From c65a8eec7f6975587cfc014a14661ba1ad80a521 Mon Sep 17 00:00:00 2001 From: Masahiro Tanaka Date: Sat, 5 Sep 2026 17:03:09 +0900 Subject: [PATCH 1/3] Partition AutoEP experts per layer under ZeRO-3 Signed-off-by: Masahiro Tanaka --- deepspeed/module_inject/auto_ep.py | 5 +- deepspeed/runtime/engine.py | 29 +++- .../runtime/zero/test_autoep_zero3_init.py | 131 ++++++++++++++++++ 3 files changed, 163 insertions(+), 2 deletions(-) create mode 100644 tests/unit/runtime/zero/test_autoep_zero3_init.py diff --git a/deepspeed/module_inject/auto_ep.py b/deepspeed/module_inject/auto_ep.py index 277ae833fc04..95e2e29ac863 100644 --- a/deepspeed/module_inject/auto_ep.py +++ b/deepspeed/module_inject/auto_ep.py @@ -13,7 +13,7 @@ import math import re from collections import OrderedDict -from typing import Literal +from typing import Callable, Literal import torch import torch.nn as nn @@ -544,11 +544,14 @@ def replace_moe_layers( specs: list[MoELayerSpec], ep_size: int, ep_rank: int, + on_moe_layer_replaced: Callable[[nn.Module], None] | None = None, ) -> None: """Replace multiple MoE modules and batch post-replacement recorder retargeting.""" replacements: list[tuple[MoELayerSpec, nn.Module]] = [] for spec in specs: replacement = self._replace_moe_layer_without_retarget(spec, ep_size, ep_rank) + if on_moe_layer_replaced is not None: + on_moe_layer_replaced(replacement) replacements.append((spec, replacement)) logger.info(f"AutoEP: replaced '{spec.moe_module_name}' with AutoEPMoELayer " f"(ep_size={ep_size}, ep_rank={ep_rank}, " diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 86918bd71c5a..93d0bfde1ca5 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -721,13 +721,40 @@ def _configure_expert_parallel(self, model): if specs: validate_autoep_post_detection(autoep_config, specs) - auto_ep.replace_moe_layers(specs, ep_size=ep_size, ep_rank=ep_rank) + convert_to_zero_parameters = self._autoep_zero3_param_converter(model) + on_moe_layer_replaced = None + if convert_to_zero_parameters is not None: + + def on_moe_layer_replaced(replacement): + self._partition_autoep_zero3_experts(replacement, convert_to_zero_parameters) + + auto_ep.replace_moe_layers(specs, + ep_size=ep_size, + ep_rank=ep_rank, + on_moe_layer_replaced=on_moe_layer_replaced) logger.info(f"AutoEP: replaced {len(specs)} MoE layer(s) with ep_size={ep_size}") # Re-tag optimizer flags for newly created AutoEP parameters from deepspeed import set_optimizer_flags set_optimizer_flags(self._config, model) + def _autoep_zero3_param_converter(self, model): + if not self.zero_optimization_partition_weights(): + return None + return next((param.convert_to_zero_parameters + for param in model.parameters() if hasattr(param, "convert_to_zero_parameters")), None) + + @staticmethod + def _partition_autoep_zero3_experts(replacement, convert_to_zero_parameters): + expert_params = list(replacement.experts.named_parameters()) + for name, param in expert_params: + group_name = getattr(param, "ds_zero_partition_group_name", None) + if group_name is None: + raise AssertionError(f"AutoEP replacement expert parameter '{name}' is missing a ZeRO partition " + "group name.") + param.ds_zero_partition_process_group = groups._get_expert_data_parallel_group(group_name) + convert_to_zero_parameters(param_list=[param for _, param in expert_params]) + def _autoep_sequence_parallel_world_size(self): if self.mpu is not None and hasattr(self.mpu, 'get_sequence_parallel_world_size'): return self.mpu.get_sequence_parallel_world_size() diff --git a/tests/unit/runtime/zero/test_autoep_zero3_init.py b/tests/unit/runtime/zero/test_autoep_zero3_init.py new file mode 100644 index 000000000000..b05cb63ecf77 --- /dev/null +++ b/tests/unit/runtime/zero/test_autoep_zero3_init.py @@ -0,0 +1,131 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +from types import SimpleNamespace + +import torch +import torch.nn as nn + +import deepspeed.runtime.engine as ds_engine +from deepspeed.module_inject.auto_ep import AutoEP +from deepspeed.runtime.engine import DeepSpeedEngine + + +class _Experts(nn.Module): + + def __init__(self, layer_index): + super().__init__() + self.w1 = nn.Parameter(torch.full((2, 2), float(layer_index + 1))) + self.w2 = nn.Parameter(torch.full((2, 2), float(layer_index + 2)), requires_grad=False) + for param in self.parameters(): + param.ds_zero_placement_family = "autoep_expert" + param.ds_zero_partition_group_name = f"expert_group_{layer_index}" + + +class _Replacement(nn.Module): + + def __init__(self, layer_index): + super().__init__() + self.layer_index = layer_index + self.experts = _Experts(layer_index) + self.router = nn.Linear(2, 2, bias=False) + self.shared_experts = nn.Linear(2, 2, bias=False) + self.dense = nn.Linear(2, 2, bias=False) + self.num_local_experts = 1 + + +def _spec(layer_index): + return SimpleNamespace( + moe_module_name=f"layers.{layer_index}.mlp", + preset_adapter="test", + model_family="test", + ) + + +def test_autoep_zero3_partitions_each_replacement_before_next_allocation(monkeypatch): + auto_ep = object.__new__(AutoEP) + events = [] + replacements = [] + + def construct(spec, ep_size, ep_rank): + layer_index = len(replacements) + if layer_index: + assert events[-1] == ("converted", layer_index - 1) + events.append(("construct", layer_index)) + replacement = _Replacement(layer_index) + replacements.append(replacement) + return replacement + + auto_ep._replace_moe_layer_without_retarget = construct + auto_ep._retarget_transformers_output_recorders = lambda spec, replacement: events.append( + ("retarget", replacement.layer_index)) + + resolved_groups = {} + + def resolve_group(group_name): + return resolved_groups.setdefault(group_name, object()) + + converted_batches = [] + + def convert_to_zero_parameters(param_list): + layer_index = len(converted_batches) + replacement = replacements[layer_index] + expected_params = list(replacement.experts.parameters()) + assert [id(param) for param in param_list] == [id(param) for param in expected_params] + assert all(param.ds_zero_placement_family == "autoep_expert" for param in param_list) + assert [param.requires_grad for param in param_list] == [True, False] + assert [param.tolist() for param in param_list] == [ + [[float(layer_index + 1)] * 2] * 2, + [[float(layer_index + 2)] * 2] * 2, + ] + assert all(param.ds_zero_partition_process_group is resolved_groups[param.ds_zero_partition_group_name] + for param in param_list) + converted_batches.append(param_list) + events.append(("converted", layer_index)) + + monkeypatch.setattr(ds_engine.groups, "_get_expert_data_parallel_group", resolve_group) + + def on_moe_layer_replaced(replacement): + events.append(("callback", replacement.layer_index)) + DeepSpeedEngine._partition_autoep_zero3_experts(replacement, convert_to_zero_parameters) + + auto_ep.replace_moe_layers([_spec(0), _spec(1)], ep_size=2, ep_rank=0, on_moe_layer_replaced=on_moe_layer_replaced) + + assert events == [ + ("construct", 0), + ("callback", 0), + ("converted", 0), + ("construct", 1), + ("callback", 1), + ("converted", 1), + ("retarget", 0), + ] + converted_ids = {id(param) for batch in converted_batches for param in batch} + expected_ids = {id(param) for replacement in replacements for param in replacement.experts.parameters()} + excluded_ids = { + id(param) + for replacement in replacements + for module in (replacement.router, replacement.shared_experts, replacement.dense) + for param in module.parameters() + } + assert converted_ids == expected_ids + assert converted_ids.isdisjoint(excluded_ids) + assert set(resolved_groups) == {"expert_group_0", "expert_group_1"} + + +def test_autoep_zero3_eager_conversion_gates(): + engine = object.__new__(DeepSpeedEngine) + source = nn.Linear(2, 2, bias=False) + converter = lambda param_list: None + source.weight.convert_to_zero_parameters = converter + + engine.zero_optimization_partition_weights = lambda: False + assert engine._autoep_zero3_param_converter(source) is None + + engine.zero_optimization_partition_weights = lambda: True + assert engine._autoep_zero3_param_converter(source) is converter + + ordinary_source = nn.Linear(2, 2, bias=False) + assert engine._autoep_zero3_param_converter(ordinary_source) is None From 58b52ba37b484b2896e6c69db13c68703c933d75 Mon Sep 17 00:00:00 2001 From: Masahiro Tanaka Date: Sat, 5 Sep 2026 17:42:51 +0900 Subject: [PATCH 2/3] Preserve Muon flags for partitioned AutoEP experts Signed-off-by: Masahiro Tanaka --- deepspeed/__init__.py | 5 ++++- .../runtime/zero/test_autoep_zero3_init.py | 21 +++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index 26bfc7c77df6..32aa201981fb 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -84,7 +84,10 @@ def _parse_version(version_str): def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) -> None: if config_class.optimizer_name == MUON_OPTIMIZER: for name, p in model.named_parameters(): - 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 + if ndim >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")): setattr(p, "use_muon", True) else: setattr(p, "use_muon", False) diff --git a/tests/unit/runtime/zero/test_autoep_zero3_init.py b/tests/unit/runtime/zero/test_autoep_zero3_init.py index b05cb63ecf77..0ca126edc16f 100644 --- a/tests/unit/runtime/zero/test_autoep_zero3_init.py +++ b/tests/unit/runtime/zero/test_autoep_zero3_init.py @@ -9,6 +9,7 @@ import torch.nn as nn import deepspeed.runtime.engine as ds_engine +from deepspeed import set_optimizer_flags from deepspeed.module_inject.auto_ep import AutoEP from deepspeed.runtime.engine import DeepSpeedEngine @@ -129,3 +130,23 @@ def test_autoep_zero3_eager_conversion_gates(): ordinary_source = nn.Linear(2, 2, bias=False) assert engine._autoep_zero3_param_converter(ordinary_source) is None + + +def test_autoep_zero3_partitioned_experts_keep_muon_assignment(monkeypatch): + replacement = _Replacement(0) + replacement.ordinary_partition = nn.Parameter(torch.empty(0)) + replacement.ordinary_partition.ds_shape = torch.Size((2, 2)) + monkeypatch.setattr(ds_engine.groups, "_get_expert_data_parallel_group", lambda group_name: object()) + + def convert_to_zero_parameters(param_list): + for param in param_list: + param.ds_shape = param.shape + param.data = torch.empty(0, dtype=param.dtype, device=param.device) + + DeepSpeedEngine._partition_autoep_zero3_experts(replacement, convert_to_zero_parameters) + assert all(param.ndim == 1 for param in replacement.experts.parameters()) + + set_optimizer_flags(SimpleNamespace(optimizer_name="muon"), replacement) + + assert all(param.use_muon for param in replacement.experts.parameters()) + assert not replacement.ordinary_partition.use_muon From 03d2ffbfd63a2bf1c7fa80704b317e7510bdd409 Mon Sep 17 00:00:00 2001 From: Masahiro Tanaka Date: Sat, 5 Sep 2026 17:56:04 +0900 Subject: [PATCH 3/3] Preserve AutoEP expert optimizer grouping Signed-off-by: Masahiro Tanaka --- deepspeed/moe/utils.py | 7 ++--- .../runtime/zero/test_autoep_zero3_init.py | 26 +++++++++++++++++++ 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/deepspeed/moe/utils.py b/deepspeed/moe/utils.py index 20866378efac..c8fc0ef176d1 100644 --- a/deepspeed/moe/utils.py +++ b/deepspeed/moe/utils.py @@ -127,13 +127,14 @@ def split_params_into_different_moe_groups_for_optimizer( size_of_cur_group = 0 for param in cast(List[nn.Parameter], param_group['params']): - if size_of_cur_group + param.numel() <= max_group_size: + param_numel = param.ds_numel if hasattr(param, "ds_numel") else param.numel() + if size_of_cur_group + param_numel <= max_group_size: cur_group.append(param) - size_of_cur_group += param.numel() + size_of_cur_group += param_numel else: all_groups.append(cur_group) cur_group = [param] - size_of_cur_group = param.numel() + size_of_cur_group = param_numel if cur_group: all_groups.append(cur_group) diff --git a/tests/unit/runtime/zero/test_autoep_zero3_init.py b/tests/unit/runtime/zero/test_autoep_zero3_init.py index 0ca126edc16f..c0faf82c2cc9 100644 --- a/tests/unit/runtime/zero/test_autoep_zero3_init.py +++ b/tests/unit/runtime/zero/test_autoep_zero3_init.py @@ -11,6 +11,7 @@ import deepspeed.runtime.engine as ds_engine from deepspeed import set_optimizer_flags from deepspeed.module_inject.auto_ep import AutoEP +from deepspeed.moe.utils import split_params_into_different_moe_groups_for_optimizer from deepspeed.runtime.engine import DeepSpeedEngine @@ -150,3 +151,28 @@ def convert_to_zero_parameters(param_list): assert all(param.use_muon for param in replacement.experts.parameters()) assert not replacement.ordinary_partition.use_muon + + +def test_autoep_zero3_partitioned_experts_keep_optimizer_grouping(): + + def make_experts(partitioned): + params = [] + for _ in range(4): + param = nn.Parameter(torch.empty(0) if partitioned else torch.ones(4)) + param.allreduce = False + param.group_name = "ep_size_2" + if partitioned: + param.ds_numel = 4 + params.append(param) + return params + + def group_lengths(params): + groups = split_params_into_different_moe_groups_for_optimizer({ + "name": "dense-params", + "params": params, + }, + max_group_size=10) + return [len(group["params"]) for group in groups if group.get("moe")] + + assert group_lengths(make_experts(partitioned=False)) == [2, 2] + assert group_lengths(make_experts(partitioned=True)) == [2, 2]