[ZhiyuLi-Nvidia]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

TEGroupedLinear.is_first_microbatch resets to True after checkpoint resume, causing the first shared-layer call to overwrite main_grad instead of accumulating into it. When an MTP layer is reused across depths, this discards one depth’s gradient and makes resumed training differ from uninterrupted training.

• Before fix

Uninterrupted training retains:
is_first_microbatch = False

forward:   depth 1 [False] → depth 2 [False]
backward:  depth 2 accumulates → depth 1 accumulates

main_grad = grad(depth 2) + grad(depth 1)  ✓

• What is not expected after ckpt resume:

After checkpoint resume, the process-local flag resets:
is_first_microbatch = True

forward:   depth 1 [True] → depth 2 [False]
backward:  depth 2 accumulates → depth 1 overwrites (beta=0)

main_grad = grad(depth 1)  ✗
                         └─ grad(depth 2) is discarded

This only becomes destructive when the same layer is invoked multiple times, such as MTP repeated-layer sharing.

• Fixed
  Restore or initialize the state deterministically so resume matches uninterrupted training:

is_first_microbatch = False

forward:   depth 1 [False] → depth 2 [False]
backward:  depth 2 accumulates → depth 1 accumulates

main_grad = grad(depth 2) + grad(depth 1)  ✓

Why no existing test catches it

A run-to-run A/A check cannot detect this. Both fresh processes make the same mistake and agree. It only appears between a process that has trained N steps and one that has just loaded a checkpoint, so it takes three separate processes to observe.

Fix

Two independent reasons the flag must not reach TE, both resolved to None:

1. Nobody maintains it. is_first_microbatch_tracked(config), added next to set_is_first_microbatch, returns whether the flag is actually re-armed for this config. Outside fp8/fp4/kitchen/quant_recipe it is set once at construction and cleared by the first forward, so it is stale for the rest of the process.
2. It is maintained, but it means the wrong thing. A repeated MTP layer is invoked once per depth, so its first forward is paired with its last backward. The True lands on the last wgrad GEMM and overwrites main_grad with beta=0, discarding what the later depths accumulated. mark_is_first_microbatch_unsafe tags the shared layer's TE submodules at construction (only when mtp_num_layers > 1).

Configuration	Before	After
BF16 / non-quantized	stale True/False	None — always accumulate
FP8 / FP4 / Kitchen, ordinary layer	maintained True/False	unchanged; the cache optimization is kept
FP8 / FP4 / Kitchen, repeated MTP layer	maintained True/False	None — always accumulate
Transpose cache disabled	None	None

None rather than False: both accumulate, but False also tells TE to reuse its cached quantized weights, which would go stale across optimizer steps now that nothing re-arms the flag. The cost is re-quantizing one MTP layer per forward, and --fp8-param-gather removes even that by quantizing in the optimizer instead.

The three copies of the gate are now one _resolve_is_first_microbatch helper so they cannot drift apart.

Thanks @zhongbozhu for catching that the first version stopped at reason 1.

Verification

Three-leg check — train-through to step 10 / save at 5 / resume from 5 — all legs in one allocation, comparing checkpoint shards byte for byte:

	result
2 nodes	A/A identical; RESUME identical, all 7 state shards byte-identical at iter_10
64 nodes	255/256 shards byte-identical (the 256th is rank 0's args blob, which differs by construction)
without this fix	diverges in grad-norm one step after the resume, loss still matching

Confirming arm: disabling gradient-accumulation fusion also makes the divergence vanish, which is what the mechanism predicts — with fusion off the and short-circuits and is_first_microbatch can no longer select overwrite.

The BF16 leg above is what the cluster runs; the quantized repeated-MTP leg is covered by unit test rather than a resume run. test_repeated_layer_opts_out_of_is_first_microbatch builds a TE-spec MTP block both ways under fp8, and asserts the resolver answers None for every flag-carrying submodule when the layer is shared and True when it is not.

Issue tracking

For PRs from open-source community contributors:

• New features: a linked issue is required. Please open a feature request and reference it here before submitting the PR.
• Small updates (bug fixes, minor improvements): a linked issue is recommended and will accelerate the PR review process.

Linked issue:

Contribution process

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!

All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.

Step 1: Mark PR as "Ready for Review"

1. When your PR is ready, click Ready for Review.
2. An oncall reviewer is auto-assigned and expert reviewers are notified based on your changes.
   • Some PRs may jump straight to step 2. This is determined by .github/CODEOWNERS.

⚠️ Only mark as ready once merge-conflicts are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

Step 2: Final Review

For PRs that change megatron/core, once all expert reviewers have approved, the Final Review label is applied automatically and final reviewers are assigned.

For PRs outside megatron/core, this step is skipped.

Step 3: Approved

Once all required reviewers have approved, the Approved label is applied automatically.

Merge

Any member of mcore-engineers will be able to merge your PR.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[ZhiyuLi-Nvidia]

/ok to test 293a85b

[zhongbozhu]

Good Catch~

With mtp_num_layers=2 and mtp_use_repeated_layer=True, logical MTP layers 0 and 1 invoke the same physical TE layer. Their forwards receive True then False, while backward runs in reverse, causing layer 1’s gradient to accumulate first and then be overwritten by layer 0. Currently, is_first_microbatch controls both whether TE refreshes/reuses the quantized weight cache during forward and whether fused wgrad overwrites or accumulates into main_grad during backward.

This issue only occurs when we are reusing one layer for 2 layers in mtp_use_repeated_layer.

In this PR, it's changed to the following condition, which is still not fully fixed:

Configuration	Before PR	After PR	Status
BF16/non-quantized	Stale True/False state passed to TE	is_first_microbatch=None	Fixed: Resume and repeated-layer gradients always accumulate correctly
FP8/FP4/Kitchen	Maintained True/False passed to TE	Maintained True/False still passed to TE	Not fixed: Quantized repeated-layer MTP can still overwrite an earlier depth's gradient
Transpose cache disabled	is_first_microbatch=None	is_first_microbatch=None	Already safe: No behavioral change

To summarize:

• For a regular layer, the first microbatch’s backward can safely overwrite main_grad because it is the first backward pass to write to that buffer.
• For repeated MTP, however, the first forward corresponds to the last backward, so it must not overwrite main_grad. If this optimization is retained, the last MTP depth’s backward should overwrite, while earlier depths should accumulate. The simplest and safest approach is to accumulate for every depth.

Adding yet another toggle to TE just for this is not recommended, so I suggest just passing None to those repeated MTP layers even if it means hurting performance because of repetitive weight quantization. One WAR is to just turn on --fp8-param-gather so that we are quantizing in optimizer instead of TE fwd pass, so then the TE weight quantization cache is essentially doing nothing.

[zhongbozhu]

Left commens~ please take a look

[ZhiyuLi-Nvidia]

Thanks @zhongbozhu — you were right that the first version only closed the BF16 face of this. Taking your suggestion: repeated MTP layers now get None unconditionally, including under fp8/fp4/kitchen.

MultiTokenPredictionBlock.__init__ tags the shared layer's TE submodules with is_repeated_layer, and the three TE call sites now go through one _resolve_is_first_microbatch helper instead of three copies of the condition:

if (
    module.disable_parameter_transpose_cache
    or not is_first_microbatch_tracked(module.config)
    or getattr(module, 'is_repeated_layer', False)
):
    return None
return module.is_first_microbatch

No new TE toggle, as you asked.

Why None and not False. Both accumulate, but False also tells TE to reuse its cached quantized weights (cache_weight = is_first_microbatch is not None, update_ws = is_first_microbatch is None or is_first_microbatch in grouped_linear.py). If we simply stopped re-arming these modules they would sit at False forever and serve a cache that goes stale across optimizer steps — correct gradients computed from the wrong weights. None is the only value that is both accumulate-always and cache-free.

Perf cost is re-quantizing one MTP layer per forward, and --fp8-param-gather removes it, as you noted.

Coverage. test_repeated_layer_opts_out_of_is_first_microbatch builds a TE-spec MTP block both ways under fp8 and asserts the resolver answers None for every flag-carrying submodule when the layer is shared and True when it is not.

One thing I checked while scoping this, in case it saves you the question: the CUDA-graph path is unaffected. CudaGraphManager keeps its own is_first_microbatch and drives the fp8 cache through _set_skip_fp8_weight_update_tensor (cuda_graphs.py:869), and TE forces is_first_microbatch = False whenever fp8_graph_capturing() supplies that tensor — so the captured graph always accumulates and the overwrite cannot occur there regardless of this marker.

[ZhiyuLi-Nvidia]

/ok to test 0ab7f8a

[ZhiyuLi-Nvidia]

/ok to test 1d1ecae

[nemo-automation-bot]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/34568330019