From 98caec79c30ad88ac4630c4bc76676d0109d5957 Mon Sep 17 00:00:00 2001 From: hazelduan Date: Mon, 19 Oct 2026 22:00:10 +0800 Subject: [PATCH 1/4] fix(autoep): preserve Qwen3.5 routing and rank-local checkpoints Signed-off-by: hazelduan --- deepspeed/module_inject/auto_ep.py | 4 + deepspeed/module_inject/auto_ep_layer.py | 2 + .../module_inject/auto_ep_presets/base.py | 4 + .../auto_ep_presets/qwen3_5_moe.py | 13 +- deepspeed/moe/ep_router.py | 9 +- deepspeed/runtime/engine.py | 5 + docs/code-docs/source/autoep.rst | 6 +- tests/unit/v1/moe/test_autoep_unit.py | 114 +++++++++++++++++- 8 files changed, 143 insertions(+), 14 deletions(-) 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 013dd21a38b5..cc1249207ca7 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -594,6 +594,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 8315ef9c80bc..f7c764ef25e0 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/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 3e230e061367..b0d5e96ef3c4 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -4271,6 +4271,7 @@ def load_module_state_dict(self, if checkpoint.get(FROZEN_PARAM_FRAGMENTS, None) is not None: saved_frozen_params = checkpoint[FROZEN_PARAM_FRAGMENTS] + restored_experts = DeepSpeedEngine._autoep_expert_parameter_names(None, self.module) for param in self.module.parameters(): if is_optimized_parameter(param): continue @@ -4279,6 +4280,10 @@ def load_module_state_dict(self, name = self.param_names[param] if hasattr(param, 'ds_id'): param.ds_tensor.data.copy_(saved_frozen_params[name].data) + elif name in restored_experts and name in module_state_dict: + # AutoEP already loaded this rank's expert files above. The + # shared checkpoint's frozen fragments belong to EP rank 0. + continue else: param.data.copy_(saved_frozen_params[name].data) diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index 1859b02cf595..1a85309035d4 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -1,3 +1,7 @@ +.. Copyright (c) DeepSpeed Team. +.. SPDX-License-Identifier: Apache-2.0 +.. DeepSpeed Team + AutoEP (Automatic Expert Parallelism) ===================================== @@ -33,7 +37,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 312a946851ba..81fd21ed594b 100644 --- a/tests/unit/v1/moe/test_autoep_unit.py +++ b/tests/unit/v1/moe/test_autoep_unit.py @@ -855,6 +855,34 @@ def test_load_module_state_dict_nonstrict_keeps_nonstrict_semantics_with_allowed with pytest.raises(RuntimeError, match="outside AutoEP expert"): engine.load_module_state_dict(checkpoint, strict=True, allowed_missing_keys=["weight"]) + @pytest.mark.parametrize("ep_rank", [1, 2, 3]) + def test_autoep_resume_keeps_rank_local_frozen_experts(self, ep_rank): + # The shared model checkpoint contains rank-0 frozen fragments. Those + # must not overwrite expert tensors restored from this rank's files. + layer = object.__new__(AutoEPMoELayer) + nn.Module.__init__(layer) + layer.experts = GroupedExperts(4, 8, 2, use_grouped_mm=False) + model = nn.ModuleDict({"moe": layer, "dense": nn.Linear(4, 4, bias=False)}) + model.requires_grad_(False) + engine = object.__new__(DeepSpeedEngine) + object.__setattr__(engine, "module", model) + object.__setattr__(engine, "param_names", {p: name for name, p in model.named_parameters()}) + checkpoint = { + "module": { + name: torch.full_like(p, ep_rank + 1) + for name, p in model.named_parameters() + }, + ds_engine.FROZEN_PARAM_FRAGMENTS: { + name: torch.ones_like(p) + for name, p in model.named_parameters() + }, + } + engine.load_module_state_dict(checkpoint) + for parameter in layer.experts.parameters(): + torch.testing.assert_close(parameter, torch.full_like(parameter, ep_rank + 1)) + # Non-expert frozen-fragment restoration retains its existing behavior. + torch.testing.assert_close(model["dense"].weight, torch.ones_like(model["dense"].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) @@ -1029,9 +1057,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 +2200,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 +2839,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) From 9a87125d42b170a511556eb1d21cff31697b79ce Mon Sep 17 00:00:00 2001 From: hazelduan Date: Mon, 19 Oct 2026 22:00:10 +0800 Subject: [PATCH 2/4] feat(npu): expose native grouped GEMM through accelerator capabilities Signed-off-by: hazelduan --- accelerator/abstract_accelerator.py | 4 + accelerator/npu_accelerator.py | 5 + deepspeed/moe/ep_experts.py | 47 ++++++- deepspeed/ops/npu/__init__.py | 4 + deepspeed/ops/npu/grouped_mm.py | 75 ++++++++++ docs/code-docs/source/autoep.rst | 3 +- tests/unit/v1/moe/test_hifloat8_experts.py | 152 +++++++++++++++++++++ 7 files changed, 288 insertions(+), 2 deletions(-) create mode 100644 deepspeed/ops/npu/__init__.py create mode 100644 deepspeed/ops/npu/grouped_mm.py create mode 100644 tests/unit/v1/moe/test_hifloat8_experts.py diff --git a/accelerator/abstract_accelerator.py b/accelerator/abstract_accelerator.py index 9f33b4f75bbc..c843ddabb9c1 100644 --- a/accelerator/abstract_accelerator.py +++ b/accelerator/abstract_accelerator.py @@ -215,6 +215,10 @@ def is_triton_supported(self): def prefer_triton_grouped_mm(self): return False + def get_grouped_mm(self): + """Return a native differentiable grouped GEMM override, if available.""" + return None + # Graph operations @abc.abstractmethod def create_graph(self): diff --git a/accelerator/npu_accelerator.py b/accelerator/npu_accelerator.py index 82008e2bab33..07262ea43a6c 100644 --- a/accelerator/npu_accelerator.py +++ b/accelerator/npu_accelerator.py @@ -60,6 +60,11 @@ def is_synchronized_device(self): def use_host_timers(self): return self.is_synchronized_device() + def get_grouped_mm(self): + from deepspeed.ops.npu.grouped_mm import grouped_mm + + return grouped_mm + def resolves_data_dependency(self): return self.is_synchronized_device() diff --git a/deepspeed/moe/ep_experts.py b/deepspeed/moe/ep_experts.py index 3bb3db5dd435..4028d1ba49df 100644 --- a/deepspeed/moe/ep_experts.py +++ b/deepspeed/moe/ep_experts.py @@ -300,6 +300,35 @@ def _run_experts_triton_grouped_mm( # --------------------------------------------------------------------------- +def _run_experts_with_grouped_mm(w1, + w2, + w3, + x, + num_tokens_per_expert, + activation="swiglu", + alpha=1.702, + limit=7.0, + *, + grouped_mm): + """Keep routing, padding and activation independent of the GEMM backend.""" + offsets = num_tokens_per_expert.cumsum(0).to(torch.int64) + rows = int(offsets[-1].item()) + if rows < 0 or rows > x.shape[0]: + raise ValueError("Expert token counts exceed the available rows") + inputs = x[:rows] + + def mm(lhs, weight): + return grouped_mm(lhs, weight.to(x.dtype), offsets, trans_b=True).to(lhs.dtype) + + gate = mm(inputs, w1) + up = mm(inputs, w3) + hidden = apply_expert_activation(gate, up, activation, alpha, limit, fused=False) + output = mm(hidden, w2) + if rows == x.shape[0]: + return output + return torch.cat((output, output.new_zeros((x.shape[0] - rows, output.shape[-1])))) + + class GroupedExperts(nn.Module): """Grouped expert computation for MoE layers. @@ -356,6 +385,7 @@ def __init__( self.w3.is_expert_group = True self.use_triton_grouped_mm = False self.use_grouped_mm = use_grouped_mm + self.grouped_mm = None # Resolve the Triton path. The device-specific decision is delegated to # the accelerator backend (e.g. the CUDA backend prefers Triton on @@ -364,7 +394,8 @@ def __init__( if use_grouped_mm and not disable_triton_grouped_mm: self.use_triton_grouped_mm = get_accelerator().prefer_triton_grouped_mm() - if use_grouped_mm and not hasattr(torch, "_grouped_mm") and not self.use_triton_grouped_mm: + if (use_grouped_mm and not hasattr(torch, "_grouped_mm") and not self.use_triton_grouped_mm + and get_accelerator().get_grouped_mm() is None): raise RuntimeError("GroupedExperts was constructed with use_grouped_mm=True but " "torch._grouped_mm is not available in this PyTorch build. " "Upgrade PyTorch to a build that provides torch._grouped_mm, install " @@ -392,6 +423,20 @@ def forward( """ act = (self.activation, self.activation_alpha, self.activation_limit) + grouped_mm = self.grouped_mm + accelerator = get_accelerator() + if grouped_mm is None and self.use_grouped_mm and accelerator.on_accelerator(x): + grouped_mm = accelerator.get_grouped_mm() + if grouped_mm is not None: + if getattr(self, "gate_up_impl", "auto") == "fused": + raise RuntimeError("The accelerator grouped GEMM override requires separate gate/up projections") + return _run_experts_with_grouped_mm(self.w1, + self.w2, + self.w3, + x, + num_tokens_per_expert, + *act, + grouped_mm=grouped_mm) if self.use_triton_grouped_mm: return _run_experts_triton_grouped_mm(self.w1, self.w2, self.w3, x, num_tokens_per_expert, *act) elif self.use_grouped_mm: diff --git a/deepspeed/ops/npu/__init__.py b/deepspeed/ops/npu/__init__.py new file mode 100644 index 000000000000..bc010555389b --- /dev/null +++ b/deepspeed/ops/npu/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Native Ascend operator adapters.""" diff --git a/deepspeed/ops/npu/grouped_mm.py b/deepspeed/ops/npu/grouped_mm.py new file mode 100644 index 000000000000..72e1cb41f0b1 --- /dev/null +++ b/deepspeed/ops/npu/grouped_mm.py @@ -0,0 +1,75 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Native BF16/FP16 grouped GEMM with explicit list-input autograd.""" + +import torch + + +def _native_mm(x, weight, offsets, group_type=0): + import torch_npu + + return torch_npu.npu_grouped_matmul([x], [weight], + group_list=offsets, + group_type=group_type, + group_list_type=0, + split_item=3, + output_dtype=x.dtype)[0] + + +class _GroupedMatmul(torch.autograd.Function): + + @staticmethod + def forward(ctx, x, weight, offsets, trans_b): + ctx.trans_b = trans_b + ctx.input_shape = x.shape + ctx.weight_shape = weight.shape + ctx.empty = x.shape[0] == 0 + saved = [offsets.to(torch.int64)] + if ctx.needs_input_grad[0]: + saved.append(weight) + if ctx.needs_input_grad[1]: + saved.append(x) + ctx.save_for_backward(*saved) + if ctx.empty: + columns = weight.shape[-2] if trans_b else weight.shape[-1] + return x.new_empty((0, columns)) + right = weight.transpose(-1, -2) if trans_b else weight + return _native_mm(x, right, saved[0]) + + @staticmethod + def backward(ctx, grad): + if ctx.empty: + return (grad.new_zeros(ctx.input_shape) if ctx.needs_input_grad[0] else None, + grad.new_zeros(ctx.weight_shape) if ctx.needs_input_grad[1] else None, None, None) + saved = iter(ctx.saved_tensors) + offsets = next(saved) + weight = next(saved) if ctx.needs_input_grad[0] else None + x = next(saved) if ctx.needs_input_grad[1] else None + grad = grad.contiguous() + dx = dw = None + if ctx.needs_input_grad[0]: + right = weight if ctx.trans_b else weight.transpose(-1, -2) + dx = _native_mm(grad, right, offsets) + if ctx.needs_input_grad[1]: + # Produce dW directly in the parameter's packed storage layout. + left, right = (grad.transpose(0, 1), x) if ctx.trans_b else (x.transpose(0, 1), grad) + dw = _native_mm(left, right, offsets, group_type=2) + return dx, dw, None, None + + +def grouped_mm(x, weight, offsets, *, trans_b=False): + if not isinstance(trans_b, bool): + raise TypeError("trans_b must be a boolean") + if x.ndim != 2 or weight.ndim != 3 or offsets.ndim != 1: + raise ValueError("Expected token matrix, packed expert weights and cumulative offsets") + inner = weight.shape[-1] if trans_b else weight.shape[-2] + if not offsets.numel() or weight.shape[0] != offsets.numel() or x.shape[1] != inner: + raise ValueError("Grouped matmul shapes do not match") + if x.device.type != "npu" or weight.device != x.device or offsets.device != x.device: + raise ValueError("Grouped matmul inputs must share one NPU device") + if offsets.dtype not in (torch.int32, torch.int64): + raise TypeError("Grouped offsets must be int32 or int64") + if x.dtype not in (torch.bfloat16, torch.float16) or weight.dtype != x.dtype: + raise TypeError("Native grouped matmul requires matching BF16 or FP16 operands") + return _GroupedMatmul.apply(x, weight, offsets, trans_b) diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index 1a85309035d4..3bcd4da8f1dc 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -148,7 +148,8 @@ value other than ``False`` instead of silently changing the requested behavior. 1. During ``deepspeed.initialize()``, AutoEP scans the model for MoE layers using preset-defined patterns (router name, expert name, weight shapes). 2. Detected MoE blocks are replaced with ``AutoEPMoELayer``, which uses - TorchTitan's grouped GEMM kernels and AllToAll token dispatch. + grouped GEMM and AllToAll token dispatch. The accelerator may supply a + native grouped operation; otherwise the existing PyTorch/Triton paths apply. 3. EP/EDP process groups are created automatically based on ``autoep_size``. 4. Expert parameters are marked for expert-data-parallel gradient reduction; router and shared-expert parameters use standard data-parallel reduction. diff --git a/tests/unit/v1/moe/test_hifloat8_experts.py b/tests/unit/v1/moe/test_hifloat8_experts.py new file mode 100644 index 000000000000..18f0d816ed6a --- /dev/null +++ b/tests/unit/v1/moe/test_hifloat8_experts.py @@ -0,0 +1,152 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Expert precision selection must preserve the base model and gradient contract.""" + +import pytest +import torch + +from deepspeed.accelerator import get_accelerator +from deepspeed.moe.ep_experts import GroupedExperts + + +@pytest.mark.parametrize("counts", [[2, 1], [0, 3], [3, 0], [0, 0]]) +def test_reference_experts_match_independent_linear_modules(counts): + # Catches wrong gate/up/down orientation and lost empty-expert gradients. + torch.manual_seed(42) + experts = GroupedExperts(8, 16, 2, use_grouped_mm=False) + references = [] + with torch.no_grad(): + for weight in experts.parameters(): + weight.normal_(std=0.1) + for i in range(2): + projections = [ + torch.nn.Linear(8, 16, bias=False), + torch.nn.Linear(8, 16, bias=False), + torch.nn.Linear(16, 8, bias=False) + ] + for linear, weight in zip(projections, (experts.w1, experts.w3, experts.w2)): + linear.weight.copy_(weight[i]) + references.append(projections) + x = torch.randn(sum(counts), 8, requires_grad=True) + ref_x = x.detach().clone().requires_grad_() + actual = experts(x, torch.tensor(counts)) + ref_outputs = [] + for rows, (gate, up, down) in zip(ref_x.split(counts), references): + ref_outputs.append(down(torch.nn.functional.silu(gate(rows)) * up(rows))) + expected = torch.cat(ref_outputs) + torch.testing.assert_close(actual, expected) + actual.sum().backward() + expected.sum().backward() + torch.testing.assert_close(x.grad, ref_x.grad) + for i, projections in enumerate(references): + for weight, linear in zip((experts.w1, experts.w3, experts.w2), projections): + torch.testing.assert_close(weight.grad[i], linear.weight.grad) + if counts[i] == 0: + assert torch.count_nonzero(weight.grad[i]) == 0 + + +def test_npu_expert_helper_preserves_upstream_activation(monkeypatch): + """The NPU GEMM dispatch must use AutoEP's resolved expert activation.""" + # Pin the merge regression without NPU hardware: the stub implements the + # native grouped GEMM contract while the helper supplies activation policy. + import deepspeed.moe.ep_experts as ep_experts + + def fake_grouped_mm(lhs, weight, offsets, *, trans_b=False): + weight = weight.transpose(-1, -2) if trans_b else weight + starts = [0] + offsets.tolist()[:-1] + return torch.cat([lhs[start:end] @ weight[i] for i, (start, end) in enumerate(zip(starts, offsets.tolist()))]) + + torch.manual_seed(7) + counts = torch.tensor([2, 1]) + x = torch.randn(4, 8, requires_grad=True) + weights = [torch.randn(shape, requires_grad=True) for shape in ((2, 16, 8), (2, 8, 16), (2, 16, 8))] + w1, w2, w3 = weights + actual = ep_experts._run_experts_with_grouped_mm(w1, + w2, + w3, + x, + counts, + "swiglu_clamped", + limit=0.5, + grouped_mm=fake_grouped_mm) + + x_ref = x.detach().clone().requires_grad_() + w1_ref, w2_ref, w3_ref = [w.detach().clone().requires_grad_() for w in weights] + pieces = [] + start = 0 + for i, count in enumerate(counts.tolist()): + rows = x_ref[start:start + count] + gate = rows @ w1_ref[i].T + up = rows @ w3_ref[i].T + hidden = torch.nn.functional.silu(gate.clamp(max=0.5)) * up.clamp(min=-0.5, max=0.5) + pieces.append(hidden @ w2_ref[i].T) + start += count + expected = torch.cat((torch.cat(pieces), x_ref.new_zeros((1, 8)))) + torch.testing.assert_close(actual, expected) + actual.sum().backward() + expected.sum().backward() + for value, reference in ((x, x_ref), (w1, w1_ref), (w2, w2_ref), (w3, w3_ref)): + torch.testing.assert_close(value.grad, reference.grad) + + +@pytest.mark.skipif(get_accelerator().device_name() != "npu", reason="requires an available NPU") +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.parametrize("trans_b", [False, True]) +@pytest.mark.parametrize("input_grad,weight_grad", [(True, False), (False, True), (True, True)]) +@pytest.mark.parametrize("counts", [[0, 8, 0, 4], [0, 0, 0, 0]]) +def test_npu_native_grouped_mm_matches_matmul_with_partial_gradients(dtype, trans_b, input_grad, weight_grad, counts): + # Catches saved-tensor ordering errors, packed dW orientation and empty-group leakage. + from deepspeed.ops.npu.grouped_mm import grouped_mm + + torch.npu.set_device(0) + torch.manual_seed(71) + inner, columns = 64, 32 + shape = (len(counts), columns, inner) if trans_b else (len(counts), inner, columns) + x = (torch.randn(sum(counts), inner, device="npu", dtype=dtype) * 0.1).requires_grad_(input_grad) + weight = (torch.randn(shape, device="npu", dtype=dtype) * 0.1).requires_grad_(weight_grad) + ref_x = x.detach().clone().requires_grad_(input_grad) + ref_weight = weight.detach().clone().requires_grad_(weight_grad) + offsets = torch.tensor(counts, device="npu", dtype=torch.int64).cumsum(0) + actual = grouped_mm(x, weight, offsets, trans_b=trans_b) + + pieces = [] + for index, rows in enumerate(ref_x.split(counts)): + right = ref_weight[index].T if trans_b else ref_weight[index] + pieces.append(rows @ right) + expected = torch.cat(pieces) + grad = torch.randn_like(expected) * 0.1 + actual.backward(grad) + expected.backward(grad) + torch.npu.synchronize() + + pairs = [(actual, expected)] + for value, reference, needs_grad in ((x, ref_x, input_grad), (weight, ref_weight, weight_grad)): + if needs_grad: + assert value.grad is not None and reference.grad is not None + pairs.append((value.grad, reference.grad)) + else: + assert value.grad is None and reference.grad is None + rtol, atol = (0.02, 0.0005) if dtype == torch.bfloat16 else (0.003, 0.00005) + for value, reference in pairs: + assert value.dtype == reference.dtype == dtype + assert bool(torch.isfinite(value).all()) + torch.testing.assert_close(value, reference, rtol=rtol, atol=atol) + if weight_grad: + for index, count in enumerate(counts): + if count == 0: + assert torch.count_nonzero(weight.grad[index]) == 0 + + +def test_grouped_override_rejects_explicit_fused_gate_up(): + # Upstream can request a fused gate/up path; a three-GEMM override cannot + # silently ignore that public choice after this integration is merged. + experts = GroupedExperts(8, 16, 2, use_grouped_mm=False) + experts.gate_up_impl = "fused" + + def unexpected_kernel(*args, **kwargs): + raise AssertionError("Unsupported execution must fail before a GEMM") + + experts.grouped_mm = unexpected_kernel + with pytest.raises(RuntimeError, match="separate gate/up"): + experts(torch.randn(2, 8), torch.tensor([1, 1])) From b49b298f2cf27642cecd4d9aaf50f3e8ac5ca8a5 Mon Sep 17 00:00:00 2001 From: hazelduan Date: Mon, 19 Oct 2026 22:00:10 +0800 Subject: [PATCH 3/4] perf(autoep): generate large permutation indices on NPU Signed-off-by: hazelduan --- deepspeed/module_inject/auto_ep_layer.py | 10 +-- deepspeed/moe/ep_kernels.py | 50 +++++++++-- docs/code-docs/source/autoep.rst | 6 ++ tests/unit/v1/moe/test_autoep_permute_rows.py | 88 +++++++++++++++++++ 4 files changed, 143 insertions(+), 11 deletions(-) diff --git a/deepspeed/module_inject/auto_ep_layer.py b/deepspeed/module_inject/auto_ep_layer.py index cc1249207ca7..778421d554c3 100644 --- a/deepspeed/module_inject/auto_ep_layer.py +++ b/deepspeed/module_inject/auto_ep_layer.py @@ -382,7 +382,7 @@ def permute_by_local_expert( n_tokens: original token count before padding (for unpermute) """ from deepspeed.moe.ep_kernels import (generate_permute_indices, permute_rows, permute_rows_supported, - TOKEN_GROUP_ALIGN_SIZE_M) + prefer_cpu_permutation_indices, TOKEN_GROUP_ALIGN_SIZE_M) if local_counts.ndim == 1: # [E_local]: already aggregated over sources (ep_degree=1) @@ -404,9 +404,9 @@ def permute_by_local_expert( x_padded_per_expert = n_tokens + num_local_experts * alignment padded_max_len = ((x_padded_per_expert + alignment - 1) // alignment) * alignment - # Use the pure-PyTorch path for host tensors. The CPU accelerator reports - # CPU tensors as "on accelerator", but Triton still requires a GPU driver. - use_cpu = tokens.device.type == "cpu" + # Small NPU groups retain the lower-latency host reference path. + use_cpu = (local_counts_flat.device.type == "cpu" + or prefer_cpu_permutation_indices(tokens.device.type, local_counts_flat.numel())) counts_for_permute = local_counts_flat.cpu() if use_cpu else local_counts_flat with torch.no_grad(): permuted_indices, m_sizes, _offsets = generate_permute_indices( @@ -417,7 +417,7 @@ def permute_by_local_expert( alignment, use_cpu=use_cpu, ) - if not use_cpu: + if tokens.device.type != "cpu": permuted_indices = permuted_indices.to(tokens.device) m_sizes = m_sizes.to(tokens.device) diff --git a/deepspeed/moe/ep_kernels.py b/deepspeed/moe/ep_kernels.py index d23b23f60e3a..18a5fdde974f 100644 --- a/deepspeed/moe/ep_kernels.py +++ b/deepspeed/moe/ep_kernels.py @@ -37,8 +37,7 @@ _TRITON_AVAILABLE = True except ImportError: - logger.info("Triton not available; using pure-PyTorch CPU fallback for " - "permutation index generation.") + logger.info("Triton not available; using the available PyTorch permutation-index backend.") # --------------------------------------------------------------------------- # Alignment constant @@ -52,6 +51,14 @@ - mxfp8: 32 (scaling block size) """ +NPU_DEVICE_PERMUTE_MIN_GROUPS = 64 + + +def prefer_cpu_permutation_indices(device_type: str, group_count: int) -> bool: + """Use host indexing when device launch overhead dominates small NPU groups.""" + return device_type == "cpu" or (device_type == "npu" and group_count < NPU_DEVICE_PERMUTE_MIN_GROUPS) + + # --------------------------------------------------------------------------- # Utility: round up # --------------------------------------------------------------------------- @@ -115,10 +122,16 @@ def fill_indices_wrapper( block_size: int = 128, max_blocks: int = 1024, ) -> torch.Tensor: - """Launch the Triton kernel to fill permutation indices. - - Falls back to :func:`fill_indices_cpu` when Triton is unavailable. - """ + """Fill permutation indices with the available device implementation.""" + if tokens_per_expert_group.device.type == "npu": + return fill_indices_torch( + tokens_per_expert_group, + start_index_values, + write_offsets, + experts_per_rank, + num_ranks, + max_len, + ) if not _TRITON_AVAILABLE: return fill_indices_cpu( tokens_per_expert_group, @@ -182,6 +195,31 @@ def fill_indices_cpu( return permuted_indices +def fill_indices_torch( + tokens_per_expert_group: torch.Tensor, + start_index_values: torch.Tensor, + write_offsets: torch.Tensor, + experts_per_rank: int, + num_ranks: int, + max_len: int, +) -> torch.Tensor: + """Build the padded permutation without host counts; AutoEP supplies ``max_len >= token count``.""" + counts = tokens_per_expert_group.to(torch.int64) + source_ends = start_index_values + counts + source_prefix = (counts.view(num_ranks, experts_per_rank).cumsum(0) - + counts.view(num_ranks, experts_per_rank)).reshape(-1) + source_rows = torch.arange(max_len, device=counts.device, dtype=torch.int64) + valid = source_rows < source_ends[-1] + group_ids = torch.searchsorted(source_ends, source_rows, right=True).clamp_max(counts.numel() - 1) + expert_ids = group_ids.remainder(experts_per_rank) + destinations = (write_offsets[expert_ids] + source_prefix[group_ids] + source_rows - start_index_values[group_ids]) + # Extra slot absorbs all rows beyond the real token count; it is discarded. + destinations = torch.where(valid, destinations, max_len).to(torch.int64) + result = torch.full((max_len + 1, ), -1, dtype=torch.int32, device=counts.device) + result.scatter_(0, destinations, source_rows.to(torch.int32)) + return result[:-1] + + # =================================================================== # generate_permute_indices # =================================================================== diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index 3bcd4da8f1dc..0706be82f77b 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -154,6 +154,12 @@ value other than ``False`` instead of silently changing the requested behavior. 4. Expert parameters are marked for expert-data-parallel gradient reduction; router and shared-expert parameters use standard data-parallel reduction. +On Ascend NPU, AutoEP generates permutation indices on device when there are +at least 64 source/expert groups. Smaller groups use the CPU reference path +because device launch overhead can dominate. The ordering and padding contract +is identical; this internal choice does not change the AutoEP configuration. +Other dispatch steps may still synchronize with the host. + **Router outputs and activation checkpointing:** Models using Hugging Face's model-level router-logit recording capture the diff --git a/tests/unit/v1/moe/test_autoep_permute_rows.py b/tests/unit/v1/moe/test_autoep_permute_rows.py index b58a92b03fa9..e17f4ccc0dfb 100644 --- a/tests/unit/v1/moe/test_autoep_permute_rows.py +++ b/tests/unit/v1/moe/test_autoep_permute_rows.py @@ -41,6 +41,94 @@ def _grads(function, inputs, upstream): return output.detach(), grad +@pytest.mark.parametrize("rows_by_source", [ + [[0, 0], [0, 0]], + [[3, 0, 2], [1, 4, 0]], + [[0, 5, 0, 1], [7, 0, 0, 0]], +]) +def test_torch_indexer_matches_cpu_reference(rows_by_source): + counts = torch.tensor(rows_by_source, dtype=torch.int32) + ranks, experts = counts.shape + flat = counts.flatten() + starts = flat.cumsum(0) - flat + alignment = ep_kernels.TOKEN_GROUP_ALIGN_SIZE_M + sizes = ((counts.sum(0).clamp_min(alignment) + alignment - 1) // alignment) * alignment + write_offsets = sizes.cumsum(0) - sizes + max_len = int(flat.sum()) + experts * alignment + alignment + expected = ep_kernels.fill_indices_cpu(flat, starts, write_offsets, experts, ranks, max_len) + actual = ep_kernels.fill_indices_torch(flat, starts, write_offsets, experts, ranks, max_len) + assert torch.equal(actual, expected) + + +@pytest.mark.parametrize("device_type,groups,expected", [ + ("cpu", 256, True), + ("cuda", 256, False), + ("npu", 32, True), + ("npu", 48, True), + ("npu", 64, False), + ("npu", 256, False), +]) +def test_permutation_index_backend_keeps_small_npu_groups_on_cpu(device_type, groups, expected): + assert ep_kernels.prefer_cpu_permutation_indices(device_type, groups) is expected + + +@pytest.mark.skipif(get_accelerator().device_name() != "npu", reason="requires NPU kernels") +@pytest.mark.parametrize("strided", [False, True]) +@pytest.mark.parametrize("rows_by_source", [ + [[0, 0], [0, 0]], + [[3, 0, 2], [1, 4, 0]], + [[0, 5, 0, 1], [7, 0, 0, 0]], + [[0 if expert % 5 == 0 else (rank + expert) % 3 for expert in range(16)] for rank in range(4)], + [[0] * 16 for _ in range(4)], +]) +def test_npu_permutation_matches_cpu_indices_and_gradients(rows_by_source, strided): + counts_cpu = torch.tensor(rows_by_source, dtype=torch.int32) + ranks, experts = counts_cpu.shape + rows = int(counts_cpu.sum()) + alignment = ep_kernels.TOKEN_GROUP_ALIGN_SIZE_M + max_len = ep_kernels._round_up(rows + experts * alignment, alignment) + expected_indices, expected_sizes, _ = ep_kernels.generate_permute_indices(counts_cpu.flatten(), + experts, + ranks, + max_len, + alignment, + use_cpu=True) + counts = counts_cpu.to("npu:0") + if strided: + tokens = torch.randn((rows, 32), device="npu:0", dtype=torch.bfloat16)[:, ::2].detach().requires_grad_() + else: + tokens = torch.randn((rows, 16), device="npu:0", dtype=torch.bfloat16, requires_grad=True) + actual, indices, sizes, actual_rows = permute_by_local_expert(tokens, counts) + assert actual_rows == rows + assert torch.equal(indices.cpu(), expected_indices) + assert torch.equal(sizes.cpu(), expected_sizes) + + reference_tokens = tokens.detach().clone().requires_grad_() + reference = _reference_permute(reference_tokens, expected_indices.to("npu:0")) + upstream = torch.randn_like(actual) + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + (actual * upstream).sum().backward() + (reference * upstream).sum().backward() + torch.testing.assert_close(tokens.grad, reference_tokens.grad, rtol=0, atol=0) + + +@pytest.mark.skipif(get_accelerator().device_name() != "npu", reason="requires NPU kernels") +def test_npu_permutation_accepts_host_counts(): + counts_cpu = torch.ones((4, 16), dtype=torch.int32) + tokens = torch.randn((64, 16), device="npu:0", dtype=torch.bfloat16, requires_grad=True) + reference_tokens = tokens.detach().clone().requires_grad_() + actual, indices, sizes, _ = permute_by_local_expert(tokens, counts_cpu) + reference, reference_indices, reference_sizes, _ = permute_by_local_expert(reference_tokens, + counts_cpu.to("npu:0")) + assert torch.equal(indices, reference_indices) + assert torch.equal(sizes, reference_sizes) + assert torch.equal(actual, reference) + upstream = torch.randn_like(actual) + (actual * upstream).sum().backward() + (reference * upstream).sum().backward() + assert torch.equal(tokens.grad, reference_tokens.grad) + + @pytest.mark.skipif(get_accelerator().device_name() != "cuda" or not ep_kernels._TRITON_AVAILABLE, reason="the gather-based reorder needs CUDA and Triton") class TestAutoEPPermuteRows(DistributedTest): From 95dd780ce4e2f534b484610f273d3ef21c505baf Mon Sep 17 00:00:00 2001 From: hazelduan Date: Mon, 19 Oct 2026 22:24:33 +0800 Subject: [PATCH 4/4] feat(hifloat8): integrate AO-owned quantize_ training handlers Signed-off-by: hazelduan --- accelerator/abstract_accelerator.py | 4 + accelerator/npu/__init__.py | 4 + accelerator/npu/hifloat8.py | 29 + accelerator/npu_accelerator.py | 5 + deepspeed/runtime/config.py | 52 ++ deepspeed/runtime/engine.py | 99 ++++ docs/_tutorials/hifloat8_training.md | 117 ++++ docs/code-docs/source/autoep.rst | 19 + .../runtime/hifloat8_resume_validation.py | 169 ++++++ tests/unit/runtime/test_hifloat8.py | 518 ++++++++++++++++++ tests/unit/v1/moe/test_hifloat8_experts.py | 87 +++ 11 files changed, 1103 insertions(+) create mode 100644 accelerator/npu/__init__.py create mode 100644 accelerator/npu/hifloat8.py create mode 100644 docs/_tutorials/hifloat8_training.md create mode 100644 tests/unit/runtime/hifloat8_resume_validation.py create mode 100644 tests/unit/runtime/test_hifloat8.py diff --git a/accelerator/abstract_accelerator.py b/accelerator/abstract_accelerator.py index c843ddabb9c1..bff7866def30 100644 --- a/accelerator/abstract_accelerator.py +++ b/accelerator/abstract_accelerator.py @@ -215,6 +215,10 @@ def is_triton_supported(self): def prefer_triton_grouped_mm(self): return False + def get_hifloat8_provider(self): + """Return optional public HiFloat8 APIs for this accelerator.""" + return None + def get_grouped_mm(self): """Return a native differentiable grouped GEMM override, if available.""" return None diff --git a/accelerator/npu/__init__.py b/accelerator/npu/__init__.py new file mode 100644 index 000000000000..4f9abc2536dd --- /dev/null +++ b/accelerator/npu/__init__.py @@ -0,0 +1,4 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Optional Ascend capability discovery.""" diff --git a/accelerator/npu/hifloat8.py b/accelerator/npu/hifloat8.py new file mode 100644 index 000000000000..dfef2cd5f948 --- /dev/null +++ b/accelerator/npu/hifloat8.py @@ -0,0 +1,29 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Resolve the optional public TorchAO HiFloat8 APIs on NPU.""" + +import importlib +from importlib.metadata import PackageNotFoundError, version + +from packaging.version import Version + + +def get_hifloat8_provider(): + hint = ("Install torchao>=0.18,<0.19 and a compatible torchao-npu build " + "providing torchao.hifloat8; keep automatic activation enabled.") + try: + installed_version = Version(version("torchao")) + except PackageNotFoundError as error: + raise ImportError(f"TorchAO HiFloat8 is unavailable. {hint}") from error + if not Version("0.18") <= installed_version < Version("0.19"): + raise RuntimeError(f"Unsupported torchao version {installed_version}. {hint}") + try: + api = importlib.import_module("torchao.hifloat8") + except ImportError as error: + raise ImportError(f"TorchAO HiFloat8 API is unavailable. {hint}") from error + required = ("HiFloat8Config", "assert_hifloat8_training_available") + missing = [name for name in required if not hasattr(api, name)] + if missing: + raise ImportError(f"TorchAO HiFloat8 API is missing {missing}. {hint}") + return api diff --git a/accelerator/npu_accelerator.py b/accelerator/npu_accelerator.py index 07262ea43a6c..7fecd6e2e2b8 100644 --- a/accelerator/npu_accelerator.py +++ b/accelerator/npu_accelerator.py @@ -60,6 +60,11 @@ def is_synchronized_device(self): def use_host_timers(self): return self.is_synchronized_device() + def get_hifloat8_provider(self): + from .npu.hifloat8 import get_hifloat8_provider + + return get_hifloat8_provider() + def get_grouped_mm(self): from deepspeed.ops.npu.grouped_mm import grouped_mm diff --git a/deepspeed/runtime/config.py b/deepspeed/runtime/config.py index 390a0be5c559..07bc5188693f 100755 --- a/deepspeed/runtime/config.py +++ b/deepspeed/runtime/config.py @@ -372,6 +372,57 @@ def get_hybrid_engine_config(param_dict): return hybrid_engine_config +HIFLOAT8 = "hifloat8" +HIFLOAT8_DEFAULT = { + "enabled": False, + "backend": "torchao", + "module_name_patterns": (), + "min_numel": 0, + "expected_module_count": None, + "config": None, +} + + +def get_hifloat8_config(param_dict): + if "low_precision" in param_dict: + raise DeepSpeedConfigError("Use the existing hifloat8 section, not low_precision") + config = param_dict.get(HIFLOAT8, {}) + if not isinstance(config, dict): + raise DeepSpeedConfigError(f"'{HIFLOAT8}' must be an object, got {type(config).__name__}") + + unknown = sorted(set(config) - set(HIFLOAT8_DEFAULT)) + if unknown: + raise DeepSpeedConfigError(f"Unknown '{HIFLOAT8}' fields: {unknown}") + + result = dict(HIFLOAT8_DEFAULT) + result.update(config) + if not isinstance(result["enabled"], bool): + raise DeepSpeedConfigError("'hifloat8.enabled' must be a boolean") + if result["backend"] == "torchao_npu": + result["backend"] = "torchao" + if result["backend"] != "torchao": + raise DeepSpeedConfigError("'hifloat8.backend' must be 'torchao'; the legacy torch_npu helper was removed") + policy = result["config"] + if policy is not None: + if not isinstance(policy, dict): + raise DeepSpeedConfigError("'hifloat8.config' must be an object or null") + patterns = result["module_name_patterns"] + if not isinstance(patterns, + (list, tuple)) or not all(isinstance(pattern, str) and pattern for pattern in patterns): + raise DeepSpeedConfigError("'hifloat8.module_name_patterns' must be a list of non-empty strings") + if result["enabled"] and not patterns: + raise DeepSpeedConfigError("'hifloat8.module_name_patterns' cannot be empty when HiFloat8 is enabled") + min_numel = result["min_numel"] + if isinstance(min_numel, bool) or not isinstance(min_numel, int) or min_numel < 0: + raise DeepSpeedConfigError("'hifloat8.min_numel' must be a non-negative integer") + expected_count = result["expected_module_count"] + if expected_count is not None and (isinstance(expected_count, bool) or not isinstance(expected_count, int) + or expected_count <= 0): + raise DeepSpeedConfigError("'hifloat8.expected_module_count' must be a positive integer or null") + result["module_name_patterns"] = tuple(patterns) + return result + + def get_expert_data_topo_config(param_dict): return get_scalar_param(param_dict, USE_DATA_BEFORE_EXPERT_PARALLEL, USE_DATA_BEFORE_EXPERT_PARALLEL_DEFAULT) @@ -536,6 +587,7 @@ def _initialize_params(self, param_dict): self.use_data_before_expert_parallel_ = get_expert_data_topo_config(param_dict) self.hybrid_engine = get_hybrid_engine_config(param_dict) + self.hifloat8_config = get_hifloat8_config(param_dict) self.pipeline = get_pipeline_config(param_dict) diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index b0d5e96ef3c4..5588505449b0 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -586,6 +586,7 @@ def __init__(self, self._configure_with_arguments(args, mpu) self.pipeline_parallelism = isinstance(model, PipelineModule) self._do_sanity_check() + self._validate_hifloat8_configuration() if self.log_level() is not None: set_log_level_from_string(self.log_level()) eager_model_parameters = _model_parameters_need_remap(model_parameters) @@ -619,6 +620,7 @@ def __init__(self, # Configure distributed model self._configure_distributed_model(model) + self._configure_hifloat8() # These hooks should be disabled later if DeepCompile is not active. self.module_forward_pre_hook = self._create_module_forward_pre_hook() @@ -1662,6 +1664,9 @@ def fp16_enabled(self): def bfloat16_enabled(self): return self._config.bfloat16_config.enabled + def hifloat8_enabled(self): + return self._config.hifloat8_config["enabled"] + def fp16_master_weights_and_gradients(self): return self._config.float16_config.fp16_master_weights_and_grads @@ -2047,6 +2052,100 @@ def _set_client_model(self, model): # register module attribute in engine but avoid getattr self.__dict__['module'] = model + def _validate_hifloat8_configuration(self): + if not self.hifloat8_enabled(): + return + if not self.bfloat16_enabled(): + raise RuntimeError("HiFloat8 training requires BF16 model parameters") + if self.pipeline_parallelism: + raise RuntimeError("HiFloat8 training has not been validated with DeepSpeed PipelineModule") + if self.autotp_size() > 1: + raise RuntimeError("HiFloat8 training has not been validated with DeepSpeed AutoTP") + if self.mpu is not None: + raise RuntimeError("HiFloat8 training has not been validated with a custom model-parallel unit") + if self.has_moe_layers and any(isinstance(module, MoE) for module in self.module.modules()): + raise RuntimeError("HiFloat8 experts require AutoEP, not the explicit DeepSpeed MoE layer") + if self.zero_optimization_stage() >= ZeroStageEnum.weights: + raise RuntimeError("HiFloat8 phase-1 supports DeepSpeed ZeRO stages 0, 1, and 2 only") + + def _configure_hifloat8(self): + if not self.hifloat8_enabled(): + return + self._validate_hifloat8_configuration() + + from fnmatch import fnmatchcase + from deepspeed.moe.ep_experts import GroupedExperts + + config = self._config.hifloat8_config + patterns = config["module_name_patterns"] + min_numel = config["min_numel"] + selected_names = [] + selected_experts = [] + total_numel = 0 + for name, module in self.module.named_modules(): + matrix_numel = getattr(module, "in_features", 0) * getattr(module, "out_features", 0) + if (type(module) is torch.nn.Linear and matrix_numel >= min_numel + and any(fnmatchcase(name, pattern) for pattern in patterns)): + selected_names.append(name) + total_numel += matrix_numel + if (isinstance(module, GroupedExperts) + and min(module.w1[0].numel(), module.w2[0].numel(), module.w3[0].numel()) >= min_numel + and any(fnmatchcase(name, pattern) for pattern in patterns)): + selected_experts.append((name, module)) + if not selected_names and not selected_experts: + raise RuntimeError("HiFloat8 module selection matched no nn.Linear or AutoEP GroupedExperts modules; " + f"patterns={list(patterns)}, min_numel={min_numel}") + + eligible_names = selected_names + [name for name, _ in selected_experts] + expected_count = config.get("expected_module_count") + if expected_count is not None and len(eligible_names) != expected_count: + raise RuntimeError( + f"HiFloat8 expected_module_count={expected_count}, selected {len(eligible_names)} modules") + unmatched = [pattern for pattern in patterns if not any(fnmatchcase(name, pattern) for name in eligible_names)] + if unmatched: + raise RuntimeError(f"HiFloat8 patterns matched no eligible modules: {unmatched}") + if any(getattr(module, "gate_up_impl", "auto") == "fused" for _name, module in selected_experts): + raise RuntimeError("HiFloat8 experts require separate gate/up projections") + + logger.info( + "HiFloat8 selected %d Linear modules (%d matrix elements); probing native kernels before conversion", + len(selected_names), + total_numel, + ) + try: + provider = get_accelerator().get_hifloat8_provider() + if provider is None: + raise RuntimeError(f"No HiFloat8 provider on {get_accelerator().device_name()}") + policy = provider.HiFloat8Config.from_dict(config.get("config")) + grouped_policy = None + if selected_experts: + if not hasattr(provider, "HiFloat8GroupedConfig"): + raise ImportError("Install a compatible TorchAO extension providing HiFloat8GroupedConfig") + grouped_policy = provider.HiFloat8GroupedConfig.from_dict(config.get("config")) + provider.assert_hifloat8_training_available(device=self.device, + config=policy, + linear=bool(selected_names), + grouped=bool(selected_experts)) + except (ImportError, RuntimeError, ValueError, TypeError) as error: + raise RuntimeError( + f"HiFloat8 selected {len(selected_names)} Linear modules ({total_numel} matrix elements), " + f"but provider/native kernel validation failed before conversion: {error}") from error + from torchao.quantization import ModuleFqnToConfig, quantize_ + + module_configs = {name: policy for name in selected_names} + module_configs.update({name: grouped_policy for name, _module in selected_experts}) + quantize_(self.module, ModuleFqnToConfig(module_configs), filter_fn=None) + self.hifloat8_converted_module_names = tuple(selected_names) + self.hifloat8_grouped_module_names = tuple(name for name, _ in selected_experts) + logger.info("HiFloat8 enabled %d AutoEP grouped expert modules: %s", len(selected_experts), + self.hifloat8_grouped_module_names) + logger.info( + "HiFloat8 training converted %d Linear modules (%d matrix elements): %s", + len(selected_names), + total_numel, + selected_names, + ) + def _configure_distributed_model(self, model): self._set_client_model(model) apply_zero_leaf_module_config(self.module, getattr(self._config.zero_config, "leaf_module", None)) diff --git a/docs/_tutorials/hifloat8_training.md b/docs/_tutorials/hifloat8_training.md new file mode 100644 index 000000000000..d92684864fee --- /dev/null +++ b/docs/_tutorials/hifloat8_training.md @@ -0,0 +1,117 @@ +--- +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +title: "TorchAO HiFloat8 training" +--- + +DeepSpeed selects ordinary Linear and AutoEP GroupedExperts modules, probes the +active accelerator's optional HiFloat8 APIs, and converts both kinds through one +public `torchao.quantization.quantize_` call. DeepSpeed owns selection, activation, +routing/padding, communication, optimizer and checkpoint lifecycle. TorchAO owns +numerical policy and the quantization/GEMM implementations. + +## Quick start + +Add the following section to your existing DeepSpeed training configuration. +AutoEP topology, batch size and optimizer settings remain in that configuration. + +```json +{ + "bf16": {"enabled": true}, + "zero_optimization": {"stage": 2}, + "hifloat8": { + "enabled": true, + "backend": "torchao", + "module_name_patterns": ["*.experts", "*.shared_experts.*_proj"], + "min_numel": 65536 + } +} +``` + +Pass the configuration to the normal initialization call; DeepSpeed applies +HiFloat8 after AutoEP replacement and device placement, before optimizer setup: + +```python +import deepspeed + +engine, optimizer, _, _ = deepspeed.initialize( + model=model, + model_parameters=model.parameters(), + config=ds_config, +) +``` + +For a Dense-only model, omit AutoEP and select the actual Linear module names. +Set `hifloat8.enabled=false` to return to the high-precision path. + +AutoEP topology and trainable-parameter selection are configured separately. +Every pattern must match an eligible module. Optional `expected_module_count` +checks the complete selection before backend loading. The existing schema is +retained; legacy `backend="torchao_npu"` normalizes to `"torchao"`. The former +`torch_npu` training-helper backend and the temporary `low_precision` section are +rejected with a migration instruction. + +## Conversion and device ownership + +`DeepSpeedAccelerator.get_hifloat8_provider()` defaults to `None`. The NPU +implementation lazily returns the public `torchao.hifloat8` APIs after checking +optional dependency/version availability. Unsupported accelerators fail before +conversion. No NPU import or device-specific kernel selection occurs in Engine. + +`ModuleFqnToConfig` maps selected Linear modules to `HiFloat8Config` and selected +experts to AO's `HiFloat8GroupedConfig`. Both parse the same numerical settings +through their common `HiFloat8Config` base. Each recipe is immutable and reused +across its selected modules; no nested policy wrapper is needed. +Both handlers are registered by AO. The grouped handler binds AO's existing GEMM +through the module's `grouped_mm` slot; it does not import a DeepSpeed container. +Engine selects configurations and calls `quantize_`, without defining a handler +or selecting a grouped function. Both preserve Parameter identity, names, metadata, `requires_grad`, +mode and state-dict keys. There is no second model-conversion API or registry. +The expert container remains owned by DeepSpeed and retains its activation and +routing/padding behavior. + +The separate `get_grouped_mm()` accelerator hook supplies ordinary high-precision +GEMM. Its NPU implementation calls TorchNPU with explicit list-input autograd, +packed `[E,N,K]` weights and empty-group gradients. The operator implementation +lives in `deepspeed/ops/npu/grouped_mm.py`; the accelerator only resolves the +capability. HiFloat8 API discovery lives in `accelerator/npu/hifloat8.py`. +BF16/FP16 grouped training does not require TorchAO. CPU/CUDA dispatch retains +its existing implementation. + +## Optional TorchAO policy + +Install `torchao>=0.18,<0.19` and the paired Ascend extension with public HiFloat8 +policy/probe APIs and registered handlers. Selected experts additionally require +`HiFloat8GroupedConfig`; Dense-only selection does not require grouped exports. +With the extension's automatic activation enabled, DeepSpeed imports +`torchao.hifloat8`; users do not need to import `torchao_npu` or convert the model +manually. Upstream TorchAO alone does not provide this HiFloat8 implementation. +Missing APIs, versions, +invalid policy or unavailable requested forward/dX/dW kernels fail before +conversion; policy parsing uses `HiFloat8Config.from_dict`. + +Missing, null or empty `config` uses the established input/weight/gradient scale +denominators `15/15/224`. Explicit overrides remain supported. `compute_dtype` +selects high-precision staging and GEMM output precision (`null`, `"bfloat16"` +or `"float16"`); operands remain HiFloat8 payloads and accumulation is +kernel-owned. Converted Linear and expert modules return their caller's dtype, +so FP16 working GEMMs compose with BF16 residuals. Raw tensor helpers retain +their working/output dtype contract. Scaling currently supports per-tensor only. + +## Validation boundary + +The current NPU HiFloat8 path requires BF16 model parameters and accepts ZeRO +stages 0–2. ZeRO-3, PipelineModule, AutoTP, custom MPU and explicit DeepSpeed MoE +layers are rejected, as is an explicit fused gate/up override on selected experts. +Quantized checkpoint ingestion into ZeRO-3 is unsupported; +checkpoints retain high-precision weights and their original state names. +Reconstruct the same policy for resume. The opt-in DP2/GAS2 lifecycle entry point +is `tests/unit/runtime/hifloat8_resume_validation.py`, whose experts are replicated. +It does not establish AutoEP checkpoint recovery. + +The integration includes Qwen3.5 text-wrapper routing compatibility and NPU +permutation index generation. These reuse AutoEP's existing preset and kernel +extension points without adding HiFloat8 configuration. Full-parameter convergence, +LoRA, offload, compile, EP8, multi-node and CUDA low-precision hardware validation +remain separate gates. diff --git a/docs/code-docs/source/autoep.rst b/docs/code-docs/source/autoep.rst index 0706be82f77b..a929fff23264 100644 --- a/docs/code-docs/source/autoep.rst +++ b/docs/code-docs/source/autoep.rst @@ -154,6 +154,25 @@ value other than ``False`` instead of silently changing the requested behavior. 4. Expert parameters are marked for expert-data-parallel gradient reduction; router and shared-expert parameters use standard data-parallel reduction. +**Optional HiFloat8 training:** + +The accelerator's ``get_hifloat8_provider()`` hook supplies public TorchAO +policy, capability-probe and recipe APIs, or ``None`` when unavailable. +With the existing ``hifloat8`` configuration, selected Linear and GroupedExperts +modules are converted together by ``torchao.quantization.quantize_`` using +``ModuleFqnToConfig`` and registered module handlers. Policy parsing belongs to +TorchAO; selection and training lifecycle checks belong to DeepSpeed. +The engine creates AO's ``HiFloat8GroupedConfig`` from the same numerical +settings as Linear. Both configurations inherit the common policy. AO +registers both module handlers; its grouped handler binds the GEMM through the +expert container's ``grouped_mm`` slot without importing DeepSpeed. The +accelerator supplies public configuration/probe APIs, not a GEMM callable. + +The separate ``get_grouped_mm()`` hook supplies high-precision grouped GEMM. +On NPU it uses TorchNPU autograd without a TorchAO dependency. HiFloat8 requires +BF16 parameters and currently supports ZeRO stages 0–2; quantized checkpoint +ingestion into ZeRO-3 remains unsupported. + On Ascend NPU, AutoEP generates permutation indices on device when there are at least 64 source/expert groups. Smaller groups use the CPU reference path because device launch overhead can dominate. The ordering and padding contract diff --git a/tests/unit/runtime/hifloat8_resume_validation.py b/tests/unit/runtime/hifloat8_resume_validation.py new file mode 100644 index 000000000000..f806d5c44512 --- /dev/null +++ b/tests/unit/runtime/hifloat8_resume_validation.py @@ -0,0 +1,169 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team +"""Opt-in NPU DP2 ZeRO-2 save/load validation for Dense and grouped HiFloat8. + +Run with a valid two-device Ascend rank table and the optional TorchAO HiFloat8 extension installed:: + + torchrun --nproc_per_node=2 tests/unit/runtime/hifloat8_resume_validation.py \ + --output /tmp/hifloat8-resume + +GroupedExperts are replicated for this DP2 lifecycle check. AutoEP expert +parallel routing is validated separately with a model integration run. +""" +import argparse +import copy +from dataclasses import asdict +import json +import os +from pathlib import Path + +import torch +import torch_npu # noqa: F401 +import deepspeed +import deepspeed.comm as dist +from deepspeed.moe.ep_experts import GroupedExperts +from deepspeed.runtime.fp16.loss_scaler import LossScaler +from torchao.hifloat8 import HiFloat8Config +import torchao.hifloat8.hifloat8_grouped as hifloat8_api + + +class Model(torch.nn.Module): + + def __init__(self): + super().__init__() + self.dense = torch.nn.Linear(128, 128, bias=False) + self.experts = GroupedExperts(128, 128, 2, use_grouped_mm=False) + with torch.no_grad(): + for p in self.parameters(): + p.normal_(std=0.02) + + def forward(self, x): + counts = torch.tensor((0, x.shape[0]), device=x.device, dtype=torch.int64) + return self.experts(self.dense(x), counts) + + +def snapshot(value): + if isinstance(value, torch.Tensor): + return value.detach().cpu().clone() + if isinstance(value, dict): + return {k: snapshot(v) for k, v in value.items()} + if isinstance(value, tuple): + return tuple(snapshot(v) for v in value) + if isinstance(value, list): + return [snapshot(v) for v in value] + return copy.deepcopy(value) + + +def compare_nested(a, b): + if isinstance(a, torch.Tensor): + assert torch.equal(a.cpu(), b.cpu()) + elif isinstance(a, LossScaler): + assert type(a) is type(b) + compare_nested(vars(a), vars(b)) + elif isinstance(a, dict): + assert a.keys() == b.keys() + for key in a: + compare_nested(a[key], b[key]) + elif isinstance(a, (tuple, list)): + assert len(a) == len(b) + for x, y in zip(a, b): + compare_nested(x, y) + else: + assert a == b, (a, b) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--output', type=Path, required=True) + args = parser.parse_args() + rank = int(os.environ['RANK']) + torch.npu.set_device(int(os.environ['LOCAL_RANK'])) + deepspeed.init_distributed(dist_backend='hccl') + torch.manual_seed(42) + initial = Model().bfloat16().state_dict() + config = { + 'train_micro_batch_size_per_gpu': 1, + 'gradient_accumulation_steps': 2, + 'bf16': { + 'enabled': True + }, + 'zero_optimization': { + 'stage': 2, + 'overlap_comm': False + }, + 'hifloat8': { + 'enabled': True, + 'backend': 'torchao', + 'module_name_patterns': ['dense', 'experts'], + 'expected_module_count': 2 + }, + 'steps_per_print': 10000, + } + generator = torch.Generator().manual_seed(123 + rank) + batches = [torch.randn(32, 128, generator=generator).bfloat16().to('npu') for _ in range(8)] + targets = [torch.randn(32, 128, generator=generator).bfloat16().to('npu') * 0.01 for _ in range(8)] + + observed_policies = [] + original_grouped_mm = hifloat8_api.hifloat8_grouped_mm + + def recorded_grouped_mm(*args, **kwargs): + observed_policies.append(kwargs["config"]) + return original_grouped_mm(*args, **kwargs) + + hifloat8_api.hifloat8_grouped_mm = recorded_grouped_mm + + def create(): + model = Model().bfloat16() + model.load_state_dict(initial) + optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3) + engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config) + assert isinstance(engine.module.dense.config, HiFloat8Config) + return engine + + def step(engine, i): + observed_policies.clear() + loss = (engine(batches[i]).float() - targets[i].float()).square().mean() + assert observed_policies and all(policy is observed_policies[0] for policy in observed_policies) + assert asdict(observed_policies[0]) == asdict(engine.module.dense.config) + engine.backward(loss) + engine.step() + return float(loss.detach()) + + engine = create() + losses = [] + for i in range(8): + losses.append(step(engine, i)) + if i == 3: + assert engine.global_steps == 2 + engine.save_checkpoint(str(args.output / 'checkpoint'), tag='step2', client_state={'marker': 2}) + final = snapshot(engine.module.state_dict()) + optimizer_final = snapshot(engine.optimizer.state_dict()) + resumed = create() + path, state = resumed.load_checkpoint(str(args.output / 'checkpoint'), tag='step2') + assert path and state['marker'] == 2 and resumed.global_steps == 2 + resumed_losses = [step(resumed, i) for i in range(4, 8)] + assert resumed_losses == losses[4:] + compare_nested(final, resumed.module.state_dict()) + compare_nested(optimizer_final, resumed.optimizer.state_dict()) + assert resumed.global_steps == 4 + args.output.mkdir(parents=True, exist_ok=True) + result = { + 'rank': rank, + 'world_size': dist.get_world_size(), + 'zero_stage': 2, + 'gradient_accumulation_steps': 2, + 'optimizer_steps': 4, + 'losses': losses, + 'resume_losses': resumed_losses, + 'parameters_exact': True, + 'optimizer_exact': True, + 'same_policy_values': True + } + (args.output / f'result_rank{rank}.json').write_text(json.dumps(result, indent=2)) + dist.barrier() + dist.destroy_process_group() + + +if __name__ == '__main__': + main() diff --git a/tests/unit/runtime/test_hifloat8.py b/tests/unit/runtime/test_hifloat8.py new file mode 100644 index 000000000000..2abdabb52d46 --- /dev/null +++ b/tests/unit/runtime/test_hifloat8.py @@ -0,0 +1,518 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 +# DeepSpeed Team + +from dataclasses import asdict +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from deepspeed.runtime.config import DeepSpeedConfigError, get_hifloat8_config +from deepspeed.runtime.engine import DeepSpeedEngine + +# Native tests require the public API; a missing extension must fail there. +# Gate only device availability, without importing TorchAO during collection. +requires_npu = pytest.mark.skipif(not hasattr(torch, "npu") or not torch.npu.is_available(), reason="NPU unavailable") + + +def test_hifloat8_config_defaults_disabled(): + assert get_hifloat8_config({}) == { + "enabled": False, + "backend": "torchao", + "module_name_patterns": (), + "min_numel": 0, + "expected_module_count": None, + "config": None, + } + + +@pytest.mark.parametrize( + "config", + [ + { + "hifloat8": { + "enabled": True + } + }, + { + "hifloat8": { + "enabled": "yes", + "module_name_patterns": ["*"] + } + }, + { + "hifloat8": { + "enabled": True, + "module_name_patterns": [""], + "min_numel": 0 + } + }, + { + "hifloat8": { + "enabled": True, + "module_name_patterns": ["*"], + "min_numel": -1 + } + }, + { + "hifloat8": { + "enabled": True, + "module_name_patterns": ["*"], + "probe_kernel": False + } + }, + { + "hifloat8": { + "enabled": True, + "module_name_patterns": ["*"], + "unknown": 1 + } + }, + ], +) +def test_hifloat8_config_rejects_invalid_values(config): + with pytest.raises(DeepSpeedConfigError): + get_hifloat8_config(config) + + +@pytest.mark.parametrize("count", [0, -1, True, 1.5, "84"]) +def test_hifloat8_config_rejects_invalid_expected_count(count): + with pytest.raises(DeepSpeedConfigError): + get_hifloat8_config({"hifloat8": {"expected_module_count": count}}) + + +def test_hifloat8_config_accepts_dense_selection_contract(): + # Existing Swift Dense configurations require exactly 28 * 3 projections. + config = get_hifloat8_config({"hifloat8": {"expected_module_count": 84}}) + assert config["expected_module_count"] == 84 + + +@pytest.mark.parametrize("backend", ["", "unknown", None, 1]) +def test_hifloat8_config_rejects_unknown_backend(backend): + with pytest.raises(DeepSpeedConfigError, match="hifloat8.backend"): + get_hifloat8_config({"hifloat8": {"backend": backend}}) + + +def test_hifloat8_config_selects_torchao_npu(): + assert get_hifloat8_config({"hifloat8": {"backend": "torchao"}})["backend"] == "torchao" + + +@requires_npu +@pytest.mark.parametrize("policy", [ + None, {}, { + "input_dst_type_max": 15, + "weight_dst_type_max": 15, + "grad_dst_type_max": 224, + "scale_policy": "pertensor", + "compute_dtype": "bfloat16" + }, { + "input_dst_type_max": 31, + "weight_dst_type_max": 31, + "grad_dst_type_max": 127, + "compute_dtype": "bfloat16" + } +]) +def test_torchao_backend_preserves_dense_and_grouped_parameters(policy, monkeypatch): + if not hasattr(torch, "npu") or not torch.npu.is_available(): + pytest.skip("NPU unavailable") + from deepspeed.moe.ep_experts import GroupedExperts + from torchao.hifloat8 import HiFloat8Linear + import torchao.hifloat8.hifloat8_grouped as api + + policies = [] + original = api.hifloat8_grouped_mm + + def recorded(*args, **kwargs): + policies.append(kwargs["config"]) + return original(*args, **kwargs) + + monkeypatch.setattr(api, "hifloat8_grouped_mm", recorded) + model = nn.ModuleDict({ + "dense": nn.Linear(16, 16, bias=False), + "experts": GroupedExperts(16, 32, 2, use_grouped_mm=True), + }).to("npu").bfloat16() + with torch.no_grad(): + for parameter in model["experts"].parameters(): + parameter.normal_(std=0.02) + parameters_before = dict(model.named_parameters()) + keys_before = tuple(model.state_dict()) + engine = _make_engine(model, ["dense", "experts"]) + engine.device = torch.device("npu:0") + engine._config.hifloat8_config["backend"] = "torchao" + engine._config.hifloat8_config["config"] = policy + engine._configure_hifloat8() + + assert isinstance(model["dense"], HiFloat8Linear) + assert model["experts"].grouped_mm is not None + x = torch.randn(16, 16, device="npu", dtype=torch.bfloat16) + model["experts"](model["dense"](x), torch.tensor([0, 16], device="npu", dtype=torch.int64)) + assert policies and all(value is policies[0] for value in policies) + assert asdict(policies[0]) == asdict(model["dense"].config) + assert model["dense"].config.grad_dst_type_max == (policy.get("grad_dst_type_max", 224) if policy else 224) + assert tuple(model.state_dict()) == keys_before + assert all(dict(model.named_parameters())[name] is parameter for name, parameter in parameters_before.items()) + + +class _ToyModel(nn.Module): + + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.mlp = nn.ModuleDict({ + "gate_proj": nn.Linear(4, 8, bias=False), + "up_proj": nn.Linear(4, 8, bias=False), + "down_proj": nn.Linear(8, 4, bias=False), + }) + self.model.self_attn = nn.ModuleDict({"q_proj": nn.Linear(4, 4, bias=False)}) + + +def _make_engine(model, patterns, min_numel=0): + engine = object.__new__(DeepSpeedEngine) + nn.Module.__init__(engine) + engine._config = SimpleNamespace( + hifloat8_config={ + "enabled": True, + "module_name_patterns": tuple(patterns), + "min_numel": min_numel, + }, + bfloat16_config=SimpleNamespace(enabled=True), + tensor_parallel_config=SimpleNamespace(autotp_size=1), + zero_optimization_stage=2, + ) + engine._config.hifloat8_config = get_hifloat8_config({"hifloat8": engine._config.hifloat8_config}) + engine._config.float16_config = SimpleNamespace(enabled=False) + engine.pipeline_parallelism = False + engine.has_moe_layers = False + engine.mpu = None + engine.device = torch.device("cpu") + engine._set_client_model(model) + return engine + + +@requires_npu +def test_engine_selects_routed_experts_without_replacing_parameters(monkeypatch): + if not hasattr(torch, "npu"): + pytest.skip("NPU accelerator required") + # Catches silently selecting the router or orphaning packed expert weights. + from deepspeed.moe.ep_experts import GroupedExperts + import torchao.hifloat8 as implementation + + model = nn.ModuleDict({ + "experts": GroupedExperts(8, 16, 2, use_grouped_mm=False), + "router": nn.Linear(8, 2, bias=False), + }) + model.experts.w1.requires_grad_(False) + parameters = dict(model.named_parameters()) + keys = tuple(model.state_dict()) + monkeypatch.setattr(implementation, "assert_hifloat8_training_available", lambda **kwargs: None) + engine = _make_engine(model, ["experts"]) + engine._configure_hifloat8() + assert engine.hifloat8_grouped_module_names == ("experts", ) + assert engine.hifloat8_converted_module_names == () + assert model.experts.grouped_mm is not None + assert type(model.router) is nn.Linear + assert not model.experts.w1.requires_grad + assert tuple(model.state_dict()) == keys + assert all(dict(model.named_parameters())[name] is parameter for name, parameter in parameters.items()) + + +@requires_npu +def test_engine_converts_only_selected_modules_and_preserves_parameters(monkeypatch): + if not hasattr(torch, "npu"): + pytest.skip("NPU accelerator required") + model = _ToyModel() + before = dict(model.named_parameters()) + probe_calls = [] + import torchao.hifloat8 as implementation + + HiFloat8Linear = implementation.HiFloat8Linear + + monkeypatch.setattr( + implementation, + "assert_hifloat8_training_available", + lambda **kwargs: probe_calls.append(kwargs), + ) + engine = _make_engine(model, ["*.mlp.gate_proj", "*.mlp.up_proj", "*.mlp.down_proj"]) + engine._config.hifloat8_config["expected_module_count"] = 3 + engine._configure_hifloat8() + + assert len(probe_calls) == 1 + assert probe_calls[0]["device"] == torch.device("cpu") + assert probe_calls[0]["linear"] and not probe_calls[0]["grouped"] + assert probe_calls[0]["config"] is model.model.mlp["gate_proj"].config + assert engine.hifloat8_converted_module_names == ( + "model.mlp.gate_proj", + "model.mlp.up_proj", + "model.mlp.down_proj", + ) + assert all(isinstance(model.model.mlp[name], HiFloat8Linear) for name in model.model.mlp) + assert type(model.model.self_attn["q_proj"]) is nn.Linear + after = dict(model.named_parameters()) + assert before.keys() == after.keys() + assert all(after[name] is parameter for name, parameter in before.items()) + + +def test_engine_fails_when_module_selection_is_empty(): + engine = _make_engine(_ToyModel(), ["*.does_not_exist"]) + with pytest.raises(RuntimeError, match="matched no nn.Linear"): + engine._configure_hifloat8() + + +def test_engine_rejects_partially_unmatched_patterns_before_conversion(): + # A typo must not silently leave part of the requested model in BF16. + model = _ToyModel() + engine = _make_engine(model, ["*.mlp.gate_proj", "*.missing_projection"]) + with pytest.raises(RuntimeError, match="patterns matched no eligible modules"): + engine._configure_hifloat8() + assert type(model.model.mlp["gate_proj"]) is nn.Linear + + +def test_engine_rejects_wrong_module_count_before_kernel_probe(monkeypatch): + # A partially selected model must fail before native initialization or mutation. + from deepspeed.accelerator import get_accelerator + + def unexpected_provider(backend, format): + pytest.fail("module-count mismatch must fail before resolving the optional backend") + + monkeypatch.setattr(get_accelerator(), "get_hifloat8_provider", unexpected_provider) + model = _ToyModel() + engine = _make_engine(model, ["*.mlp.gate_proj"]) + engine._config.hifloat8_config["expected_module_count"] = 3 + with pytest.raises(RuntimeError, match="expected_module_count"): + engine._configure_hifloat8() + assert type(model.model.mlp["gate_proj"]) is nn.Linear + + +def test_hifloat8_validation_rejects_autotp(): + engine = _make_engine(_ToyModel(), ["*"]) + engine._config.tensor_parallel_config.autotp_size = 2 + + with pytest.raises(RuntimeError, match="AutoTP"): + engine._validate_hifloat8_configuration() + + +def test_missing_optional_backend_reports_selection_before_model_mutation(monkeypatch): + # Provider policy parsing may report a missing optional dependency on any + # accelerator; the engine must retain its diagnostic and parameter contract. + from deepspeed.accelerator import get_accelerator + + error = ImportError("optional low-precision backend is not installed") + + def missing_policy(values): + raise error + + provider = SimpleNamespace(HiFloat8Config=SimpleNamespace(from_dict=missing_policy)) + monkeypatch.setattr(get_accelerator(), "get_hifloat8_provider", lambda: provider) + model = _ToyModel() + parameters = dict(model.named_parameters()) + engine = _make_engine(model, ["*.mlp.gate_proj"]) + with pytest.raises(RuntimeError, + match=r"selected 1 Linear modules \(32 matrix elements\).+before conversion") as caught: + engine._configure_hifloat8() + assert caught.value.__cause__ is error + assert type(model.model.mlp["gate_proj"]) is nn.Linear + assert all(dict(model.named_parameters())[name] is parameter for name, parameter in parameters.items()) + + +def test_hifloat8_validation_rejects_custom_model_parallel_unit(): + engine = _make_engine(_ToyModel(), ["*"]) + engine.mpu = object() + + with pytest.raises(RuntimeError, match="model-parallel unit"): + engine._validate_hifloat8_configuration() + + +@requires_npu +@pytest.mark.parametrize("error", [RuntimeError("unsupported device"), ImportError("missing torchao_npu backend")]) +def test_engine_kernel_failure_reports_selection_and_does_not_mutate(monkeypatch, error): + if not hasattr(torch, "npu"): + pytest.skip("NPU accelerator required") + import torchao.hifloat8 as implementation + + model = _ToyModel() + before = dict(model.named_parameters()) + + def fail_probe(**_kwargs): + raise error + + monkeypatch.setattr(implementation, "assert_hifloat8_training_available", fail_probe) + engine = _make_engine(model, ["*.mlp.gate_proj"]) + + with pytest.raises(RuntimeError, match=r"selected 1 Linear modules \(32 matrix elements\).+before conversion"): + engine._configure_hifloat8() + + assert type(model.model.mlp["gate_proj"]) is nn.Linear + after = dict(model.named_parameters()) + assert before.keys() == after.keys() + assert all(after[name] is parameter for name, parameter in before.items()) + + +@pytest.mark.parametrize("policy", [True, 1, "bf16", []]) +def test_hifloat8_config_rejects_non_object_policy(policy): + with pytest.raises(DeepSpeedConfigError, match="hifloat8.config"): + get_hifloat8_config({"hifloat8": {"backend": "torchao_npu", "config": policy}}) + + +def test_hifloat8_legacy_backend_is_rejected_with_migration_instruction(): + with pytest.raises(DeepSpeedConfigError, match="legacy torch_npu helper was removed"): + get_hifloat8_config({"hifloat8": {"backend": "torch_npu"}}) + policy = {"grad_dst_type_max": 15, "compute_dtype": "bfloat16"} + assert get_hifloat8_config({"hifloat8": {"config": policy}})["config"] == policy + + +def test_minimal_hifloat8_json_requires_no_numerical_fields(): + config = get_hifloat8_config({ + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": 2 + }, + "hifloat8": { + "enabled": True, + "backend": "torchao_npu", + "module_name_patterns": ["*.experts", "*.shared_experts.*_proj"], + "min_numel": 65536, + }, + }) + assert config["enabled"] and config["config"] is None + assert config["backend"] == "torchao" + assert config["min_numel"] == 65536 + + +def test_hifloat8_config_keeps_policy_opaque_and_normalizes_legacy_backend(): + policy = {"backend_specific_option": 7} + assert get_hifloat8_config({"hifloat8": {"backend": "torchao_npu", "config": policy}})["backend"] == "torchao" + assert get_hifloat8_config({"hifloat8": {"config": policy}})["config"] == policy + + +def test_accelerator_without_hifloat8_fails_before_conversion(monkeypatch): + from deepspeed.accelerator import get_accelerator + + monkeypatch.setattr(get_accelerator(), "get_hifloat8_provider", lambda: None) + model = _ToyModel() + engine = _make_engine(model, ["*.mlp.gate_proj"]) + with pytest.raises(RuntimeError, match="No HiFloat8 provider"): + engine._configure_hifloat8() + assert type(model.model.mlp["gate_proj"]) is nn.Linear + + +@requires_npu +@pytest.mark.parametrize("grouped", [False, True]) +def test_engine_requires_grouped_config_only_when_experts_are_selected(monkeypatch, grouped): + # A Linear-only backend must work; a missing grouped handler must fail before any conversion. + import torchao.hifloat8 as api + from deepspeed.accelerator import get_accelerator + from deepspeed.moe.ep_experts import GroupedExperts + + provider = SimpleNamespace(HiFloat8Config=api.HiFloat8Config, + assert_hifloat8_training_available=api.assert_hifloat8_training_available) + monkeypatch.setattr(get_accelerator(), "get_hifloat8_provider", lambda: provider) + model = nn.ModuleDict({"dense": nn.Linear(16, 16)}) + if grouped: + model["experts"] = GroupedExperts(16, 32, 2, use_grouped_mm=False) + model = model.to("npu").bfloat16() + parameters = dict(model.named_parameters()) + engine = _make_engine(model, list(model)) + engine.device = torch.device("npu:0") + if grouped: + with pytest.raises(RuntimeError, match="HiFloat8GroupedConfig"): + engine._configure_hifloat8() + assert type(model["dense"]) is nn.Linear + assert model["experts"].grouped_mm is None + else: + engine._configure_hifloat8() + assert isinstance(model["dense"], api.HiFloat8Linear) + output = model["dense"](torch.ones(16, 16, device="npu", dtype=torch.bfloat16)) + output.float().sum().backward() + assert bool(torch.isfinite(output).all()) + assert all(dict(model.named_parameters())[name] is value for name, value in parameters.items()) + + +def test_npu_provider_missing_dependency_has_install_hint(monkeypatch): + from importlib.metadata import PackageNotFoundError + from deepspeed.accelerator.npu import hifloat8 + + def missing_version(name): + raise PackageNotFoundError(name) + + monkeypatch.setattr(hifloat8, "version", missing_version) + with pytest.raises(ImportError, match="Install torchao"): + hifloat8.get_hifloat8_provider() + + +def test_removed_configuration_extension_reports_existing_schema(): + # Do not silently disable a run using the short-lived generic section. + with pytest.raises(DeepSpeedConfigError, match="hifloat8"): + get_hifloat8_config({"low_precision": {"enabled": False}}) + + +@requires_npu +@pytest.mark.parametrize("kind", ["mixed", "linear", "grouped"]) +def test_public_quantize_entry_converts_dense_and_grouped_in_place(monkeypatch, kind): + # A split expert path would bypass the public API; wrap the real collaborator. + import torchao.quantization as quantization + from torchao.hifloat8 import HiFloat8Linear + from deepspeed.moe.ep_experts import GroupedExperts + + torch.manual_seed(59) + if kind == "mixed": + model = nn.ModuleDict({"dense": nn.Linear(16, 16), "experts": GroupedExperts(16, 32, 2, use_grouped_mm=False)}) + patterns = ["dense", "experts"] + elif kind == "linear": + model = nn.Linear(16, 16) + patterns = ["*"] + else: + model = GroupedExperts(16, 32, 2, use_grouped_mm=False) + patterns = ["*"] + model = model.to("npu").bfloat16() + with torch.no_grad(): + for parameter in model.parameters(): + parameter.normal_(std=0.02) + before = dict(model.named_parameters()) + keys = tuple(model.state_dict()) + calls = [] + policies = [] + import torchao.hifloat8.hifloat8_grouped as api + original_grouped = api.hifloat8_grouped_mm + + def recorded_grouped(*args, **kwargs): + policies.append(kwargs["config"]) + return original_grouped(*args, **kwargs) + + monkeypatch.setattr(api, "hifloat8_grouped_mm", recorded_grouped) + original = quantization.quantize_ + + def counted(actual_model, config, **kwargs): + calls.append(actual_model) + return original(actual_model, config, **kwargs) + + monkeypatch.setattr(quantization, "quantize_", counted) + engine = _make_engine(model, patterns) + engine.device = torch.device("npu:0") + engine._configure_hifloat8() + assert calls == [model] + assert engine.module is model + assert tuple(model.state_dict()) == keys + assert all(dict(model.named_parameters())[name] is value for name, value in before.items()) + x = torch.randn(16, 16, device="npu", dtype=torch.bfloat16, requires_grad=True) + counts = torch.tensor([0, 16], device="npu", dtype=torch.int64) + if kind == "mixed": + assert isinstance(model["dense"], HiFloat8Linear) + output = model["experts"](model["dense"](x), counts) + elif kind == "linear": + assert isinstance(model, HiFloat8Linear) + output = model(x) + else: + assert model.grouped_mm is not None + output = model(x, counts) + if kind == "mixed": + assert policies and all(value is policies[0] for value in policies) + assert asdict(policies[0]) == asdict(model["dense"].config) + assert output.dtype == x.dtype + output.float().square().mean().backward() + assert bool(torch.isfinite(x.grad).all()) + assert all(parameter.grad is not None and bool(torch.isfinite(parameter.grad).all()) + for parameter in model.parameters()) diff --git a/tests/unit/v1/moe/test_hifloat8_experts.py b/tests/unit/v1/moe/test_hifloat8_experts.py index 18f0d816ed6a..6757a6c84990 100644 --- a/tests/unit/v1/moe/test_hifloat8_experts.py +++ b/tests/unit/v1/moe/test_hifloat8_experts.py @@ -90,6 +90,54 @@ def fake_grouped_mm(lhs, weight, offsets, *, trans_b=False): torch.testing.assert_close(value.grad, reference.grad) +@pytest.mark.skipif(get_accelerator().device_name() != "npu", reason="requires an available NPU") +@pytest.mark.parametrize("counts", [[64, 64], [128, 0], [0, 128], [1, 127], [0, 0], [32, 32, 32, 32], [0, 127, 0, 1], + [1, 0, 0, 127], [0, 0, 0, 128], [0, 0, 0, 0], [2] * 64]) +@pytest.mark.parametrize("dim,hidden_dim", [(128, 128), (2048, 512)]) +@pytest.mark.parametrize("backend", ["bf16", "torchao"]) +def test_npu_hifloat8_expert_forward_and_gradients(counts, dim, hidden_dim, backend): + # Catches unsupported native dispatch, inaccurate dX/dW, and empty-expert leakage. + import torch_npu + + torch_npu.npu.set_device(0) + torch.manual_seed(42) + baseline = GroupedExperts(dim, hidden_dim, len(counts), use_grouped_mm=False).to(device="npu", + dtype=torch.bfloat16) + with torch.no_grad(): + for weight in baseline.parameters(): + weight.normal_(std=0.02) + candidate = GroupedExperts(dim, hidden_dim, len(counts)).to(device="npu", dtype=torch.bfloat16) + candidate.load_state_dict(baseline.state_dict()) + if backend == "torchao": + from torchao.quantization import quantize_ + from torchao.hifloat8 import HiFloat8GroupedConfig + + quantize_(candidate, HiFloat8GroupedConfig(), filter_fn=lambda module, name: name == "") + x = torch.randn(sum(counts), dim, device="npu", dtype=torch.bfloat16, requires_grad=True) + candidate_x = x.detach().clone().requires_grad_() + groups = torch.tensor(counts, device="npu", dtype=torch.int64) + reference = baseline(x, groups) + actual = candidate(candidate_x, groups) + grad = torch.randn_like(reference) + reference.backward(grad) + actual.backward(grad) + torch_npu.npu.synchronize() + pairs = [(actual, reference), (candidate_x.grad, x.grad)] + pairs.extend((actual_w.grad, ref_w.grad) for actual_w, ref_w in zip(candidate.parameters(), baseline.parameters())) + for actual_t, reference_t in pairs: + assert torch.isfinite(actual_t).all() + assert torch.isfinite(reference_t).all() + if not reference_t.numel() or reference_t.float().norm() == 0: + assert torch.count_nonzero(actual_t) == 0 + continue + actual_flat, reference_flat = actual_t.float().flatten(), reference_t.float().flatten() + assert torch.nn.functional.cosine_similarity(actual_flat, reference_flat, dim=0) >= 0.99 + assert (actual_flat - reference_flat).norm() / reference_flat.norm() <= 0.15 + for index, count in enumerate(counts): + if count == 0: + assert all(torch.count_nonzero(weight.grad[index]) == 0 for weight in candidate.parameters()) + + @pytest.mark.skipif(get_accelerator().device_name() != "npu", reason="requires an available NPU") @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.parametrize("trans_b", [False, True]) @@ -138,6 +186,45 @@ def test_npu_native_grouped_mm_matches_matmul_with_partial_gradients(dtype, tran assert torch.count_nonzero(weight.grad[index]) == 0 +@pytest.mark.skipif(get_accelerator().device_name() != "npu", reason="requires an available NPU") +def test_npu_provider_experts_preserve_bfloat16_residual_with_float16_compute(): + # A working-dtype leak promotes the residual and breaks the BF16-facing module contract. + torch.npu.set_device(0) + torch.manual_seed(73) + experts = GroupedExperts(64, 64, 4).to(device="npu", dtype=torch.bfloat16) + projection = torch.nn.Linear(64, 32).to(device="npu", dtype=torch.bfloat16) + with torch.no_grad(): + for weight in experts.parameters(): + weight.normal_(std=0.02) + parameters = dict(experts.named_parameters()) + state_keys = tuple(experts.state_dict()) + from torchao.quantization import quantize_ + from torchao.hifloat8 import HiFloat8GroupedConfig + + policy = HiFloat8GroupedConfig.from_dict({"compute_dtype": "float16"}) + quantize_(experts, policy, filter_fn=lambda module, name: name == "") + counts = torch.tensor([0, 8, 0, 4], device="npu", dtype=torch.int64) + x = torch.randn(15, 64, device="npu", dtype=torch.bfloat16, requires_grad=True) + branch = experts(x, counts) + residual = x + branch + output = projection(residual) + assert output.dtype == branch.dtype == residual.dtype == x.dtype + assert torch.count_nonzero(branch[12:]) == 0 + output.float().square().mean().backward() + torch.npu.synchronize() + + assert x.grad.dtype == torch.bfloat16 and bool(torch.isfinite(x.grad).all()) + assert tuple(experts.state_dict()) == state_keys + for name, parameter in experts.named_parameters(): + assert parameter is parameters[name] + assert parameter.dtype == parameter.grad.dtype == torch.bfloat16 + assert bool(torch.isfinite(parameter.grad).all()) + assert torch.count_nonzero(parameter.grad[0]) == torch.count_nonzero(parameter.grad[2]) == 0 + for parameter in projection.parameters(): + assert parameter.dtype == parameter.grad.dtype == torch.bfloat16 + assert bool(torch.isfinite(parameter.grad).all()) + + def test_grouped_override_rejects_explicit_fused_gate_up(): # Upstream can request a fused gate/up path; a three-GEMM override cannot # silently ignore that public choice after this integration is merged.