What does this PR do ?

PR for main: #3797
Based on: #3766

Support Megatron-FSDP with EP overlap 1F1B schedule

Supported Feature Combinations

# without param sharding
--use-megatron-fsdp --data-parallel-sharding-strategy "optim_grads" --overlap-moe-expert-parallel-comm
--use-megatron-fsdp --data-parallel-sharding-strategy "optim_grads" --overlap-moe-expert-parallel-comm --delay-wgrad-compute


# param sharding
--use-megatron-fsdp --data-parallel-sharding-strategy "optim_grads_params" --overlap-moe-expert-parallel-comm
--use-megatron-fsdp --data-parallel-sharding-strategy "optim_grads_params" --overlap-moe-expert-parallel-comm --delay-wgrad-compute

Problem

The EP overlap schedule (overlap_moe_expert_parallel_comm) interleaves forward and backward passes at the sub-module level (attn, mlp, moe_dispatch, moe_combine) to overlap Expert Parallel A2A communication with compute. This schedule bypasses TransformerLayer.forward() entirely, calling sub-modules directly. As a result, FSDP hooks registered on TransformerLayer — including parameter all-gather, parameter release, gradient reduction, and root pre/post backward — are never triggered, making FSDP with optim_grads_params incompatible with this schedule.

FSDP Hook Triggering: Standard vs 1F1B Overlap

Summary: The root cause of all differences is that the 1F1B overlap schedule decomposes TransformerLayer.forward() into sub-operations (attn, mlp, moe_dispatch, moe_combine) and interleaves them across microbatches. This means:

1. MegatronFSDP.forward() is never called → _replace_param_with_raw_if_needed must be explicit.
2. TransformerLayer.forward() is never called as a whole → module-level register_forward_hook / register_forward_pre_hook on the FSDP unit are not fired → _post_forward and _post_backward_release_module must be wired manually via set_fsdp_reshard_hooks.
3. No root-level output tensor in autograd graph → _root_pre_backward / _root_post_backward must be called manually.
4. Hooks that operate at the parameter level (_process_post_backward_gradients) or sub-module level (_pre_forward_param_unshard, _pre_backward_param_unshard with enable_fine_grained=True) continue to work automatically because the schedule still calls sub-module forwards, and autograd still writes to parameter .grad.

Solution

Bridge the gap between Megatron-FSDP's hook-based lifecycle and the EP overlap schedule's direct sub-module invocation by:

1. Enabling fine-grained per-sub-module all-gather hooks so each sub-module can unshard its own parameters.
2. Exposing manual handles for FSDP lifecycle events that the schedule must drive explicitly.

Code changes

combined_1f1b.py — Schedule-level FSDP integration

• Add find_megatron_fsdp() helper to locate the FSDP wrapper from the model chain.
• Call fsdp_wrapper._replace_param_with_raw_if_needed() before the schedule starts (normally done inside MegatronFSDP.forward()).
• Manually invoke fsdp_wrapper.pre_backward() / post_backward() around the combined forward+backward step since the schedule bypasses the autograd hooks that normally trigger them.
• Wire per-layer FSDP release callbacks via set_fsdp_reshard_hooks().

megatron_fsdp.py — Expose manual hook handles & support delayed wgrad

• Add skip_backward_hook parameter to _root_pre_backward so the manual caller can prevent the redundant queue_callback(_root_post_backward).
• Expose pre_backward, post_backward, post_forward_release_module, and post_backward_release_module as public attributes for external schedule use.
• Add is_delayed parameter to _process_post_backward_gradients to distinguish delayed wgrad gradient processing from the standard path, preventing deferred parameters from being filtered out by the skip_backward_post_hook guard.
• Register fine-grained backward all-gather hooks (_register_pre_backward_param_unshard_hook) on every sub-module when overlap_moe_expert_parallel_comm + optim_grads_params, so backward all-gather works at the sub-module level.
• In per-parameter register_post_accumulate_grad_hook, skip parameters with skip_backward_post_hook=True so their gradient processing is deferred to the delayed wgrad path.

mcore_fsdp_adapter.py — Auto-configure FSDP for EP overlap

• Enable enable_fine_grained_param_gather_hook when overlap_moe_expert_parallel_comm is set (for all sharding strategies), so every sub-module gets its own forward all-gather/unshard hooks.
• Assert that fsdp_double_buffer and partial CUDA graph scopes are disabled when overlap_moe_expert_parallel_comm is set.
• Assert that fsdp_unit_modules must be [TransformerLayer] when combined with optim_grads_params.

model_chunk_schedule_plan.py — Per-layer FSDP hook wiring

• Add set_fsdp_reshard_hooks() to TransformerLayerSchedulePlan that attaches post-forward and post-backward release callbacks to the correct schedule nodes (last forward node, first backward node).

fine_grained_callables.py — Node-level FSDP hooks

• Override forward() / backward() in TransformerLayerNode to invoke FSDP reshard hooks at layer boundaries.
• In backward_dw(), collect and execute delayed post_wgrad_grad_acc_hook for FSDP gradient processing, then invoke the reshard hook.
• Add set_fsdp_post_forward_reshard_hook() / set_fsdp_post_backward_reshard_hook() registration methods.
• Add parameters() iterator to _BackwardDWWrapper so FSDP can discover parameters owned by the attn and shared-expert backward-dw callables.

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.