feat(quantization): PTQ support for Step-3.7 MoE checkpoints

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 <dir>`` 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``.

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,42 +1921,77 @@ 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)
 
     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,

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

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

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