fix: Do not pass is_first_microbatch to TE when it is not maintained

[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

Add is_first_microbatch_tracked(config) next to set_is_first_microbatch, returning whether the flag is actually maintained for this config, and pass None to TE when it is not. TE then falls back to its own accumulate decision rather than acting on a value nobody updates.

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.

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.