Fix TE grouped MLP fused main-grad setup

[Wong4j]

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

What does this PR do ?

This PR exposes Megatron-FSDP main-grad buffers to the TE grouped MLP op-fuser path before fused execution.

Issue tracking

Linked issue: N/A. This is a small compatibility bug fix for Megatron-FSDP with the TE grouped MLP op-fuser path.

Changes

• Initialize param.main_grad from param.get_main_grad() for grouped MLP fused parameters when needed.
• Keep overwrite_main_grad enabled for the fused path.
• Add focused unit coverage for the hook behavior and real TE construction path.

Validation

• python3 -m compileall -q megatron/core/transformer/moe/experts.py tests/unit_tests/transformer/moe/test_grouped_mlp.py
• On a GB200 node:
  • targeted main-grad / TE GroupedMLP unit tests passed
  • 1-node / 4-GPU mock-data MoE + TE op-fuser + Megatron-FSDP e2e training reached iter 3/3
  • skipped iterations: 0; NaN iterations: 0; Slurm exit: 0:0
• CHECK_ONLY=true tools/autoformat.sh passed for this PR branch

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

[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.

[Phlip79]

/ok to test 05236b6

[yaox12]

/claude strict-review

[claude]

Strict review passed — no significant issues found. LGTM

Traced the full data/dtype flow for this fix:

Correctness (verified): The change fills a genuine gap. Megatron-FSDP's _bucket_forward sets overwrite_main_grad=True but does not pre-populate param.main_grad; the non-fused wgrad paths (tensor_parallel/layers.py:586, core/utils.py:1195) set it lazily inside wgrad_compute, but the TE op-fuser path has no equivalent hook, so param.main_grad would be None when TE's fused wgrad accumulation writes into it. Exposing it in the pre-forward hook mirrors the reference FSDP pattern (megatron_fsdp.py:652,670).

Backward compatibility (verified safe): get_main_grad/overwrite_main_grad exist only on FSDP-managed params, so the getattr(..., None) is not None and hasattr(...) guards leave the DDP + fuse_wgrad_accumulation path completely untouched — no default or behavior change for existing users.

Details that are correct:

• is None identity check (not truthiness) avoids the ambiguous-tensor-bool error.
• Under gradient accumulation, main_grad is reset to None only at grad-reduce time (reset_param_main_grad), so the guard is a no-op on later microbatches and the same bucket is reused — consistent.
• parameters(recurse=False) is correct: TEGroupedLinear registers its weights/biases directly on the module.
• No new direct parallel_state.get_*_group() reads; all new identifiers have live use paths.

The unit test cleanly validates the guard logic (uses pytest.fail to assert get_main_grad is skipped when main_grad is already set, and that non-fuse_wgrad params are left alone).

Risk level: low.

[kvareddy]

@fanshiqing can you please take a look at these changes?

[Wong4j]

@aklife97 @yashaswikarnati Could you please help review these changes?

[ericharper]

/ok to test d000dad

[Wong4j]

@kvareddy @lmcafee-nvidia Could you please help review these changes?

[yaox12]

/ok to test fba3816

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

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