diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 7c526d8255f..200ec6e5dae 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -108,12 +108,13 @@ def _is_quant_moe_linear(module: nn.Module) -> bool: name gets uniquified — would bypass this handler and export without the input-amax fallback. The name check is kept as a fallback so stand-in modules match too. - The wrapper lives in the optional transformers plugin, hence the lazy import. + The wrapper lives in the Step PTQ modeling module, imported lazily like other export + modules do to avoid an import cycle with ``modelopt.torch.quantization``. """ if not hasattr(module, "experts"): return False try: - from modelopt.torch.quantization.plugins.huggingface import _QuantMoELinear + from modelopt.torch.models.step3p5.modeling_ptq import _QuantMoELinear if isinstance(module, _QuantMoELinear): return True diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index 8bd9e6cc3f8..92f31011ea3 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/modelopt/torch/export/layerwise_export.py @@ -286,8 +286,8 @@ def export_layer( a fusing format can rediscover which modules share an input; omit them only when nothing fuses. """ - # Local, as in every other export module: the plugin imports transformers. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Local, as in every other export module: see pack_quantized_weights. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear assert self._bound, "export_layer() before bind()" assert not self._finalized, "export_layer() called after finalize()" diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 4c39674f585..f6116f4746a 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1027,10 +1027,9 @@ def _prepare_model_for_export(model, dtype, is_modelopt_qlora): def pack_quantized_weights(model, dtype, is_modelopt_qlora: bool = False) -> None: """Quantize every module's weight in place, then rebuild the fused MoE linears.""" - # Deferred: modelopt.torch.quantization.plugins.huggingface imports transformers at module - # scope, and transformers is an optional extra -- importing it here keeps - # ``import modelopt.torch.export`` working without it. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the + # per-model PTQ module at module scope here would risk an import cycle. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear _process_quantized_modules(model, dtype, is_modelopt_qlora) _reconstruct_fused_moe_linear(model) diff --git a/modelopt/torch/export/unified_export_hf_streaming.py b/modelopt/torch/export/unified_export_hf_streaming.py index 7f692194fa6..9188bb5780e 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -346,9 +346,9 @@ def _export_transformers_checkpoint_streaming( NotImplementedError: if the model's conversion mapping regroups tensors. RuntimeError: if decoder layers cannot be discovered for layer-wise materialization. """ - # Deferred: the huggingface plugin imports transformers at module scope, and transformers - # is an optional extra -- keep ``import modelopt.torch.export`` working without it. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the + # per-model PTQ module at module scope here would risk an import cycle. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear export_dir = Path(export_dir) # Materialization dispatch walks the module tree from the root; without these maps each @@ -545,9 +545,9 @@ def collect_export_tensors( All the gathers finish before this returns, so the caller can write or postprocess without stalling anyone. Returning a list rather than a generator is what guarantees that. """ - # Deferred: the huggingface plugin imports transformers at module scope, and transformers - # is an optional extra -- keep ``import modelopt.torch.export`` working without it. - from modelopt.torch.quantization.plugins.huggingface import _reconstruct_fused_moe_linear + # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the + # per-model PTQ module at module scope here would risk an import cycle. + from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear _assert_fsdp2_owns_every_mesh_dim(model) my_rank, world = _dist.rank(), _dist.size() diff --git a/modelopt/torch/models/README.md b/modelopt/torch/models/README.md index 287fd7f7da1..c0caf7f6e01 100644 --- a/modelopt/torch/models/README.md +++ b/modelopt/torch/models/README.md @@ -53,18 +53,19 @@ an intra-family detail, not a package API. ## PTQ modeling Quantized-module wrappers and registrations that only one model needs (e.g. Llama4's fused -BMM experts) live in `/modeling_ptq.py`; generic +BMM experts, Step's expert-indexed `MoELinear`) live in `/modeling_ptq.py`; generic ones (fused/sequential MoE auto-detection, attention, `FP8Linear`) stay in `modelopt/torch/quantization/plugins/huggingface.py`. The HF plugin imports every `modeling_ptq` from an explicit list at its end, so add the new model type there. Leave the model's `__init__.py` importing only `specs`: `modeling_ptq` pulls in quantization and transformers, which `import modelopt.torch.models` must not. -A `modeling_ptq.py` registers on import — static classes via `QuantModuleRegistry.register` -behind a `try`/`except ImportError`, remote-code classes via a callback added to +A `modeling_ptq.py` registers on import — static classes via `QuantModuleRegistry.register`, +imported from transformers at module scope; remote-code classes via a callback added to `CUSTOM_MODEL_PLUGINS` — and exports nothing (`__all__ = []`). Every package still has a `specs.py`, even when its spec only records `model_type` and where the modeling code comes -from. +from. A family revision that shares another type's PTQ code (`step3p7` uses `step3p5`'s) still +gets a `modeling_ptq.py` that imports the shared one, so loading it by model type works. ## Sections diff --git a/modelopt/torch/models/step3p5/modeling_ptq.py b/modelopt/torch/models/step3p5/modeling_ptq.py new file mode 100644 index 00000000000..9e3d40d49b8 --- /dev/null +++ b/modelopt/torch/models/step3p5/modeling_ptq.py @@ -0,0 +1,264 @@ +# 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-family PTQ modeling (HF model types ``step3p5``, ``step3p7``). + +The package is named for the first Step release; later revisions share its expert-indexed +``MoELinear`` layout and are matched by :func:`_is_step_family_model`. +""" + +import inspect +import re + +import torch +import torch.nn as nn +from torch import Tensor + +from modelopt.torch.models import hf_model_type +from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry +from modelopt.torch.quantization.plugins.custom import CUSTOM_MODEL_PLUGINS + +__all__: list[str] = [] + + +class _QuantMoELinear(QuantModule): + """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 + each expert gets its own weight_quantizer and input_quantizer, calibrated + only on tokens actually routed to that expert. + + On export, _reconstruct_fused_moe_linear() stacks the per-expert quantized + weights and scales back into the original 3D format. + + Note: we use expansion-then-reconstruction rather than the add_module() approach + because vLLM requires stacked 3D scaling factors; per-expert expanded keys are + not accepted by the downstream serving engine. + """ + + def _setup(self): + # Deferred: accelerate is a separate optional dependency, and the HF plugin (which + # imports this module) must stay importable without it. + from accelerate import init_empty_weights + + # Accelerate's CPU/disk offload (`device_map="auto"`, `--offload_folder`) leaves + # `weight` as a meta tensor and keeps the real value in the module's offload hook, + # keyed on the original `weight` name. Expanding that would copy meta storage into + # every expert and then delete the key the hook restores into, silently producing a + # checkpoint of zeros. Refuse instead of corrupting. + if self.weight.is_meta or getattr(getattr(self, "_hf_hook", None), "offload", False): + raise NotImplementedError( + f"{type(self).__name__}: expert-indexed MoELinear weights cannot be quantized " + "while offloaded by Accelerate (the weight is a meta tensor whose value lives " + "in the offload hook). Load the model without CPU/disk offload — more GPUs, or " + "a device_map that keeps the MoE layers resident — and re-run." + ) + + dtype, device = self.weight.dtype, self.weight.device + + with init_empty_weights(): + experts = nn.ModuleList( + [ + nn.Linear(self.in_features, self.out_features, bias=False) + for _ in range(self.num_experts) + ] + ) + + for i in range(self.num_experts): + experts[i].to_empty(device=device) + with torch.no_grad(): + experts[i].weight.data = self.weight[i].detach().to(dtype=dtype, device=device) + + delattr(self, "weight") + self.experts = experts + + def forward(self, x, expert_id): + # experts[expert_id] is a _QuantLinear after quantization wrapping, providing + # per-expert input_quantizer and weight_quantizer. + # + # MoELinear.forward always promotes to fp32 for the matmul regardless of storage + # dtype (`F.linear(x.float(), self.weight[expert_id].float())`), so leaving the + # expert's weight at its native storage dtype (e.g. bf16) and downcasting x to + # match before the matmul would compute in bf16 and change the model's output even + # with every quantizer disabled -- the bf16 rounding this class exists to quantize + # *past*, not to reintroduce as a side effect of conversion. + # + # A prior version of this fix instead expanded every expert's weight in fp32 + # permanently in `_setup`. That reproduces Step's fp32 compute but turns a per-call + # transient promotion into persistent model state: on Step-3.7's full routed-expert + # set (42 layers x 3 projections x 288 experts x 4096 x 1280), doubling from bf16 to + # fp32 adds roughly 354 GiB held throughout calibration, on top of device placement + # already sized for bf16 -- a model that loaded successfully can then OOM. It also + # left disabled/unquantized experts reconstructed at fp32 in the exported checkpoint. + # + # Instead, only the one expert actually being called is promoted, transiently, for + # the duration of this one call -- matching Step's own per-call `.float()` memory + # profile instead of Step-3.7's full expert set. `expert.weight` is read here + # outside any `quantize_weight()` context, so `_get_quantized_weight` passes it + # through unchanged and this is the real underlying nn.Parameter (the same pattern + # `_setup` above uses), not a value computed by the quantizer -- so reassigning its + # `.data` genuinely mutates the persisted storage, not a transient wrapper. + # + # This must keep calling `expert(x)` (`__call__`, not `.forward()`) rather than + # reimplementing the input/weight-quantize/output-quantize sequence inline: some + # calibration algorithms (e.g. `local_hessian_calibrate`) register a + # `forward_pre_hook` directly on the quantized Linear module, which only fires + # through standard `nn.Module.__call__` dispatch. + expert = self.experts[expert_id] + original_weight = expert.weight.data + with torch.no_grad(): + expert.weight.data = original_weight.float() + try: + out = expert(x.float()) + finally: + with torch.no_grad(): + expert.weight.data = original_weight + return out.float() + + +def _is_expert_indexed_moe_linear(module: nn.Module) -> bool: + """Whether ``module`` packs one projection's experts into an expert-indexed 3-D weight. + + 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 are picked + up without another hardcoded name, but the shape alone is not a sufficient contract: the + replacement forward indexes ``self.experts[expert_id]``, so it only works for callers that + pass a **scalar expert index**. Grouped-GEMM MoE layers share the exact same 3-D weight and + attribute set while passing a per-expert token-count *tensor* instead (e.g. Moondream3's + ``MoeFusedLinear.forward(input, m_sizes)``, which would raise ``TypeError: only integer + tensors of a single element can be converted to an index`` on the first calibration + forward). The second parameter must therefore be named ``expert_id``, which is the + scalar-index contract both Step revisions declare. + """ + 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 + # The wrapper rebuilds the weight as `num_experts` Linears of (out_features, in_features), + # so a 3-D weight laid out any other way would silently copy the wrong slices. + if tuple(weight.shape) != (module.num_experts, module.out_features, module.in_features): + return False + try: + params = list(inspect.signature(type(module).forward).parameters.values())[1:] + except (TypeError, ValueError): + return False + # The replacement forward is exactly `(x, expert_id)`, so anything the caller could pass + # beyond those two — a keyword-only `router_state`, *args, **kwargs — would raise once + # converted. Require the signature to match what the wrapper can honour. + return ( + len(params) == 2 + and all(p.kind is p.POSITIONAL_OR_KEYWORD and p.default is p.empty for p in params) + and params[1].name == "expert_id" + ) + + +_STEP_FAMILY_RE = re.compile(r"(?i)^step\d") + + +def _is_step_family_model(model: nn.Module) -> bool: + """Whether ``model`` is a Step-family root model (Step-3.5, Step-3.7, or a future revision). + + Matched against the ``step`` convention shared by ``model_type`` (``"step3p5"``, + ``"step3p7"``) and the remote-code class name (``Step3p5ForCausalLM``, + ``Step3p7ForConditionalGeneration``), not an exact revision, so a new Step release is + still picked up without another hardcoded name. This is deliberately narrower than the + shape/signature check in :func:`_is_expert_indexed_moe_linear` alone: that check accepts + any module with a matching 3-D weight and an ``(x, expert_id)`` forward, which is a + coincidence risk on its own -- an unrelated architecture happening to reuse the parameter + name ``expert_id`` with different semantics (a per-expert bias or post-scale, say) would + be claimed and have that behavior silently dropped by the replacement wrapper. Gating on + the model family keeps the shape check doing what it is actually good at: telling + Step revisions apart without a class-name allowlist, rather than distinguishing Step + from arbitrary third-party MoE code. + """ + model_type = hf_model_type(model) or "" + return bool(_STEP_FAMILY_RE.match(model_type) or _STEP_FAMILY_RE.match(type(model).__name__)) + + +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``. + """ + if not _is_step_family_model(model): + return + 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 :class:`_QuantMoELinear` per-expert weights back to the 3-D MoELinear format. + + 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) + + 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]). + + 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(): + if not isinstance(module, _QuantMoELinear): + continue + + n = module.num_experts + experts = module.experts + + # Reconstruct 3D weight: [num_experts, out_features, in_features] + module.weight = nn.Parameter( + torch.stack([experts[i].weight.data for i in range(n)]), + requires_grad=False, + ) + + # Stack per-expert scales back under the original attribute names. + # Check all experts: some may lack input_scale if they were never routed + # during calibration, so only stack when every expert has the attribute. + for attr in ("weight_scale", "weight_scale_2", "input_scale"): + if all(hasattr(experts[i], attr) for i in range(n)): + module.register_buffer( + attr, + torch.stack([getattr(experts[i], attr) for i in range(n)]), + ) + + # Remove expanded experts — the reconstructed 3D tensors replace them + del module.experts + + +CUSTOM_MODEL_PLUGINS.add(register_moe_linear_on_the_fly) diff --git a/modelopt/torch/models/step3p7/modeling_ptq.py b/modelopt/torch/models/step3p7/modeling_ptq.py new file mode 100644 index 00000000000..e43c5237278 --- /dev/null +++ b/modelopt/torch/models/step3p7/modeling_ptq.py @@ -0,0 +1,26 @@ +# 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 modeling (HF model type ``step3p7``). + +Step-3.7 shares Step-3.5's expert-indexed ``MoELinear`` layout, so its PTQ support is the +Step-family code in ``step3p5/modeling_ptq.py``, which matches every Step revision. This +module imports it, so that loading the PTQ modeling for model type ``step3p7`` by name +registers that support too. +""" + +from ..step3p5 import modeling_ptq as _step3p5_modeling_ptq # noqa: F401 + +__all__: list[str] = [] diff --git a/modelopt/torch/models/step3p7/specs.py b/modelopt/torch/models/step3p7/specs.py index f38c0339989..74db68aed4b 100644 --- a/modelopt/torch/models/step3p7/specs.py +++ b/modelopt/torch/models/step3p7/specs.py @@ -15,8 +15,8 @@ """Step-3.7 specs (HF model type ``step3p7``). -Shares Step-3.5's architecture; its PTQ modeling (in the HF quantization plugin) -matches every Step revision. +Shares Step-3.5's architecture; its PTQ modeling lives in ``step3p5/modeling_ptq.py``, +which matches every Step revision. """ from ..specs import ModelSpec, register diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 05e589835ac..4110a95e5bd 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -18,7 +18,6 @@ import importlib import inspect import logging -import re import warnings from contextlib import contextmanager from functools import partial @@ -38,7 +37,6 @@ validate_triton_attention_envelope, ) from modelopt.torch.kernels.quantization.gemm import IS_AVAILABLE as IS_TRITON_AVAILABLE -from modelopt.torch.models import hf_model_type from modelopt.torch.opt.dynamic import DynamicModule from modelopt.torch.utils.distributed import ParallelState @@ -1484,6 +1482,8 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): "gpt_oss", "llama4", "nemotron_h", + "step3p5", + "step3p7", ): importlib.import_module(f"modelopt.torch.models.{_model_type}.modeling_ptq") @@ -1492,235 +1492,8 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model): ) -class _QuantMoELinear(QuantModule): - """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 - each expert gets its own weight_quantizer and input_quantizer, calibrated - only on tokens actually routed to that expert. - - On export, _reconstruct_fused_moe_linear() stacks the per-expert quantized - weights and scales back into the original 3D format. - - Note: we use expansion-then-reconstruction rather than the add_module() approach - because vLLM requires stacked 3D scaling factors; per-expert expanded keys are - not accepted by the downstream serving engine. - """ - - def _setup(self): - from accelerate import init_empty_weights - - # Accelerate's CPU/disk offload (`device_map="auto"`, `--offload_folder`) leaves - # `weight` as a meta tensor and keeps the real value in the module's offload hook, - # keyed on the original `weight` name. Expanding that would copy meta storage into - # every expert and then delete the key the hook restores into, silently producing a - # checkpoint of zeros. Refuse instead of corrupting. - if self.weight.is_meta or getattr(getattr(self, "_hf_hook", None), "offload", False): - raise NotImplementedError( - f"{type(self).__name__}: expert-indexed MoELinear weights cannot be quantized " - "while offloaded by Accelerate (the weight is a meta tensor whose value lives " - "in the offload hook). Load the model without CPU/disk offload — more GPUs, or " - "a device_map that keeps the MoE layers resident — and re-run." - ) - - dtype, device = self.weight.dtype, self.weight.device - - with init_empty_weights(): - experts = nn.ModuleList( - [ - nn.Linear(self.in_features, self.out_features, bias=False) - for _ in range(self.num_experts) - ] - ) - - for i in range(self.num_experts): - experts[i].to_empty(device=device) - with torch.no_grad(): - experts[i].weight.data = self.weight[i].detach().to(dtype=dtype, device=device) - - delattr(self, "weight") - self.experts = experts - - def forward(self, x, expert_id): - # experts[expert_id] is a _QuantLinear after quantization wrapping, providing - # per-expert input_quantizer and weight_quantizer. - # - # MoELinear.forward always promotes to fp32 for the matmul regardless of storage - # dtype (`F.linear(x.float(), self.weight[expert_id].float())`), so leaving the - # expert's weight at its native storage dtype (e.g. bf16) and downcasting x to - # match before the matmul would compute in bf16 and change the model's output even - # with every quantizer disabled -- the bf16 rounding this class exists to quantize - # *past*, not to reintroduce as a side effect of conversion. - # - # A prior version of this fix instead expanded every expert's weight in fp32 - # permanently in `_setup`. That reproduces Step's fp32 compute but turns a per-call - # transient promotion into persistent model state: on Step-3.7's full routed-expert - # set (42 layers x 3 projections x 288 experts x 4096 x 1280), doubling from bf16 to - # fp32 adds roughly 354 GiB held throughout calibration, on top of device placement - # already sized for bf16 -- a model that loaded successfully can then OOM. It also - # left disabled/unquantized experts reconstructed at fp32 in the exported checkpoint. - # - # Instead, only the one expert actually being called is promoted, transiently, for - # the duration of this one call -- matching Step's own per-call `.float()` memory - # profile instead of Step-3.7's full expert set. `expert.weight` is read here - # outside any `quantize_weight()` context, so `_get_quantized_weight` passes it - # through unchanged and this is the real underlying nn.Parameter (the same pattern - # `_setup` above uses), not a value computed by the quantizer -- so reassigning its - # `.data` genuinely mutates the persisted storage, not a transient wrapper. - # - # This must keep calling `expert(x)` (`__call__`, not `.forward()`) rather than - # reimplementing the input/weight-quantize/output-quantize sequence inline: some - # calibration algorithms (e.g. `local_hessian_calibrate`) register a - # `forward_pre_hook` directly on the quantized Linear module, which only fires - # through standard `nn.Module.__call__` dispatch. - expert = self.experts[expert_id] - original_weight = expert.weight.data - with torch.no_grad(): - expert.weight.data = original_weight.float() - try: - out = expert(x.float()) - finally: - with torch.no_grad(): - expert.weight.data = original_weight - return out.float() - - -def _is_expert_indexed_moe_linear(module: nn.Module) -> bool: - """Whether ``module`` packs one projection's experts into an expert-indexed 3-D weight. - - 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 are picked - up without another hardcoded name, but the shape alone is not a sufficient contract: the - replacement forward indexes ``self.experts[expert_id]``, so it only works for callers that - pass a **scalar expert index**. Grouped-GEMM MoE layers share the exact same 3-D weight and - attribute set while passing a per-expert token-count *tensor* instead (e.g. Moondream3's - ``MoeFusedLinear.forward(input, m_sizes)``, which would raise ``TypeError: only integer - tensors of a single element can be converted to an index`` on the first calibration - forward). The second parameter must therefore be named ``expert_id``, which is the - scalar-index contract both Step revisions declare. - """ - 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 - # The wrapper rebuilds the weight as `num_experts` Linears of (out_features, in_features), - # so a 3-D weight laid out any other way would silently copy the wrong slices. - if tuple(weight.shape) != (module.num_experts, module.out_features, module.in_features): - return False - try: - params = list(inspect.signature(type(module).forward).parameters.values())[1:] - except (TypeError, ValueError): - return False - # The replacement forward is exactly `(x, expert_id)`, so anything the caller could pass - # beyond those two — a keyword-only `router_state`, *args, **kwargs — would raise once - # converted. Require the signature to match what the wrapper can honour. - return ( - len(params) == 2 - and all(p.kind is p.POSITIONAL_OR_KEYWORD and p.default is p.empty for p in params) - and params[1].name == "expert_id" - ) - - -_STEP_FAMILY_RE = re.compile(r"(?i)^step\d") - - -def _is_step_family_model(model: nn.Module) -> bool: - """Whether ``model`` is a Step-family root model (Step-3.5, Step-3.7, or a future revision). - - Matched against the ``step`` convention shared by ``model_type`` (``"step3p5"``, - ``"step3p7"``) and the remote-code class name (``Step3p5ForCausalLM``, - ``Step3p7ForConditionalGeneration``), not an exact revision, so a new Step release is - still picked up without another hardcoded name. This is deliberately narrower than the - shape/signature check in :func:`_is_expert_indexed_moe_linear` alone: that check accepts - any module with a matching 3-D weight and an ``(x, expert_id)`` forward, which is a - coincidence risk on its own -- an unrelated architecture happening to reuse the parameter - name ``expert_id`` with different semantics (a per-expert bias or post-scale, say) would - be claimed and have that behavior silently dropped by the replacement wrapper. Gating on - the model family keeps the shape check doing what it is actually good at: telling - Step revisions apart without a class-name allowlist, rather than distinguishing Step - from arbitrary third-party MoE code. - """ - model_type = hf_model_type(model) or "" - return bool(_STEP_FAMILY_RE.match(model_type) or _STEP_FAMILY_RE.match(type(model).__name__)) - - -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``. - """ - if not _is_step_family_model(model): - return - 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 :class:`_QuantMoELinear` per-expert weights back to the 3-D MoELinear format. - - 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) - - 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]). - - 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(): - if not isinstance(module, _QuantMoELinear): - continue - - n = module.num_experts - experts = module.experts - - # Reconstruct 3D weight: [num_experts, out_features, in_features] - module.weight = nn.Parameter( - torch.stack([experts[i].weight.data for i in range(n)]), - requires_grad=False, - ) - - # Stack per-expert scales back under the original attribute names. - # Check all experts: some may lack input_scale if they were never routed - # during calibration, so only stack when every expert has the attribute. - for attr in ("weight_scale", "weight_scale_2", "input_scale"): - if all(hasattr(experts[i], attr) for i in range(n)): - module.register_buffer( - attr, - torch.stack([getattr(experts[i], attr) for i in range(n)]), - ) - - # Remove expanded experts — the reconstructed 3D tensors replace them - del module.experts - - CUSTOM_MODEL_PLUGINS.update( [ - 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/plugins/modelopt/skills/ptq/references/checkpoint-validation.md b/plugins/modelopt/skills/ptq/references/checkpoint-validation.md index 3d755933134..c836681147f 100644 --- a/plugins/modelopt/skills/ptq/references/checkpoint-validation.md +++ b/plugins/modelopt/skills/ptq/references/checkpoint-validation.md @@ -201,6 +201,6 @@ Layers silently skipped because the quantization config patterns don't match the ## What to do when warnings appear -- **Layers should have been quantized** (e.g., MoE experts with `nvfp4_mlp_only`): the quantization config patterns missed them. Fix by adding the missing pattern to the config and re-running PTQ. Check if ModelOpt already has a plugin for the model in `modelopt/torch/quantization/plugins/huggingface.py`. +- **Layers should have been quantized** (e.g., MoE experts with `nvfp4_mlp_only`): the quantization config patterns missed them. Fix by adding the missing pattern to the config and re-running PTQ. Check if ModelOpt already has a plugin for the model in `modelopt/torch/models//modeling_ptq.py`. - **Layers are intentionally unquantized** (e.g., attention layers with `nvfp4_mlp_only`): they should be in the `exclude_modules` list but the export didn't add them. Add them manually to both `hf_quant_config.json` and `config.json` `quantization_config.ignore` in the checkpoint to prevent deployment failures. diff --git a/plugins/modelopt/skills/ptq/references/unsupported-models.md b/plugins/modelopt/skills/ptq/references/unsupported-models.md index 4ed7959c71b..09514ebf4be 100644 --- a/plugins/modelopt/skills/ptq/references/unsupported-models.md +++ b/plugins/modelopt/skills/ptq/references/unsupported-models.md @@ -66,7 +66,7 @@ Manual dequantization is only needed for **non-standard parameter names** (e.g., Read the model source to identify how weights are stored. **If all linear layers are plain `nn.Linear`, no custom code is needed** — ModelOpt quantizes them automatically. -**For HuggingFace models**, check `modelopt/torch/quantization/plugins/huggingface.py` first — it already registers patches for common non-standard modules (`Llama4TextExperts`, `FP8Linear`, `FalconLinear`, `Conv1D`, `Qwen3_5MoeExperts`, etc.). If your model's non-standard class is already registered there, no extra code is needed. +**For HuggingFace models**, check `modelopt/torch/models//modeling_ptq.py` (model-specific patches, e.g. `Llama4TextExperts`, `FalconLinear`) and `modelopt/torch/quantization/plugins/huggingface.py` (generic ones: `FP8Linear`, `Conv1D`, auto-detected MoE such as `Qwen3_5MoeExperts`) first. If your model's non-standard class is already registered, no extra code is needed. Custom patches are required when: @@ -94,7 +94,7 @@ Compare against the `enable`/`disable` patterns in the config. Add custom overri After Steps A-D: - **No patches needed** (all standard modules) → run `hf_ptq.py` with a smoke test (`--calib_size 4`). If it succeeds, proceed with full calibration. If it fails, read the error and revisit Steps C/D. -- **Patches needed** → patch ModelOpt directly using the patterns below (add `QuantModule` in `modelopt/torch/quantization/plugins/huggingface.py`, update `modelopt/torch/export/` if needed), then run `hf_ptq.py` with a smoke test. This is preferred over writing a standalone script because it reuses all existing `hf_ptq.py` logic. Debug failures iteratively — quantization errors often reveal additional modules that need patching. +- **Patches needed** → patch ModelOpt directly using the patterns below (add `QuantModule` in `modelopt/torch/models//modeling_ptq.py`, update `modelopt/torch/export/` if needed), then run `hf_ptq.py` with a smoke test. This is preferred over writing a standalone script because it reuses all existing `hf_ptq.py` logic. Debug failures iteratively — quantization errors often reveal additional modules that need patching. --- @@ -138,7 +138,7 @@ class QuantCustomModule(OriginalModule): - **Fused experts** (`gate_up_proj` + `down_proj` 3D tensors, all native transformers MoEs) → auto-detected by `register_fused_experts_on_the_fly`, handled by `_QuantFusedExperts`. Covers Mixtral, Qwen, DeepSeek, Jamba, OlMoE, etc. - **Sequential per-expert `nn.Linear`** with `gate` + `experts` (e.g. remote-code models) → auto-detected by `register_sparse_moe_on_the_fly`. -**Custom MoE** (non-standard layout not matching auto-detection) requires patching. Find the closest pattern in the plugin (`modelopt/torch/quantization/plugins/huggingface.py`): +**Custom MoE** (non-standard layout not matching auto-detection) requires patching. Find the closest pattern in `modelopt/torch/models/*/modeling_ptq.py`: | MoE design | Strategy | Plugin example | | --- | --- | --- | @@ -163,7 +163,7 @@ for name, module in model.named_modules(): ## Pattern 3: Registering with ModelOpt -**When patching the plugin directly** (preferred): Use `QuantModuleRegistry.register` in `modelopt/torch/quantization/plugins/huggingface.py`, following existing examples: +**When patching ModelOpt directly** (preferred): Use `QuantModuleRegistry.register` in `modelopt/torch/models//modeling_ptq.py` and add `` to the `modeling_ptq` import loop at the end of `modelopt/torch/quantization/plugins/huggingface.py`, following existing examples: ```python from modelopt.torch.quantization.nn import QuantModuleRegistry @@ -179,6 +179,8 @@ def register_my_model_on_the_fly(model): if QuantModuleRegistry.get(mod_type) is None: QuantModuleRegistry.register({mod_type: f"hf.{mod_type.__name__}"})(QuantCustomModule) break + +CUSTOM_MODEL_PLUGINS.add(register_my_model_on_the_fly) # from ...quantization.plugins.custom ``` **When writing a standalone script** (fallback): Use `mtq.register()`: diff --git a/tests/unit/torch/models/test_modeling_ptq_registration.py b/tests/unit/torch/models/test_modeling_ptq_registration.py index fbaa761d38d..6782b64c344 100644 --- a/tests/unit/torch/models/test_modeling_ptq_registration.py +++ b/tests/unit/torch/models/test_modeling_ptq_registration.py @@ -35,6 +35,8 @@ "modelopt.torch.models.gpt_oss.modeling_ptq", "modelopt.torch.models.llama4.modeling_ptq", "modelopt.torch.models.nemotron_h.modeling_ptq", + "modelopt.torch.models.step3p5.modeling_ptq", + "modelopt.torch.models.step3p7.modeling_ptq", "modelopt.torch.quantization", "modelopt.torch.export", ] @@ -48,6 +50,7 @@ from modelopt.torch.models.falcon.modeling_ptq import register_falcon_linears_on_the_fly from modelopt.torch.models.nemotron_h.modeling_ptq import is_nemotron_h_model +from modelopt.torch.models.step3p5.modeling_ptq import register_moe_linear_on_the_fly from modelopt.torch.quantization.nn import QuantModuleRegistry from modelopt.torch.quantization.plugins.custom import CUSTOM_MODEL_PLUGINS from modelopt.torch.quantization.plugins.huggingface import is_homogeneous_hf_model @@ -74,6 +77,7 @@ assert GptOssExperts in QuantModuleRegistry, "GptOssExperts not registered" assert register_falcon_linears_on_the_fly in CUSTOM_MODEL_PLUGINS, "Falcon callback missing" +assert register_moe_linear_on_the_fly in CUSTOM_MODEL_PLUGINS, "Step callback missing" # The first matching decoder discoverer wins, so Nemotron-H's must precede the generic one. predicates = [is_supported for is_supported, _ in LayerActivationCollector._decoder_layer_support] diff --git a/tests/unit/torch/quantization/plugins/test_moe_linear.py b/tests/unit/torch/models/test_step3p5_modeling_ptq.py similarity index 99% rename from tests/unit/torch/quantization/plugins/test_moe_linear.py rename to tests/unit/torch/models/test_step3p5_modeling_ptq.py index cceaacd44d5..66b52325c24 100644 --- a/tests/unit/torch/quantization/plugins/test_moe_linear.py +++ b/tests/unit/torch/models/test_step3p5_modeling_ptq.py @@ -27,13 +27,13 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.hf_export_handlers import _export_moe_linear from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry -from modelopt.torch.quantization.nn import QuantModuleRegistry -from modelopt.torch.quantization.plugins.huggingface import ( +from modelopt.torch.models.step3p5.modeling_ptq import ( _is_expert_indexed_moe_linear, _QuantMoELinear, _reconstruct_fused_moe_linear, register_moe_linear_on_the_fly, ) +from modelopt.torch.quantization.nn import QuantModuleRegistry NUM_EXPERTS = 4 HIDDEN_SIZE = 32