Repository navigation
Support independent expert optimizer instances in MFSDP v2 - #7341
Autumn1998 wants to merge 1 commit into
Conversation
|
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:
See the contribution guide for more details. |
wujingyue
left a comment
wujingyue
left a comment
There was a problem hiding this comment.
Are you sure it's needed? This changes the mesh and the mesh is typically aligned with network topology -- most likely [inter-nvl, intra-nvl] for production settings.
Oh, maybe this is for supporting something like HFSDP for dense and Zero3+EP for MoE? That would make sense to me. |
Yes, I recently found that some customers require using HSDP, with the outer dimension set to ZeRO-1 and the inner dimension set to ZeRO-3 for dense part, while using a fully ZeRO-1 configuration plus EP for MoE part. Alternatively, we could add a single switch to control whether HSDP is enabled for the expert, for example: --enable-expert-hsdp 🤔 |
…zer-instances Cherry-pick NVIDIA#7341: support independent expert optimizer instances in MFSDP v2
|
cc @wujingyue for the earlier discussion about the HSDP integration . |
This use case makes sense to me.
I'm not sure about the best API. Could you rebase the PR so I can review it again? In general, we've been using DataParallelConfig as a domain-specific language to dictate how to fully_shard. However, the friction has always been that DataParallelConfig is less powerful. Have you considered letting users (instead of the mcore adapter) fully_shard directly and still running that model in Megatron's training loop? That could be the right solution going forward. |
8cb509e to
9507679
Compare
That sounds useful for advanced users, but I think most users still want MCore to handle the sharding setup. I’d keep the config-based path and offer direct fully_shard for cases that need more control. |
330a790 to
26e0d41
Compare
|
/ok to test 26e0d41 |
Dense and expert parameters can now use different numbers of distributed optimizer instances, so each domain gets its own outer (HSDP) group. Add expert_num_distributed_optimizer_instances to DistributedDataParallelConfig, initialize_model_parallel and the --expert-num-distributed-optimizer-instances argument. It defaults to num_distributed_optimizer_instances, which keeps the existing behavior, and a value of 1 gives experts no outer axis. Only MFSDP v2 accepts differing counts; v1 DDP and DistOpt reject them. This replaces the expert_hsdp flag. expert_outer_dp_sharding_strategy still defaults to outer_dp_sharding_strategy, or to no_shard when experts have no outer axis. parallel_state keeps the expert inter-instance group and adds a dense one. get_inter_distributed_optimizer_instance_group returns the dense group by default and the expert group with expert=True, exposed as ProcessGroupCollection.inter_dist_opt and inter_expt_dist_opt. The dense group reuses the expert group when the two counts match and is built only when they differ. Extend the hybrid placements test to cover independent dense and expert outer sizes. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Signed-off-by: tongliu <tongliu@nvidia.com>
26e0d41 to
8ebc58e
Compare
|
/ok to test 8ebc58e |
Summary
Building on #7277, this change adds
expert_num_distributed_optimizer_instancesso dense and expert parameters can use independent instance counts.num_distributed_optimizer_instances, preserving existing behavior.