fix(mtp): skip disabled loss and handle empty targets
Summary
This follows up on the MTP loss handling in the remote-code model. Passing labels currently runs MTP even when mtp_loss_scaling_factor is zero. If the MTP targets are all ignored after shifting, its loss becomes NaN and 0 * NaN contaminates the main loss.
This change:
- Requires a positive loss scaling factor before running MTP.
- Uses the existing loss function's summed-loss normalization with a clamped target count, so an empty MTP target set contributes zero. Caller-provided
num_items_in_batchis preserved. - Iterates directly over the MTP hidden states, avoiding overwriting the list when more than one MTP layer is configured.
No changes to config.json, weights, cache handling, or the inference output policy. When MTP is disabled by a zero scaling factor, mtp_loss and mtp_logits are None, including when labels are supplied. The tuple layout is unchanged: supplying labels adds the usual leading loss, not a separate MTP field.
Why use num_items_in_batch for MTP?
The positive-scaling check skips disabled MTP computation. This normalization handles a different case: MTP is enabled, but shifting the labels leaves no supervised targets at a particular depth.
ForCausalLMLoss normally uses mean cross entropy. When all effective targets are ignored, that mean is undefined and produces NaN. Supplying num_items_in_batch selects the existing summed-cross-entropy path and divides by that value instead. The fallback therefore computes:
MTP loss = summed cross entropy / max(number of effective target tokens, 1)
With valid targets, this is mathematically the same token-mean loss as before. With no valid targets, the summed loss is zero and the denominator is one, so the auxiliary loss and its gradients are zero rather than NaN. Here the count refers to supervised target tokens, not the number of sequences in the batch.
For example, a single sequence with labels [-100, 3, -100, -100] has a valid main next-token target. After the first MTP roll, the labels become [3, -100, -100, -100]; the additional causal shift performed by the existing loss function leaves only ignored targets. Counting before that final shift would incorrectly report one valid MTP target.
The fallback counts the targets actually consumed by the existing loss function, or uses the explicit shift_labels if supplied. It does not change target alignment. It only supplies a denominator when num_items_in_batch is absent or None; an explicit caller-provided denominator is left untouched to preserve Trainer/custom-loop normalization. The kwargs are copied for the auxiliary call, so neither the main loss nor the caller's kwargs are changed. Counting and clamping use tensor operations, without branching in Python on whether the target set is empty.
This is limited to the MTP auxiliary loss; it does not redefine the main loss for a batch with no valid main targets.
Validation
Tested on AIStudio with Transformers 4.56.2 and PyTorch 2.8.0a0+34c6371d24.nv25.08, using the actual remote-code model with small random weights on CPU (two MLA decoder layers and an MTP layer).
- Reproduced the unnecessary MTP forward and NaN total loss on the current main before applying the patch.
- Verified that zero scaling skips the MTP forward and returns exactly the main loss.
- Verified finite loss and backward gradients when positive scaling encounters empty MTP targets.
- Compared valid-target losses and gradients against main, including mixed masked batches, explicit shifted labels, and caller-provided normalization.
- Checked bit-exact inference logits for prefill and cached decoding.
- Checked 32 combinations of cache, hidden-state, attention, and router-output flags across zero and positive scaling; tuple output only gains the standard leading loss.
- Checked configurations with zero and two MTP layers, including targets exhausted at the second MTP depth.
These are focused control-flow/loss regression tests, not a full-checkpoint or KDA-kernel validation.
shift_labels in kwargs. If a caller passes shift_labels, it is forwarded into the MTP loss call. ForCausalLMLoss, it then skips its own shifting and uses those labels instead of the rolled MTP labels. Every MTP depth would then train against the main head’s targets. This predates the PR, but the new count now treats it as valid and your tests list “explicit shifted labels” as a covered case. Consider mtp_loss_kwargs.pop("shift_labels", None), or we can confirm it’s intentional.
The rest part LGTM.
Thanks for catching this — you’re right. The baseline already forwards the main head’s shift_labels into each MTP loss call, so this is a pre-existing target-alignment bug, not a regression introduced by this patch. Reusing the main head’s targets at every MTP depth is not intentional.
The intended alignment is that the main head at position t predicts token t+1, the first MTP depth predicts t+2, the second predicts t+3, and so on, with out-of-range and masked targets ignored. Without explicit shift_labels, the cumulative label rolls followed by the causal loss function’s own shift are intended to produce that alignment. When the main head’s explicit shift_labels is forwarded, however, it overrides the rolled labels and breaks that alignment.
I should also clarify the validation wording: the explicit-shift-label test checked loss/gradient parity with the baseline, not correctness against independently constructed MTP targets. That was insufficient to establish correctness for this case.
To keep this PR minimal, I’d like to retain its scope: skip disabled MTP computation, make the auxiliary loss finite when no targets remain, and fix iteration over multiple MTP layers. I’ll examine the inherited target-alignment issue in a separate follow-up PR, including explicit shifted labels, multiple depths, masking/boundaries, and loss/gradient checks against independently constructed targets. This PR should not be read as resolving all existing MTP training issues.
Would that scope work for you, so we can land the minimal fixes first? Thanks again for the careful review!
LGTM