[bug] gpt_step does not pass loss_mask to the model, so the MTP loss (and its gradient) is computed over padding

Repo: NVIDIA-NeMo/Megatron-Bridge Component: src/megatron/bridge/training/gpt_step.py
Affects: every recipe trained through gpt_step.forward_step (and deepseek_v4_step) with mtp_num_layers > 0. Observed on 26.08 (Megatron-Bridge v0.6.0), still present on main.

Problem

_forward_step_common builds

forward_args = {
    "input_ids": tokens,
    "position_ids": position_ids,
    "attention_mask": attention_mask,
    "labels": labels,
}
...
output_tensor = model(**forward_args)

and only uses loss_mask later in _create_loss_function. GPTModel.forward accepts loss_mask and hands it to megatron.core.transformer.multi_token_prediction.process_mtp_loss, which does

if loss_mask is None:
    loss_mask = torch.ones_like(mtp_labels)

So the language-model loss is masked (Bridge's own loss_func) but the MTP loss and the MTP gradient are computed over every position, including padding and, for answer-only SFT, the prompt. Only the return_schedule_plan branch passes loss_mask (model.build_schedule_plan(..., loss_mask=loss_mask)). Megatron-LM's own pretrain_gpt.py passes loss_mask to the model.

Logs / observed effect (DeepSeek-V4-Pro SFT, SQuAD, 32 nodes H100, TP1/PP32/EP8, MBS 1)

The SFT dataset pads each micro-batch to the longest sample (pad id 1 = eos). A 189-token sample padded to 304:

iteration        1/      10 | lm loss: 3.050079E+00 | mtp_1 loss: 1.580068E+01 | grad norm: 85.956

Per-position cross-entropy of the MTP head recomputed in fp32 from the logits (probe):

mtp per-token CE: mean=21.606 median=1.950 p90=54.812 p99=72.253 max=85.188 frac(CE>20)=0.380
mtp worst pos=189 ce=85.19 label=1 pred=0 logit_true=9.31 logit_max=94.50
mtp worst pos=190 ce=78.44 label=1 pred=0 logit_true=8.56 logit_max=87.00

All positions with CE > 20 are padding (label = pad/eos, prediction = token 0 with logit ≈ 90). 96.5 % of the MTP loss mass came from padding. With answer-only loss the real loss_mask had 2–5 valid tokens per sample, so the padded positions dominated completely. The gradient of those positions flows into the shared embedding and output layer.

Expected behavior and verified fix

Pass the mask through:

forward_args = {
    "input_ids": tokens,
    "position_ids": position_ids,
    "attention_mask": attention_mask,
    "labels": labels,
    "loss_mask": loss_mask,
}

Same run with the mask injected (wrapper that stashes loss_mask in get_batch and injects it into process_mtp_loss when None):

iteration        1/      10 | lm loss: 3.042851E+00 | mtp_1 loss: 4.416957E+00 | grad norm: 67.528
validation loss at iteration 10 | lm loss value: 3.512455E-01 (was 3.586473E-01 without the fix)

lm loss unchanged (3.043 vs 3.050), mtp_1 loss 15.8 → 4.42, step-1 grad norm 86 → 67.5, validation loss improves. A 64-node run (GBS 128) reproduces: mtp_1 loss 4.13 → 0.36 over 10 steps, validation 0.266.

Notes for the fix:

• deepseek_v4_step.forward_step reuses _forward_step_common, so one change covers both.
• Bridge's loss_func should keep applying the mask to the LM loss as today. process_mtp_loss rolls the mask itself for the MTP depth, so the input-aligned mask returned by get_batch is the right tensor to pass.
• For packed sequences (packed_seq_params), process_mtp_loss also handles the roll per segment. We ran CP=2 + offline-packed THD with the mask injected without issues.

Environment

nvcr.io/nvidia/nemo:26.08 (Megatron-Bridge v0.6.0 c932511, Megatron-LM .dev.commit 43e45e13 for DSv4), 8×H100 80 GB nodes. Full write-up with logs: internal README §5.4-② (can be shared on request).

Affected area

training (gpt_step.py; deepseek_v4_step.py reuses _forward_step_common)

Regression?

Not a regression as far as we know: main builds forward_args the same way today.

A PR with the one-line change and a unit test follows.

Start in src/megatron/bridge/training/gpt_step.py at _forward_step_common and inspect how forward_args reaches GPTModel.forward; deepseek_v4_step reuses this path. Add or update the unit test mentioned in the issue, then verify that the MTP loss respects loss_mask while the existing language-model loss remains unchanged.