SFT Loss Masking

The problem, from first principles

Supervised fine-tuning shows a pretrained model examples of a prompt followed by a good response and trains it with the ordinary next-token loss. One detail decides what the model actually learns: the loss must be counted only on the response tokens. If prompt tokens are counted too, the model spends capacity learning to reproduce questions, and it gets credit for copying text that was handed to it. Padding must be excluded for a different reason, because it carries no information at all. This question builds the mask and the masked loss, including the one-position shift between a token and the position that predicts it.

From theory to code

Implement response_mask and sft_loss.

Constraints

• A batch has B sequences padded to length T. Sequence b has prompt_lens[b] prompt tokens followed by response tokens up to seq_lens[b], then padding.
• response_mask(prompt_lens, seq_lens, T) returns a (B, T) boolean array that is True exactly at token positions t with prompt_lens[b] <= t < seq_lens[b].
• logits has shape (B, T, V): logits[b, t] predicts token t + 1. sft_loss(logits, token_ids, prompt_lens, seq_lens) therefore compares logits[:, :-1] with token_ids[:, 1:] and the mask mask[:, 1:].
• Return the mean negative log-likelihood over the masked target tokens, using a numerically stable log-softmax.
• Every sequence has at least one response token.

Hints