Leverage TE wgrad accumulation fusion in MFSDP v2

[PACTHEMAN123]

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

What does this PR do?

Leverage Transformer Engine wgrad accumulation fusion in MFSDP v2 so supported TE/MCore linear layers write weight gradients directly into MFSDP v2 reduce-scatter input buffers.

The implementation follows the MFSDP runtime schedule described in megatron/core/distributed/fsdp/src/docs/runtime_schedule.md.

Issue tracking

Linked issue: Fixes #6942

Contribution process

Pre-checks

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

The test implementations used for validation are intentionally kept outside this production-only PR. The validation results are reported below.

Implementation

• Allow supported TE/MCore fused linear layers to write directly into preallocated MFSDP v2 reduce-scatter input buffers.
• Allocate and zero fused-wgrad destinations on the reduce-scatter stream.
• Preserve accumulation semantics when the same parameter is used by multiple fused backward paths, including MTP and the main GPT output layer.
• Pass explicit accumulation behavior through MCore and Transformer Engine wgrad GEMM paths.
• Avoid TE first-microbatch overwrite behavior for MFSDP-managed parameters.
• Keep ordinary autograd gradients and tied input/output contributions supported.
• Preserve backward all-gather prefetch ordering and stream/event dependencies from the MFSDP runtime schedule.

MFSDP v2 fused wgrad accumulation rejects unsupported combinations:

• Transformer Engine older than 2.10
• CPU offloading
• overlap_moe_expert_parallel_comm
• overlap_dispatch_backward_with_experts_wgrad
• use_transformer_engine_op_fuser
• NCCL-UB/symmetric-memory communication

The delayed expert-wgrad overlap combination is intentionally rejected because MFSDP v2 does not currently implement the delayed post_wgrad_grad_acc_hook callback required by MoELayer.backward_dw().

Validation

• Native 2-layer GPT with 2 MTP layers completed baseline and fused training runs.
• Main LM loss, both MTP losses, and gradient norms matched at printed precision.

Benchmark

Configuration:

• 8x NVIDIA H20
• 8.03B-parameter GPT model
• 32 layers, hidden size 4096, FFN size 14336
• Sequence length 2048
• MFSDP v2 with pure data parallelism
• Two microbatches per rank
• Global batch size 16
• 15 iterations, with the first 5 excluded as warmup
• Reported step time is the trimmed mean of iterations 6-15

Metric	Baseline	Fused	Delta
Step time	1674.150 ms	1653.875 ms	-1.211%
Aggregate throughput	19572.9 tok/s	19812.9 tok/s	+1.226%
Peak allocated memory	24518.26 MB	24746.81 MB	+0.932%
Peak reserved memory	30656 MB	30072 MB	-1.905%
Allocated after an iteration	11586.54 MB	11588.01 MB	+1.47 MB

The first-step language-model loss was 12.58627 for both paths. Both runs completed all 15 iterations without OOM or NCCL errors, and all eight GPUs were released after completion.

Limitations

MFSDP v2 continues to use the static backward-order prefetch design described in runtime_schedule.md.

Delayed expert wgrad overlap and CPU offloading remain unsupported when fused wgrad accumulation is enabled.

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

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[wujingyue]

Thanks for the PR!

For the first PR, don't touch mcore_fsdp_adapter or things outside megatron_fsdp.

Add unit tests.

Simplify -- quite some slops at this moment.