diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ec9e7ffdf5e..2a26d9c181e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -10,6 +10,8 @@ Changelog - Add the ``nvfp4_act_headroom`` calibration algorithm for NVFP4 **activation** global scales. Instead of setting the global scale from the largest per-block amax seen during calibration (plain ``max``, which leaves no room above it so any larger activation saturates), it anchors the scale to a low percentile of the per-block amax distribution, leaving the rest of the FP8 block-scale range as headroom: ``amax = max(rho * anchor, upper)``, where ``anchor`` and ``upper`` are the per-block amaxes at ``anchor_percentile`` (default 1) and ``upper_percentile`` (default 99.99; set to 100 to never clip calibration data), and ``rho`` (default 16384) is the headroom factor. Applies only to NVFP4 dynamic-block input quantizers; ``SequentialQuantizer`` activation quantizers raise. Weight scales are an orthogonal axis selected by a nested ``weight_scale_algorithm`` (``max`` by default, or ``mse`` / ``local_hessian``), so one recipe can combine a weight calibration with this activation policy in a single pass. Ships ``modelopt_recipes/general/ptq/nvfp4_act_headroom-kv_fp8_cast.yaml``, which mirrors ``nvfp4_default-kv_fp8_cast`` with only the calibration algorithm swapped and exports a standard NVFP4 checkpoint. +- Add PTQ support for Step-3.7 (``stepfun-ai/Step-3.7-Flash``). The expert-indexed ``MoELinear`` modules these checkpoints ship via ``trust_remote_code`` are now detected structurally instead of by model class name, so the routed experts get per-expert quantizers on any Step revision. Quantize with the new ``huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast`` or ``huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8`` recipes — the general recipes select experts by module names Step does not use and would leave the model unquantized. + *Megatron Framework (M-LM / M-Bridge)* - Add SFT-masked data support to ``examples/megatron_bridge/distill.py``: ``--sft --sft_dataset_root `` distills on raw prompt-completion JSONL (``{"input", "output"}`` records) with the loss masked to the response tokens, using Megatron-Bridge's ``FinetuningDatasetConfig`` and the model's own HuggingFace tokenizer instead of the pretraining ``GPTDataset`` and ``NullTokenizer``. diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index dde24513086..28602d4eb69 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -1875,7 +1875,7 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): class _QuantMoELinear(QuantModule): - """Quantization wrapper for Step3p5 MoELinear modules (fused expert weights). + """Quantization wrapper for expert-indexed MoELinear modules (fused expert weights). MoELinear has weight shape [num_experts, out_features, in_features] with forward(x, expert_id). We expand it into per-expert nn.Linear modules so @@ -1921,29 +1921,64 @@ def forward(self, x, expert_id): return expert(x).float() -def register_step3p5_moe_on_the_fly(model): - """Register Step3p5 MoELinear for quantization. +def _is_expert_indexed_moe_linear(module: nn.Module) -> bool: + """Whether ``module`` packs one projection's experts into an expert-indexed 3-D weight. - Step3p5 uses a custom MoELinear class (loaded via trust_remote_code) with - weight shape [num_experts, out_features, in_features] and forward(x, expert_id). - We detect it by model class name, then grab the type from the first MoE layer. + The Step family (``stepfun-ai/Step-3.5-Flash``, ``stepfun-ai/Step-3.7-Flash``) ships a + custom ``MoELinear`` via ``trust_remote_code``: a plain ``nn.Module`` holding a single + ``weight`` of shape ``[num_experts, out_features, in_features]``, whose + ``forward(x, expert_id)`` runs ``F.linear`` against the selected expert's slice. The + weights therefore live on the projection submodule rather than on the expert container, + which is what :func:`_fused_experts_wrapper_class` looks for, and the module is not an + ``nn.Linear``, so neither the fused-experts path nor the plain linear path claims it. + + Detection is structural rather than keyed on class names so new Step revisions (or any + other model shipping the same layout) are picked up without another hardcoded name. """ - if type(model).__name__ not in ("Step3p5ForCausalLM", "Step3p5Model"): - return - for module in model.modules(): - if type(module).__name__ == "Step3p5MoEMLP": - moe_linear_type = type(module.up_proj) - if QuantModuleRegistry.get(moe_linear_type) is None: - QuantModuleRegistry.register({moe_linear_type: f"hf.{moe_linear_type.__name__}"})( - _QuantMoELinear - ) - break + weight = getattr(module, "weight", None) + if not isinstance(weight, (nn.Parameter, Tensor)) or weight.dim() != 3: + return False + if not all(hasattr(module, attr) for attr in ("num_experts", "in_features", "out_features")): + return False + # `_QuantMoELinear.forward` takes (x, expert_id), so only claim modules whose callers + # already drive them that way. + try: + params = list(inspect.signature(type(module).forward).parameters.values())[1:] + except (TypeError, ValueError): + return False + positional = [ + p + for p in params + if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) and p.default is p.empty + ] + return len(positional) == 2 + + +def register_moe_linear_on_the_fly(model): + """Register expert-indexed ``MoELinear`` modules (Step-3.5 / Step-3.7) for quantization. + + Without this the routed experts carry no quantizer at all: an experts-only recipe matches + nothing and the export writes a checkpoint with ``quant_algo: null``. + """ + visited_types = set() + for name, module in model.named_modules(): + mod_type = type(module) + if mod_type in visited_types or QuantModuleRegistry.get(mod_type) is not None: + continue + visited_types.add(mod_type) + + if _is_expert_indexed_moe_linear(module): + print( + f"\033[1mDetected expert-indexed MoE linear '{name}' of type " + f"{mod_type.__name__}, registering with _QuantMoELinear.\033[0m" + ) + QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(_QuantMoELinear) def _reconstruct_fused_moe_linear(model: nn.Module) -> None: - """Reconstruct QuantMoELinear per-expert weights back to original 3D MoELinear format. + """Reconstruct :class:`_QuantMoELinear` per-expert weights back to the 3-D MoELinear format. - After _process_quantized_modules, each expert's nn.Linear inside QuantMoELinear has: + After _process_quantized_modules, each expert's nn.Linear inside the wrapper has: - weight: fp4-quantized tensor [out_features, in_features] - weight_scale, weight_scale_2: per-block / global scales - input_scale: activation scale (if calibrated) @@ -1951,12 +1986,12 @@ def _reconstruct_fused_moe_linear(model: nn.Module) -> None: This stacks them back into the original MoELinear layout so the exported state_dict uses the original key names (e.g. moe.up_proj.weight with shape [N, out, in]). - Note: QuantMoELinear is the dynamically generated class name (Quant + MoELinear), - not _QuantMoELinear which is the implementation class. + Matched by wrapper type rather than by the dynamically generated class name (``Quant`` + + the model's own class name): a model whose class is not spelled ``MoELinear`` would + otherwise quantize normally but export unusable per-expert keys. """ for _name, module in model.named_modules(): - # Match QuantMoELinear (dynamically generated name) not _QuantMoELinear (implementation class) - if type(module).__name__ != "QuantMoELinear": + if not isinstance(module, _QuantMoELinear): continue n = module.num_experts @@ -1986,7 +2021,7 @@ def _reconstruct_fused_moe_linear(model: nn.Module) -> None: [ register_falcon_linears_on_the_fly, register_dbrx_moe_on_the_fly, - register_step3p5_moe_on_the_fly, + register_moe_linear_on_the_fly, register_fused_experts_on_the_fly, force_eager_experts_impl_on_the_fly, register_sparse_moe_on_the_fly, diff --git a/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml new file mode 100644 index 00000000000..251f807b53b --- /dev/null +++ b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast.yaml @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Model-specific PTQ recipe: Step-3.7 routed-experts-only NVFP4 with FP8 KV-cache cast. +# +# Step's remote code names the MoE block `moe` (not `experts` / `block_sparse_moe`), so the +# general experts-only recipe matches nothing on this architecture. The routed experts are +# `moe.{gate,up,down}_proj`, the router is `moe.gate`, and each MoE layer also has a dense +# `share_expert` that stays in BF16. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8_cast: configs/ptq/units/kv_fp8_cast + +metadata: + recipe_type: ptq + description: >- + Applies dynamic NVFP4 to Step-3.7 routed-expert weight and input quantizers, plus FP8 + KV-cache cast mode using constant amax; the router, the shared expert and the dense MLPs + stay unquantized. Uses max calibration. + +quantize: + algorithm: + method: max + layerwise: false + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*moe*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8_cast + - $import: default_disabled_quantizers + # Router and shared expert are matched by `*moe*` above; disable them last (later wins). + - quantizer_name: '*moe.gate.*' + enable: false + - quantizer_name: '*share_expert*' + enable: false diff --git a/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml new file mode 100644 index 00000000000..5298adabf22 --- /dev/null +++ b/modelopt_recipes/huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8.yaml @@ -0,0 +1,59 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Model-specific PTQ recipe: Step-3.7 MLP/MoE NVFP4 with calibrated FP8 KV cache. +# +# Same module-naming caveat as the experts-only recipe: Step's MoE block is `moe`, so the +# general `*mlp*` / `*.experts.*` patterns reach only the dense layers' `mlp` submodules and +# leave the routed experts — the bulk of the model — in BF16. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + kv_fp8: configs/ptq/units/kv_fp8 + +metadata: + recipe_type: ptq + description: >- + Applies dynamic NVFP4 to Step-3.7 routed-expert and dense-MLP weight and input + quantizers, plus calibrated FP8 KV-cache quantization; the router and the shared expert + stay unquantized. Uses max calibration. + +quantize: + algorithm: + method: max + layerwise: false + quant_cfg: + - $import: base_disable_all + - quantizer_name: '*moe*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*moe*input_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp*weight_quantizer' + cfg: + $import: nvfp4 + - quantizer_name: '*mlp*input_quantizer' + cfg: + $import: nvfp4 + - $import: kv_fp8 + - $import: default_disabled_quantizers + # Router and shared expert are matched by `*moe*` above; disable them last (later wins). + - quantizer_name: '*moe.gate.*' + enable: false + - quantizer_name: '*share_expert*' + enable: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index a1fb2d94c57..9ef749824a0 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -295,6 +295,15 @@ A lighter case: **`step3p5/Step3.5-Flash/ptq/nvfp4-mlp-only`** is close to to one released checkpoint and carrying instance-specific disables (`share_expert`, `moe.gate`, the conv1d branches). +**`step3p7/ptq/{nvfp4_experts_only-kv_fp8_cast,nvfp4_mlp_only-kv_fp8}`** are the +Step-3.7 equivalents, and the reason they exist is **module naming**: Step calls +the MoE block `moe` and the dense sibling `share_expert`, so the general +recipes' `*.experts.*`, `*block_sparse_moe*` and `*mlp*` patterns match nothing +on the routed experts — the general recipe would quantize *nothing* and export a +checkpoint with `quant_algo: null`. These select `*moe*` instead and disable the +router (`moe.gate`) and `share_expert` on top. Use them, not the general +recipes, for any Step checkpoint. + ### Algorithm overrides — `gemma`, `gemma4`, `mpt` These quantize the **same layers** as the general recipes; only the diff --git a/tests/unit/recipe/test_step3p7_recipes.py b/tests/unit/recipe/test_step3p7_recipes.py new file mode 100644 index 00000000000..2dc2f0c7339 --- /dev/null +++ b/tests/unit/recipe/test_step3p7_recipes.py @@ -0,0 +1,172 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Step-3.7 PTQ recipes: what the `moe` / `share_expert` naming does and does not match.""" + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch.quantization.nn import QuantModuleRegistry + +HIDDEN_SIZE = 32 +MOE_INTERMEDIATE_SIZE = 16 +NUM_EXPERTS = 2 + + +class _MoELinear(nn.Module): + """Step's expert-indexed projection: one 3-D weight, ``forward(x, expert_id)``.""" + + def __init__(self, num_experts, in_features, out_features): + super().__init__() + self.num_experts = num_experts + self.in_features = in_features + self.out_features = out_features + self.weight = nn.Parameter(torch.randn(num_experts, out_features, in_features) * 0.02) + + def forward(self, x, expert_id): + return F.linear(x.float(), self.weight[expert_id].float()) + + +class _StepMoEMLP(nn.Module): + def __init__(self): + super().__init__() + self.num_experts = NUM_EXPERTS + self.gate = nn.Linear(HIDDEN_SIZE, NUM_EXPERTS, bias=False) # router + self.up_proj = _MoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.gate_proj = _MoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.down_proj = _MoELinear(NUM_EXPERTS, MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE) + + +class _StepMLP(nn.Module): + """Dense FFN — used both as the non-MoE layers' ``mlp`` and as ``share_expert``.""" + + def __init__(self): + super().__init__() + self.gate_proj = nn.Linear(HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE, bias=False) + self.up_proj = nn.Linear(HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE, bias=False) + self.down_proj = nn.Linear(MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE, bias=False) + + +class _StepMoELayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False) + self.moe = _StepMoEMLP() + self.share_expert = _StepMLP() + + +class _StepDenseLayer(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = nn.Module() + self.self_attn.q_proj = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False) + self.mlp = _StepMLP() + + +class _StepModel(nn.Module): + """Mirrors Step3p7ForConditionalGeneration's module paths (one MoE, one dense layer).""" + + def __init__(self): + super().__init__() + self.model = nn.Module() + self.model.language_model = nn.Module() + self.model.language_model.layers = nn.ModuleList([_StepMoELayer(), _StepDenseLayer()]) + self.lm_head = nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE, bias=False) + + +@pytest.fixture(autouse=True) +def _unregister_moe_linear(): + yield + if QuantModuleRegistry.get(_MoELinear) is not None: + QuantModuleRegistry.unregister(_MoELinear) + + +def _quantize_with_recipe(name): + """Convert (no calibration) with a built-in recipe and return the model.""" + model = _StepModel() + config = load_recipe(name).quantize.model_dump() + config["algorithm"] = None + # `mtq.quantize` runs the custom-model plugins itself, which is what registers Step's + # `MoELinear` — no explicit registration here, so this also covers that hook firing. + mtq.quantize(model, config) + return model + + +def _enabled(module, quantizer="weight_quantizer"): + return getattr(module, quantizer).is_enabled + + +@pytest.mark.parametrize( + "recipe_name", + [ + "huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast", + "huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8", + ], +) +def test_routed_experts_are_quantized(recipe_name): + """The routed experts — the bulk of the model — must end up quantized, per expert.""" + model = _quantize_with_recipe(recipe_name) + moe = model.model.language_model.layers[0].moe + + for proj in ("up_proj", "gate_proj", "down_proj"): + experts = getattr(moe, proj).experts + assert len(experts) == NUM_EXPERTS + for expert in experts: + assert _enabled(expert) + assert _enabled(expert, "input_quantizer") + # Dynamic NVFP4: 16-element blocks along the input dim. + assert expert.weight_quantizer.block_sizes[-1] == 16 + assert expert.weight_quantizer.num_bits == (2, 1) + + +@pytest.mark.parametrize( + "recipe_name", + [ + "huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast", + "huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8", + ], +) +def test_router_shared_expert_and_head_stay_bf16(recipe_name): + """`*moe*` also matches the router; the shared expert and lm_head stay unquantized too.""" + model = _quantize_with_recipe(recipe_name) + moe_layer = model.model.language_model.layers[0] + + assert not _enabled(moe_layer.moe.gate) + for proj in ("gate_proj", "up_proj", "down_proj"): + assert not _enabled(getattr(moe_layer.share_expert, proj)) + assert not _enabled(model.lm_head) + + +def test_experts_only_leaves_dense_mlp_bf16(): + model = _quantize_with_recipe("huggingface/step3p7/ptq/nvfp4_experts_only-kv_fp8_cast") + dense_mlp = model.model.language_model.layers[1].mlp + + for proj in ("gate_proj", "up_proj", "down_proj"): + assert not _enabled(getattr(dense_mlp, proj)) + + +def test_mlp_only_also_quantizes_dense_mlp(): + model = _quantize_with_recipe("huggingface/step3p7/ptq/nvfp4_mlp_only-kv_fp8") + dense_mlp = model.model.language_model.layers[1].mlp + + for proj in ("gate_proj", "up_proj", "down_proj"): + assert _enabled(getattr(dense_mlp, proj)) + # Attention projections are out of scope for both recipes. + assert not _enabled(model.model.language_model.layers[1].self_attn.q_proj) diff --git a/tests/unit/torch/quantization/plugins/test_moe_linear.py b/tests/unit/torch/quantization/plugins/test_moe_linear.py new file mode 100644 index 00000000000..659e65a864f --- /dev/null +++ b/tests/unit/torch/quantization/plugins/test_moe_linear.py @@ -0,0 +1,181 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for _QuantMoELinear: expert-indexed MoE weights (Step-3.5 / Step-3.7 remote code).""" + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +pytest.importorskip("transformers") + +import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.nn import QuantModuleRegistry +from modelopt.torch.quantization.plugins.huggingface import ( + _is_expert_indexed_moe_linear, + _QuantMoELinear, + _reconstruct_fused_moe_linear, + register_moe_linear_on_the_fly, +) + +NUM_EXPERTS = 4 +HIDDEN_SIZE = 32 +MOE_INTERMEDIATE_SIZE = 16 +TOP_K = 2 + + +class _SyntheticMoELinear(nn.Module): + """Mimics Step-3.5 / Step-3.7 ``MoELinear`` (verbatim layout from their remote code).""" + + def __init__(self, num_experts, in_features, out_features): + super().__init__() + self.num_experts = num_experts + self.in_features = in_features + self.out_features = out_features + self.weight = nn.Parameter(torch.randn(num_experts, out_features, in_features) * 0.02) + + def forward(self, x, expert_id): + return F.linear(x.float(), self.weight[expert_id].float()) + + +class _SyntheticStepMoEMLP(nn.Module): + """Mimics ``Step3p7MoEMLP``: a router plus three expert-indexed projections.""" + + def __init__(self): + super().__init__() + self.num_experts = NUM_EXPERTS + self.top_k = TOP_K + self.gate = nn.Linear(HIDDEN_SIZE, NUM_EXPERTS, bias=False) + self.act_fn = nn.SiLU() + self.up_proj = _SyntheticMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.gate_proj = _SyntheticMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + self.down_proj = _SyntheticMoELinear(NUM_EXPERTS, MOE_INTERMEDIATE_SIZE, HIDDEN_SIZE) + + def forward(self, hidden_states): + tokens = hidden_states.view(-1, HIDDEN_SIZE) + routing = F.softmax(self.gate(tokens).float(), dim=-1) + weights, indices = torch.topk(routing, self.top_k, dim=-1) + out = torch.zeros_like(tokens) + for expert_id in range(self.num_experts): + pos, token_idx = torch.where(indices == expert_id) + if token_idx.numel() == 0: + continue + current = tokens[pos] + gate = self.act_fn(self.gate_proj(current, expert_id)) + up = self.up_proj(current, expert_id) + expert_out = self.down_proj(gate * up, expert_id) + out.index_add_(0, pos, (expert_out * weights[pos, token_idx, None]).to(out.dtype)) + return out.view_as(hidden_states) + + +class _TinyStepModel(nn.Module): + def __init__(self): + super().__init__() + self.moe = _SyntheticStepMoEMLP() + + def forward(self, x): + return self.moe(x) + + +@pytest.fixture(autouse=True) +def _unregister_synthetic_moe_linear(): + """Keep the on-the-fly registration from leaking into other tests.""" + yield + if QuantModuleRegistry.get(_SyntheticMoELinear) is not None: + QuantModuleRegistry.unregister(_SyntheticMoELinear) + + +def _moe_quant_cfg(): + """Per-tensor INT8 on the expert projections only — CPU-friendly, no kernels needed.""" + return { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*moe*weight_quantizer", "cfg": {"num_bits": 8, "axis": 0}}, + {"quantizer_name": "*moe*input_quantizer", "cfg": {"num_bits": 8, "axis": None}}, + {"quantizer_name": "*moe.gate.*", "enable": False}, + ], + "algorithm": "max", + } + + +def test_expert_indexed_moe_linear_is_detected(): + assert _is_expert_indexed_moe_linear( + _SyntheticMoELinear(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + ) + + +@pytest.mark.parametrize( + "module", + [ + pytest.param(nn.Linear(HIDDEN_SIZE, HIDDEN_SIZE), id="plain_linear_2d_weight"), + pytest.param(nn.LayerNorm(HIDDEN_SIZE), id="norm_1d_weight"), + ], +) +def test_unrelated_modules_are_not_claimed(module): + assert not _is_expert_indexed_moe_linear(module) + + +def test_module_with_3d_weight_but_other_forward_is_not_claimed(): + """A 3-D weight alone is not enough — the forward must take ``(x, expert_id)``.""" + + class _NotExpertIndexed(_SyntheticMoELinear): + def forward(self, x, top_k_index, top_k_weights): + return x + + assert not _is_expert_indexed_moe_linear( + _NotExpertIndexed(NUM_EXPERTS, HIDDEN_SIZE, MOE_INTERMEDIATE_SIZE) + ) + + +def test_registration_is_not_gated_on_model_class_name(): + """Detection is structural, so a Step-3.7-style model registers as readily as Step-3.5.""" + model = _TinyStepModel() + assert QuantModuleRegistry.get(_SyntheticMoELinear) is None + + register_moe_linear_on_the_fly(model) + + assert issubclass(QuantModuleRegistry.get(_SyntheticMoELinear), _QuantMoELinear) + + +def test_expert_indexed_moe_is_quantized_and_reconstructed(): + """Each expert gets its own quantizers, and export folds them back to the 3-D layout.""" + torch.manual_seed(0) + model = _TinyStepModel() + reference_weight = model.moe.up_proj.weight.detach().clone() + + def forward_loop(m): + m(torch.randn(2, 8, HIDDEN_SIZE)) + + mtq.quantize(model, _moe_quant_cfg(), forward_loop=forward_loop) + + # Every expert of every projection carries its own calibrated quantizer pair. + for proj in ("up_proj", "gate_proj", "down_proj"): + experts = getattr(model.moe, proj).experts + assert len(experts) == NUM_EXPERTS + for expert in experts: + assert expert.weight_quantizer.is_enabled + assert expert.weight_quantizer.amax is not None + # The router stays untouched. + assert not model.moe.gate.weight_quantizer.is_enabled + + _reconstruct_fused_moe_linear(model) + + # Back to the original ``[num_experts, out_features, in_features]`` parameter, so the + # exported keys match the hub checkpoint instead of per-expert names. + up_proj = model.moe.up_proj + assert not hasattr(up_proj, "experts") + assert up_proj.weight.shape == reference_weight.shape + assert torch.equal(up_proj.weight, reference_weight)