[OMNIML-3817] Move model-specific PTQ modeling into modelopt/torch/models [2/4]

modelopt/torch/models/README.md

@@ -50,6 +50,22 @@ star-import. A value shared between two model types is private and imported expl
 (`gemma4_text` reuses `gemma4`'s `_GEMMA4_MOE_SPEC`); sub-model types of one family are
 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 `<model_type>/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
+`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.
+
 ## Sections
 
 A `ModelSpec` holds one attribute per section, each `None` unless the model fills it:

modelopt/torch/models/__init__.py

@@ -26,6 +26,9 @@
 
 The per-model file is named for what it holds, not for who reads it: a model's spec is
 general model data, and export is only its first consumer.
+
+A package may also hold ``modeling_ptq.py``, the model-specific PTQ wrappers. It is not imported
+here; the HF quantization plugin imports it (see ``README.md``).
 """
 
 from .moe import is_moe

modelopt/torch/models/falcon/modeling_ptq.py

@@ -0,0 +1,44 @@
+# 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.
+
+"""Falcon PTQ modeling (HF model type ``falcon``)."""
+
+from transformers.models.falcon.modeling_falcon import FalconLinear
+
+from modelopt.torch.quantization.nn import QuantModuleRegistry
+from modelopt.torch.quantization.nn.modules.quant_linear import _QuantLinear
+from modelopt.torch.quantization.plugins.custom import CUSTOM_MODEL_PLUGINS
+
+__all__: list[str] = []
+
+
+if FalconLinear not in QuantModuleRegistry:
+    QuantModuleRegistry.register({FalconLinear: "hf.FalconLinear"})(_QuantLinear)
+
+
+def register_falcon_linears_on_the_fly(model):
+    """Register Falcon linear modules as a QUANT_MODULE.
+
+    Certain falcon models (for example, falcon 40b) use remote code, which are loaded dynamically, to build their model.
+    Therefore, we need to register the linear on the fly before quantization.
+    """
+    if type(model).__name__ in ["RWForCausalLM", "FalconForCausalLM"]:
+        linear_type = type(model.transformer.h[0].self_attention.dense)
+        # Create a QuantFalconLinear class on the fly
+        if QuantModuleRegistry.get(linear_type) is None:
+            QuantModuleRegistry.register({linear_type: linear_type.__name__})(_QuantLinear)
+
+
+CUSTOM_MODEL_PLUGINS.add(register_falcon_linears_on_the_fly)

modelopt/torch/models/falcon/specs.py

@@ -20,5 +20,5 @@
 __all__: list[str] = []
 
 # Dense, and nothing model-specific for export. Older remote-code Falcon checkpoints
-# (RWForCausalLM) report other model types and are handled by the HF quantization plugin alone.
-register(ModelSpec(model_type="falcon", min_transformers_version="4.57"))
+# (RWForCausalLM) report other model types and are handled by modeling_ptq.py alone.
+register(ModelSpec(model_type="falcon", min_transformers_version="5.5"))

modelopt/torch/models/llama4/modeling_ptq.py

@@ -0,0 +1,55 @@
+# 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.
+
+"""Llama4 PTQ modeling (HF model type ``llama4``)."""
+
+import torch
+from transformers.models.llama4.modeling_llama4 import Llama4TextExperts
+
+from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry, TensorQuantizer
+from modelopt.torch.quantization.plugins.huggingface import (
+    _transposed_quantize,
+    _TransposedExpertsCalibMixin,
+)
+
+__all__: list[str] = []
+
+
+class _QuantLlama4TextExperts(_TransposedExpertsCalibMixin, QuantModule):
+    def _setup(self):
+        self.gate_up_proj_input_quantizer = TensorQuantizer()
+        self.gate_up_proj_weight_quantizer = TensorQuantizer()
+        self.down_proj_input_quantizer = TensorQuantizer()
+        self.down_proj_weight_quantizer = TensorQuantizer()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = hidden_states.view(self.num_experts, -1, self.hidden_size)
+        gate_up = torch.bmm(
+            self.gate_up_proj_input_quantizer(hidden_states),
+            _transposed_quantize(self.gate_up_proj, self.gate_up_proj_weight_quantizer),
+        )
+        gate, up = gate_up.chunk(2, dim=-1)  # not supported for DTensors
+        next_states = torch.bmm(
+            self.down_proj_input_quantizer(up * self.act_fn(gate)),
+            _transposed_quantize(self.down_proj, self.down_proj_weight_quantizer),
+        )
+        next_states = next_states.view(-1, self.hidden_size)
+        return next_states
+
+
+if Llama4TextExperts not in QuantModuleRegistry:
+    QuantModuleRegistry.register({Llama4TextExperts: "hf.Llama4TextExperts"})(
+        _QuantLlama4TextExperts
+    )

modelopt/torch/models/llama4/specs.py

@@ -26,7 +26,7 @@
 register(
     ModelSpec(
         model_type="llama4",
-        min_transformers_version="4.57",
+        min_transformers_version="5.5",
         moe_spec=MoESpec(
             block_names=("Llama4TextMoe",),
             expert_linear_names=("gate_up_proj", "down_proj"),

modelopt/torch/models/nemotron_h/modeling_ptq.py

@@ -0,0 +1,52 @@
+# 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.
+
+"""Nemotron-H PTQ modeling (HF model type ``nemotron_h``)."""
+
+import torch.nn as nn
+
+from modelopt.torch.quantization.plugins.huggingface import _is_supported_hf_model
+from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
+
+__all__: list[str] = []
+
+
+def is_nemotron_h_model(model: nn.Module) -> bool:
+    return get_nemotron_h_decoder_layers(model) is not None
+
+
+def get_nemotron_h_decoder_layers(model: nn.Module) -> nn.ModuleList | None:
+    if not _is_supported_hf_model(model):
+        return None
+
+    # Custom remote-code checkpoint uses model.backbone.layers;
+    # native transformers NemotronHModel uses model.model.layers.
+    for container_attr in ("backbone", "model"):
+        container = getattr(model, container_attr, None)
+        if container is not None and hasattr(container, "layers"):
+            layers = container.layers
+            if layers and hasattr(layers[0], "block_type"):
+                return layers
+
+    return None
+
+
+# Order matters: more specific predicates must be registered first because
+# the first matching entry wins.  Nemotron-H must precede the generic
+# homogeneous HF discoverer (which explicitly rejects Nemotron-H); the HF plugin
+# imports this module before registering that discoverer.
+LayerActivationCollector.register_decoder_layer_support(
+    is_nemotron_h_model, get_nemotron_h_decoder_layers
+)

modelopt/torch/models/qwen3_vl_moe/specs.py

@@ -20,13 +20,12 @@
 __all__: list[str] = []
 
 # Qwen3VLMoeTextExperts is fused on every supported transformers: 3-D gate_up_proj and
-# down_proj parameters (transposed before 5.12, standard layout from 5.12 on). Before 5.12
-# the PTQ wrapper unrolls them into gate_proj/up_proj/down_proj ModuleLists (see
-# quantization/plugins/huggingface.py); that rewrite is not iterable per expert, so no export path groups it.
+# down_proj parameters (transposed before 5.12, standard layout from 5.12 on), so there is no
+# (gate, up) pair left to fuse and no per-expert grouping for export.
 register(
     ModelSpec(
         model_type="qwen3_vl_moe",
-        min_transformers_version="4.57",
+        min_transformers_version="5.5",
         moe_spec=MoESpec(
             block_names=("Qwen3VLMoeTextSparseMoeBlock",),
             expert_linear_names=("gate_up_proj", "down_proj"),

modelopt/torch/quantization/plugins/huggingface.py

@@ -15,6 +15,7 @@
 
 """Support quantization for huggingface layers."""
 
+import importlib
 import inspect
 import logging
 import re
@@ -29,9 +30,7 @@
 from torch import Tensor
 from torch.nn.functional import linear
 from transformers.integrations.finegrained_fp8 import FP8Linear
-from transformers.models.falcon.modeling_falcon import FalconLinear
 from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts
-from transformers.models.llama4.modeling_llama4 import Llama4TextExperts
 from transformers.models.t5.modeling_t5 import T5Attention
 
 from modelopt.torch.kernels.common.attention import IS_AVAILABLE as TRITON_FA_AVAILABLE
@@ -769,28 +768,6 @@ def layer_sync_moe_local_experts_amax(self, sync_weight_amax=False):
         sync_moe_expert_amax(self.experts, sync_weight_amax=sync_weight_amax)
 
 
-class _QuantLlama4TextExperts(_TransposedExpertsCalibMixin, QuantModule):
-    def _setup(self):
-        self.gate_up_proj_input_quantizer = TensorQuantizer()
-        self.gate_up_proj_weight_quantizer = TensorQuantizer()
-        self.down_proj_input_quantizer = TensorQuantizer()
-        self.down_proj_weight_quantizer = TensorQuantizer()
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = hidden_states.view(self.num_experts, -1, self.hidden_size)
-        gate_up = torch.bmm(
-            self.gate_up_proj_input_quantizer(hidden_states),
-            _transposed_quantize(self.gate_up_proj, self.gate_up_proj_weight_quantizer),
-        )
-        gate, up = gate_up.chunk(2, dim=-1)  # not supported for DTensors
-        next_states = torch.bmm(
-            self.down_proj_input_quantizer(up * self.act_fn(gate)),
-            _transposed_quantize(self.down_proj, self.down_proj_weight_quantizer),
-        )
-        next_states = next_states.view(-1, self.hidden_size)
-        return next_states
-
-
 def _get_fused_expert_intermediate_dim(module):
     """Resolve the intermediate (expert) dimension from a fused expert module.
 
@@ -1215,14 +1192,6 @@ def unpack_weight(self):
             del self.weight_scale_inv
 
 
-if Llama4TextExperts not in QuantModuleRegistry:
-    QuantModuleRegistry.register({Llama4TextExperts: "hf.Llama4TextExperts"})(
-        _QuantLlama4TextExperts
-    )
-
-if FalconLinear not in QuantModuleRegistry:
-    QuantModuleRegistry.register({FalconLinear: "hf.FalconLinear"})(_QuantLinear)
-
 try:
     from compressed_tensors.linear.compressed_linear import CompressedLinear
 
@@ -1327,19 +1296,6 @@ def forward(
     QuantModuleRegistry.register({GptOssExperts: "hf.GptOssExperts"})(_QuantGptOssExperts)
 
 
-def register_falcon_linears_on_the_fly(model):
-    """Register Falcon linear modules as a QUANT_MODULE.
-
-    Certain falcon models (for example, falcon 40b) use remote code, which are loaded dynamically, to build their model.
-    Therefore, we need to register the linear on the fly before quantization.
-    """
-    if type(model).__name__ in ["RWForCausalLM", "FalconForCausalLM"]:
-        linear_type = type(model.transformer.h[0].self_attention.dense)
-        # Create a QuantFalconLinear class on the fly
-        if QuantModuleRegistry.get(linear_type) is None:
-            QuantModuleRegistry.register({linear_type: linear_type.__name__})(_QuantLinear)
-
-
 def _has_num_experts(obj):
     # n_routed_experts: NemotronH-style MoE
     return hasattr(obj, "num_experts") or hasattr(obj, "n_routed_experts")
@@ -1529,27 +1485,9 @@ def _is_supported_hf_model(model):
     return isinstance(model, tuple(supported_models))
 
 
-def is_nemotron_h_model(model: nn.Module) -> bool:
-    return get_nemotron_h_decoder_layers(model) is not None
-
-
-def get_nemotron_h_decoder_layers(model: nn.Module) -> nn.ModuleList | None:
-    if not _is_supported_hf_model(model):
-        return None
-
-    # Custom remote-code checkpoint uses model.backbone.layers;
-    # native transformers NemotronHModel uses model.model.layers.
-    for container_attr in ("backbone", "model"):
-        container = getattr(model, container_attr, None)
-        if container is not None and hasattr(container, "layers"):
-            layers = container.layers
-            if layers and hasattr(layers[0], "block_type"):
-                return layers
-
-    return None
-
-
 def is_homogeneous_hf_model(model: nn.Module) -> bool:
+    from modelopt.torch.models.nemotron_h.modeling_ptq import is_nemotron_h_model
+
     if is_nemotron_h_model(model):
         return False
     decoder_layers = get_homogeneous_hf_decoder_layers(model)
@@ -1627,12 +1565,17 @@ def _is_param_grad_enabled_for_auto_quantize(pname, model):
     _is_param_grad_enabled_for_auto_quantize,
 )
 
-# Order matters: more specific predicates must be registered first because
-# the first matching entry wins.  Nemotron-H must precede the generic
-# homogeneous HF discoverer (which explicitly rejects Nemotron-H).
-LayerActivationCollector.register_decoder_layer_support(
-    is_nemotron_h_model, get_nemotron_h_decoder_layers
-)
+# Model-specific PTQ support lives with its model in
+# ``modelopt/torch/models/<model_type>/modeling_ptq.py`` and registers itself on import. It is
+# imported here: after the generic wrappers it builds on are defined, and before the
+# homogeneous decoder discoverer below, since the first matching discoverer wins and
+# Nemotron-H's is the more specific one.
+for _model_type in (
+    "falcon",
+    "llama4",
+    "nemotron_h",
+):
+    importlib.import_module(f"modelopt.torch.models.{_model_type}.modeling_ptq")
 
 LayerActivationCollector.register_decoder_layer_support(
     is_homogeneous_hf_model, get_homogeneous_hf_decoder_layers
@@ -1867,7 +1810,6 @@ def _reconstruct_fused_moe_linear(model: nn.Module) -> None:
 
 CUSTOM_MODEL_PLUGINS.update(
     [
-        register_falcon_linears_on_the_fly,
         register_moe_linear_on_the_fly,
         register_fused_experts_on_the_fly,
         force_eager_experts_impl_on_the_fly,