Next-Token Loss with Padding and Prompt Masks

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