Mask the intended assistant prediction targets

For a causal language model, a target token is predicted from preceding visible context. User and system tokens can remain in the attention context while their own prediction targets contribute no supervised loss. Padding commonly contributes neither useful context nor targets under the collator's defined behavior. Keep attention masking and label masking separate in the data structure. Masking an input position without considering the model's causal shift can accidentally supervise the wrong boundary token.

Define the mask over actual prediction targets after the trainer's alignment convention. Do not blindly search for the printed word assistant in text: it can appear inside user content, and control markers vary across templates. TRL 0.24.0 supports assistant-only loss when the conversational template provides assistant-token-mask information through its documented generation regions. That is a version-specific capability with template requirements, not a guarantee that every tokenizer or custom template yields a correct mask automatically.

Decide whether the assistant's end-of-turn token is supervised. Including it can teach the intended stopping behavior; excluding it defines a different objective and may be intentional in a specialized setup. Role-prefix tokens can also be inside or outside the supervised region depending on the exact template. Document the chosen target policy and test it. Do not assume a single universal EOS rule across model families or let padding share an ID in a way that silently masks real stopping targets.

A global mean over all selected tokens gives longer responses more weight than an equal mean of per-example means. Gradient accumulation and distributed averaging can introduce further denominator differences when microbatches contain unequal target counts. If the intended objective is a global token mean, sum compatible loss numerators and selected-token denominators under the trainer's supported reduction rather than averaging arbitrary local means. Zero-target examples need rejection or explicit handling so division by zero is not hidden.

Print or inspect a small batch with token IDs, labels, and masks at the exact trainer boundary. Check every assistant turn, user/system spans, padding, and terminal tokens. Run a tiny manually computed loss case before a long training job. During training, record selected target tokens and loss by task or response length so a lower aggregate loss cannot merely reflect a changed mixture or denominator. A fully masked answer may appear in context without teaching the intended response.

The function lab is a limited normalization contrast. Reset and set both a=1, x=1, predict y=1 and slope 2, then inspect. Reset and re-set a=0.5, x=1 to predict y=0.5. This exact quadratic changes a coefficient, not a language-model loss mask. It has no token probabilities, causal shift, or batches. Use the comparison to explain that scaling an objective changes optimization magnitude, while correct supervision also requires selecting the intended target positions rather than merely shrinking every term.