SFT Examples with Loss Masking

The problem

In supervised fine-tuning (SFT) the model is trained on conversations: a prompt (the user's message) and an answer (what we want the model to say). It reads the whole conversation, but it must learn only the answer: it should not be trained to write the user's message. You will build the training arrays for one example and compute its loss.

Part 1: the arrays

Concatenate the prompt and the answer, cut at max_len tokens (cut the end), and pad with pad_id up to exactly max_len. Two arrays come out:

• input_ids: the tokens, padding included;
• labels: the same numbers, except -100 on every prompt position and every padding position. -100 is the value that PyTorch's cross-entropy ignores.

Example. prompt_ids = [5, 6], answer_ids = [7, 8], max_len = 6, pad_id = 0.

• input_ids = [5, 6, 7, 8, 0, 0];
• labels = [-100, -100, 7, 8, -100, -100].

With max_len = 3 the end is cut: input_ids = [5, 6, 7] and labels = [-100, -100, 7]. The answer is only partly kept, and there is no padding.

Part 2: the loss, with the shift

The model reads token $t$ and predicts token $t + 1$. So the logits at position $t$ are scored against the label of position $t + 1$: the logits without the last row (logits[:-1]) against the labels without the first (labels[1:]). The loss is the mean cross-entropy over the positions whose label is not -100.

Example. The first example above, 6 positions. labels[1:] is [-100, 7, 8, -100, -100]: position 0 has no target (it would predict a prompt token), position 1 must predict 7, position 2 must predict 8, the rest is padding. If the model gives token 7 probability $0.5$ at position 1 and token 8 probability $0.25$ at position 2, the loss is $(-\log 0.5 - \log 0.25)/2 = 1.04$.

The classic bug: comparing logits with labels without the shift scores position 2 against 7 and position 3 against 8: wrong tokens, and no error is raised.

What to implement

• build_example(prompt_ids, answer_ids, max_len, pad_id): return (input_ids, labels), two integer arrays of length max_len.
• sft_loss(logits, input_ids, labels): the loss of part 2. logits has shape $(T, V)$.

Requirements

• Shift by one as described: the label of position $t$ is what the model must predict after reading token $t$.
• If no position is left (everything masked), return 0.0.
• Stable log-softmax; no Python loops over positions.

Hint

Build labels with a boolean mask of "answer positions": len(prompt) <= i < len(prompt) + len(answer), cut to max_len, and use np.where(mask, input_ids, -100).