What does this PR do ?

dev PR: #2800

This MR adds HSDP support for Expert Parallelism. With this change, DeepSeek‑v3 and other MoE models can be trained using M‑FSDP as well as HSDP+EP. HSDP enables better use of local bandwidth within a larger nvlink domain.

We ran convergence tests on the DeepSeek Proxy model. model. The experiments were based on the latest dev branch, with an additional patch that fixes MoE checkpoint save/load.

To validate the correctness of HSDP+EP, we loaded an FSDP+EP checkpoint at step 100 and resumed training with HSDP+EP. The loss curves and gradient norms matched closely.

For example, to use HSDP=2 and EP size=2, you need to add these argument to your script

--num-distributed-optimizer-instances 2 \
--outer-dp-sharding-strategy optim
--expert-model-parallel-size 2 

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share a design doc with the team. If you're unsure what's the best way to do so, contact the @megatron-oncall.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• 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

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.