[Reference, not for merge] Move model-specific PTQ modeling into modelopt/torch/models

CHANGELOG.rst

@@ -72,6 +72,7 @@ Changelog
 - Rename the architecture-specific recipe tier from ``modelopt_recipes/huggingface/`` to ``modelopt_recipes/model_type/`` to clarify that it holds recipes shared across every checkpoint of a Hugging Face ``model_type``. Saved ``--recipe huggingface/<model_type>/...`` paths still resolve via a backward-compatibility alias but now emit a ``FutureWarning``, so update them to ``model_type/<model_type>/...`` as the ``huggingface/`` prefix is deprecated.
 - The single-format quantization CLI flags are deprecated in favour of ``--recipe`` and will be removed in a future release; passing one now emits a ``FutureWarning``. ``examples/hf_ptq``: ``--qformat`` and ``--kv_cache_qformat``. ``examples/megatron_bridge/quantize.py``: ``--quant_cfg``, ``--kv_cache_quant`` and ``--weight_only``. ``examples/torch_onnx/torch_quant_to_onnx.py``: ``--qformat``. A recipe carries the quantization config, the calibration algorithm and the KV-cache setting in one file, so they cannot drift apart the way separate flags can -- and ``--recipe`` already took precedence over all six, silently on ``hf_ptq`` and with a warning on ``megatron_bridge`` -- with one gap the recipe closes rather than inherits: a weight AutoQuantize recipe that omits ``kv_cache`` still falls back to ``--kv_cache_qformat``, so set ``kv_cache`` in the recipe when migrating. Use a recipe from ``modelopt_recipes/general/ptq/``, an architecture-specific one under ``modelopt_recipes/model_type/<model_type>/``, or a checkpoint-specific one under ``modelopt_recipes/models/``. The warning fires only when a flag is passed explicitly: ``--qformat`` defaults to ``fp8`` and ``--kv_cache_qformat`` to ``fp8_cast``, so warning on the defaults would fire on every run, including runs that correctly use ``--recipe``. ``examples/speculative_decoding/scripts/quantize_drafter.py`` keeps ``--qformat`` undeprecated: it has no ``--recipe`` alternative yet.
 - The TensorRT-LLM checkpoint export format is deprecated and will be removed in 0.49.0: ``export_tensorrt_llm_checkpoint`` and ``torch_to_tensorrt_llm_checkpoint`` now emit a ``DeprecationWarning`` on use. Use ``export_hf_checkpoint``, which exports a unified Hugging Face checkpoint deployable on TensorRT-LLM, vLLM and SGLang. Its implementation moved to ``modelopt.torch.export.trtllm``, so import those two functions from there and the ``ModelConfig`` dataclasses from ``modelopt.torch.export.trtllm.model_config``; both functions remain importable from ``modelopt.torch.export`` for this release only.
+- ``modelopt.torch.export.get_model_type`` and ``modelopt.torch.export.model_utils.MODEL_NAME_TO_TYPE`` are deprecated and will be removed in 0.49.0: they return TensorRT-LLM model names and now emit a ``DeprecationWarning``. Use the Hugging Face model type, ``model.config.model_type``, instead. ``export_tensorrt_llm_checkpoint`` now detects its ``decoder_type`` when it is omitted. ``examples/hf_ptq`` now uses Hugging Face model types throughout.
 - The ``evaluation`` agent skill no longer supports GDPVal: its task recipe, example config and Apptainer SIF helper are removed, and the shared NeMo Gym machinery it carried now lives at ``references/gym.md`` with the launcher wrapper renamed ``scripts/nel-gym.sh``. GDPVal is an AA-suite member, so an "AA" request now generates the ``aa/`` tasks only -- report per-task scores rather than an aggregate compared against a published AA Index.
 - Deprecate ``metadata.recipe_type`` in recipe YAML, for both single-file recipes and a directory recipe's ``metadata.yml``. A recipe now says what kind it is with a ``# modelopt-schema:`` comment naming its schema class, or by delegating to a recipe that does (single-file recipes only -- a directory recipe has no body to delegate through); ``recipe_type`` is still read and still honoured, so a recipe outside this repo keeps working unchanged, but new recipes should leave it out and every recipe shipped here has been converted. Where both are present they must agree, and so must a recipe and the recipe it delegates to -- a disagreement is an error rather than a silent preference.
 

docs/source/guides/_customized_model_quantization.rst

@@ -18,7 +18,7 @@ As ModelOpt cannot detect these linear ops out-of-the-box, a HugggingFace plugin
 #. Try quantize the DBRX model after the plugin is implemented, feel free to follow the `hf_ptq example <https://github.com/NVIDIA/Model-Optimizer/tree/main/examples/hf_ptq>`_.
 #. Export the quantized model with :meth:`export_hf_checkpoint <modelopt.torch.export.unified_export_hf.export_hf_checkpoint>`. If the customized model is not supported by TensorRT-LLM, add support in its PyTorch backend and adapt the HF exporter if needed. See the :doc:`unified HF export guide <../deployment/3_unified_hf>` or :doc:`contact us <../support/1_contact>` for help.
 
-The following code snippet is excerpted from ``modelopt/torch/quantization/plugins/huggingface.py``
+The following code snippet is excerpted from ``modelopt/torch/models/dbrx/modeling_ptq.py``
 
 .. code-block:: python
 

examples/hf_ptq/README.md

@@ -696,7 +696,7 @@ from modelopt.torch.export.trtllm import export_tensorrt_llm_checkpoint
 with torch.inference_mode():
     export_tensorrt_llm_checkpoint(
         model,  # The quantized model.
-        decoder_type,  # The type of the model, e.g gpt, gptj, or llama.
+        decoder_type,  # The TensorRT-LLM model type, e.g gpt, gptj, or llama. Detected from the model class if None.
         dtype,  # The exported weights data type.
         export_dir,  # The directory where the exported files will be stored.
         inference_tensor_parallel,  # The number of GPUs used in the inference time tensor parallel.

examples/hf_ptq/example_utils.py

@@ -892,15 +892,15 @@ def is_model_on_gpu(model) -> bool:
     return all("cuda" in str(param.device) for param in model.parameters())
 
 
-def is_enc_dec(model_type) -> bool:
-    """Return whether the model_type uses encoder-decoder-style preview decode.
+def is_enc_dec(model_type: str | None) -> bool:
+    """Return whether the Hugging Face model_type is an encoder-decoder model.
 
-    Controls whether ``hf_ptq.py`` slices off the prompt prefix from
-    ``.generate()`` output. ``diffusion_gemma`` is structurally encoder-decoder
-    but returns prompt+canvas concatenated, so it stays OFF this list (AR-style
-    decode applies).
+    These models use encoder-decoder-style preview decode (``hf_ptq.py`` does not slice off
+    the prompt prefix from ``.generate()`` output) and still export TensorRT-LLM checkpoints.
+    ``diffusion_gemma`` is structurally encoder-decoder but returns prompt+canvas concatenated,
+    so it stays OFF this list (AR-style decode applies).
     """
-    return model_type in ["t5", "bart", "whisper"]
+    return model_type in ["t5", "mt5", "umt5", "bart", "mbart", "whisper"]
 
 
 def _resolve_model_path(model_name_or_path: str, trust_remote_code: bool = False) -> str:

examples/hf_ptq/hf_ptq.py

@@ -78,13 +78,13 @@
     export_hf_checkpoint,
     export_hf_vllm_fq_checkpoint,
     export_speculative_decoding,
-    get_model_type,
     has_spec_opt,
     save_expert_token_count_table,
 )
 from modelopt.torch.export.layerwise_export import LayerwiseExporter
 from modelopt.torch.export.model_utils import get_language_model_from_vl, is_multimodal_model
 from modelopt.torch.export.trtllm import export_tensorrt_llm_checkpoint
+from modelopt.torch.models import hf_model_type
 from modelopt.torch.quantization.config import need_calibration
 from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights
 from modelopt.torch.quantization.utils import is_quantized
@@ -151,14 +151,16 @@ def _kv_cfg_uses_constant_amax(kv_quant_cfg: list[dict[str, Any]]) -> bool:
 mto.enable_huggingface_checkpointing()
 
 
-def extract_and_prepare_language_model_from_vl(full_model):
+def extract_and_prepare_language_model_from_vl(
+    full_model: torch.nn.Module,
+) -> torch.nn.Module | None:
     """Extract language model from VL model and disable quantization for non-language components.
 
     Args:
         full_model: The full VLM model
 
     Returns:
-        tuple: (language_model, model_type) or (None, None) if not a VLM
+        The language model, or None if not a VLM
     """
     language_model_lineage = get_language_model_from_vl(full_model, strict=True)
     if language_model_lineage is not None:
@@ -178,10 +180,9 @@ def extract_and_prepare_language_model_from_vl(full_model):
                     mtq.quantize(module, disabled_quant_cfg, forward_loop=None)
                     memo.add(module)
 
-        model_type = get_model_type(language_model)
-        return language_model, model_type
+        return language_model
 
-    return None, None
+    return None
 
 
 class _DeviceDataLoader:
@@ -655,7 +656,9 @@ def load_model(args: argparse.Namespace):
             )
         calibration_only = True
 
-    model_type = get_model_type(full_model)
+    # The checkpoint's Hugging Face model_type. Always taken from the root model: a VLM's language
+    # model reports its own sub-config's type, which varies across transformers versions.
+    model_type = hf_model_type(full_model)
 
     if args.use_fsdp2:
         device = args.dist_state.device
@@ -713,12 +716,9 @@ def load_model(args: argparse.Namespace):
         # Plain PTQ quantizes only the language model. Recipes keep the complete VLM so their
         # quantizer rules can target vision and language components in one state.
         if args.recipe is None:
-            extracted_lm, extracted_model_type = extract_and_prepare_language_model_from_vl(
-                full_model
-            )
+            extracted_lm = extract_and_prepare_language_model_from_vl(full_model)
             if extracted_lm is not None:
                 language_model = extracted_lm
-                model_type = extracted_model_type
     else:
         if args.specdec_offline_dataset is not None:
             language_model = full_model
@@ -739,12 +739,9 @@ def load_model(args: argparse.Namespace):
             # would leave modelopt state on the ancestors and make auto_quantize() fail with
             # "multiple modelopt states".
             if args.recipe is None:
-                extracted_lm, extracted_model_type = extract_and_prepare_language_model_from_vl(
-                    full_model
-                )
+                extracted_lm = extract_and_prepare_language_model_from_vl(full_model)
                 if extracted_lm is not None:
                     language_model = extracted_lm
-                    model_type = extracted_model_type
 
         tokenizer = get_tokenizer(args.pyt_ckpt_path, trust_remote_code=args.trust_remote_code)
 
@@ -863,7 +860,11 @@ def mono_quantize(
         warnings.warn("Skipping quantization: model is already quantized.")
 
 
-def assert_layerwise_export_compatible(args, full_model, algorithm) -> None:
+def assert_layerwise_export_compatible(
+    args: argparse.Namespace,
+    full_model: torch.nn.Module,
+    algorithm: str | dict | list | None,
+) -> None:
     """Refuse layerwise export before calibration starts, not after the run is paid for.
 
     Layerwise export writes each layer's shard during calibration and finishes the checkpoint
@@ -906,7 +907,7 @@ def assert_layerwise_export_compatible(args, full_model, algorithm) -> None:
         ),
         (
             "an encoder-decoder model_type (t5/bart/whisper)",
-            getattr(full_model.config, "model_type", None) in ("t5", "bart", "whisper"),
+            is_enc_dec(hf_model_type(full_model)),
             "export_tensorrt_llm_checkpoint()",
         ),
     ):
@@ -924,16 +925,12 @@ def export_quantized(
     language_model: torch.nn.Module,
     model_type: str | None,
     tokenizer: PreTrainedTokenizerBase | None,
-    default_padding_side,
-    default_pad_token,
+    default_padding_side: str | None,
+    default_pad_token: str | None,
 ):
     # Not inference_mode: the FSDP2 path gathers full params in this context and
     # inference tensors break the subsequent state_dict() -> param.detach().
     with torch.no_grad():
-        if model_type is None:
-            print(f"Unknown model type {type(language_model).__name__}. Continue exporting...")
-            model_type = f"unknown:{type(language_model).__name__}"
-
         export_path = args.export_path
 
         # Early exit for speculative decoding checkpoints
@@ -953,7 +950,7 @@ def export_quantized(
 
         start_time = time.time()
         is_tensorrt_llm_export = (
-            model_type in ["t5", "bart", "whisper"]
+            is_enc_dec(model_type)
             or args.sparsity_fmt != "dense"
             or "int8_smoothquant" in args.qformat
         )
@@ -969,9 +966,9 @@ def export_quantized(
             # Move meta tensor back to device before exporting.
             remove_hook_from_module(language_model, recurse=True)
 
+            # The deprecated exporter detects its TensorRT-LLM decoder_type from the model class.
             export_tensorrt_llm_checkpoint(
                 language_model,
-                model_type,
                 export_dir=export_path,
                 inference_tensor_parallel=args.inference_tensor_parallel,
                 inference_pipeline_parallel=args.inference_pipeline_parallel,
@@ -1082,8 +1079,9 @@ def pre_quantize(
     # Generate preview before quantization
     if args.skip_generate:
         generated_ids_before_ptq = None
-    elif model_type == "deepseek":
-        # DeepSeek generation may go OOM, so we skip it
+    elif model_type is not None and model_type.startswith(("deepseek", "kimi_k2")):
+        # DeepSeek generation may go OOM, so we skip it. Kimi-K2 (kimi_k2) and Kimi-K2.5
+        # (kimi_k25) are DeepSeek-V3 architectures under their own model_type.
         generated_ids_before_ptq = None
     elif model_type == "nemotron_h":
         # NemotronH (SSM/Mamba hybrid) modeling code does not work with accelerate's big model inference
@@ -1117,13 +1115,13 @@ def post_quantize(
     model_type: str | None,
     tokenizer: PreTrainedTokenizerBase | None,
     processor: ProcessorMixin | None,
-    preview_input_ids,
-    preview_attention_mask,
-    generated_ids_before_ptq,
-    is_nemotron_vl_model,
-    first_text_speech_dataset,
-    default_padding_side,
-    default_pad_token,
+    preview_input_ids: torch.Tensor | None,
+    preview_attention_mask: torch.Tensor | None,
+    generated_ids_before_ptq: Any,
+    is_nemotron_vl_model: bool,
+    first_text_speech_dataset: str | None,
+    default_padding_side: str | None,
+    default_pad_token: str | None,
     calib_dataloader: DataLoader,
 ):
     """
@@ -1147,6 +1145,7 @@ def post_quantize(
             default_pad_token,
         )
         return
+    assert preview_input_ids is not None, "pre_quantize returns preview inputs for this path"
 
     if args.verbose and args.dist_state.is_main:
         try:
@@ -1161,7 +1160,7 @@ def post_quantize(
     generated_ids_after_ptq = None
     if generated_ids_before_ptq is None:
         pass
-    elif model_type != "llama4" and not is_nemotron_vl_model:
+    elif model_type not in ("llama4", "llama4_text") and not is_nemotron_vl_model:
         # Our fake quantizer may not be fully compatible with torch.compile.
         # This is a best-effort sanity check: e.g. a `device_map="auto"` load that offloads
         # part of the model to CPU (seen on unified-memory single-GPU hosts) can make a

examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb

@@ -79,7 +79,7 @@
     "- `EFFECTIVE_BITS` is the average precision target across quantized layers.\n",
     "- `Q_FORMATS` defines the list of quantization formats used during the AutoQuant search.\n",
     "- `KV_FORMAT` allows optional quantization of KV cache, applied after main quant.\n",
-    "- `EXPORT_FMT` supports exporting to Hugging Face (`hf`) or TensorRT-LLM (`tensorrt_llm`).\n",
+    "- `EXPORT_FMT` supports exporting to Hugging Face (`hf`) or the deprecated TensorRT-LLM checkpoint format (`tensorrt_llm`).\n",
     "- Also logs into Hugging Face\n",
     "\n",
     "💡 Try adding formats like `\"nvfp4\"`, `\"w4a8_awq\"` to explore tradeoffs."
@@ -99,7 +99,7 @@
     "Q_FORMATS = \"fp8,int4_awq\"  # search space\n",
     "KV_FORMAT = \"none\"  # or \"none\" to skip\n",
     "EXPORT_DIR = \"llama3_8b_autoq\"  # output folder\n",
-    "EXPORT_FMT = \"tensorrt_llm\"  # or \"hf\"\n",
+    "EXPORT_FMT = \"hf\"  # or \"tensorrt_llm\" (deprecated)\n",
     "# ----------------------------\n",
     "DEVICE = \"cuda\"\n",
     "DTYPE = torch.float16  # keep default for faster search\n",
@@ -372,7 +372,6 @@
     "if EXPORT_FMT == \"tensorrt_llm\":\n",
     "    export_tensorrt_llm_checkpoint(\n",
     "        model,\n",
-    "        model_type=\"llama\",\n",
     "        export_dir=EXPORT_DIR,\n",
     "        inference_tensor_parallel=1,\n",
     "        inference_pipeline_parallel=1,\n",

examples/llm_qat/llama_factory/llama_factory.py

@@ -35,6 +35,7 @@
 from wrapt import register_post_import_hook
 
 import modelopt.torch.opt as mto
+from modelopt.torch.models import hf_model_type
 from modelopt.torch.quantization.plugins.transformers_trainer import QADTrainer, QATTrainer
 
 logger = logging.get_logger(__name__)
@@ -144,7 +145,7 @@ def load_model(
                 )
             else:
                 model = load_class.from_pretrained(**init_kwargs)
-                if getattr(model.config, "model_type", None) == "qwen2_5_omni":
+                if hf_model_type(model) == "qwen2_5_omni":
                     model = model.thinker  # use part of Omni model
         else:
             raise ValueError(f"Model type {type(config)} is not supported.")

modelopt/torch/export/hf_export_handlers.py

@@ -125,12 +125,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

modelopt/torch/export/layerwise_export.py

@@ -284,8 +284,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()"