wehub-resource-sync
1188 lines
53 KiB
Python
1188 lines
53 KiB
Python
# Copyright (c) DeepSpeed Team.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
"""Compact critical-path tests for AutoEP."""
|
|
|
|
import ast
|
|
import inspect
|
|
from collections import OrderedDict
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
import deepspeed.runtime.engine as ds_engine
|
|
import deepspeed.runtime.zero.stage3 as zero_stage3
|
|
import deepspeed.moe.ep_repack as ep_repack
|
|
from deepspeed.module_inject.auto_ep import AutoEP, _resolve_route_scale
|
|
from deepspeed.module_inject.auto_ep_config import (
|
|
AutoEPConfig,
|
|
MoELayerSpec,
|
|
PRESET_MODELS,
|
|
fill_autoep_config_from_hf,
|
|
parse_autoep_config,
|
|
validate_autoep_config,
|
|
validate_autoep_post_detection,
|
|
)
|
|
from deepspeed.module_inject.auto_ep_layer import (
|
|
AutoEPMoELayer,
|
|
apply_scores_before_experts_if_enabled,
|
|
combine_from_routed,
|
|
resolve_score_apply_mode,
|
|
)
|
|
from deepspeed.module_inject.auto_ep_preset_adapters import get_preset_adapter
|
|
from deepspeed.module_inject.auto_ep_presets.registry import (
|
|
preset_name_for_hf_model_type,
|
|
unsupported_preset_for_hf_model_type,
|
|
)
|
|
from deepspeed.moe.layer import MoE
|
|
from deepspeed.moe.ep_experts import GroupedExperts
|
|
from deepspeed.moe.ep_kernels import TokenReorderer
|
|
from deepspeed.moe.ep_repack import repack_expert_weights
|
|
from deepspeed.moe.ep_router import TokenChoiceTopKRouter
|
|
from deepspeed.runtime.engine import DeepSpeedEngine
|
|
from deepspeed.runtime.zero.stage3 import DeepSpeedZeroOptimizer_Stage3
|
|
from deepspeed.utils import groups
|
|
from unit.v1.moe.autoep_test_utils import (
|
|
MockMoEBlock,
|
|
MockMoETransformer,
|
|
UNSUPPORTED_LOAD_BALANCE_VALUES,
|
|
assert_causal_lm_outputs_close,
|
|
assert_load_balance_coeff_rejection_message,
|
|
replace_autoep_layers,
|
|
skip_unless_transformers_has,
|
|
state_matched_models,
|
|
tiny_mixtral_config,
|
|
)
|
|
|
|
|
|
def _runtime_config(**kwargs):
|
|
kwargs.setdefault("use_grouped_mm", False)
|
|
return AutoEPConfig(**kwargs)
|
|
|
|
|
|
def _make_spec(**kwargs):
|
|
defaults = dict(
|
|
moe_module_name="model.layers.0.mlp",
|
|
model_family="mixtral",
|
|
router_name="gate",
|
|
experts_name="experts",
|
|
expert_storage="fused_3d",
|
|
expert_w1_name="gate_up_proj",
|
|
expert_w2_name="down_proj",
|
|
expert_w3_name=None,
|
|
num_experts=4,
|
|
top_k=2,
|
|
hidden_size=64,
|
|
ffn_hidden_size=128,
|
|
score_func="softmax",
|
|
score_apply="post",
|
|
route_norm=True,
|
|
gate_bias=False,
|
|
return_router_logits=False,
|
|
router_logits_capture_target="none",
|
|
router_logits_capture_index=None,
|
|
router_logits_capture_layer_name=None,
|
|
has_shared_experts=False,
|
|
shared_experts_name="",
|
|
shared_experts_gate_name="",
|
|
)
|
|
defaults.update(kwargs)
|
|
return MoELayerSpec(**defaults)
|
|
|
|
|
|
def _assert_same_dtype_device(actual, expected):
|
|
assert actual.dtype == expected.dtype
|
|
assert actual.device == expected.device
|
|
|
|
|
|
def _mark_fake_zero_param(param, full_data, partition_data=None, ds_id=0, name="param"):
|
|
param.ds_id = ds_id
|
|
param.ds_shape = torch.Size(full_data.shape)
|
|
param._autoep_test_full_data = full_data.detach().clone()
|
|
param._autoep_test_name = name
|
|
if partition_data is None:
|
|
partition_data = torch.zeros(1, dtype=full_data.dtype, device=full_data.device)
|
|
param.data = partition_data.detach().clone()
|
|
return param
|
|
|
|
|
|
class FakeGatheredParameters:
|
|
calls = []
|
|
|
|
def __init__(self, params, modifier_rank=None, fwd_module=None, enabled=True):
|
|
self.params = list(params)
|
|
self.modifier_rank = modifier_rank
|
|
self.enabled = enabled
|
|
self._saved_data = []
|
|
FakeGatheredParameters.calls.append({
|
|
"names": [getattr(param, "_autoep_test_name", f"param{param.ds_id}") for param in self.params],
|
|
"modifier_rank":
|
|
modifier_rank,
|
|
"enabled":
|
|
enabled,
|
|
})
|
|
|
|
def __enter__(self):
|
|
if not self.enabled:
|
|
return
|
|
for param in self.params:
|
|
self._saved_data.append((param, param.data))
|
|
param.data = param._autoep_test_full_data.detach().clone()
|
|
|
|
def __exit__(self, *exc):
|
|
if not self.enabled:
|
|
return
|
|
for param, data in self._saved_data:
|
|
param.data = data
|
|
|
|
|
|
class MockSharedExpert(nn.Module):
|
|
|
|
def __init__(self, hidden_size=64):
|
|
super().__init__()
|
|
self.up_proj = nn.Linear(hidden_size, hidden_size, bias=False)
|
|
self.gate_proj = nn.Linear(hidden_size, hidden_size, bias=False)
|
|
self.down_proj = nn.Linear(hidden_size, hidden_size, bias=False)
|
|
|
|
|
|
class MockDeepSeekV3Config:
|
|
model_type = "deepseek_v3"
|
|
n_routed_experts = 8
|
|
num_experts_per_tok = 2
|
|
hidden_size = 64
|
|
moe_intermediate_size = 128
|
|
n_group = 4
|
|
topk_group = 2
|
|
routed_scaling_factor = 2.5
|
|
|
|
|
|
class MockDeepSeekV3Expert(nn.Module):
|
|
|
|
def __init__(self, hidden_size=64, ffn_hidden=128):
|
|
super().__init__()
|
|
self.gate_proj = nn.Linear(hidden_size, ffn_hidden, bias=False)
|
|
self.up_proj = nn.Linear(hidden_size, ffn_hidden, bias=False)
|
|
self.down_proj = nn.Linear(ffn_hidden, hidden_size, bias=False)
|
|
|
|
|
|
class MockDeepSeekV3MoEBlock(nn.Module):
|
|
|
|
def __init__(self, num_experts=8, ffn_hidden=128, hidden_size=64):
|
|
super().__init__()
|
|
self.gate = nn.Linear(hidden_size, num_experts, bias=False)
|
|
self.experts = nn.ModuleList([MockDeepSeekV3Expert(hidden_size, ffn_hidden) for _ in range(num_experts)])
|
|
self.shared_experts = MockSharedExpert(hidden_size)
|
|
|
|
|
|
class MockDeepSeekV3Transformer(nn.Module):
|
|
|
|
def __init__(self, num_layers=2, num_experts=8):
|
|
super().__init__()
|
|
self.config = MockDeepSeekV3Config()
|
|
self.config.n_routed_experts = num_experts
|
|
self.model = nn.Module()
|
|
self.model.layers = nn.ModuleList([self._make_layer(num_experts) for _ in range(num_layers)])
|
|
|
|
@staticmethod
|
|
def _make_layer(num_experts):
|
|
layer = nn.Module()
|
|
layer.mlp = MockDeepSeekV3MoEBlock(num_experts)
|
|
return layer
|
|
|
|
|
|
class TestAutoEPConfig:
|
|
|
|
def test_parse_and_validate_enabled_size_contract(self):
|
|
disabled = parse_autoep_config({})
|
|
assert disabled.enabled is False
|
|
assert disabled.autoep_size == 1
|
|
assert disabled.validate_folding_routing is False
|
|
assert disabled.load_balance_coeff is None
|
|
assert disabled._load_balance_coeff_explicit is False
|
|
|
|
config = parse_autoep_config({
|
|
"enabled": True,
|
|
"autoep_size": 4,
|
|
"preset_model": "mixtral",
|
|
"load_balance_coeff": None,
|
|
"score_apply": "pre",
|
|
"route_scale": 2.0,
|
|
"validate_folding_routing": True,
|
|
})
|
|
|
|
assert config.enabled is True
|
|
assert config.autoep_size == 4
|
|
assert config.preset_model == "mixtral"
|
|
assert config.validate_folding_routing is True
|
|
assert config.load_balance_coeff is None
|
|
assert config._load_balance_coeff_explicit is True
|
|
assert config.score_apply == "pre"
|
|
assert config.route_scale == 2.0
|
|
validate_autoep_config(config, world_size=4, pp_size=1, tp_size=1, sp_size=1)
|
|
|
|
def test_validate_folding_routing_requires_boolean(self):
|
|
with pytest.raises(ValueError, match="validate_folding_routing"):
|
|
validate_autoep_config(AutoEPConfig(enabled=True, validate_folding_routing="true"),
|
|
world_size=1,
|
|
pp_size=1,
|
|
tp_size=1,
|
|
sp_size=1)
|
|
|
|
@pytest.mark.parametrize("value", UNSUPPORTED_LOAD_BALANCE_VALUES)
|
|
def test_load_balance_coeff_rejected_at_parse(self, value):
|
|
with pytest.raises(ValueError) as exc_info:
|
|
parse_autoep_config({"enabled": True, "load_balance_coeff": value})
|
|
assert_load_balance_coeff_rejection_message(exc_info.value, value)
|
|
|
|
@pytest.mark.parametrize("enabled", [True, False])
|
|
@pytest.mark.parametrize("value", [0.01, False, "0.01"])
|
|
def test_load_balance_coeff_rejected_by_validate(self, enabled, value):
|
|
config = AutoEPConfig(enabled=enabled, load_balance_coeff=value)
|
|
|
|
with pytest.raises(ValueError) as exc_info:
|
|
validate_autoep_config(config, world_size=1, pp_size=1, tp_size=1, sp_size=1)
|
|
assert_load_balance_coeff_rejection_message(exc_info.value, value)
|
|
|
|
def test_ep_size_validation_rejects_invalid_topology(self):
|
|
validate_autoep_config(AutoEPConfig(enabled=True, autoep_size=2),
|
|
world_size=8,
|
|
pp_size=1,
|
|
tp_size=2,
|
|
sp_size=1)
|
|
with pytest.raises(ValueError, match="must divide the stage size"):
|
|
validate_autoep_config(AutoEPConfig(enabled=True, autoep_size=3),
|
|
world_size=8,
|
|
pp_size=1,
|
|
tp_size=1,
|
|
sp_size=1)
|
|
with pytest.raises(ValueError, match="exceeds num_experts"):
|
|
validate_autoep_post_detection(AutoEPConfig(enabled=True, autoep_size=16), [_make_spec(num_experts=8)])
|
|
|
|
def test_expert_tensor_parallel_size_is_parsed_but_limited_to_one(self):
|
|
config = parse_autoep_config({
|
|
"enabled": True,
|
|
"autoep_size": 2,
|
|
"expert_tensor_parallel_size": 1,
|
|
})
|
|
assert config.expert_tensor_parallel_size == 1
|
|
|
|
config.expert_tensor_parallel_size = 2
|
|
with pytest.raises(ValueError, match="expert_tensor_parallel_size=1"):
|
|
validate_autoep_config(config, world_size=4, pp_size=1, tp_size=1, sp_size=1)
|
|
|
|
def test_configure_expert_parallel_uses_engine_mpu_sequence_parallel_size(self, monkeypatch):
|
|
|
|
class SequenceParallelMPU:
|
|
|
|
def get_model_parallel_world_size(self):
|
|
return 1
|
|
|
|
def get_sequence_parallel_world_size(self):
|
|
return 2
|
|
|
|
class EmptyAutoEP:
|
|
|
|
def __init__(self, model, config):
|
|
pass
|
|
|
|
def ep_parser(self):
|
|
return []
|
|
|
|
observed = {}
|
|
|
|
def record_validate(config, world_size, pp_size, tp_size, sp_size):
|
|
observed["validate"] = {
|
|
"world_size": world_size,
|
|
"pp_size": pp_size,
|
|
"tp_size": tp_size,
|
|
"sp_size": sp_size,
|
|
}
|
|
|
|
def record_create(**kwargs):
|
|
observed["create"] = kwargs
|
|
|
|
monkeypatch.setattr(groups, "mpu", None)
|
|
monkeypatch.setattr(groups, "_get_sequence_parallel_world_size", lambda: 1)
|
|
monkeypatch.setattr(groups, "_create_expert_and_data_parallel", record_create)
|
|
monkeypatch.setattr(groups, "_get_expert_parallel_group", lambda name: object())
|
|
monkeypatch.setattr(ds_engine.dist, "get_world_size", lambda: 4)
|
|
monkeypatch.setattr(ds_engine.dist, "get_rank", lambda group=None: 0)
|
|
monkeypatch.setattr("deepspeed.module_inject.auto_ep.AutoEP", EmptyAutoEP)
|
|
monkeypatch.setattr("deepspeed.module_inject.auto_ep_config.validate_autoep_config", record_validate)
|
|
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.mpu = SequenceParallelMPU()
|
|
engine._config = SimpleNamespace(
|
|
expert_parallel_config=AutoEPConfig(enabled=True, autoep_size=2),
|
|
tensor_parallel_config=SimpleNamespace(autotp_size=1),
|
|
use_data_before_expert_parallel_=False,
|
|
zero_config=SimpleNamespace(offload_optimizer=None, offload_param=None),
|
|
zero_optimization_stage=0,
|
|
)
|
|
|
|
engine._configure_expert_parallel(model=nn.Module())
|
|
|
|
assert groups.mpu is None
|
|
assert observed["validate"]["sp_size"] == 2
|
|
assert observed["create"]["mp_size"] == 2
|
|
assert observed["create"]["mp_mode"] == "sp"
|
|
|
|
def test_configure_expert_parallel_rejects_bwc_tensor_model_parallel_mpu(self, monkeypatch):
|
|
|
|
class TensorParallelMPU:
|
|
|
|
def get_tensor_model_parallel_world_size(self):
|
|
return 2
|
|
|
|
monkeypatch.setattr(groups, "_get_sequence_parallel_world_size", lambda: 1)
|
|
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.mpu = TensorParallelMPU()
|
|
engine._config = SimpleNamespace(
|
|
expert_parallel_config=AutoEPConfig(enabled=True, autoep_size=2),
|
|
tensor_parallel_config=SimpleNamespace(autotp_size=1),
|
|
use_data_before_expert_parallel_=False,
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="bwc_tensor_model_parallel_world_size=2"):
|
|
engine._configure_expert_parallel(model=nn.Module())
|
|
|
|
def test_autoep_sequence_parallel_size_falls_back_to_groups_helper(self, monkeypatch):
|
|
monkeypatch.setattr(groups, "_get_sequence_parallel_world_size", lambda: 3)
|
|
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.mpu = object()
|
|
|
|
assert engine._autoep_sequence_parallel_world_size() == 3
|
|
|
|
def test_zero3_compatibility_gate_rejects_native_moe(self):
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.__dict__["module"] = nn.Sequential(MoE(hidden_size=4, expert=nn.Linear(4, 4), num_experts=1))
|
|
engine.has_moe_layers = True
|
|
engine.sequence_parallel_size = 1
|
|
engine._config = SimpleNamespace(
|
|
tensor_parallel_config=SimpleNamespace(autotp_size=1),
|
|
expert_parallel_config=AutoEPConfig(enabled=True, autoep_size=1),
|
|
)
|
|
|
|
with pytest.raises(AssertionError, match="Native DeepSpeed MoE"):
|
|
engine._validate_zero3_moe_compatibility()
|
|
|
|
def test_zero3_compatibility_gate_allows_constrained_autoep(self):
|
|
model = MockMoETransformer(num_layers=1)
|
|
replace_autoep_layers(model, "mixtral")
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.__dict__["module"] = model
|
|
engine.has_moe_layers = True
|
|
engine.sequence_parallel_size = 1
|
|
engine.zero_quantized_gradients = lambda: False
|
|
engine._config = SimpleNamespace(
|
|
tensor_parallel_config=SimpleNamespace(autotp_size=1),
|
|
expert_parallel_config=AutoEPConfig(enabled=True, autoep_size=1),
|
|
)
|
|
|
|
engine._validate_zero3_moe_compatibility()
|
|
|
|
def test_zero3_compatibility_gate_rejects_sequence_parallel(self):
|
|
model = MockMoETransformer(num_layers=1)
|
|
replace_autoep_layers(model, "mixtral")
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.__dict__["module"] = model
|
|
engine.has_moe_layers = True
|
|
engine.sequence_parallel_size = 2
|
|
engine.zero_quantized_gradients = lambda: False
|
|
engine._config = SimpleNamespace(
|
|
tensor_parallel_config=SimpleNamespace(autotp_size=1),
|
|
expert_parallel_config=AutoEPConfig(enabled=True, autoep_size=1),
|
|
)
|
|
|
|
with pytest.raises(AssertionError, match="sequence parallelism"):
|
|
engine._validate_zero3_moe_compatibility()
|
|
|
|
def test_zero3_compatibility_gate_rejects_active_autotp(self):
|
|
model = MockMoETransformer(num_layers=1)
|
|
replace_autoep_layers(model, "mixtral")
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.__dict__["module"] = model
|
|
engine.has_moe_layers = True
|
|
engine.sequence_parallel_size = 1
|
|
engine.zero_quantized_gradients = lambda: False
|
|
engine._config = SimpleNamespace(
|
|
tensor_parallel_config=SimpleNamespace(autotp_size=2),
|
|
expert_parallel_config=AutoEPConfig(enabled=True, autoep_size=1),
|
|
)
|
|
|
|
with pytest.raises(AssertionError, match="AutoTP"):
|
|
engine._validate_zero3_moe_compatibility()
|
|
|
|
def test_zero3_compatibility_gate_rejects_quantized_gradients(self):
|
|
model = MockMoETransformer(num_layers=1)
|
|
replace_autoep_layers(model, "mixtral")
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.__dict__["module"] = model
|
|
engine.has_moe_layers = True
|
|
engine.sequence_parallel_size = 1
|
|
engine.zero_quantized_gradients = lambda: True
|
|
engine._config = SimpleNamespace(
|
|
tensor_parallel_config=SimpleNamespace(autotp_size=1),
|
|
expert_parallel_config=AutoEPConfig(enabled=True, autoep_size=1),
|
|
)
|
|
|
|
with pytest.raises(AssertionError, match="zero_quantized_gradients"):
|
|
engine._validate_zero3_moe_compatibility()
|
|
|
|
def test_zero3_compatibility_gate_rejects_mics(self):
|
|
model = MockMoETransformer(num_layers=1)
|
|
replace_autoep_layers(model, "mixtral")
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.__dict__["module"] = model
|
|
engine.has_moe_layers = True
|
|
engine.sequence_parallel_size = 1
|
|
engine.zero_quantized_gradients = lambda: False
|
|
engine._config = SimpleNamespace(
|
|
mics_shard_size=2,
|
|
zero_config=SimpleNamespace(zero_hpz_partition_size=1),
|
|
tensor_parallel_config=SimpleNamespace(autotp_size=1),
|
|
expert_parallel_config=AutoEPConfig(enabled=True, autoep_size=1),
|
|
)
|
|
|
|
with pytest.raises(AssertionError, match="MiCS"):
|
|
engine._validate_zero3_moe_compatibility()
|
|
|
|
def test_zero3_compatibility_gate_rejects_hpzero(self):
|
|
model = MockMoETransformer(num_layers=1)
|
|
replace_autoep_layers(model, "mixtral")
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.__dict__["module"] = model
|
|
engine.has_moe_layers = True
|
|
engine.sequence_parallel_size = 1
|
|
engine.zero_quantized_gradients = lambda: False
|
|
engine._config = SimpleNamespace(
|
|
mics_shard_size=0,
|
|
zero_config=SimpleNamespace(zero_hpz_partition_size=2),
|
|
tensor_parallel_config=SimpleNamespace(autotp_size=1),
|
|
expert_parallel_config=AutoEPConfig(enabled=True, autoep_size=1),
|
|
)
|
|
|
|
with pytest.raises(AssertionError, match="hpZeRO"):
|
|
engine._validate_zero3_moe_compatibility()
|
|
|
|
def test_autoep_layer_marks_zero3_param_placement_families(self):
|
|
model = MockMoETransformer(num_layers=1)
|
|
replace_autoep_layers(model, "mixtral")
|
|
autoep_layer = next(module for module in model.modules() if isinstance(module, AutoEPMoELayer))
|
|
|
|
for param in autoep_layer.experts.parameters():
|
|
assert param.ds_zero_placement_family == "autoep_expert"
|
|
assert param.ds_zero_partition_group_name == autoep_layer.ep_group_name
|
|
|
|
for param in autoep_layer.router.parameters():
|
|
assert param.ds_zero_placement_family == "replicated"
|
|
|
|
def test_zero3_checkpoint_metadata_includes_partition_group_ranks(self):
|
|
optimizer = object.__new__(DeepSpeedZeroOptimizer_Stage3)
|
|
param = nn.Parameter(torch.empty(1))
|
|
param.ds_zero_placement_family = "autoep_expert"
|
|
param.ds_zero_partition_group_name = "ep_size_2"
|
|
optimizer.fp16_groups = [[param]]
|
|
optimizer._get_sub_group_partition_count = lambda _: 2
|
|
optimizer._get_sub_group_partition_rank = lambda _: 1
|
|
optimizer._get_sub_group_partition_ranks = lambda _: [1, 3]
|
|
|
|
metadata = optimizer._zero3_partition_group_metadata()
|
|
|
|
assert metadata == [{
|
|
"sub_group": 0,
|
|
"partition_count": 2,
|
|
"partition_rank": 1,
|
|
"partition_ranks": [1, 3],
|
|
"families": ["autoep_expert"],
|
|
"group_names": ["ep_size_2"],
|
|
}]
|
|
|
|
param.ds_zero_placement_family = "replicated"
|
|
param.ds_zero_partition_group_name = None
|
|
assert optimizer._zero3_partition_group_metadata() is None
|
|
|
|
def test_zero3_cpu_offload_grad_norm_reduces_autoep_expert_parallel_group(self, monkeypatch):
|
|
optimizer = object.__new__(DeepSpeedZeroOptimizer_Stage3)
|
|
param = nn.Parameter(torch.empty(1))
|
|
param.ds_zero_placement_family = "autoep_expert"
|
|
param.ds_zero_partition_group_name = "ep_size_2"
|
|
optimizer.model_parallel_rank = 0
|
|
optimizer.norm_for_param_grads = {7: 3.0}
|
|
optimizer.get_param_id = lambda _: 7
|
|
optimizer._assert_same_partition_group = lambda _: None
|
|
optimizer._get_param_partition_group = lambda _: "expert_data_parallel"
|
|
optimizer._model_parallel_all_reduce = lambda tensor, op: None
|
|
optimizer._autoep_expert_parallel_group = lambda _: "expert_parallel"
|
|
calls = []
|
|
|
|
def fake_all_reduce(tensor, op=None, group=None):
|
|
calls.append(group)
|
|
|
|
class FakeAccelerator:
|
|
|
|
def FloatTensor(self, values):
|
|
return torch.FloatTensor(values)
|
|
|
|
monkeypatch.setattr(zero_stage3, "get_accelerator", lambda: FakeAccelerator())
|
|
monkeypatch.setattr(zero_stage3.dist, "all_reduce", fake_all_reduce)
|
|
|
|
norm = optimizer.complete_grad_norm_calculation_for_cpu_offload([param])
|
|
|
|
assert calls == ["expert_data_parallel", "expert_parallel"]
|
|
assert torch.isfinite(norm)
|
|
|
|
def test_zero3_autoep_reduce_scatter_grads_average_by_global_dp(self, monkeypatch):
|
|
optimizer = object.__new__(DeepSpeedZeroOptimizer_Stage3)
|
|
optimizer.dp_process_group = "global_data_parallel"
|
|
optimizer.dtype = torch.float32
|
|
optimizer.gradient_accumulation_dtype = torch.float32
|
|
optimizer.postscale_gradients = True
|
|
optimizer.gradient_predivide_factor = 1.0
|
|
optimizer.all2all_process_group = None
|
|
optimizer._assert_same_partition_group = lambda _: None
|
|
optimizer._get_param_partition_group = lambda _: "expert_data_parallel"
|
|
optimizer._autoep_expert_parallel_group = lambda _: "expert_parallel"
|
|
param = nn.Parameter(torch.ones(4))
|
|
param.grad = torch.ones(4)
|
|
|
|
class FakeAccelerator:
|
|
|
|
def device_count(self):
|
|
return 4
|
|
|
|
def fake_get_world_size(group=None):
|
|
return 2 if group == "expert_data_parallel" else 4
|
|
|
|
def fake_reduce_scatter(grads, process_group):
|
|
assert process_group == "expert_data_parallel"
|
|
return [torch.full((2, ), 8.0)]
|
|
|
|
monkeypatch.setattr(zero_stage3, "get_accelerator", lambda: FakeAccelerator())
|
|
monkeypatch.setattr(zero_stage3.dist, "get_world_size", fake_get_world_size)
|
|
monkeypatch.setattr(zero_stage3, "reduce_scatter_coalesced", fake_reduce_scatter)
|
|
|
|
grad_partitions = optimizer._DeepSpeedZeroOptimizer_Stage3__avg_scatter_grads([param], torch.float32)
|
|
|
|
torch.testing.assert_close(grad_partitions[0], torch.full((2, ), 4.0))
|
|
|
|
def test_zero3_autoep_contiguous_grads_average_by_global_dp(self, monkeypatch):
|
|
optimizer = object.__new__(DeepSpeedZeroOptimizer_Stage3)
|
|
optimizer.dp_process_group = "global_data_parallel"
|
|
optimizer.ipg_buckets = {torch.float32: SimpleNamespace(params=[], process_group="expert_data_parallel")}
|
|
optimizer.postscale_gradients = True
|
|
optimizer.gradient_predivide_factor = 1.0
|
|
optimizer.sequence_parallel_size = 1
|
|
optimizer.gradient_accumulation_dtype = torch.float32
|
|
optimizer._assert_same_partition_group = lambda _: None
|
|
optimizer._autoep_expert_parallel_group = lambda _: "expert_parallel"
|
|
optimizer._apply_distributed_muon_update = lambda communication_data_type, buffer: None
|
|
param = nn.Parameter(torch.empty(2))
|
|
param.grad = torch.zeros(2)
|
|
param.partition_numel = lambda: 1
|
|
optimizer.ipg_buckets[torch.float32].params = [param]
|
|
|
|
def fake_get_world_size(group=None):
|
|
return 2 if group == "expert_data_parallel" else 4
|
|
|
|
def fake_all_reduce(tensor, group=None):
|
|
assert group == "expert_data_parallel"
|
|
tensor.mul_(2)
|
|
|
|
monkeypatch.setattr(zero_stage3.dist, "get_world_size", fake_get_world_size)
|
|
monkeypatch.setattr(zero_stage3.dist, "get_rank", lambda group=None: 0)
|
|
monkeypatch.setattr(zero_stage3.dist, "all_reduce", fake_all_reduce)
|
|
|
|
grad_partitions = optimizer._DeepSpeedZeroOptimizer_Stage3__avg_scatter_contiguous_grads(
|
|
torch.tensor([4.0, 8.0]), torch.float32)
|
|
|
|
torch.testing.assert_close(grad_partitions[0], torch.tensor([2.0]))
|
|
|
|
def test_pipeline_load_module_state_dict_accepts_autoep_zero3_fetch_kwarg(self):
|
|
from deepspeed.runtime.pipe.engine import PipelineEngine
|
|
|
|
signature = inspect.signature(PipelineEngine.load_module_state_dict)
|
|
|
|
assert "z3_params_to_fetch" in signature.parameters
|
|
assert "allowed_missing_keys" in signature.parameters
|
|
|
|
def test_load_module_state_dict_nonstrict_keeps_nonstrict_semantics_with_allowed_missing_keys(self):
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
# bypass nn.Module.__setattr__, which requires Module.__init__
|
|
object.__setattr__(engine, "module", nn.Linear(2, 2))
|
|
checkpoint = {"module": {"unexpected_key": torch.zeros(1)}}
|
|
|
|
# strict=False must keep the documented non-strict load semantics even
|
|
# when AutoEP expert keys are allowed to be missing.
|
|
engine.load_module_state_dict(checkpoint, strict=False, allowed_missing_keys=["weight"])
|
|
|
|
with pytest.raises(RuntimeError, match="outside AutoEP expert"):
|
|
engine.load_module_state_dict(checkpoint, strict=True, allowed_missing_keys=["weight"])
|
|
|
|
def test_resolve_zero3_param_placement_rejects_pre_partitioned_expert_on_wrong_group(self, monkeypatch):
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
model = nn.Linear(2, 2, bias=False)
|
|
# bypass nn.Module.__setattr__, which requires Module.__init__
|
|
object.__setattr__(engine, "module", model)
|
|
|
|
expert_group = object()
|
|
other_group = object()
|
|
monkeypatch.setattr(ds_engine.groups, "_get_expert_data_parallel_group", lambda name: expert_group)
|
|
monkeypatch.setattr(ds_engine.dist, "get_rank", lambda group=None: 0)
|
|
monkeypatch.setattr(ds_engine.dist, "get_world_size", lambda group=None: 1)
|
|
monkeypatch.setattr(ds_engine.dist,
|
|
"get_all_ranks_from_group",
|
|
lambda group: [0] if group is expert_group else [0, 1],
|
|
raising=False)
|
|
|
|
param = model.weight
|
|
param.ds_zero_placement_family = "autoep_expert"
|
|
param.ds_zero_partition_group_name = "ep_size_2"
|
|
param.ds_id = 0
|
|
param.ds_process_group = other_group
|
|
|
|
with pytest.raises(AssertionError, match="already ZeRO-partitioned over a non-expert process group"):
|
|
engine._resolve_zero3_param_placement()
|
|
|
|
# A pre-partitioned expert param over the matching group is accepted
|
|
# and keeps metadata derived from its actual partition group.
|
|
param.ds_process_group = expert_group
|
|
engine._resolve_zero3_param_placement()
|
|
assert param.ds_zero_partition_process_group is expert_group
|
|
|
|
def test_autoep_zero3_16bit_export_guard_directs_to_universal_conversion(self):
|
|
engine = object.__new__(DeepSpeedEngine)
|
|
engine.zero_optimization_partition_weights = lambda: True
|
|
engine._has_autoep_layers = lambda: True
|
|
|
|
with pytest.raises(NotImplementedError, match="ds_to_universal.py"):
|
|
engine._raise_if_autoep_zero3_consolidated_export("save_16bit_model")
|
|
|
|
def test_universal_converter_detects_zero3_partitioned_autoep_model_state(self, tmp_path):
|
|
from deepspeed.checkpoint.constants import (
|
|
AUTOEP_LAYERS_KEY,
|
|
AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION,
|
|
AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY,
|
|
AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY,
|
|
AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT,
|
|
)
|
|
from deepspeed.checkpoint.ds_to_universal import (
|
|
_autoep_expert_param_names_by_rank,
|
|
_get_zero3_model_state_files,
|
|
_uses_zero3_partitioned_autoep_metadata,
|
|
)
|
|
|
|
zero3_model_file = tmp_path / "zero_pp_rank_0_mp_rank_00_model_states.pt"
|
|
expert_file = tmp_path / "layer_0_expert_0_mp_rank_00_model_states.pt"
|
|
metadata = [{
|
|
"moe_layer_id": 0,
|
|
"module_path": "model.layers.0.mlp",
|
|
"num_experts": 4,
|
|
"num_local_experts": 2,
|
|
"ep_size": 2,
|
|
"expert_key_prefix": "model.layers.0.mlp.experts",
|
|
AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY: AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT,
|
|
AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY: AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION,
|
|
"ep_group_name": "ep_size_2",
|
|
"ep_rank": 0,
|
|
"expert_data_parallel_rank": 0,
|
|
"expert_data_parallel_world_size": 1,
|
|
"global_expert_start": 0,
|
|
"global_expert_end": 2,
|
|
}]
|
|
torch.save({AUTOEP_LAYERS_KEY: metadata}, zero3_model_file)
|
|
torch.save({"expert": torch.empty(1)}, expert_file)
|
|
|
|
model_files = _get_zero3_model_state_files(str(tmp_path))
|
|
expert_param_names, metadata_by_rank = _autoep_expert_param_names_by_rank(model_files)
|
|
|
|
assert model_files == [str(zero3_model_file)]
|
|
assert expert_param_names == {
|
|
"model.layers.0.mlp.experts.w1",
|
|
"model.layers.0.mlp.experts.w2",
|
|
"model.layers.0.mlp.experts.w3",
|
|
}
|
|
assert _uses_zero3_partitioned_autoep_metadata(metadata_by_rank[0])
|
|
|
|
def test_universal_stage3_extract_accepts_tuple_param_shapes(self, tmp_path):
|
|
from deepspeed.checkpoint.constants import OPTIMIZER_STATE_DICT
|
|
from deepspeed.checkpoint.ds_to_universal import extract_zero_shards_stage3
|
|
|
|
optim_file = tmp_path / "zero_pp_rank_0_mp_rank_00_optim_states.pt"
|
|
torch.save(
|
|
{
|
|
OPTIMIZER_STATE_DICT: {
|
|
"optimizer_state_dict": {
|
|
"state": [{
|
|
"exp_avg": torch.arange(6, dtype=torch.float32),
|
|
"exp_avg_sq": torch.arange(6, dtype=torch.float32) + 10,
|
|
}]
|
|
},
|
|
"fp32_flat_groups": [torch.arange(6, dtype=torch.float32) + 20],
|
|
}
|
|
},
|
|
optim_file,
|
|
)
|
|
|
|
temp_dir = tmp_path / "tmp"
|
|
extract_zero_shards_stage3([str(optim_file)], [OrderedDict([("dense.weight", (2, 3))])], 1, str(temp_dir), 0)
|
|
|
|
fp32_fragment = torch.load(temp_dir / "dense.weight" / "0" / "fp32.00", weights_only=False)
|
|
exp_avg_fragment = torch.load(temp_dir / "dense.weight" / "0" / "exp_avg.00", weights_only=False)
|
|
torch.testing.assert_close(fp32_fragment, torch.arange(6, dtype=torch.float32) + 20)
|
|
torch.testing.assert_close(exp_avg_fragment, torch.arange(6, dtype=torch.float32))
|
|
|
|
def test_zero_to_fp32_rejects_zero3_partitioned_autoep_checkpoint(self, tmp_path):
|
|
from deepspeed.checkpoint.constants import (
|
|
AUTOEP_LAYERS_KEY,
|
|
AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION,
|
|
AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY,
|
|
AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY,
|
|
AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT,
|
|
BUFFER_NAMES,
|
|
PARAM_SHAPES,
|
|
)
|
|
from deepspeed.utils.zero_to_fp32 import _raise_if_autoep_zero3_partitioned_checkpoint
|
|
|
|
model_file = tmp_path / "zero_pp_rank_0_mp_rank_00_model_states.pt"
|
|
torch.save(
|
|
{
|
|
BUFFER_NAMES: [],
|
|
PARAM_SHAPES: [],
|
|
"module": {},
|
|
"shared_params": {},
|
|
AUTOEP_LAYERS_KEY: [{
|
|
"moe_layer_id": 0,
|
|
"module_path": "model.layers.0.mlp",
|
|
"num_experts": 4,
|
|
"num_local_experts": 2,
|
|
"ep_size": 2,
|
|
"expert_key_prefix": "model.layers.0.mlp.experts",
|
|
AUTOEP_ZERO3_EXPERT_STATE_FORMAT_KEY: AUTOEP_ZERO3_PARTITIONED_EXPERT_STATE_FORMAT,
|
|
AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION_KEY: AUTOEP_ZERO3_EXPERT_STATE_FORMAT_VERSION,
|
|
"ep_group_name": "ep_size_2",
|
|
"ep_rank": 0,
|
|
"expert_data_parallel_rank": 0,
|
|
"expert_data_parallel_world_size": 1,
|
|
"global_expert_start": 0,
|
|
"global_expert_end": 2,
|
|
}],
|
|
},
|
|
model_file,
|
|
)
|
|
|
|
with pytest.raises(NotImplementedError, match="ds_to_universal.py"):
|
|
_raise_if_autoep_zero3_partitioned_checkpoint([str(model_file)])
|
|
|
|
# parse_model_states is the guard point used by
|
|
# _get_fp32_state_dict_from_zero_checkpoint, which loads each model
|
|
# state file only once.
|
|
from deepspeed.utils.zero_to_fp32 import parse_model_states
|
|
with pytest.raises(NotImplementedError, match="ds_to_universal.py"):
|
|
parse_model_states([str(model_file)])
|
|
|
|
def test_preset_registry_core_contracts(self):
|
|
assert set(PRESET_MODELS) == {"mixtral", "qwen3_moe", "qwen3_5_moe", "deepseek_v2", "deepseek_v3"}
|
|
assert preset_name_for_hf_model_type("mixtral") == "mixtral"
|
|
assert preset_name_for_hf_model_type("qwen2_moe") == "qwen3_moe"
|
|
assert preset_name_for_hf_model_type("llama4_text") is None
|
|
|
|
qwen35 = unsupported_preset_for_hf_model_type("qwen3_5_moe")
|
|
assert qwen35 is not None
|
|
assert "qwen3_5_moe_text" in qwen35[1].unsupported_hf_model_type_notes["qwen3_5_moe"]
|
|
assert PRESET_MODELS["deepseek_v2"].supports_expert_bias is False
|
|
assert PRESET_MODELS["deepseek_v3"].unsupported_router_bias_names == ()
|
|
|
|
def test_fill_autoep_config_from_hf_defaults(self):
|
|
config = AutoEPConfig(enabled=True, autoep_size=2)
|
|
|
|
fill_autoep_config_from_hf(config, MockDeepSeekV3Config())
|
|
|
|
assert config.num_expert_groups == 4
|
|
assert config.num_limited_groups == 2
|
|
assert config.route_scale == pytest.approx(2.5)
|
|
|
|
def test_fill_autoep_config_from_hf_preserves_explicit_values(self):
|
|
config = AutoEPConfig(enabled=True,
|
|
autoep_size=2,
|
|
num_expert_groups=8,
|
|
num_limited_groups=1,
|
|
routed_scaling_factor=3.0,
|
|
route_scale=3.0)
|
|
|
|
fill_autoep_config_from_hf(config, MockDeepSeekV3Config())
|
|
|
|
assert config.num_expert_groups == 8
|
|
assert config.num_limited_groups == 1
|
|
assert config.route_scale == pytest.approx(3.0)
|
|
|
|
@pytest.mark.parametrize("value", ["2.5", True, float("nan"), float("inf")])
|
|
def test_invalid_routed_scaling_factor_rejected(self, value):
|
|
with pytest.raises(ValueError, match="routed_scaling_factor"):
|
|
_resolve_route_scale(AutoEPConfig(enabled=True, routed_scaling_factor=value), None)
|
|
|
|
|
|
class TestRoutingAndLayerSemantics:
|
|
|
|
def test_router_route_scale_and_group_limited_routing(self):
|
|
base = TokenChoiceTopKRouter(64, 8, 4, 2, 2, "softmax", False, 1.0, False)
|
|
scaled = TokenChoiceTopKRouter(64, 8, 4, 2, 2, "softmax", False, 2.5, False)
|
|
scaled.load_state_dict(base.state_dict())
|
|
x = torch.randn(50, 64)
|
|
|
|
base_scores, base_experts, base_counts = base(x)
|
|
scaled_scores, scaled_experts, scaled_counts = scaled(x)
|
|
|
|
assert torch.equal(scaled_experts, base_experts)
|
|
assert torch.allclose(scaled_scores, base_scores * 2.5, atol=1e-5)
|
|
assert torch.equal(scaled_counts, base_counts)
|
|
assert base_counts.shape == (8, )
|
|
|
|
def test_grouped_experts_and_token_reorderer(self):
|
|
experts = GroupedExperts(dim=64, hidden_dim=128, num_experts=4, use_grouped_mm=False)
|
|
nn.init.normal_(experts.w1, std=0.02)
|
|
nn.init.normal_(experts.w2, std=0.02)
|
|
nn.init.normal_(experts.w3, std=0.02)
|
|
out = experts(torch.randn(8, 64), torch.tensor([2, 2, 2, 2]))
|
|
assert out.shape == (8, 64)
|
|
assert not torch.isnan(out).any()
|
|
|
|
top_scores = torch.randn(20, 2)
|
|
selected_experts = torch.randint(0, 4, (20, 2))
|
|
scores_sorted, indices_sorted, counts = TokenReorderer(num_experts=4, top_k=2)(top_scores, selected_experts)
|
|
assert scores_sorted.shape == (40, )
|
|
assert set(indices_sorted.tolist()) == set(range(40))
|
|
assert torch.equal(counts, torch.bincount(selected_experts.reshape(-1), minlength=4).to(counts.dtype))
|
|
|
|
def test_score_application_and_combine(self):
|
|
x = torch.randn(4, 8)
|
|
scores = torch.tensor([0.25, 0.5, 0.75, 1.0])
|
|
expected = x.float() * scores.reshape(-1, 1)
|
|
torch.testing.assert_close(apply_scores_before_experts_if_enabled(x, scores, "pre"), expected.to(x.dtype))
|
|
|
|
spec = _make_spec(score_apply="post")
|
|
assert resolve_score_apply_mode(spec, "auto") == "post"
|
|
expert_output = torch.ones(4, 8)
|
|
top_scores = torch.tensor([[0.6, 0.4], [0.7, 0.3]])
|
|
out = combine_from_routed(expert_output, top_scores, torch.arange(4), 2, "post", "weighted_sum", (1, 2, 8))
|
|
torch.testing.assert_close(out[0, 0], torch.ones(8))
|
|
|
|
def test_autoep_layer_forward_and_expert_bias_rejection(self):
|
|
source = MockMoEBlock(num_experts=4, ffn_hidden=128, hidden_size=64)
|
|
layer = AutoEPMoELayer(_make_spec(route_scale=2.5),
|
|
source,
|
|
ep_size=1,
|
|
ep_rank=0,
|
|
config=_runtime_config(enabled=True, autoep_size=1))
|
|
out = layer(torch.randn(2, 8, 64))
|
|
assert layer._is_autoep_layer is True
|
|
assert layer.num_experts == 4
|
|
assert layer.router.route_scale == pytest.approx(2.5)
|
|
assert out.shape == (2, 8, 64)
|
|
assert not torch.isnan(out).any()
|
|
|
|
with pytest.raises(ValueError, match="load_balance_coeff/expert_bias"):
|
|
AutoEPMoELayer(_make_spec(model_family="no_bias_family", supports_expert_bias=False),
|
|
source,
|
|
ep_size=1,
|
|
ep_rank=0,
|
|
config=AutoEPConfig(enabled=True, autoep_size=1, load_balance_coeff=0.02))
|
|
|
|
|
|
class TestModelDetectionAndReplacement:
|
|
|
|
def test_mixtral_detect_replace_and_mock_forward(self):
|
|
model = MockMoETransformer(num_layers=2, moe_every_n=1)
|
|
auto_ep = AutoEP(model, _runtime_config(enabled=True, autoep_size=1, preset_model="mixtral"))
|
|
specs = auto_ep.ep_parser()
|
|
|
|
assert len(specs) == 2
|
|
assert specs[0].model_family == "mixtral"
|
|
auto_ep.replace_moe_layer(specs[0], ep_size=1, ep_rank=0)
|
|
assert isinstance(model.model.layers[0].mlp, AutoEPMoELayer)
|
|
assert model(torch.randn(1, 4, 64)).shape == (1, 4, 100)
|
|
|
|
def test_fused_replacement_preserves_frozen_experts_and_trainable_router(self):
|
|
model = MockMoETransformer(num_layers=1, num_experts=4, moe_every_n=1).to(dtype=torch.bfloat16)
|
|
source = model.model.layers[0].mlp
|
|
source.experts.gate_up_proj.requires_grad_(False)
|
|
source.experts.down_proj.requires_grad_(False)
|
|
source.gate.weight.requires_grad_(True)
|
|
|
|
auto_ep = AutoEP(model, _runtime_config(enabled=True, autoep_size=1, preset_model="mixtral"))
|
|
spec = auto_ep.ep_parser()[0]
|
|
auto_ep.replace_moe_layer(spec, ep_size=1, ep_rank=0)
|
|
|
|
replaced = model.model.layers[0].mlp
|
|
assert isinstance(replaced, AutoEPMoELayer)
|
|
assert replaced.experts.w1.requires_grad is False
|
|
assert replaced.experts.w2.requires_grad is False
|
|
assert replaced.experts.w3.requires_grad is False
|
|
assert replaced.router.gate.weight.requires_grad is True
|
|
_assert_same_dtype_device(replaced.router.gate.weight, source.gate.weight)
|
|
_assert_same_dtype_device(replaced.experts.w1, source.experts.gate_up_proj)
|
|
_assert_same_dtype_device(replaced.experts.w2, source.experts.down_proj)
|
|
_assert_same_dtype_device(replaced.experts.w3, source.experts.gate_up_proj)
|
|
|
|
def test_zero_init_source_gathered_for_parser_router_and_fused_repack(self, monkeypatch):
|
|
FakeGatheredParameters.calls = []
|
|
monkeypatch.setattr(ep_repack, "GatheredParameters", FakeGatheredParameters)
|
|
|
|
model = MockMoETransformer(num_layers=1, num_experts=4, moe_every_n=1)
|
|
source = model.model.layers[0].mlp
|
|
expected_gate = source.gate.weight.detach().clone()
|
|
expected_gate_up = source.experts.gate_up_proj.detach().clone()
|
|
expected_down = source.experts.down_proj.detach().clone()
|
|
|
|
_mark_fake_zero_param(source.gate.weight, expected_gate, ds_id=1, name="router.weight")
|
|
_mark_fake_zero_param(source.experts.gate_up_proj, expected_gate_up, ds_id=2, name="experts.gate_up_proj")
|
|
_mark_fake_zero_param(source.experts.down_proj, expected_down, ds_id=3, name="experts.down_proj")
|
|
|
|
auto_ep = AutoEP(model, _runtime_config(enabled=True, autoep_size=1, preset_model="mixtral"))
|
|
specs = auto_ep.ep_parser()
|
|
assert len(specs) == 1
|
|
assert specs[0].expert_storage == "fused_3d"
|
|
assert specs[0].num_experts == 4
|
|
assert specs[0].hidden_size == 64
|
|
|
|
auto_ep.replace_moe_layer(specs[0], ep_size=1, ep_rank=0)
|
|
|
|
replaced = model.model.layers[0].mlp
|
|
torch.testing.assert_close(replaced.router.gate.weight, expected_gate)
|
|
torch.testing.assert_close(replaced.experts.w1, expected_gate_up[:, :128, :])
|
|
torch.testing.assert_close(replaced.experts.w3, expected_gate_up[:, 128:, :])
|
|
torch.testing.assert_close(replaced.experts.w2, expected_down)
|
|
assert [call["names"] for call in FakeGatheredParameters.calls] == [
|
|
["router.weight"],
|
|
["experts.gate_up_proj", "experts.down_proj"],
|
|
]
|
|
assert all(call["modifier_rank"] is None for call in FakeGatheredParameters.calls)
|
|
|
|
def test_module_list_replacement_preserves_frozen_experts_and_trainable_router(self, monkeypatch):
|
|
monkeypatch.setattr(get_preset_adapter("deepseek_v3"), "_installed_transformers_version", lambda: "5.0.0")
|
|
model = MockDeepSeekV3Transformer(num_layers=1, num_experts=4).to(dtype=torch.bfloat16)
|
|
source = model.model.layers[0].mlp
|
|
for expert in source.experts:
|
|
for param in expert.parameters():
|
|
param.requires_grad_(False)
|
|
source.gate.weight.requires_grad_(True)
|
|
source.gate.e_score_correction_bias = nn.Parameter(torch.zeros(4,
|
|
dtype=source.gate.weight.dtype,
|
|
device=source.gate.weight.device),
|
|
requires_grad=True)
|
|
|
|
auto_ep = AutoEP(model, _runtime_config(enabled=True, autoep_size=2))
|
|
spec = auto_ep.ep_parser()[0]
|
|
auto_ep.replace_moe_layer(spec, ep_size=2, ep_rank=0)
|
|
|
|
replaced = model.model.layers[0].mlp
|
|
assert isinstance(replaced, AutoEPMoELayer)
|
|
assert replaced.experts.w1.requires_grad is False
|
|
assert replaced.experts.w2.requires_grad is False
|
|
assert replaced.experts.w3.requires_grad is False
|
|
assert replaced.router.gate.weight.requires_grad is True
|
|
assert replaced.router.e_score_correction_bias.requires_grad is True
|
|
_assert_same_dtype_device(replaced.router.gate.weight, source.gate.weight)
|
|
_assert_same_dtype_device(replaced.router.e_score_correction_bias, source.gate.e_score_correction_bias)
|
|
_assert_same_dtype_device(replaced.experts.w1, source.experts[0].gate_proj.weight)
|
|
_assert_same_dtype_device(replaced.experts.w2, source.experts[0].down_proj.weight)
|
|
_assert_same_dtype_device(replaced.experts.w3, source.experts[0].up_proj.weight)
|
|
|
|
def test_module_list_zero_source_gathers_all_experts_in_global_order(self, monkeypatch):
|
|
FakeGatheredParameters.calls = []
|
|
monkeypatch.setattr(ep_repack, "GatheredParameters", FakeGatheredParameters)
|
|
monkeypatch.setattr(get_preset_adapter("deepseek_v3"), "_installed_transformers_version", lambda: "5.0.0")
|
|
|
|
model = MockDeepSeekV3Transformer(num_layers=1, num_experts=4)
|
|
source = model.model.layers[0].mlp
|
|
for expert_idx, expert in enumerate(source.experts):
|
|
for offset, (suffix, param) in enumerate((
|
|
("w1", expert.gate_proj.weight),
|
|
("w2", expert.down_proj.weight),
|
|
("w3", expert.up_proj.weight),
|
|
)):
|
|
full_data = param.detach().clone()
|
|
_mark_fake_zero_param(param,
|
|
full_data,
|
|
ds_id=10 + 3 * expert_idx + offset,
|
|
name=f"e{expert_idx}.{suffix}")
|
|
|
|
auto_ep = AutoEP(model, _runtime_config(enabled=True, autoep_size=2))
|
|
spec = auto_ep.ep_parser()[0]
|
|
w1, w2, w3 = repack_expert_weights(source.experts, spec, ep_rank=1, ep_size=2)
|
|
|
|
expected_w1 = torch.stack([
|
|
source.experts[2].gate_proj.weight._autoep_test_full_data,
|
|
source.experts[3].gate_proj.weight._autoep_test_full_data
|
|
])
|
|
expected_w2 = torch.stack([
|
|
source.experts[2].down_proj.weight._autoep_test_full_data,
|
|
source.experts[3].down_proj.weight._autoep_test_full_data
|
|
])
|
|
expected_w3 = torch.stack([
|
|
source.experts[2].up_proj.weight._autoep_test_full_data,
|
|
source.experts[3].up_proj.weight._autoep_test_full_data
|
|
])
|
|
|
|
torch.testing.assert_close(w1, expected_w1)
|
|
torch.testing.assert_close(w2, expected_w2)
|
|
torch.testing.assert_close(w3, expected_w3)
|
|
assert [call["names"] for call in FakeGatheredParameters.calls] == [
|
|
["e0.w1", "e0.w2", "e0.w3"],
|
|
["e1.w1", "e1.w2", "e1.w3"],
|
|
["e2.w1", "e2.w2", "e2.w3"],
|
|
["e3.w1", "e3.w2", "e3.w3"],
|
|
]
|
|
|
|
def test_module_list_mixed_expert_requires_grad_flags_are_rejected(self, monkeypatch):
|
|
monkeypatch.setattr(get_preset_adapter("deepseek_v3"), "_installed_transformers_version", lambda: "5.0.0")
|
|
model = MockDeepSeekV3Transformer(num_layers=1, num_experts=4)
|
|
source = model.model.layers[0].mlp
|
|
source.experts[0].gate_proj.weight.requires_grad_(False)
|
|
source.experts[1].gate_proj.weight.requires_grad_(True)
|
|
|
|
auto_ep = AutoEP(model, _runtime_config(enabled=True, autoep_size=2))
|
|
spec = auto_ep.ep_parser()[0]
|
|
with pytest.raises(ValueError, match="mixed requires_grad flags"):
|
|
auto_ep.replace_moe_layer(spec, ep_size=2, ep_rank=0)
|
|
|
|
model = MockDeepSeekV3Transformer(num_layers=1, num_experts=4)
|
|
source = model.model.layers[0].mlp
|
|
source.experts[1].gate_proj.to(dtype=torch.float64)
|
|
|
|
auto_ep = AutoEP(model, _runtime_config(enabled=True, autoep_size=2))
|
|
spec = auto_ep.ep_parser()[0]
|
|
with pytest.raises(ValueError, match="mixed dtype/device"):
|
|
auto_ep.replace_moe_layer(spec, ep_size=2, ep_rank=0)
|
|
|
|
def test_hf_mixtral_causal_lm_matches_autoep_with_router_logits(self):
|
|
transformers = pytest.importorskip("transformers")
|
|
skip_unless_transformers_has(transformers,
|
|
"MixtralConfig",
|
|
"MixtralForCausalLM",
|
|
min_version="5.0.0",
|
|
reason="Mixtral AutoEP router-logit capture")
|
|
|
|
torch.manual_seed(1234)
|
|
config = tiny_mixtral_config(transformers)
|
|
native_model, autoep_model = state_matched_models(transformers.MixtralForCausalLM, config)
|
|
replace_autoep_layers(autoep_model, "mixtral")
|
|
assert_causal_lm_outputs_close(native_model,
|
|
autoep_model,
|
|
output_router_logits=True,
|
|
compare_router_logits=True,
|
|
compare_aux_loss=True,
|
|
compare_logits=False)
|
|
|
|
def test_qwen_adapter_guards(self, monkeypatch):
|
|
monkeypatch.setattr(get_preset_adapter("qwen3_moe"), "_installed_transformers_version", lambda: "5.0.0")
|
|
model = MockMoETransformer(num_layers=1, num_experts=4, moe_every_n=1)
|
|
model.config.model_type = "qwen2_moe"
|
|
model.config.num_experts = model.config.num_local_experts
|
|
|
|
specs = AutoEP(model, _runtime_config(enabled=True, autoep_size=1)).ep_parser()
|
|
|
|
assert len(specs) == 1
|
|
assert specs[0].model_family == "qwen3_moe"
|
|
|
|
model.config.model_type = "qwen3_5_moe"
|
|
with pytest.raises(ValueError, match="qwen3_5_moe_text"):
|
|
AutoEP(model, _runtime_config(enabled=True, autoep_size=1))._resolve_presets()
|
|
|
|
def test_deepseek_v3_detection_and_score_correction_bias_copy(self, monkeypatch):
|
|
FakeGatheredParameters.calls = []
|
|
monkeypatch.setattr(ep_repack, "GatheredParameters", FakeGatheredParameters)
|
|
monkeypatch.setattr(get_preset_adapter("deepseek_v3"), "_installed_transformers_version", lambda: "5.0.0")
|
|
model = MockDeepSeekV3Transformer(num_layers=1, num_experts=8)
|
|
auto_ep = AutoEP(model, _runtime_config(enabled=True, autoep_size=2))
|
|
specs = auto_ep.ep_parser()
|
|
|
|
assert len(specs) == 1
|
|
assert specs[0].model_family == "deepseek_v3"
|
|
assert specs[0].expert_storage == "module_list"
|
|
assert specs[0].expert_w1_name == "gate_proj"
|
|
assert specs[0].has_shared_experts is True
|
|
|
|
source_bias = torch.arange(8, dtype=torch.float32)
|
|
model.model.layers[0].mlp.gate.e_score_correction_bias = nn.Parameter(source_bias.clone())
|
|
_mark_fake_zero_param(model.model.layers[0].mlp.gate.e_score_correction_bias,
|
|
source_bias,
|
|
ds_id=100,
|
|
name="router.e_score_correction_bias")
|
|
|
|
auto_ep.replace_moe_layer(specs[0], ep_size=2, ep_rank=0)
|
|
|
|
replaced = model.model.layers[0].mlp
|
|
assert isinstance(replaced, AutoEPMoELayer)
|
|
assert replaced.router.e_score_correction_bias is not None
|
|
torch.testing.assert_close(replaced.router.e_score_correction_bias, source_bias)
|
|
assert ["router.e_score_correction_bias"] in [call["names"] for call in FakeGatheredParameters.calls]
|
|
|
|
|
|
def _eager_pep604_lines(module):
|
|
"""Line numbers where a module evaluates PEP 604 unions at import time."""
|
|
tree = ast.parse(inspect.getsource(module))
|
|
defers_annotations = any(
|
|
isinstance(node, ast.ImportFrom) and node.module == "__future__" and any(alias.name == "annotations"
|
|
for alias in node.names)
|
|
for node in tree.body)
|
|
if defers_annotations:
|
|
return []
|
|
offending_lines = []
|
|
for node in ast.walk(tree):
|
|
annotations = []
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
|
|
arguments = node.args
|
|
for arg in arguments.args + arguments.posonlyargs + arguments.kwonlyargs + [
|
|
arguments.vararg, arguments.kwarg
|
|
]:
|
|
if arg is not None and arg.annotation is not None:
|
|
annotations.append(arg.annotation)
|
|
if node.returns is not None:
|
|
annotations.append(node.returns)
|
|
elif isinstance(node, ast.AnnAssign):
|
|
annotations.append(node.annotation)
|
|
for annotation in annotations:
|
|
for sub_node in ast.walk(annotation):
|
|
if isinstance(sub_node, ast.BinOp) and isinstance(sub_node.op, ast.BitOr):
|
|
offending_lines.append(sub_node.lineno)
|
|
return sorted(set(offending_lines))
|
|
|
|
|
|
class TestPy39AnnotationSafety:
|
|
|
|
def test_autoep_import_chain_defers_pep604_annotations(self):
|
|
"""PEP 604 unions (``int | None``) in def signatures or class-level
|
|
annotations are evaluated at import time, so on Python 3.9 they raise
|
|
TypeError while the module is imported; that escapes the engine's
|
|
``except ImportError`` guards around AutoEP and breaks every
|
|
``deepspeed.initialize()`` (issue #8102). Every module in the AutoEP
|
|
import chain must defer annotation evaluation with
|
|
``from __future__ import annotations``."""
|
|
import deepspeed.moe.ep_count as ep_count
|
|
import deepspeed.moe.ep_experts as ep_experts
|
|
import deepspeed.moe.ep_kernels as ep_kernels
|
|
import deepspeed.moe.ep_router as ep_router
|
|
import deepspeed.module_inject.auto_ep as auto_ep
|
|
import deepspeed.module_inject.auto_ep_config as auto_ep_config
|
|
import deepspeed.module_inject.auto_ep_layer as auto_ep_layer
|
|
import deepspeed.module_inject.auto_ep_preset_adapters as preset_adapters
|
|
import deepspeed.module_inject.auto_ep_presets.base as presets_base
|
|
import deepspeed.module_inject.auto_ep_presets.registry as presets_registry
|
|
|
|
autoep_import_chain = [
|
|
ep_count, ep_experts, ep_kernels, ep_router, ep_repack, auto_ep, auto_ep_config, auto_ep_layer,
|
|
preset_adapters, presets_base, presets_registry
|
|
]
|
|
for module in autoep_import_chain:
|
|
offending_lines = _eager_pep604_lines(module)
|
|
assert not offending_lines, (
|
|
f"{module.__name__} evaluates PEP 604 unions at import time (lines {offending_lines}); "
|
|
f"on Python 3.9 this raises TypeError during import and escapes the engine's "
|
|
f"except-ImportError guards (issue #8102). Add 'from __future__ import annotations'.")
|