Partition AutoEP expert parameters per layer under ZeRO-3 - #8424
Conversation
Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
Signed-off-by: Masahiro Tanaka <tanaka.masahiro@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 03d2ffbfd6
ℹ️ 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".
| @@ -0,0 +1,178 @@ | |||
| # Copyright (c) DeepSpeed Team. | |||
There was a problem hiding this comment.
This is a non-merge commit, but its message contains no Signed-off-by trailer, so it violates the repository's mandatory DCO/CI requirement and may be rejected before the code can merge. Recreate the commit using git commit --signoff with the configured author identity.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
| if p.ndim >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")): | ||
| is_partitioned_autoep_expert = (getattr(p, "ds_zero_placement_family", None) == "autoep_expert" | ||
| and hasattr(p, "ds_shape")) | ||
| ndim = len(p.ds_shape) if is_partitioned_autoep_expert else p.ndim |
There was a problem hiding this comment.
this can be a bigger bug. regardless of this work: under zero.Init, models are already partitioned and p.ndim = 1 even for 2+ dimension tensors causing all parameters to not use muon. Do we want to take this opportunity to fix this?
2 participants
Problem
When a model created under
zero.Initis passed todeepspeed.initialize, AutoEP creates new local-expert tensors while replacing MoE layers. ZeRO-3 previously converted those tensors only after every layer had been replaced, so full replacement-expert storage accumulated across layers and increased the initialization peak.This addresses the initialization-memory behavior described as problem 1 in #8353. It does not change the independent debug-map retention behavior covered by #8356.
Approach
After each AutoEP layer is repacked, copied, tagged, and installed, resolve its expert-data-parallel group and immediately pass only its newly created expert parameters through the existing ZeRO conversion helper. Optimizer classification and grouping use the experts' ZeRO logical shape and size after partitioning. Existing source gathers, router/shared/dense placement, and batched output-recorder retargeting remain unchanged. The later ZeRO setup already skips parameters converted by this step.