Fine-Tuning Loss Masks: Verify Which Tokens Train Before Spending Compute

Before launching supervised fine-tuning, inspect the labels produced by the actual tokenizer, chat template and collator. A falling loss does not prove that the intended responses are being trained. Require role-level token counts, preserved reply endings and deliberate failure tests before allocating training compute.

Consider a maintenance assistant that receives a system policy, a pump-status request, a tool-like assistant instruction, a tool result and a final recommendation. You want it to learn the assistant’s approved behaviour, not reproduce the operator’s questions or the tool’s measurements. That objective is a data-processing contract. Neither a role field in the original JSON nor an attention mask proves that the contract survives tokenization and batching.

This article tests that narrow contract with synthetic English and German conversations. It does not report customer data, a production training run or an improvement in model quality. If the decision to train is still open, start with the fine-tuning, RAG and prompting decision framework: a correct objective cannot fix a use case that actually needs fresher retrieval.

Attention and supervision answer different questions

For a causal language model, input_ids are the tokens presented to the model. The attention mechanism controls which context positions can be used; an ordinary padding attention_mask distinguishes real input from batch padding. Labels identify the prediction targets. In the tested path, labels set to -100 are ignored by cross-entropy. User text can remain visible as context while contributing no direct target loss.

Do not remove all user tokens from attention merely because you do not want to train their reproduction. The model needs the request to predict the response. Conversely, an attention value of one does not imply that a token should be a supervised target. Inspect the post-collation batch, not only the source examples.

Causal alignment matters too. In the ordinary Transformers causal-LM contract used here, labels initially align with input positions; the model’s loss pairs logits at position t with the label at t+1. Do not pre-shift labels and then let the model shift them again. Count usable targets after that shift, and fail closed when the denominator is zero.

Choose the objective before choosing a flag

TRL’s versioned SFTTrainer documentation distinguishes conversational language-modeling datasets from prompt-completion datasets. assistant_only_loss selects assistant spans when the template can return their masks. completion_only_loss selects the completion portion of a prompt-completion record. With both enabled, the effective target set is their intersection, not their union.

That distinction changes a multi-turn training example. A full conversation can supervise two assistant replies. A prompt-completion record whose prompt contains the earlier turns and whose completion contains only the final reply does not train the earlier assistant response. In this fixture, completion-only also includes the final assistant role header; assistant-only excludes role headers by explicit template policy. Do not generalize that header choice to every model template.

The table is a selection guide, not a recommendation to replace a pretrained model’s native control tokens. The fixture uses an intentionally small custom template so every boundary is inspectable. For a real checkpoint, preserve its expected format and test a compatible training template. Generation markers are template instructions, not additional tokens that should appear in the model input.

What was actually executed

The CPU fixture pins TRL 0.23.1, Transformers 4.56.2, tokenizers 0.22.2 and PyTorch 2.9.0+cpu. It ran on Python 3.14.9. These are tested versions, not a claim that they are the latest releases. A local WordLevel tokenizer has 54 vocabulary entries, explicit role tokens and an [EOS] token deliberately reused as padding. A randomly initialized GPT-2 model has one layer, two attention heads and a 32-dimensional embedding. No pretrained weights were downloaded and no optimizer step was run.

Each language has six messages: system → user → assistant → tool → user → assistant. The English messages are “Follow approved maintenance policy”; “Check pump status”; “CALL inspect pump”; “pressure normal vibration high”; “What should I do next”; and “Request operator review before intervention”. The German equivalents are “Befolgen Sie freigegebene Wartungsregeln”; “Prüfen Sie den Pumpenstatus”; “AUFRUF prüfen Pumpe”; “Druck normal Schwingung erhöht”; “Was soll ich als Nächstes tun”; and “Fordern Sie vor Eingriffen eine Bedienerprüfung an”. These are plain synthetic strings, not real tool-call JSON or an industrial operating procedure.

The policy supervises assistant content and its genuine EOS, but not role headers, system/user/tool text or padding. The tokenizer uses WhitespaceSplit and a vocabulary constructed from the corpus plus special tokens. The assistant branch places content and EOS inside the generation span; its following separating space is outside. add_generation_prompt is false. The trainer prepares both the full-conversation and final-response prompt-completion datasets, and its actual collator constructs the labels.

Expected versus actual supervised tokens

The table reports positions in the padded batch, including role headers and genuine message-ending tokens in their respective role totals. “Expected” comes from the independent role/header policy; “actual” counts labels different from -100. Padding is counted separately. These are two synthetic examples, not language-wide token-efficiency measurements.

The full-conversation assistant-only objective yields 10 English and 12 German targets. Final-response completion-only yields 7 and 9; intersecting it with assistant-only yields 6 and 8. The difference is observable, even though all three configurations receive the same conversation content. Compare objectives using their target sets and denominators, not a lower aggregate loss alone.

A tiny forward pass produced loss 4.020492. An independently computed cross-entropy over logits excluding the final position and labels excluding the first position matched it, with 22 supervised targets. This checks alignment and masking arithmetic only. A finite loss from a random model is not evidence of learning, safety or business value.

A boundary bug the fixture actually caught

An initial template put the separating space after the assistant EOS inside the generation span. With this pinned tokenizer/template path, the returned assistant mask included the following tool role header. The role-level assertion failed even though the rendered conversation looked correct. Moving the separator outside the generation span restored the intended mask; a retained negative test recreates the leak while confirming that input_ids remain identical.

This is a fixture-specific observation, not a claim that every Hugging Face chat template has this defect. It demonstrates why comparing rendered text or token IDs alone is insufficient: two templates can produce identical model input and different supervised targets. Another regression test moves EOS outside the intended assistant span and requires the mask change to be detected.

The Transformers chat-template guide also warns against duplicating special tokens. When formatting to text first, tokenize with add_special_tokens=False; the fixture checks equivalence with direct template tokenization and requires exactly one genuine EOS per message. Adapt that invariant to the checkpoint’s actual end-of-turn and end-of-sequence semantics rather than assuming every architecture uses one shared ending token.

Negative tests that must survive refactoring

All-masked example: replace the assistant mask with zeros. The fixture’s validation rejects the example because no shifted target remains. Do not rely on a backend to produce a useful error; an empty supervised set can create undefined or otherwise uninformative loss behaviour.

Truncation before the reply: cut the token sequence before the first supervised token and require rejection. Truncation inside a reply: a one-token reply prefix still has a nonzero denominator, so the fixture separately detects the missing ending boundary. Decide whether partial responses are allowed by policy; do not let maximum sequence length decide silently. Track rejected and truncated examples by language and source.

EOS equals padding: the tested TRL collator preserves supervised genuine EOS while ignoring batch padding. The generic Transformers causal language-modeling collator, used with the same shared token ID, masks genuine EOS as well. That is an executed contrast, not proof that every collator has this behaviour. The data-collator reference describes distinct task-specific collators; a shared class name is not an interchangeable contract.

Packing: a separate padding-free collator check concatenates the two examples with explicit sequence lengths. It verifies position resets and ignored boundary-start labels. It does not run packed attention through a model and does not prove cross-example attention isolation. A position reset is metadata, not an isolation mechanism by itself. Test the actual attention implementation and packing strategy before enabling packing in production.

Turn the fixture into a preflight gate

First, freeze the model revision, tokenizer files, template hash, dataset snapshot, library lockfile, role policy, truncation policy and packing configuration. Store a token trace containing position, decoded token, role, input ID, attention value, label and expected supervision for synthetic or appropriately protected samples. Do not export customer prompts into unrestricted CI logs.

Second, assert the target contract after preprocessing and collation. Require nonzero shifted targets for accepted examples; no supervised system/user/tool positions under this policy; exact agreement between expected and actual labels; preserved approved reply endings; ignored padding; and no unexpected special-token duplication. Aggregate accepted, rejected and supervised-token counts by language and task. A zero-error aggregate can still hide a missing segment if its denominator is zero.

Third, run mutation tests on every tokenizer, template, collator or library change. Include the all-masked, cut-reply, EOS/padding and same-input/different-mask cases. Pinning a package name is not enough: retain the resolved environment and tokenizer/template artifacts. If the production template cannot return reliable assistant spans, stop and resolve that incompatibility rather than claiming assistant-only training from a configuration flag.

Fourth, run the tiny aligned-loss check and then the intended runtime’s packing/attention checks. Only afterward start a bounded training pilot. The next gate is independent held-out task evaluation, followed by the quality and serving controls for quantisation and adapters. Mask correctness is upstream of those release decisions, not a substitute for them.

What this gate cannot prove

Correct labels do not establish training rights, representative data, useful learning or safe actions. The synthetic WordLevel tokenizer does not exercise a production subword vocabulary, native tool serialization or every multilingual boundary. The small CPU forward pass does not measure GPU throughput, memory use or training cost. Legal and privacy questions about the proposed training data require the responsible reviewers.

If your team is preparing an expensive fine-tuning run, begin with a short objective-validation review: one approved role policy, a pinned batch trace, mutation tests and a clear rejection report. AI implementation consulting can help turn those artifacts into a repeatable preflight gate. The useful outcome is a training job whose target tokens are explained before compute is spent—not merely a curve that slopes downward.