diff --git a/deepspeed/module_inject/auto_ep.py b/deepspeed/module_inject/auto_ep.py index b119a218620a..8b3abcf58cad 100644 --- a/deepspeed/module_inject/auto_ep.py +++ b/deepspeed/module_inject/auto_ep.py @@ -126,11 +126,13 @@ def _has_3d_expert_params(module: nn.Module, preset: MoEModelPreset) -> bool: def _get_num_experts_from_config(model_config, preset: MoEModelPreset) -> int | None: """Extract num_experts from model.config using the preset's attribute name.""" + model_config = getattr(model_config, "text_config", model_config) return getattr(model_config, preset.num_experts_attr, None) def _get_top_k_from_config(model_config, preset: MoEModelPreset) -> int | None: """Extract top_k from model.config using the preset's attribute name.""" + model_config = getattr(model_config, "text_config", model_config) return getattr(model_config, preset.top_k_attr, None) @@ -611,6 +613,8 @@ def ep_parser(self) -> list[MoELayerSpec]: router_logits_capture_mode=forward_contract.router_logits_capture_mode, moe_output_shape=forward_contract.moe_output_shape, e_score_correction_bias_path=e_score_correction_bias_path, + router_topk_sorted=preset.router_topk_sorted, + router_scores_in_gate_dtype=preset.router_scores_in_gate_dtype, expert_activation=expert_activation, expert_activation_alpha=activation_alpha, expert_activation_limit=activation_limit, diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index 3c64b6370606..a90d722cbe4e 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -595,6 +595,8 @@ def __init__( route_scale=spec.route_scale, gate_bias=spec.gate_bias, group_score_func=spec.group_score_func, + topk_sorted=spec.router_topk_sorted, + scores_in_gate_dtype=spec.router_scores_in_gate_dtype, ) # Copy gate weights _copy_parameter_data(self.router.gate.weight, source_gate.weight) diff --git a/deepspeed/module_inject/auto_ep_presets/base.py b/deepspeed/module_inject/auto_ep_presets/base.py index d5904ad38bf9..1ea7ec8cf04e 100644 --- a/deepspeed/module_inject/auto_ep_presets/base.py +++ b/deepspeed/module_inject/auto_ep_presets/base.py @@ -58,6 +58,8 @@ class MoEModelPreset: min_transformers_version: str | None = None validated_transformers_versions: str = "" docs_support_notes: str = "" + router_topk_sorted: bool = False + router_scores_in_gate_dtype: bool = False @dataclass @@ -97,6 +99,8 @@ class MoELayerSpec: router_logits_capture_mode: Literal["raw", "post_score"] = "post_score" moe_output_shape: Literal["batched", "flat"] = "batched" e_score_correction_bias_path: str | None = None + router_topk_sorted: bool = False + router_scores_in_gate_dtype: bool = False #: Expert MLP activation resolved for this layer (see MoEModelPreset.expert_activation). expert_activation: str = "swiglu" expert_activation_alpha: float = 1.702 diff --git a/deepspeed/module_inject/auto_ep_presets/qwen3_5_moe.py b/deepspeed/module_inject/auto_ep_presets/qwen3_5_moe.py index 0b17e9c27de5..4ebf3e50893f 100644 --- a/deepspeed/module_inject/auto_ep_presets/qwen3_5_moe.py +++ b/deepspeed/module_inject/auto_ep_presets/qwen3_5_moe.py @@ -20,7 +20,7 @@ PRESET_NAME = "qwen3_5_moe" PRESET = MoEModelPreset( - moe_layer_pattern=r"model\.layers\.\d+\.mlp", + moe_layer_pattern=r"model\.(?:language_model\.)?layers\.\d+\.mlp", router_pattern="gate", experts_pattern="experts", expert_storage="fused_3d", @@ -37,13 +37,11 @@ shared_experts_pattern="shared_expert", shared_experts_gate_pattern="shared_expert_gate", preset_adapter="qwen3_5_moe", - hf_model_types=("qwen3_5_moe_text", ), - unsupported_hf_model_type_notes={ - "qwen3_5_moe": ("AutoEP supports the Qwen3.5 text backbone preset path; pass the " - "text-backbone model/config with model_type='qwen3_5_moe_text'.") - }, + hf_model_types=("qwen3_5_moe_text", "qwen3_5_moe"), min_transformers_version="5.2.0", - docs_support_notes="Requires the Qwen3.5 text-backbone qwen3_5_moe_text model type.", + docs_support_notes="Text backbone and the language-model path of the Qwen3.5 MoE wrapper.", + router_topk_sorted=True, + router_scores_in_gate_dtype=True, ) @@ -61,6 +59,7 @@ def adjust_forward_contract(self, contract: ForwardContract) -> ForwardContract: return_router_logits=True, capture_target="router", capture_index=1, + router_logits_capture_mode="raw", ) def retarget_transformers_output_recorders( diff --git a/deepspeed/moe/ep_router.py b/deepspeed/moe/ep_router.py index 57c6df682445..a0e528aedd57 100644 --- a/deepspeed/moe/ep_router.py +++ b/deepspeed/moe/ep_router.py @@ -60,6 +60,8 @@ def __init__( route_scale: float, gate_bias: bool, group_score_func: str = "top2_sum", + topk_sorted: bool = False, + scores_in_gate_dtype: bool = False, ): super().__init__() self.gate = nn.Linear(dim, num_experts, bias=gate_bias) @@ -71,6 +73,8 @@ def __init__( self.route_norm = route_norm self.route_scale = route_scale self.group_score_func = group_score_func + self.topk_sorted = topk_sorted + self.scores_in_gate_dtype = scores_in_gate_dtype # Trainable expert score correction bias (e.g. DeepSeek-V3/Moonlight noaux_tc). # Separate from the dynamic load-balancing expert_bias passed in forward(). self.e_score_correction_bias = None @@ -152,6 +156,7 @@ def forward( """ # Gate projection -> (T, num_experts) scores = self.gate(x) + gate_dtype = scores.dtype # Scoring in float32 to avoid loss explosion if self.score_func == "sigmoid": @@ -172,7 +177,7 @@ def forward( scores_for_choice = self._get_node_limited_routing_scores(scores_for_choice) # Select top-k experts per token - _, selected_experts_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False) + _, selected_experts_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=self.topk_sorted) # Gather original (unbiased) scores for selected experts top_scores = scores.gather(dim=1, index=selected_experts_indices) @@ -183,6 +188,8 @@ def forward( top_scores = top_scores / denominator top_scores = top_scores * self.route_scale + if self.scores_in_gate_dtype: + top_scores = top_scores.to(gate_dtype) num_tokens_per_expert = count_tokens_per_expert(selected_experts_indices, self.num_experts) diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index a20891eb904c..004fb3c34a6a 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -33,7 +33,7 @@ Transformers build that exposes the matching config/model classes, but the tested ``4.x`` builds do not match the validated AutoEP layout. * - ``qwen3_5_moe`` - ``5.2.0`` - - Requires the Qwen3.5 text-backbone ``qwen3_5_moe_text`` model type; + - Supports text input through the text backbone or the MoE wrapper; vision input is untested; for performance on Qwen3.5's Gated DeltaNet layers, install optimized kernels. See the `Hugging Face Transformers kernel loading docs `__ diff --git a/tests/unit/v1/moe/test_autoep_unit.py b/tests/unit/v1/moe/test_autoep_unit.py index af58fae49f98..80d90901f96f 100644 --- a/tests/unit/v1/moe/test_autoep_unit.py +++ b/tests/unit/v1/moe/test_autoep_unit.py @@ -1029,9 +1029,8 @@ def test_preset_registry_core_contracts(self): assert preset_name_for_hf_model_type("minimax_m3_vl_text") == "minimax_m3" 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_name_for_hf_model_type("qwen3_5_moe") == "qwen3_5_moe" + assert unsupported_preset_for_hf_model_type("qwen3_5_moe") is None minimax = unsupported_preset_for_hf_model_type("minimax_m3_vl") assert minimax is not None assert "minimax_m3_vl_text" in minimax[1].unsupported_hf_model_type_notes["minimax_m3_vl"] @@ -2173,8 +2172,9 @@ def test_qwen_adapter_guards(self, monkeypatch): 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() + monkeypatch.setattr(get_preset_adapter("qwen3_5_moe"), "_installed_transformers_version", lambda: "5.16.1") + specs = AutoEP(model, _runtime_config(enabled=True, autoep_size=1)).ep_parser() + assert specs[0].model_family == "qwen3_5_moe" def test_deepseek_v3_detection_and_score_correction_bias_copy(self, monkeypatch): FakeGatheredParameters.calls = [] @@ -2811,3 +2811,79 @@ def test_fused_gate_up_source_feeds_both_w1_and_w3(self): assert [id(p) for p in w1] == [id(experts_source.gate_up_proj)] assert [id(p) for p in w3] == [id(experts_source.gate_up_proj)] assert [id(p) for p in w2] == [id(experts_source.down_proj)] + + +@pytest.mark.parametrize("wrapper", [False, True]) +@pytest.mark.parametrize("autocast", [False, True]) +def test_qwen35_autoep_preserves_native_router_and_text_outputs(wrapper, autocast): + # Catches wrong wrapper paths, top-k/dtype drift, and post-softmax logit capture. + transformers = pytest.importorskip("transformers") + required = ("Qwen3_5MoeTextConfig", "Qwen3_5MoeConfig", "Qwen3_5MoeForCausalLM", + "Qwen3_5MoeForConditionalGeneration") + if not all(hasattr(transformers, name) for name in required): + pytest.skip("Transformers does not provide the Qwen3.5 MoE model classes") + Qwen3_5MoeTextConfig = transformers.Qwen3_5MoeTextConfig + Qwen3_5MoeConfig = transformers.Qwen3_5MoeConfig + Qwen3_5MoeForCausalLM = transformers.Qwen3_5MoeForCausalLM + Qwen3_5MoeForConditionalGeneration = transformers.Qwen3_5MoeForConditionalGeneration + + torch.manual_seed(42) + config = Qwen3_5MoeTextConfig(vocab_size=128, + hidden_size=128, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=64, + moe_intermediate_size=128, + shared_expert_intermediate_size=128, + num_experts=4, + num_experts_per_tok=2, + layer_types=["full_attention"], + use_cache=False) + cls = Qwen3_5MoeForCausalLM + if wrapper: + config = Qwen3_5MoeConfig(text_config=config.to_dict(), + vision_config={ + "depth": 1, + "hidden_size": 32, + "intermediate_size": 64, + "num_heads": 4, + "out_hidden_size": 128, + "num_position_embeddings": 16 + }) + cls = Qwen3_5MoeForConditionalGeneration + config._attn_implementation = "eager" + reference = cls(config).to(dtype=torch.bfloat16).eval() + candidate = copy.deepcopy(reference) + autoep = AutoEP( + candidate, + parse_autoep_config({ + "enabled": True, + "autoep_size": 1, + "preset_model": "qwen3_5_moe", + "use_grouped_mm": False + })) + specs = autoep.ep_parser() + assert len(specs) == 1 + path = "model.language_model.layers.0.mlp" if wrapper else "model.layers.0.mlp" + assert specs[0].moe_module_name == path + original = reference.get_submodule(path) + autoep.replace_moe_layer(specs[0], ep_size=1, ep_rank=0) + replaced = candidate.get_submodule(path) + x = torch.randn(13, 128, dtype=torch.float32 if autocast else torch.bfloat16) + with torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast): + logits, scores, indices = original.gate(x) + actual_scores, actual_indices, counts = replaced.router(x) + assert actual_scores.dtype == scores.dtype == torch.bfloat16 + torch.testing.assert_close(actual_indices, indices, rtol=0, atol=0) + torch.testing.assert_close(actual_scores, scores, rtol=0, atol=0) + with torch.no_grad(), torch.autocast("cpu", dtype=torch.bfloat16, enabled=autocast): + _, returned_logits = replaced(x.unsqueeze(0)) + torch.testing.assert_close(returned_logits, logits, rtol=0, atol=0) + assert counts.sum() == 26 + ids = torch.tensor([[1, 5, 7, 9, 11]]) + with torch.no_grad(): + expected = reference(input_ids=ids, labels=ids) + actual = candidate(input_ids=ids, labels=ids) + torch.testing.assert_close(actual.loss, expected.loss, rtol=0, atol=0.005) + torch.testing.assert_close(actual.logits, expected.logits, rtol=0.02, atol=0.01)