What does this PR do ?

PR for main: #3795

Problem

In MoE models, the expert weight gradient (wgrad) computation during backward is serialized on the main CUDA stream. This blocks the data gradient (dgrad) from flowing to earlier layers until the expert wgrad finishes, even though there is no data dependency between them. The result is wasted GPU cycles — earlier layers' backward pass sits idle waiting for expert wgrad to complete.

With FSDP, this is further compounded because the gradient reduce-scatter for expert parameters is also blocked on the same critical path.

Solution

This PR introduces a new flag --overlap-dispatch-backward-with-experts-wgrad that separates the expert wgrad computation from the main backward stream:

1. Two autograd functions are inserted into the MoE layer's forward graph:

   • _RecordExpertDgradCompletion — placed before the expert computation; during backward, it records a CUDA event once the expert dgrad is done.
   • _RegisterDelayedWgradForExperts — placed at the dispatch boundary; during backward, it waits on the dgrad event, then launches backward_dw() on a dedicated CUDA stream, and synchronizes back to the main stream before proceeding.

2. FSDP integration — When used with MegatronFSDP, expert parameters are marked with _fsdp_delay_grad_reduce = True so the normal post-accumulate-grad hook skips them. A callback is registered via register_process_expert_grads_fn() that triggers the FSDP reduce-scatter for expert parameters only after the delayed wgrad computation completes.

3. TE GroupedLinear is configured with delay_wgrad_compute=True, which tells Transformer Engine to skip wgrad during the normal autograd backward and instead wait for an explicit backward_dw() call.

How to enable

--overlap-dispatch-backward-with-experts-wgrad

Requirements:

• Transformer Engine >= 2.3.0
• moe_grouped_gemm enabled (not legacy grouped gemm)
• Mutually exclusive with --delay-wgrad-compute (the existing A2A-overlap-based delay)
• Mutually exclusive with --overlap-moe-expert-parallel-comm

Works with both FSDP and 3-D parallelism (TP/EP/PP).

What is achieved

The expert wgrad computation runs on a separate CUDA stream, overlapping with the EP communication within the same transformer layer. This reduces the wall-clock time of the backward pass without changing numerical results — the feature is bit-exact with the non-delayed baseline (verified by unit tests comparing per-step losses and final weights over multiple optimizer steps).

Changes

Test plan

• Unit test: test_delay_wgrad_compute_for_te_grouped_gemm — full-model training loop (forward → backward → optimizer) comparing delayed vs. non-delayed across num_layers × shared_experts × dispatcher_type × fp8_flag
• Unit test: test_delay_wgrad_compute_for_te_grouped_gemm_with_fsdp — same comparison with MegatronFSDP wrapping (fully_shard_model + fully_shard_optimizer), verifying the deferred reduce-scatter path

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 the @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.