[Feature] Enable AutoEP Compatibility with ZeRO-3 - #7928
Closed
nathon-lee wants to merge 15 commits into
Closed
Conversation
This reverts commit ff88670. Co-authored-by: nathon-lee <248585198+nathon-lee@users.noreply.github.com>
Revert "fix: update 1 file reformatted." (ff88670)
This reverts commit b90aee5.
Revert accidental Muon optimizer code re-introduction from copilot PRs
Signed-off-by: nathon-lee <leejianwoo@gmail.com> fix: move torch.distributed as dist Signed-off-by: nathon-lee <leejianwoo@gmail.com> fix: update docs _tutorials autoep.md . Signed-off-by: nathon-lee <leejianwoo@gmail.com>
nathon-lee
requested review from
hwchen2017,
loadams,
tjruwase and
tohtana
as code owners
March 28, 2026 09:09
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d231f6b3bc
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
Comment on lines
+197
to
+200
| if ep_group is not None: | ||
| self.ep_group = ep_group | ||
| self.ep_group_name = f"ep_group_{id(ep_group)}" | ||
| for param in self.experts.parameters(): |
There was a problem hiding this comment.
DeepSpeedEngine._configure_distributed_model passes a boolean (use_data_before_expert_parallel_) into every module's set_deepspeed_parallelism, but this method treats any non-None value as an EP process group and overwrites self.ep_group with True/False. Once that happens, AutoEP forward calls (all_gather/all_to_all_single) run with a boolean instead of a process group and fail at runtime as soon as the layer executes.
Useful? React with 👍 / 👎.
Comment on lines
+156
to
+158
| self.reorderer = TokenReorderer( | ||
| num_experts=self.num_local_experts, | ||
| top_k=spec.top_k, |
There was a problem hiding this comment.
The router emits expert IDs in the global range [0, num_experts), but TokenReorderer is initialized with num_local_experts. For ep_size > 1, expert IDs outside the local range are mis-bucketed/dropped by the histogram logic, so token counts no longer match the sorted token stream; this corrupts dispatch metadata and can trigger incorrect routing or downstream shape/index failures in multi-rank EP runs.
Useful? React with 👍 / 👎.
Comment on lines
+2622
to
+2626
| if not hasattr(self, '_autoep_expert_optimizer'): | ||
| optimizer_cls = type(self.optimizer) | ||
| base_group = self.optimizer.param_groups[0] | ||
| expert_group = {k: v for k, v in base_group.items() if k != 'params'} | ||
| expert_group['params'] = expert_params |
There was a problem hiding this comment.
The dedicated AutoEP expert optimizer is created once from self.optimizer.param_groups[0] and then reused without any hyperparameter refresh. If a scheduler (or manual LR/WD update) changes the main optimizer during training, expert params keep stale hyperparameters while non-expert params follow the new values, causing silent optimization drift between parameter sets.
Useful? React with 👍 / 👎.
PKUWZP
self-requested a review
March 28, 2026 15:33
Collaborator
|
Thank you, @nathon-lee! This is amazing. Maybe we should focus on merging the branch first? I have left it for a while, but I will prioritize it if you can help me. |
Contributor
Author
@tohtana Thanks for pointing this out — you’re right. I did use tohtana/DeepSpeedExamples/training/expert_parallel as a reference, and I should have acknowledged that more clearly. |
Merged
Collaborator
|
Hi @nathon-lee, |
Contributor
Author
Hi @tohtana, thanks for the heads-up and for adding the missing features on top of your branch. |
Collaborator
|
@nathon-lee #7938 is missing Z3 support. Do you think you can add it? What about creating a new PR focusing on Z3 support and merge it to #7938. |
Contributor
Author
@tohtana ok |
Contributor
Author
I’ll probably wait until your AutoEP branch is merged into main before opening my PR, since my changes depend on your branch. |
Signed-off-by: nathon-lee <leejianwoo@gmail.com>
scripts/check-license.py uses 'git grep -e ^# SPDX-License-Identifier: Apache-2.0$', which doesn't match compound expressions. Split 'Apache-2.0 AND BSD-3-Clause' into a primary SPDX line plus an 'Additional license' note; THIRD_PARTY_NOTICES.md still records the BSD-3-Clause portion. Signed-off-by: nathon-lee <leejianwoo@gmail.com>
refactor(autoep-zero3): drop files already covered by deepspeedai#7938
Collaborator
|
Hi @nathon-lee, |
Contributor
Author
tohtana
added a commit
that referenced
this pull request
Jun 11, 2026
This PR adds AutoEP (Automatic Expert Parallelism) to DeepSpeed training for HuggingFace MoE models. AutoEP detects MoE blocks during `deepspeed.initialize()`, builds the required EP/EDP process groups, and replaces supported MoE blocks with an EP-enabled execution path, so expert parallelism can be enabled with DeepSpeed config only and without model code changes. Current scope in this PR is the base AutoEP feature: - ZeRO stages 0, 1, and 2 support - checkpoint save/load support - universal checkpoint conversion support ZeRO-3 extensions are intentionally left as follow-up work (#7928 should be merged for this work) Supported presets in this PR: - Mixtral - Qwen3-MoE - DeepSeek-V2 - DeepSeek-V3 For end-to-end benchmarking and testing, an AutoEP example is available in DeepSpeedExamples: - <https://github.com/tohtana/DeepSpeedExamples/tree/tohtana/add_auto_ep/training/expert_parallel> ## Attribution This implementation substantially builds on TorchTitan's MoE / expert-parallel implementation, and we want to explicitly acknowledge that prior work. The TorchTitan-derived pieces in this PR are primarily: - `deepspeed/moe/ep_router.py`: adapted from TorchTitan's `TokenChoiceTopKRouter` - `deepspeed/moe/ep_experts.py`: adapted from TorchTitan's `GroupedExperts` and grouped-GEMM expert execution path - `deepspeed/moe/ep_kernels.py`: adapted from TorchTitan's `TokenReorderer`, `generate_permute_indices`, Triton fill-indices kernel, and token-group alignment / padding helpers - `deepspeed/module_inject/auto_ep_layer.py`: adapts the same router -> reorder -> dispatch -> local expert compute -> combine structure used in TorchTitan's MoE / EP flow Relevant TorchTitan sources: - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/moe.py> - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/kernels.py> - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/utils.py> - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/expert_parallel.py> The DeepSpeed-specific work in this PR is the AutoEP integration layer around those building blocks: - HuggingFace MoE detection and structural validation - model-family presets and custom-config path - weight repacking from HF expert layouts into grouped expert tensors - DeepSpeed runtime group setup and module replacement - DeepSpeed checkpoint save/load and universal checkpoint support - DeepSpeed docs and tests ## Design The implementation is split into a few layers: - `deepspeed/module_inject/auto_ep_config.py` - user config parsing - built-in model presets - validation for EP topology and per-model constraints - `deepspeed/module_inject/auto_ep.py` - scans the model for MoE blocks - validates the detected structure - builds a `MoELayerSpec` for each supported MoE layer - replaces the original HF block with `AutoEPMoELayer` - `deepspeed/module_inject/auto_ep_layer.py` - the drop-in execution wrapper for a detected MoE block - implements router execution, token reorder, EP dispatch/combine, local expert compute, and shared-expert merge - `deepspeed/moe/ep_router.py`, `deepspeed/moe/ep_experts.py`, `deepspeed/moe/ep_kernels.py` - reusable MoE runtime pieces for routing, grouped expert compute, token permutation, and aligned grouped-GEMM execution - `deepspeed/moe/ep_repack.py` - converts HF expert weights into the grouped expert layout expected by the runtime - `deepspeed/runtime/engine.py` and checkpoint conversion code - wires AutoEP into `deepspeed.initialize()` - handles checkpoint save/load metadata and universal checkpoint integration At runtime, the execution path is: 1. detect and replace supported HF MoE blocks during initialization 2. route tokens with the EP router 3. reorder tokens by expert assignment 4. perform all-to-all dispatch across the EP group when `autoep_size > 1` 5. run local grouped expert compute 6. all-to-all combine and restore the original token order 7. merge shared experts if the model has them ## Adding new model support There are two supported ways to extend AutoEP to a new MoE model family. 1. Add a preset in `PRESET_MODELS`. This is the preferred path for a model family we want to support out of the box. A preset defines: - MoE layer pattern - router child name - experts child name - expert weight names / layout - `num_experts` and `top_k` config attributes - routing defaults - optional shared-expert structure 2. Use the custom config path. For models that are not yet built into DeepSpeed, AutoEP can be driven from config with: - `moe_layer_pattern` - `router_pattern` - `expert_pattern` - `expert_w1`, `expert_w2`, `expert_w3` - `num_experts_attr` - `top_k_attr` - optional shared-expert fields Once detection can produce a valid `MoELayerSpec`, the replacement, execution, and checkpoint paths are shared. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Ma, Guokai <guokai.ma@gmail.com> Signed-off-by: Guokai Ma <guokai.ma@intel.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com> Co-authored-by: Guokai Ma <guokai.ma@intel.com>
nathon-lee
pushed a commit
to nathon-lee/DeepSpeed_woo
that referenced
this pull request
Jul 1, 2026
This PR adds AutoEP (Automatic Expert Parallelism) to DeepSpeed training for HuggingFace MoE models. AutoEP detects MoE blocks during `deepspeed.initialize()`, builds the required EP/EDP process groups, and replaces supported MoE blocks with an EP-enabled execution path, so expert parallelism can be enabled with DeepSpeed config only and without model code changes. Current scope in this PR is the base AutoEP feature: - ZeRO stages 0, 1, and 2 support - checkpoint save/load support - universal checkpoint conversion support ZeRO-3 extensions are intentionally left as follow-up work (deepspeedai#7928 should be merged for this work) Supported presets in this PR: - Mixtral - Qwen3-MoE - DeepSeek-V2 - DeepSeek-V3 For end-to-end benchmarking and testing, an AutoEP example is available in DeepSpeedExamples: - <https://github.com/tohtana/DeepSpeedExamples/tree/tohtana/add_auto_ep/training/expert_parallel> ## Attribution This implementation substantially builds on TorchTitan's MoE / expert-parallel implementation, and we want to explicitly acknowledge that prior work. The TorchTitan-derived pieces in this PR are primarily: - `deepspeed/moe/ep_router.py`: adapted from TorchTitan's `TokenChoiceTopKRouter` - `deepspeed/moe/ep_experts.py`: adapted from TorchTitan's `GroupedExperts` and grouped-GEMM expert execution path - `deepspeed/moe/ep_kernels.py`: adapted from TorchTitan's `TokenReorderer`, `generate_permute_indices`, Triton fill-indices kernel, and token-group alignment / padding helpers - `deepspeed/module_inject/auto_ep_layer.py`: adapts the same router -> reorder -> dispatch -> local expert compute -> combine structure used in TorchTitan's MoE / EP flow Relevant TorchTitan sources: - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/moe.py> - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/kernels.py> - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/utils.py> - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/expert_parallel.py> The DeepSpeed-specific work in this PR is the AutoEP integration layer around those building blocks: - HuggingFace MoE detection and structural validation - model-family presets and custom-config path - weight repacking from HF expert layouts into grouped expert tensors - DeepSpeed runtime group setup and module replacement - DeepSpeed checkpoint save/load and universal checkpoint support - DeepSpeed docs and tests ## Design The implementation is split into a few layers: - `deepspeed/module_inject/auto_ep_config.py` - user config parsing - built-in model presets - validation for EP topology and per-model constraints - `deepspeed/module_inject/auto_ep.py` - scans the model for MoE blocks - validates the detected structure - builds a `MoELayerSpec` for each supported MoE layer - replaces the original HF block with `AutoEPMoELayer` - `deepspeed/module_inject/auto_ep_layer.py` - the drop-in execution wrapper for a detected MoE block - implements router execution, token reorder, EP dispatch/combine, local expert compute, and shared-expert merge - `deepspeed/moe/ep_router.py`, `deepspeed/moe/ep_experts.py`, `deepspeed/moe/ep_kernels.py` - reusable MoE runtime pieces for routing, grouped expert compute, token permutation, and aligned grouped-GEMM execution - `deepspeed/moe/ep_repack.py` - converts HF expert weights into the grouped expert layout expected by the runtime - `deepspeed/runtime/engine.py` and checkpoint conversion code - wires AutoEP into `deepspeed.initialize()` - handles checkpoint save/load metadata and universal checkpoint integration At runtime, the execution path is: 1. detect and replace supported HF MoE blocks during initialization 2. route tokens with the EP router 3. reorder tokens by expert assignment 4. perform all-to-all dispatch across the EP group when `autoep_size > 1` 5. run local grouped expert compute 6. all-to-all combine and restore the original token order 7. merge shared experts if the model has them ## Adding new model support There are two supported ways to extend AutoEP to a new MoE model family. 1. Add a preset in `PRESET_MODELS`. This is the preferred path for a model family we want to support out of the box. A preset defines: - MoE layer pattern - router child name - experts child name - expert weight names / layout - `num_experts` and `top_k` config attributes - routing defaults - optional shared-expert structure 2. Use the custom config path. For models that are not yet built into DeepSpeed, AutoEP can be driven from config with: - `moe_layer_pattern` - `router_pattern` - `expert_pattern` - `expert_w1`, `expert_w2`, `expert_w3` - `num_experts_attr` - `top_k_attr` - optional shared-expert fields Once detection can produce a valid `MoELayerSpec`, the replacement, execution, and checkpoint paths are shared. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Ma, Guokai <guokai.ma@gmail.com> Signed-off-by: Guokai Ma <guokai.ma@intel.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com> Co-authored-by: Guokai Ma <guokai.ma@intel.com> Signed-off-by: nathon-lee <leejianwoo@gmail.com>
nathon-lee
pushed a commit
to nathon-lee/DeepSpeed_woo
that referenced
this pull request
Jul 1, 2026
This PR adds AutoEP (Automatic Expert Parallelism) to DeepSpeed training for HuggingFace MoE models. AutoEP detects MoE blocks during `deepspeed.initialize()`, builds the required EP/EDP process groups, and replaces supported MoE blocks with an EP-enabled execution path, so expert parallelism can be enabled with DeepSpeed config only and without model code changes. Current scope in this PR is the base AutoEP feature: - ZeRO stages 0, 1, and 2 support - checkpoint save/load support - universal checkpoint conversion support ZeRO-3 extensions are intentionally left as follow-up work (deepspeedai#7928 should be merged for this work) Supported presets in this PR: - Mixtral - Qwen3-MoE - DeepSeek-V2 - DeepSeek-V3 For end-to-end benchmarking and testing, an AutoEP example is available in DeepSpeedExamples: - <https://github.com/tohtana/DeepSpeedExamples/tree/tohtana/add_auto_ep/training/expert_parallel> ## Attribution This implementation substantially builds on TorchTitan's MoE / expert-parallel implementation, and we want to explicitly acknowledge that prior work. The TorchTitan-derived pieces in this PR are primarily: - `deepspeed/moe/ep_router.py`: adapted from TorchTitan's `TokenChoiceTopKRouter` - `deepspeed/moe/ep_experts.py`: adapted from TorchTitan's `GroupedExperts` and grouped-GEMM expert execution path - `deepspeed/moe/ep_kernels.py`: adapted from TorchTitan's `TokenReorderer`, `generate_permute_indices`, Triton fill-indices kernel, and token-group alignment / padding helpers - `deepspeed/module_inject/auto_ep_layer.py`: adapts the same router -> reorder -> dispatch -> local expert compute -> combine structure used in TorchTitan's MoE / EP flow Relevant TorchTitan sources: - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/moe.py> - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/kernels.py> - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/common/moe/utils.py> - <https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/expert_parallel.py> The DeepSpeed-specific work in this PR is the AutoEP integration layer around those building blocks: - HuggingFace MoE detection and structural validation - model-family presets and custom-config path - weight repacking from HF expert layouts into grouped expert tensors - DeepSpeed runtime group setup and module replacement - DeepSpeed checkpoint save/load and universal checkpoint support - DeepSpeed docs and tests ## Design The implementation is split into a few layers: - `deepspeed/module_inject/auto_ep_config.py` - user config parsing - built-in model presets - validation for EP topology and per-model constraints - `deepspeed/module_inject/auto_ep.py` - scans the model for MoE blocks - validates the detected structure - builds a `MoELayerSpec` for each supported MoE layer - replaces the original HF block with `AutoEPMoELayer` - `deepspeed/module_inject/auto_ep_layer.py` - the drop-in execution wrapper for a detected MoE block - implements router execution, token reorder, EP dispatch/combine, local expert compute, and shared-expert merge - `deepspeed/moe/ep_router.py`, `deepspeed/moe/ep_experts.py`, `deepspeed/moe/ep_kernels.py` - reusable MoE runtime pieces for routing, grouped expert compute, token permutation, and aligned grouped-GEMM execution - `deepspeed/moe/ep_repack.py` - converts HF expert weights into the grouped expert layout expected by the runtime - `deepspeed/runtime/engine.py` and checkpoint conversion code - wires AutoEP into `deepspeed.initialize()` - handles checkpoint save/load metadata and universal checkpoint integration At runtime, the execution path is: 1. detect and replace supported HF MoE blocks during initialization 2. route tokens with the EP router 3. reorder tokens by expert assignment 4. perform all-to-all dispatch across the EP group when `autoep_size > 1` 5. run local grouped expert compute 6. all-to-all combine and restore the original token order 7. merge shared experts if the model has them ## Adding new model support There are two supported ways to extend AutoEP to a new MoE model family. 1. Add a preset in `PRESET_MODELS`. This is the preferred path for a model family we want to support out of the box. A preset defines: - MoE layer pattern - router child name - experts child name - expert weight names / layout - `num_experts` and `top_k` config attributes - routing defaults - optional shared-expert structure 2. Use the custom config path. For models that are not yet built into DeepSpeed, AutoEP can be driven from config with: - `moe_layer_pattern` - `router_pattern` - `expert_pattern` - `expert_w1`, `expert_w2`, `expert_w3` - `num_experts_attr` - `top_k_attr` - optional shared-expert fields Once detection can produce a valid `MoELayerSpec`, the replacement, execution, and checkpoint paths are shared. --------- Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: Ma, Guokai <guokai.ma@gmail.com> Signed-off-by: Guokai Ma <guokai.ma@intel.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com> Co-authored-by: Guokai Ma <guokai.ma@intel.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
[Feature] Enable AutoEP Compatibility with ZeRO-3
📌 Summary
This PR introduces compatibility between AutoEP (Expert Parallelism) and ZeRO-3.
AutoEP has historically relied on ZeRO-2 due to inherent conflicts between expert-parallel parameter partitioning and ZeRO-3’s data-parallel sharding. This PR resolves those conflicts through a minimal and targeted decoupling strategy, allowing:
This preserves AutoEP’s high-throughput execution while unlocking the memory efficiency of ZeRO-3 where applicable.
🔍 Design Overview
Instead of modifying core ZeRO-3 logic, this PR selectively bypasses ZeRO-3 mechanisms for expert parameters, while keeping the default behavior unchanged for all other parameters.
The implementation consists of four focused components:
1. Parameter Partition Bypass
Expert parameters are tagged (
_autoep_expert=True) and excluded from ZeRO-3 partitioning and gathering logic.2. Gradient Reduction Isolation
Expert gradients bypass ZeRO-3
reduce-scatterand instead useall_reducewithin the EP data-parallel group, matching AutoEP semantics.3. Optimizer State Isolation
A dedicated optimizer is introduced for expert parameters, along with FP32 master weights to ensure numerical stability during updates.
4. Checkpoint Compatibility
Expert parameters and their optimizer states are explicitly integrated into checkpoint save/load paths to ensure correct training resumption.
✅ Benefits
🧪 Testing
Verified end-to-end training correctness
Added unit tests for:
Due to limited GPU resources, validation has been performed on 2 GPUs.
If additional resources (e.g., 8 GPUs) are available, I would be very happy to further validate scalability and robustness. The additional verification should only require a few hours.
🙏 Notes
Feedback and suggestions are very welcome.
If possible, I would greatly appreciate access to larger-scale testing resources to further strengthen validation.
References
Signed-off-by: nathon-lee [leejianwoo@gmail.com]