Next-Token Loss with Padding and Prompt Masks

medium
Checkpoint challenge
Solve this to show you know the material and jump straight to Unit 19. Everything before it counts as cleared.

Training a language model means predicting token t+1 from position t, so logits and labels are shifted by one. Two masks matter:

  • padding: batches are padded to the same length, and pad positions must not count
  • prompt masking (for supervised fine-tuning): only the response tokens should be trained on, not the prompt

The convention (from Hugging Face) is a labels array with -100 at positions to ignore.

Implement, with logits: (B, T, V) and integer arrays (B, T):

  • make_labels(input_ids, attention_mask, prompt_lengths=None): a copy of input_ids with -100 wherever attention_mask == 0, and also at the first prompt_lengths[b] positions of row b if given
  • lm_loss(logits, labels): shift (use logits[:, :-1] to predict labels[:, 1:]), then return the mean cross-entropy over targets that aren't -100, using a stable log-softmax. Return 0.0 if there are none.
  • token_accuracy(logits, labels): the fraction of non-ignored shifted targets where the argmax is correct
#language-models#loss-functions#fine-tuning2 visible + 1 hidden tests
solution.py · Python 3.13
Loading editor…
Run your code to see results. ⌘↵ to run · ⌘⇧↵ to submit
Next-Token Loss with Padding and Prompt Masks · TechGradify