diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 8209ef8dd3d..cacf21a0d98 100755 --- a/CHANGELOG.rst +++ b/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//...`` paths still resolve via a backward-compatibility alias but now emit a ``FutureWarning``, so update them to ``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//``, 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. diff --git a/docs/source/guides/_customized_model_quantization.rst b/docs/source/guides/_customized_model_quantization.rst index b2fa879ecd4..c149999ca09 100644 --- a/docs/source/guides/_customized_model_quantization.rst +++ b/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 `_. #. Export the quantized model with :meth:`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 diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 27b57bdf166..4093c6b16d5 100755 --- a/examples/hf_ptq/README.md +++ b/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. diff --git a/examples/hf_ptq/example_utils.py b/examples/hf_ptq/example_utils.py index ed2f0961aa8..c6169614329 100755 --- a/examples/hf_ptq/example_utils.py +++ b/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: diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 347f6dc86c7..5d82baa9a06 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/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 diff --git a/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb b/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb index 068529d6ea9..aba5fdfce77 100644 --- a/examples/hf_ptq/notebooks/3_PTQ_AutoQuantization.ipynb +++ b/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", diff --git a/examples/llm_qat/llama_factory/llama_factory.py b/examples/llm_qat/llama_factory/llama_factory.py index 0ccf4e4d2ca..98fa47e9555 100644 --- a/examples/llm_qat/llama_factory/llama_factory.py +++ b/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.") diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 550fb0dbd6a..4fc9fc2a6d8 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/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 diff --git a/modelopt/torch/export/layerwise_export.py b/modelopt/torch/export/layerwise_export.py index 27be1c50cbb..898c28cbfdd 100644 --- a/modelopt/torch/export/layerwise_export.py +++ b/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()" diff --git a/modelopt/torch/export/model_utils.py b/modelopt/torch/export/model_utils.py index 95707639a03..11c8c8db209 100755 --- a/modelopt/torch/export/model_utils.py +++ b/modelopt/torch/export/model_utils.py @@ -21,58 +21,6 @@ from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector -MODEL_NAME_TO_TYPE = { - "GPT2": "gpt", - "Mllama": "mllama", - "Llama4": "llama4", - "Llama": "llama", - "Mistral": "llama", - "GPTJ": "gptj", - "FalconForCausalLM": "falcon", - "RWForCausalLM": "falcon", - "baichuan": "baichuan", - "MPT": "mpt", - "Bloom": "bloom", - "ChatGLM": "chatglm", - "Qwen3Moe": "qwen3moe", - "Qwen3Next": "qwen3next", - "QWen": "qwen", - "RecurrentGemma": "recurrentgemma", - # DiffusionGemma must come before "Gemma" — get_model_type substring-matches - # in order, and "gemma" is a substring of "diffusiongemma". - "DiffusionGemma": "diffusion_gemma", - "Gemma3": "gemma3", - "Gemma2": "gemma2", - "Gemma": "gemma", - "phi3small": "phi3small", - "phi3": "phi3", - "PhiMoEForCausalLM": "phi3", - "phi": "phi", - "TLGv4ForCausalLM": "phi", - "MixtralForCausalLM": "llama", - "ArcticForCausalLM": "llama", - "StarCoder": "gpt", - "Dbrx": "dbrx", - "T5": "t5", - "Bart": "bart", - "GLM": "glm", - "InternLM2ForCausalLM": "internlm", - "ExaoneForCausalLM": "exaone", - "NemotronH": "nemotron_h", - "Nemotron": "gpt", - "Deepseek": "deepseek", - "Whisper": "whisper", - "gptoss": "gptoss", - "MiniMax": "minimax", -} - -__doc__ = f"""Utility functions for model type detection and classification. - - .. code-block:: python - - {MODEL_NAME_TO_TYPE=} -""" - __all__ = [ "TiedWeightMap", "get_language_model_from_vl", @@ -80,6 +28,14 @@ "is_multimodal_model", ] +# Deprecated in 0.48.0, scheduled for removal in 0.49.0 together with the TensorRT-LLM checkpoint +# export, the only consumer of these TensorRT-LLM model names. +_MODEL_TYPE_DEPRECATION_MSG = ( + "{name} returns TensorRT-LLM model names and is deprecated as of 0.48.0; it will be removed " + "in 0.49.0. Use the Hugging Face model type, model.config.model_type, instead. " + "export_tensorrt_llm_checkpoint now detects its decoder_type when it is omitted." +) + @contextmanager def _release_exported_tensors(root: nn.Module): @@ -103,12 +59,34 @@ def _release_exported_tensors(root: nn.Module): module._buffers[buf_name] = None -def get_model_type(model): - """Try get the model type from the model name. If not found, return None.""" - for k, v in MODEL_NAME_TO_TYPE.items(): - if k.lower() in type(model).__name__.lower(): - return v - return None +def __getattr__(name: str): + if name == "MODEL_NAME_TO_TYPE": + from .trtllm.decoder_type import MODEL_NAME_TO_DECODER_TYPE + + warnings.warn( + _MODEL_TYPE_DEPRECATION_MSG.format(name="MODEL_NAME_TO_TYPE"), + DeprecationWarning, + stacklevel=2, + ) + return MODEL_NAME_TO_DECODER_TYPE + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def get_model_type(model: nn.Module): + """Try get the TensorRT-LLM model type from the model name. If not found, return None. + + .. deprecated:: 0.48.0 + Returns TensorRT-LLM model names and will be removed in 0.49.0. Use the Hugging Face + model type, ``model.config.model_type``, instead. + """ + from .trtllm.decoder_type import get_decoder_type + + warnings.warn( + _MODEL_TYPE_DEPRECATION_MSG.format(name="get_model_type"), + DeprecationWarning, + stacklevel=2, + ) + return get_decoder_type(model) def is_multimodal_model(model): diff --git a/modelopt/torch/export/trtllm/decoder_type.py b/modelopt/torch/export/trtllm/decoder_type.py new file mode 100644 index 00000000000..ef117aa000e --- /dev/null +++ b/modelopt/torch/export/trtllm/decoder_type.py @@ -0,0 +1,77 @@ +# 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. +"""TensorRT-LLM decoder type detection for the deprecated TensorRT-LLM checkpoint export. + +The decoder types here are TensorRT-LLM's own model names, not Hugging Face ``model_type`` values. +Only the TensorRT-LLM checkpoint export consumes them, so this module goes away with it in 0.49.0. +""" + +import torch.nn as nn + +__all__ = [] + +# Substring-matched against the lowercased model class name, in insertion order. +MODEL_NAME_TO_DECODER_TYPE = { + "GPT2": "gpt", + "Mllama": "mllama", + "Llama4": "llama4", + "Llama": "llama", + "Mistral": "llama", + "GPTJ": "gptj", + "FalconForCausalLM": "falcon", + "RWForCausalLM": "falcon", + "baichuan": "baichuan", + "MPT": "mpt", + "Bloom": "bloom", + "ChatGLM": "chatglm", + "Qwen3Moe": "qwen3moe", + "Qwen3Next": "qwen3next", + "QWen": "qwen", + "RecurrentGemma": "recurrentgemma", + # DiffusionGemma must come before "Gemma": "gemma" is a substring of "diffusiongemma". + "DiffusionGemma": "diffusion_gemma", + "Gemma3": "gemma3", + "Gemma2": "gemma2", + "Gemma": "gemma", + "phi3small": "phi3small", + "phi3": "phi3", + "PhiMoEForCausalLM": "phi3", + "phi": "phi", + "TLGv4ForCausalLM": "phi", + "MixtralForCausalLM": "llama", + "ArcticForCausalLM": "llama", + "StarCoder": "gpt", + "Dbrx": "dbrx", + "T5": "t5", + "Bart": "bart", + "GLM": "glm", + "InternLM2ForCausalLM": "internlm", + "ExaoneForCausalLM": "exaone", + "NemotronH": "nemotron_h", + "Nemotron": "gpt", + "Deepseek": "deepseek", + "Whisper": "whisper", + "gptoss": "gptoss", + "MiniMax": "minimax", +} + + +def get_decoder_type(model: nn.Module) -> str | None: + """Return the TensorRT-LLM decoder type for ``model``'s class name, or None if unknown.""" + name = type(model).__name__.lower() + for k, v in MODEL_NAME_TO_DECODER_TYPE.items(): + if k.lower() in name: + return v + return None diff --git a/modelopt/torch/export/trtllm/model_config_export.py b/modelopt/torch/export/trtllm/model_config_export.py index 8eae2aa94e4..3dfc35ddea7 100644 --- a/modelopt/torch/export/trtllm/model_config_export.py +++ b/modelopt/torch/export/trtllm/model_config_export.py @@ -36,6 +36,7 @@ from ..layer_utils import is_layernorm from ..quant_format import QUANTIZATION_INT4_AWQ, QUANTIZATION_W4A8_AWQ from ..quant_utils import get_quantization_format, process_layer_quant_config +from .decoder_type import get_decoder_type from .layer_utils import ( build_conv_config, build_decoder_config, @@ -98,7 +99,7 @@ def torch_to_tensorrt_llm_checkpoint( model: nn.Module, - decoder_type: str, + decoder_type: str | None = None, dtype: torch.dtype | None = None, inference_tensor_parallel: int = 0, inference_pipeline_parallel: int = 1, @@ -117,8 +118,9 @@ def torch_to_tensorrt_llm_checkpoint( Args: model: the torch model. - decoder_type: the type of the decoder, e.g. gpt, gptj, llama. - Please see :mod:`modelopt.torch.export.model_utils` for the supported models. + decoder_type: the TensorRT-LLM decoder type, e.g. gpt, gptj, llama. Detected from the + model class name if None, falling back to the generic decoder export for an + unrecognized Hugging Face model. Required for Megatron-Core models. dtype: the weights data type to export the unquantized layers or the default model data type if None. inference_tensor_parallel: The target inference time tensor parallel. We will merge or split the calibration tensor parallelism to inference. @@ -155,13 +157,20 @@ def torch_to_tensorrt_llm_checkpoint( def _torch_to_tensorrt_llm_checkpoint( model: nn.Module, - decoder_type: str, + decoder_type: str | None = None, dtype: torch.dtype | None = None, inference_tensor_parallel: int = 0, inference_pipeline_parallel: int = 1, workspace_path: Path | str | None = None, ) -> Iterator[tuple[dict[str, Any], dict[str, torch.Tensor], dict[str, Any]]]: """Generator behind :func:`torch_to_tensorrt_llm_checkpoint`; see it for the contract.""" + if decoder_type is None: + decoder_type = get_decoder_type(model) + if decoder_type is None: + # A Hugging Face model's TensorRT-LLM architecture comes from its config, so an + # unrecognized decoder_type still exports through the generic decoder path. + warn(f"Unknown decoder_type for {type(model).__name__}. Continue exporting...") + decoder_type = f"unknown:{type(model).__name__}" if dtype is None: dtype = get_dtype(model) @@ -480,7 +489,7 @@ def _torch_to_tensorrt_llm_checkpoint( def export_tensorrt_llm_checkpoint( model: nn.Module, - decoder_type: str, + decoder_type: str | None = None, dtype: torch.dtype | None = None, export_dir: Path | str = tempfile.gettempdir(), inference_tensor_parallel: int = 0, @@ -497,8 +506,9 @@ def export_tensorrt_llm_checkpoint( Args: model: the torch model. - decoder_type: the type of the decoder, e.g. gpt, gptj, llama. - Please see the model_utils.py for the supported models. + decoder_type: the TensorRT-LLM decoder type, e.g. gpt, gptj, llama. Detected from the + model class name if None, falling back to the generic decoder export for an + unrecognized Hugging Face model. Required for Megatron-Core models. dtype: the weights data type to export the unquantized layers or the default model data type if None. export_dir: the target export path. inference_tensor_parallel: The target inference time tensor parallel. diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 64e158c05bb..fe38bbfc4db 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -462,7 +462,6 @@ def requantize_resmooth_fused_llm_layers(model: torch.nn.Module): """Group modules that take the same input and register shared parameters in module.""" # TODO: Handle DBRX MoE quantization_format = get_quantization_format(model) - model_type = type(model).__name__.lower() model_hf_type = hf_model_type(model) module_names = set() # Built once: every fusion below resolves module names through it. @@ -499,7 +498,7 @@ def llm_dummy_forward(): # Check if this is a VL model that needs special input handling is_vl_model = is_multimodal_model(model) - if model_type.startswith("whisper"): + if model_hf_type == "whisper": # For Whisper models, we need to pass a fake input with the specific sequence length from transformers import AutoFeatureExtractor @@ -508,7 +507,9 @@ def llm_dummy_forward(): [1, model.config.num_mel_bins, feature_extractor.nb_max_frames], dtype=model.dtype ).to(model.device) - if is_vl_model and "nemotron" in model_type: + # Nemotron VL remote-code model_types (NemotronH_Nano_VL_V2, Llama_Nemotron_Nano_VL, + # nemotron_parse) are mixed-case, so match case-insensitively. + if is_vl_model and "nemotron" in (model_hf_type or "").lower(): # For Nemotron VL models, run optimization on just the language model/decoder. # This avoids needing pixel_values for the vision encoder. language_model_lineage = get_language_model_from_vl(model) @@ -522,7 +523,7 @@ def llm_dummy_forward(): language_model(fake_input, use_cache=False) else: raise ValueError( - f"Cannot extract language_model from Nemotron VL model (type: {model_type}). " + f"Cannot extract language_model from Nemotron VL model (type: {model_hf_type}). " "This is required for requantization/resmoothing optimization. " "Please ensure the model architecture is supported or file an issue." ) @@ -1017,10 +1018,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 cbb76489fc3..51d5b53dfb4 100644 --- a/modelopt/torch/export/unified_export_hf_streaming.py +++ b/modelopt/torch/export/unified_export_hf_streaming.py @@ -354,9 +354,9 @@ def _export_transformers_checkpoint_streaming( NotImplementedError: if the model's conversion mapping contains split rules. 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 @@ -553,9 +553,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 f1ffcf6a54a..378063fffed 100644 --- a/modelopt/torch/models/README.md +++ b/modelopt/torch/models/README.md @@ -50,6 +50,23 @@ 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. DBRX's fused +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 +`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; a family revision sharing another type's `modeling_ptq.py` (`step3p7` uses `step3p5`'s) +holds only `specs.py`. + ## Sections A `ModelSpec` holds one attribute per section, each `None` unless the model fills it: diff --git a/modelopt/torch/models/__init__.py b/modelopt/torch/models/__init__.py index c835e6b63e4..2aef8e9ddca 100644 --- a/modelopt/torch/models/__init__.py +++ b/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 @@ -71,6 +74,7 @@ deepseek, deepseek_v3, deepseek_v4, + falcon, gemma, gemma2, gemma3, @@ -78,6 +82,7 @@ gemma4_text, gpt_oss, llama, + llama4, minimax, mixtral, phimoe, @@ -88,4 +93,7 @@ qwen3_5_moe, qwen3_moe, qwen3_next, + qwen3_vl_moe, + step3p5, + step3p7, ) diff --git a/modelopt/torch/models/dbrx/modeling_ptq.py b/modelopt/torch/models/dbrx/modeling_ptq.py new file mode 100644 index 00000000000..0ef4402d6d3 --- /dev/null +++ b/modelopt/torch/models/dbrx/modeling_ptq.py @@ -0,0 +1,178 @@ +# 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. + +"""DBRX PTQ modeling (HF model type ``dbrx``).""" + +import torch +import torch.nn as nn + +from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry +from modelopt.torch.quantization.plugins.custom import CUSTOM_MODEL_PLUGINS +from modelopt.torch.quantization.plugins.huggingface import _QuantSparseSequentialMoe + +__all__: list[str] = [] + + +# For more information on DbrxExpert, see https://github.com/huggingface/transformers/blob/dcdda532/src/transformers/models/dbrx/modeling_dbrx.py#L756 +class _QuantDbrxExperts(QuantModule): + def _setup(self): + """Modify the DbrxExpert.""" + # No setup is needed for DbrxExpert, we only need to update DbrxExpertGLU + + def forward( + self, + x: torch.Tensor, + top_experts: torch.LongTensor, + top_weights: torch.Tensor, + ) -> torch.Tensor: + bsz, q_len, hidden_size = x.shape + x = x.view(-1, hidden_size) + out = torch.zeros_like(x) + + expert_mask = nn.functional.one_hot(top_experts, num_classes=self.num_experts).permute( + 2, 1, 0 + ) + for expert_idx in range(self.num_experts): + topk_idx, token_idx = torch.where(expert_mask[expert_idx]) + if token_idx.shape[0] == 0: + continue + + token_list = token_idx.tolist() + topk_list = topk_idx.tolist() + + expert_tokens = x[None, token_list].reshape(-1, hidden_size) + expert_out = ( + self.mlp(expert_tokens, expert_idx) * top_weights[token_list, topk_list, None] + ) + + out.index_add_(0, token_idx, expert_out) + + out = out.reshape(bsz, q_len, hidden_size) + return out + + +class _QuantDbrxExpertGLU(QuantModule): + def _setup(self): + """Modify the DbrxExpertGLU by using nn.Linear layers.""" + dtype, device = self.w1.dtype, self.w1.device + + def _copy_weights(modules, weights): + modules.to(dtype=dtype, device=device) + for expert_idx, module in enumerate(modules): + with torch.no_grad(): + module.weight.copy_(weights[expert_idx].detach()) + + # In transformers 5.0, DbrxExpertGLU.forward uses raw matmul: x @ w1[i] where + # w1[i] has shape (ffn_hidden_size, hidden_size). To match via F.linear (which + # computes x @ W.T), we store weights transposed: W = w1[i].T. + self.w1_linear = nn.ModuleList( + [ + nn.Linear(self.ffn_hidden_size, self.hidden_size, bias=False) + for _ in range(self.moe_num_experts) + ] + ) + _copy_weights( + self.w1_linear, + self.w1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size).transpose( + 1, 2 + ), + ) + delattr(self, "w1") + + self.v1_linear = nn.ModuleList( + [ + nn.Linear(self.ffn_hidden_size, self.hidden_size, bias=False) + for _ in range(self.moe_num_experts) + ] + ) + _copy_weights( + self.v1_linear, + self.v1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size).transpose( + 1, 2 + ), + ) + delattr(self, "v1") + + # w2: down_proj uses intermediate.matmul(w2[i].t()) = F.linear(intermediate, w2[i]) + # so W = w2[i] directly (no extra transpose needed). + self.w2_linear = nn.ModuleList( + [ + nn.Linear(self.hidden_size, self.ffn_hidden_size, bias=False) + for _ in range(self.moe_num_experts) + ] + ) + _copy_weights( + self.w2_linear, + self.w2.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size), + ) + delattr(self, "w2") + + def forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor: + x1 = self.w1_linear[expert_idx](x) + x2 = self.v1_linear[expert_idx](x) + x1 = self.activation_fn(x1) + x1 = x1 * x2 + return self.w2_linear[expert_idx](x1) + + +class _QuantDbrxFFN(_QuantSparseSequentialMoe): + @property + def num_experts(self): + return self.router.moe_num_experts + + @property + def top_k(self): + # In older transformers, top_k was stored on DbrxRouter as moe_top_k. + # In transformers 5.0, DbrxFFN stores it as a plain attribute (top_k). + if hasattr(self.router, "moe_top_k"): + return self.router.moe_top_k + return self.__dict__.get("top_k", 1) + + @top_k.setter + def top_k(self, value): + if hasattr(self.router, "moe_top_k"): + self.router.moe_top_k = value + else: + self.__dict__["top_k"] = value + + +try: + from transformers.models.dbrx.modeling_dbrx import DbrxExpertGLU, DbrxExperts, DbrxFFN + + if DbrxExperts not in QuantModuleRegistry: + QuantModuleRegistry.register({DbrxExperts: "hf.DbrxExperts"})(_QuantDbrxExperts) + + if DbrxExpertGLU not in QuantModuleRegistry: + QuantModuleRegistry.register({DbrxExpertGLU: "hf.DbrxExpertGLU"})(_QuantDbrxExpertGLU) + + if DbrxFFN not in QuantModuleRegistry: + QuantModuleRegistry.register({DbrxFFN: "hf.DbrxFFN"})(_QuantDbrxFFN) +except ImportError: + pass + + +def register_dbrx_moe_on_the_fly(model): + """Register DBRX MoE modules as QUANT_MODULE. + + The MoE class in DBRX is `transformers_modules.modeling_dbrx.DbrxExpertGLU`, which loads dynamically. + """ + if type(model).__name__ == "DbrxForCausalLM": + moe_type = type(model.transformer.blocks[0].ffn.experts.mlp) + # Create a QuantDbrxExpertGLU class on the fly + if QuantModuleRegistry.get(moe_type) is None: + QuantModuleRegistry.register({moe_type: moe_type.__name__})(_QuantDbrxExpertGLU) + + +CUSTOM_MODEL_PLUGINS.add(register_dbrx_moe_on_the_fly) diff --git a/modelopt/torch/models/dbrx/specs.py b/modelopt/torch/models/dbrx/specs.py index b114abaeba3..39997d38a2a 100644 --- a/modelopt/torch/models/dbrx/specs.py +++ b/modelopt/torch/models/dbrx/specs.py @@ -21,7 +21,7 @@ # Expert names refer to the quantized layout: _QuantDbrxExpertGLU rewrites the fused # w1/v1/w2 parameters into per-expert w1_linear/v1_linear/w2_linear ModuleLists on -# experts.mlp (see modelopt/torch/quantization/plugins/huggingface.py). +# experts.mlp (see modeling_ptq.py). register( ModelSpec( model_type="dbrx", diff --git a/modelopt/torch/models/falcon/__init__.py b/modelopt/torch/models/falcon/__init__.py new file mode 100644 index 00000000000..c08e11541cb --- /dev/null +++ b/modelopt/torch/models/falcon/__init__.py @@ -0,0 +1,19 @@ +# 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. + + +"""Per-model specs for HF model type ``falcon``.""" + +from . import specs diff --git a/modelopt/torch/models/falcon/modeling_ptq.py b/modelopt/torch/models/falcon/modeling_ptq.py new file mode 100644 index 00000000000..5844e842865 --- /dev/null +++ b/modelopt/torch/models/falcon/modeling_ptq.py @@ -0,0 +1,47 @@ +# 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 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] = [] + + +try: + from transformers.models.falcon.modeling_falcon import FalconLinear + + if FalconLinear not in QuantModuleRegistry: + QuantModuleRegistry.register({FalconLinear: "hf.FalconLinear"})(_QuantLinear) +except ImportError: + pass + + +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) diff --git a/modelopt/torch/models/falcon/specs.py b/modelopt/torch/models/falcon/specs.py new file mode 100644 index 00000000000..2b93f8f318c --- /dev/null +++ b/modelopt/torch/models/falcon/specs.py @@ -0,0 +1,24 @@ +# 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 specs (HF model type ``falcon``).""" + +from ..specs import ModelSpec, register + +__all__: list[str] = [] + +# Dense, and nothing model-specific for export. Older remote-code Falcon checkpoints +# (RWForCausalLM) report other model types and are handled by modeling_ptq.py alone. +register(ModelSpec(model_type="falcon", min_transformers_version="4.57")) diff --git a/modelopt/torch/models/gpt_oss/modeling_ptq.py b/modelopt/torch/models/gpt_oss/modeling_ptq.py new file mode 100644 index 00000000000..9bc76a9a120 --- /dev/null +++ b/modelopt/torch/models/gpt_oss/modeling_ptq.py @@ -0,0 +1,125 @@ +# 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. + +"""GPT-OSS PTQ modeling (HF model type ``gpt_oss``).""" + +from contextlib import contextmanager +from functools import partial + +import torch + +from modelopt.torch.quantization.nn import QuantModuleRegistry, TensorQuantizer +from modelopt.torch.quantization.plugins.custom import _QuantFunctionalMixin +from modelopt.torch.quantization.plugins.huggingface import ( + _transposed_quantize, + _TransposedExpertsCalibMixin, +) + +__all__: list[str] = [] + + +class _QuantGptOssExperts(_TransposedExpertsCalibMixin, _QuantFunctionalMixin): + """Quantized wrapper for `transformers.GptOssExperts`. + + Quantizes `gate_up_proj` and `down_proj` weights via dynamic attributes inside `quantize_weight()`. + Activations into `gate_up_proj` are quantized by `gate_up_proj_input_quantizer`. For `down_proj` + activation quantization, we intercept `torch.Tensor.__matmul__`/`torch.bmm` and quantize inputs + on every second call (since the first call computes `gate_up_proj` outputs and second call + computes `down_proj` outputs). + """ + + @staticmethod + def _get_quantized_weight(quantizer, module, weight): + # MoE weight is accessed for each expert in one forward pass. so lets cache it + if module._enable_weight_quantization: + if hasattr(quantizer, "_cached_quant_val"): + return getattr(quantizer, "_cached_quant_val") + quantizer._cached_quant_val = _transposed_quantize(weight, quantizer) + return quantizer._cached_quant_val + return weight + + def _setup_for_weight_quantization(self): + self._register_dynamic_attribute( + "gate_up_proj", partial(self._get_quantized_weight, self.gate_up_proj_weight_quantizer) + ) + self._register_dynamic_attribute( + "down_proj", partial(self._get_quantized_weight, self.down_proj_weight_quantizer) + ) + + def _setup(self): + assert not hasattr(self, "kernel_layer_name"), ( + "ModelOpt quantization does not support patched forward for kernel_hub" + ) + 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() + + self._register_temp_attribute("_enable_weight_quantization", False) + self._register_temp_attribute("_down_proj_mul", False) + self._setup_for_weight_quantization() + + @property + def functionals_to_replace(self): + # Use torch.ops.aten to bypass Python dispatch and avoid RecursionError + # (torch.matmul / __matmul__ can dispatch to each other) + _aten_bmm = torch.ops.aten.bmm + _aten_matmul = torch.ops.aten.matmul + + def _quantized_bmm(batch1, batch2, *, out=None): + batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1 + self._down_proj_mul = not self._down_proj_mul # toggle the flag + if out is not None: + return torch.ops.aten.bmm.out(batch1, batch2, out=out) + return _aten_bmm(batch1, batch2) + + def _tensor_matmul(self_t, other): + self_t = self.down_proj_input_quantizer(self_t) if self._down_proj_mul else self_t + self._down_proj_mul = not self._down_proj_mul + return _aten_matmul(self_t, other) + + return [ + (torch, "bmm", _quantized_bmm), + (torch.Tensor, "__matmul__", _tensor_matmul), + ] + + @contextmanager + def quantize_weight(self): + """Context in which MoE weight is quantized.""" + self._enable_weight_quantization = True + try: + yield + finally: + for module in self.modules(): + if isinstance(module, TensorQuantizer) and hasattr(module, "_cached_quant_val"): + delattr(module, "_cached_quant_val") + self._enable_weight_quantization = False + + def forward( + self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None + ) -> torch.Tensor: + """Forward method to add quantization.""" + hidden_states = self.gate_up_proj_input_quantizer(hidden_states) + with self.quantize_weight(): + return super().forward(hidden_states, router_indices, routing_weights) + + +try: + from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts + + if GptOssExperts not in QuantModuleRegistry: + QuantModuleRegistry.register({GptOssExperts: "hf.GptOssExperts"})(_QuantGptOssExperts) +except ImportError: + pass diff --git a/modelopt/torch/models/llama4/__init__.py b/modelopt/torch/models/llama4/__init__.py new file mode 100644 index 00000000000..e9fd7824c40 --- /dev/null +++ b/modelopt/torch/models/llama4/__init__.py @@ -0,0 +1,19 @@ +# 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. + + +"""Per-model specs for HF model type ``llama4``.""" + +from . import specs diff --git a/modelopt/torch/models/llama4/modeling_ptq.py b/modelopt/torch/models/llama4/modeling_ptq.py new file mode 100644 index 00000000000..9f16a4786a7 --- /dev/null +++ b/modelopt/torch/models/llama4/modeling_ptq.py @@ -0,0 +1,59 @@ +# 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 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 + + +try: + from transformers.models.llama4.modeling_llama4 import Llama4TextExperts + + if Llama4TextExperts not in QuantModuleRegistry: + QuantModuleRegistry.register({Llama4TextExperts: "hf.Llama4TextExperts"})( + _QuantLlama4TextExperts + ) +except ImportError: + pass diff --git a/modelopt/torch/models/llama4/specs.py b/modelopt/torch/models/llama4/specs.py new file mode 100644 index 00000000000..70a7779883d --- /dev/null +++ b/modelopt/torch/models/llama4/specs.py @@ -0,0 +1,36 @@ +# 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 specs (HF model type ``llama4``).""" + +from ..specs import ModelSpec, MoESpec, register + +__all__: list[str] = [] + +# Llama4TextExperts is fused: one module holding 3-D gate_up_proj and down_proj +# parameters, run through torch.bmm. There is no (gate, up) pair left to fuse, and the +# grouped-export path (get_experts_list) does not apply. The block also carries a +# shared_expert MLP, which is an ordinary dense MLP, not part of this layout. +register( + ModelSpec( + model_type="llama4", + min_transformers_version="4.57", + moe_spec=MoESpec( + block_names=("Llama4TextMoe",), + expert_linear_names=("gate_up_proj", "down_proj"), + fused_expert_names=True, + ), + ) +) diff --git a/modelopt/torch/models/nemotron_h/modeling_ptq.py b/modelopt/torch/models/nemotron_h/modeling_ptq.py new file mode 100644 index 00000000000..46670b9b7eb --- /dev/null +++ b/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 +) diff --git a/modelopt/torch/models/qwen3_vl_moe/__init__.py b/modelopt/torch/models/qwen3_vl_moe/__init__.py new file mode 100644 index 00000000000..e99407906be --- /dev/null +++ b/modelopt/torch/models/qwen3_vl_moe/__init__.py @@ -0,0 +1,19 @@ +# 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. + + +"""Per-model specs for HF model type ``qwen3_vl_moe``.""" + +from . import specs diff --git a/modelopt/torch/models/qwen3_vl_moe/modeling_ptq.py b/modelopt/torch/models/qwen3_vl_moe/modeling_ptq.py new file mode 100644 index 00000000000..9770682d225 --- /dev/null +++ b/modelopt/torch/models/qwen3_vl_moe/modeling_ptq.py @@ -0,0 +1,136 @@ +# 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. + +"""Qwen3-VL-MoE PTQ modeling (HF model type ``qwen3_vl_moe``).""" + +import torch +import torch.nn as nn + +from modelopt.torch.quantization.nn import QuantModule, QuantModuleRegistry + +__all__: list[str] = [] + + +class _QuantQwen3VLMoeTextExperts(QuantModule): + """Quantized wrapper for the pre-transformers-5.12 ``Qwen3VLMoeTextExperts`` layout. + + That layout stores ``gate_up_proj`` as (num_experts, hidden_size, 2*expert_dim) and runs + the experts through ``torch.bmm``/``@``, so it is unrolled into ``nn.Linear`` modules here. + transformers>=5.12 moved this module to the standard fused layout handled by + :class:`_QuantFusedExperts`; see the registration site below. + """ + + def _setup(self): + """Modify the Qwen3VLMoeTextExperts by using nn.Linear layers.""" + from accelerate import init_empty_weights + + dtype, device = self.gate_up_proj.dtype, self.gate_up_proj.device + + def _copy_weight(module, weight): + module.to_empty(device=device) + with torch.no_grad(): + module.weight.data = weight.detach().data.to(dtype=dtype, device=device) + + # The attribute name was changed from `intermediate_size` to `intermediate_dim` in + # https://github.com/huggingface/transformers/commit/0642963ba13f2dae0596fe489415569e1d91fbda + if hasattr(self, "intermediate_size"): + expert_dim = self.intermediate_size + elif hasattr(self, "intermediate_dim"): + expert_dim = self.intermediate_dim + else: + raise AttributeError("Could not find intermediate dimension size in model") + + with init_empty_weights(): + gate_proj = nn.ModuleList( + [ + nn.Linear(self.hidden_size, expert_dim, bias=False) + for _ in range(self.num_experts) + ] + ) + up_proj = nn.ModuleList( + [ + nn.Linear(self.hidden_size, expert_dim, bias=False) + for _ in range(self.num_experts) + ] + ) + down_proj = nn.ModuleList( + [ + nn.Linear(expert_dim, self.hidden_size, bias=False) + for _ in range(self.num_experts) + ] + ) + + for idx in range(self.num_experts): + _copy_weight(gate_proj[idx], self.gate_up_proj[idx, :, :expert_dim].T) + _copy_weight(up_proj[idx], self.gate_up_proj[idx, :, expert_dim:].T) + _copy_weight(down_proj[idx], self.down_proj[idx, :].T) + + delattr(self, "gate_up_proj") + delattr(self, "down_proj") + self.gate_proj = gate_proj + self.up_proj = up_proj + self.down_proj = down_proj + + def forward( + self, + hidden_states: torch.Tensor, + routing_weights: torch.Tensor, + router_indices: torch.Tensor, + ) -> torch.Tensor: + batch_size = hidden_states.shape[0] + hidden_states = hidden_states.reshape(-1, self.hidden_size) + next_states = torch.zeros_like(hidden_states) + with torch.no_grad(): + expert_mask = torch.nn.functional.one_hot(router_indices, num_classes=self.num_experts) + expert_mask = expert_mask.permute(2, 1, 0) + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + for expert_idx in expert_hit: + with torch.no_grad(): + _, token_idx = torch.where(expert_mask[expert_idx[0]]) + current_state = hidden_states[token_idx] + gate = self.gate_proj[expert_idx](current_state) + up = self.up_proj[expert_idx](current_state) + gated_output = up * self.act_fn(gate) + out = self.down_proj[expert_idx](gated_output) + weighted_output = out * routing_weights[token_idx, expert_idx, None] + next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype)) + next_states = next_states.view(batch_size, -1, self.hidden_size) + + return next_states + + +try: + from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeTextExperts + + # transformers>=5.12 rewrote Qwen3VLMoeTextExperts onto the standard + # ``@use_experts_implementation`` fused layout: ``hidden_size``/``expert_dim`` became + # ``hidden_dim``/``intermediate_dim``, ``gate_up_proj`` was transposed to + # (num_experts, 2*intermediate_dim, hidden_dim), and the forward now calls ``F.linear`` + # twice per expert. ``_QuantQwen3VLMoeTextExperts`` only understands the older layout, + # so registering it against the new one crashes on ``self.hidden_size`` (nvbug 6518551). + # The decorator sets ``_apply_gate`` on the class; use it to detect the new layout and + # leave those modules to ``register_fused_experts_on_the_fly``, which claims them with + # the generic ``_QuantFusedExperts``. The old layout must stay explicitly registered: + # it is structurally indistinguishable from a generic fused-experts module, yet its + # forward uses ``torch.bmm``/``@`` rather than ``F.linear``, so the generic wrapper + # would silently quantize nothing. + if Qwen3VLMoeTextExperts not in QuantModuleRegistry and not hasattr( + Qwen3VLMoeTextExperts, "_apply_gate" + ): + QuantModuleRegistry.register({Qwen3VLMoeTextExperts: "hf.Qwen3VLMoeTextExperts"})( + _QuantQwen3VLMoeTextExperts + ) +except ImportError: + pass diff --git a/modelopt/torch/models/qwen3_vl_moe/specs.py b/modelopt/torch/models/qwen3_vl_moe/specs.py new file mode 100644 index 00000000000..515ac7179ee --- /dev/null +++ b/modelopt/torch/models/qwen3_vl_moe/specs.py @@ -0,0 +1,36 @@ +# 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. + +"""Qwen3-VL-MoE specs (HF model type ``qwen3_vl_moe``).""" + +from ..specs import ModelSpec, MoESpec, register + +__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 +# modeling_ptq.py); that rewrite is not iterable per expert, so no export path groups it. +register( + ModelSpec( + model_type="qwen3_vl_moe", + min_transformers_version="4.57", + moe_spec=MoESpec( + block_names=("Qwen3VLMoeTextSparseMoeBlock",), + expert_linear_names=("gate_up_proj", "down_proj"), + fused_expert_names=True, + ), + ) +) diff --git a/modelopt/torch/models/step3p5/__init__.py b/modelopt/torch/models/step3p5/__init__.py new file mode 100644 index 00000000000..e5c436ddc2d --- /dev/null +++ b/modelopt/torch/models/step3p5/__init__.py @@ -0,0 +1,19 @@ +# 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. + + +"""Per-model specs for HF model type ``step3p5``.""" + +from . import specs diff --git a/modelopt/torch/models/step3p5/modeling_ptq.py b/modelopt/torch/models/step3p5/modeling_ptq.py new file mode 100644 index 00000000000..84dc6022261 --- /dev/null +++ b/modelopt/torch/models/step3p5/modeling_ptq.py @@ -0,0 +1,262 @@ +# 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. + +"""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): + 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/step3p5/specs.py b/modelopt/torch/models/step3p5/specs.py new file mode 100644 index 00000000000..541bdfbf42c --- /dev/null +++ b/modelopt/torch/models/step3p5/specs.py @@ -0,0 +1,26 @@ +# 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. + +"""Step-3.5 specs (HF model type ``step3p5``).""" + +from ..specs import ModelSpec, register + +__all__: list[str] = [] + +# No MoESpec: the routed experts are expert-indexed MoELinear projections on the MoE MLP +# itself, with no `experts` container, and are exported by the MoELinear handler rather +# than through the MoE-block lookups. Declaring the block would make is_moe claim it and +# send AWQ export into get_experts_list, which does not support this layout. +register(ModelSpec(model_type="step3p5", modeling_source="remote_code")) diff --git a/modelopt/torch/models/step3p7/__init__.py b/modelopt/torch/models/step3p7/__init__.py new file mode 100644 index 00000000000..1f34f24c93b --- /dev/null +++ b/modelopt/torch/models/step3p7/__init__.py @@ -0,0 +1,19 @@ +# 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. + + +"""Per-model specs for HF model type ``step3p7``.""" + +from . import specs diff --git a/modelopt/torch/models/step3p7/specs.py b/modelopt/torch/models/step3p7/specs.py new file mode 100644 index 00000000000..0846600918b --- /dev/null +++ b/modelopt/torch/models/step3p7/specs.py @@ -0,0 +1,27 @@ +# 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. + +"""Step-3.7 specs (HF model type ``step3p7``). + +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 + +__all__: list[str] = [] + +# No MoESpec, for the same reason as step3p5 (see step3p5/specs.py). +register(ModelSpec(model_type="step3p7", modeling_source="remote_code")) diff --git a/modelopt/torch/quantization/plugins/huggingface.py b/modelopt/torch/quantization/plugins/huggingface.py index 12e7c49b687..5f2ac412119 100644 --- a/modelopt/torch/quantization/plugins/huggingface.py +++ b/modelopt/torch/quantization/plugins/huggingface.py @@ -15,9 +15,9 @@ """Support quantization for huggingface layers.""" +import importlib import inspect import logging -import re import warnings from contextlib import contextmanager from functools import partial @@ -760,218 +760,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 - - -# For more information on DbrxExpert, see https://github.com/huggingface/transformers/blob/dcdda532/src/transformers/models/dbrx/modeling_dbrx.py#L756 -class _QuantDbrxExperts(QuantModule): - def _setup(self): - """Modify the DbrxExpert.""" - # No setup is needed for DbrxExpert, we only need to update DbrxExpertGLU - - def forward( - self, - x: torch.Tensor, - top_experts: torch.LongTensor, - top_weights: torch.Tensor, - ) -> torch.Tensor: - bsz, q_len, hidden_size = x.shape - x = x.view(-1, hidden_size) - out = torch.zeros_like(x) - - expert_mask = nn.functional.one_hot(top_experts, num_classes=self.num_experts).permute( - 2, 1, 0 - ) - for expert_idx in range(self.num_experts): - topk_idx, token_idx = torch.where(expert_mask[expert_idx]) - if token_idx.shape[0] == 0: - continue - - token_list = token_idx.tolist() - topk_list = topk_idx.tolist() - - expert_tokens = x[None, token_list].reshape(-1, hidden_size) - expert_out = ( - self.mlp(expert_tokens, expert_idx) * top_weights[token_list, topk_list, None] - ) - - out.index_add_(0, token_idx, expert_out) - - out = out.reshape(bsz, q_len, hidden_size) - return out - - -class _QuantDbrxExpertGLU(QuantModule): - def _setup(self): - """Modify the DbrxExpertGLU by using nn.Linear layers.""" - dtype, device = self.w1.dtype, self.w1.device - - def _copy_weights(modules, weights): - modules.to(dtype=dtype, device=device) - for expert_idx, module in enumerate(modules): - with torch.no_grad(): - module.weight.copy_(weights[expert_idx].detach()) - - # In transformers 5.0, DbrxExpertGLU.forward uses raw matmul: x @ w1[i] where - # w1[i] has shape (ffn_hidden_size, hidden_size). To match via F.linear (which - # computes x @ W.T), we store weights transposed: W = w1[i].T. - self.w1_linear = nn.ModuleList( - [ - nn.Linear(self.ffn_hidden_size, self.hidden_size, bias=False) - for _ in range(self.moe_num_experts) - ] - ) - _copy_weights( - self.w1_linear, - self.w1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size).transpose( - 1, 2 - ), - ) - delattr(self, "w1") - - self.v1_linear = nn.ModuleList( - [ - nn.Linear(self.ffn_hidden_size, self.hidden_size, bias=False) - for _ in range(self.moe_num_experts) - ] - ) - _copy_weights( - self.v1_linear, - self.v1.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size).transpose( - 1, 2 - ), - ) - delattr(self, "v1") - - # w2: down_proj uses intermediate.matmul(w2[i].t()) = F.linear(intermediate, w2[i]) - # so W = w2[i] directly (no extra transpose needed). - self.w2_linear = nn.ModuleList( - [ - nn.Linear(self.hidden_size, self.ffn_hidden_size, bias=False) - for _ in range(self.moe_num_experts) - ] - ) - _copy_weights( - self.w2_linear, - self.w2.view(self.moe_num_experts, self.ffn_hidden_size, self.hidden_size), - ) - delattr(self, "w2") - - def forward(self, x: torch.Tensor, expert_idx: int) -> torch.Tensor: - x1 = self.w1_linear[expert_idx](x) - x2 = self.v1_linear[expert_idx](x) - x1 = self.activation_fn(x1) - x1 = x1 * x2 - return self.w2_linear[expert_idx](x1) - - -class _QuantQwen3VLMoeTextExperts(QuantModule): - """Quantized wrapper for the pre-transformers-5.12 ``Qwen3VLMoeTextExperts`` layout. - - That layout stores ``gate_up_proj`` as (num_experts, hidden_size, 2*expert_dim) and runs - the experts through ``torch.bmm``/``@``, so it is unrolled into ``nn.Linear`` modules here. - transformers>=5.12 moved this module to the standard fused layout handled by - :class:`_QuantFusedExperts`; see the registration site below. - """ - - def _setup(self): - """Modify the Qwen3VLMoeTextExperts by using nn.Linear layers.""" - from accelerate import init_empty_weights - - dtype, device = self.gate_up_proj.dtype, self.gate_up_proj.device - - def _copy_weight(module, weight): - module.to_empty(device=device) - with torch.no_grad(): - module.weight.data = weight.detach().data.to(dtype=dtype, device=device) - - # The attribute name was changed from `intermediate_size` to `intermediate_dim` in - # https://github.com/huggingface/transformers/commit/0642963ba13f2dae0596fe489415569e1d91fbda - if hasattr(self, "intermediate_size"): - expert_dim = self.intermediate_size - elif hasattr(self, "intermediate_dim"): - expert_dim = self.intermediate_dim - else: - raise AttributeError("Could not find intermediate dimension size in model") - - with init_empty_weights(): - gate_proj = nn.ModuleList( - [ - nn.Linear(self.hidden_size, expert_dim, bias=False) - for _ in range(self.num_experts) - ] - ) - up_proj = nn.ModuleList( - [ - nn.Linear(self.hidden_size, expert_dim, bias=False) - for _ in range(self.num_experts) - ] - ) - down_proj = nn.ModuleList( - [ - nn.Linear(expert_dim, self.hidden_size, bias=False) - for _ in range(self.num_experts) - ] - ) - - for idx in range(self.num_experts): - _copy_weight(gate_proj[idx], self.gate_up_proj[idx, :, :expert_dim].T) - _copy_weight(up_proj[idx], self.gate_up_proj[idx, :, expert_dim:].T) - _copy_weight(down_proj[idx], self.down_proj[idx, :].T) - - delattr(self, "gate_up_proj") - delattr(self, "down_proj") - self.gate_proj = gate_proj - self.up_proj = up_proj - self.down_proj = down_proj - - def forward( - self, - hidden_states: torch.Tensor, - routing_weights: torch.Tensor, - router_indices: torch.Tensor, - ) -> torch.Tensor: - batch_size = hidden_states.shape[0] - hidden_states = hidden_states.reshape(-1, self.hidden_size) - next_states = torch.zeros_like(hidden_states) - with torch.no_grad(): - expert_mask = torch.nn.functional.one_hot(router_indices, num_classes=self.num_experts) - expert_mask = expert_mask.permute(2, 1, 0) - expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() - for expert_idx in expert_hit: - with torch.no_grad(): - _, token_idx = torch.where(expert_mask[expert_idx[0]]) - current_state = hidden_states[token_idx] - gate = self.gate_proj[expert_idx](current_state) - up = self.up_proj[expert_idx](current_state) - gated_output = up * self.act_fn(gate) - out = self.down_proj[expert_idx](gated_output) - weighted_output = out * routing_weights[token_idx, expert_idx, None] - next_states.index_add_(0, token_idx, weighted_output.to(hidden_states.dtype)) - next_states = next_states.view(batch_size, -1, self.hidden_size) - - return next_states - - def _get_fused_expert_intermediate_dim(module): """Resolve the intermediate (expert) dimension from a fused expert module. @@ -1170,27 +958,6 @@ def _is_quant_fused_experts_module(module): return isinstance(module, _QuantFusedExperts) -class _QuantDbrxFFN(_QuantSparseSequentialMoe): - @property - def num_experts(self): - return self.router.moe_num_experts - - @property - def top_k(self): - # In older transformers, top_k was stored on DbrxRouter as moe_top_k. - # In transformers 5.0, DbrxFFN stores it as a plain attribute (top_k). - if hasattr(self.router, "moe_top_k"): - return self.router.moe_top_k - return self.__dict__.get("top_k", 1) - - @top_k.setter - def top_k(self, value): - if hasattr(self.router, "moe_top_k"): - self.router.moe_top_k = value - else: - self.__dict__["top_k"] = value - - @contextmanager def patch_compressed_linear_loading(): """Context manager that patches CompressedLinear to survive custom ``_init_weights`` calls. @@ -1417,38 +1184,6 @@ def unpack_weight(self): del self.weight_scale_inv -try: - from transformers.models.llama4.modeling_llama4 import Llama4TextExperts - - if Llama4TextExperts not in QuantModuleRegistry: - QuantModuleRegistry.register({Llama4TextExperts: "hf.Llama4TextExperts"})( - _QuantLlama4TextExperts - ) -except ImportError: - pass - -try: - from transformers.models.dbrx.modeling_dbrx import DbrxExpertGLU, DbrxExperts, DbrxFFN - - if DbrxExperts not in QuantModuleRegistry: - QuantModuleRegistry.register({DbrxExperts: "hf.DbrxExperts"})(_QuantDbrxExperts) - - if DbrxExpertGLU not in QuantModuleRegistry: - QuantModuleRegistry.register({DbrxExpertGLU: "hf.DbrxExpertGLU"})(_QuantDbrxExpertGLU) - - if DbrxFFN not in QuantModuleRegistry: - QuantModuleRegistry.register({DbrxFFN: "hf.DbrxFFN"})(_QuantDbrxFFN) -except ImportError: - pass - -try: - from transformers.models.falcon.modeling_falcon import FalconLinear - - if FalconLinear not in QuantModuleRegistry: - QuantModuleRegistry.register({FalconLinear: "hf.FalconLinear"})(_QuantLinear) -except ImportError: - pass - try: from compressed_tensors.linear.compressed_linear import CompressedLinear @@ -1459,30 +1194,6 @@ def unpack_weight(self): except ImportError: pass -try: - from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeTextExperts - - # transformers>=5.12 rewrote Qwen3VLMoeTextExperts onto the standard - # ``@use_experts_implementation`` fused layout: ``hidden_size``/``expert_dim`` became - # ``hidden_dim``/``intermediate_dim``, ``gate_up_proj`` was transposed to - # (num_experts, 2*intermediate_dim, hidden_dim), and the forward now calls ``F.linear`` - # twice per expert. ``_QuantQwen3VLMoeTextExperts`` only understands the older layout, - # so registering it against the new one crashes on ``self.hidden_size`` (nvbug 6518551). - # The decorator sets ``_apply_gate`` on the class; use it to detect the new layout and - # leave those modules to ``register_fused_experts_on_the_fly``, which claims them with - # the generic ``_QuantFusedExperts``. The old layout must stay explicitly registered: - # it is structurally indistinguishable from a generic fused-experts module, yet its - # forward uses ``torch.bmm``/``@`` rather than ``F.linear``, so the generic wrapper - # would silently quantize nothing. - if Qwen3VLMoeTextExperts not in QuantModuleRegistry and not hasattr( - Qwen3VLMoeTextExperts, "_apply_gate" - ): - QuantModuleRegistry.register({Qwen3VLMoeTextExperts: "hf.Qwen3VLMoeTextExperts"})( - _QuantQwen3VLMoeTextExperts - ) -except ImportError: - pass - try: from transformers.integrations.finegrained_fp8 import FP8Linear @@ -1492,126 +1203,6 @@ def unpack_weight(self): pass -class _QuantGptOssExperts(_TransposedExpertsCalibMixin, _QuantFunctionalMixin): - """Quantized wrapper for `transformers.GptOssExperts`. - - Quantizes `gate_up_proj` and `down_proj` weights via dynamic attributes inside `quantize_weight()`. - Activations into `gate_up_proj` are quantized by `gate_up_proj_input_quantizer`. For `down_proj` - activation quantization, we intercept `torch.Tensor.__matmul__`/`torch.bmm` and quantize inputs - on every second call (since the first call computes `gate_up_proj` outputs and second call - computes `down_proj` outputs). - """ - - @staticmethod - def _get_quantized_weight(quantizer, module, weight): - # MoE weight is accessed for each expert in one forward pass. so lets cache it - if module._enable_weight_quantization: - if hasattr(quantizer, "_cached_quant_val"): - return getattr(quantizer, "_cached_quant_val") - quantizer._cached_quant_val = _transposed_quantize(weight, quantizer) - return quantizer._cached_quant_val - return weight - - def _setup_for_weight_quantization(self): - self._register_dynamic_attribute( - "gate_up_proj", partial(self._get_quantized_weight, self.gate_up_proj_weight_quantizer) - ) - self._register_dynamic_attribute( - "down_proj", partial(self._get_quantized_weight, self.down_proj_weight_quantizer) - ) - - def _setup(self): - assert not hasattr(self, "kernel_layer_name"), ( - "ModelOpt quantization does not support patched forward for kernel_hub" - ) - 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() - - self._register_temp_attribute("_enable_weight_quantization", False) - self._register_temp_attribute("_down_proj_mul", False) - self._setup_for_weight_quantization() - - @property - def functionals_to_replace(self): - # Use torch.ops.aten to bypass Python dispatch and avoid RecursionError - # (torch.matmul / __matmul__ can dispatch to each other) - _aten_bmm = torch.ops.aten.bmm - _aten_matmul = torch.ops.aten.matmul - - def _quantized_bmm(batch1, batch2, *, out=None): - batch1 = self.down_proj_input_quantizer(batch1) if self._down_proj_mul else batch1 - self._down_proj_mul = not self._down_proj_mul # toggle the flag - if out is not None: - return torch.ops.aten.bmm.out(batch1, batch2, out=out) - return _aten_bmm(batch1, batch2) - - def _tensor_matmul(self_t, other): - self_t = self.down_proj_input_quantizer(self_t) if self._down_proj_mul else self_t - self._down_proj_mul = not self._down_proj_mul - return _aten_matmul(self_t, other) - - return [ - (torch, "bmm", _quantized_bmm), - (torch.Tensor, "__matmul__", _tensor_matmul), - ] - - @contextmanager - def quantize_weight(self): - """Context in which MoE weight is quantized.""" - self._enable_weight_quantization = True - try: - yield - finally: - for module in self.modules(): - if isinstance(module, TensorQuantizer) and hasattr(module, "_cached_quant_val"): - delattr(module, "_cached_quant_val") - self._enable_weight_quantization = False - - def forward( - self, hidden_states: torch.Tensor, router_indices=None, routing_weights=None - ) -> torch.Tensor: - """Forward method to add quantization.""" - hidden_states = self.gate_up_proj_input_quantizer(hidden_states) - with self.quantize_weight(): - return super().forward(hidden_states, router_indices, routing_weights) - - -try: - from transformers.models.gpt_oss.modeling_gpt_oss import GptOssExperts - - if GptOssExperts not in QuantModuleRegistry: - QuantModuleRegistry.register({GptOssExperts: "hf.GptOssExperts"})(_QuantGptOssExperts) -except ImportError: - pass - - -def register_dbrx_moe_on_the_fly(model): - """Register DBRX MoE modules as QUANT_MODULE. - - The MoE class in DBRX is `transformers_modules.modeling_dbrx.DbrxExpertGLU`, which loads dynamically. - """ - if type(model).__name__ == "DbrxForCausalLM": - moe_type = type(model.transformer.blocks[0].ffn.experts.mlp) - # Create a QuantDbrxExpertGLU class on the fly - if QuantModuleRegistry.get(moe_type) is None: - QuantModuleRegistry.register({moe_type: moe_type.__name__})(_QuantDbrxExpertGLU) - - -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") @@ -1794,27 +1385,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) @@ -1892,249 +1465,29 @@ 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//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 ( + "dbrx", + "falcon", + "gpt_oss", + "llama4", + "nemotron_h", + "qwen3_vl_moe", + "step3p5", +): + 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 ) -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 = str(getattr(getattr(model, "config", None), "model_type", "") 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_falcon_linears_on_the_fly, - register_dbrx_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/model_type/README.md b/modelopt_recipes/model_type/README.md index 88f315ad4cb..14dda421310 100644 --- a/modelopt_recipes/model_type/README.md +++ b/modelopt_recipes/model_type/README.md @@ -25,9 +25,11 @@ Built-in recipes live in three tiers — pick the most specific that applies: ## Folder structure Recipes are categorized by the Hugging Face `model_type` string — the -value of the top-level `model_type` field in the model's `config.json` -(or, for multimodal configs, the `text_config.model_type` of the inner -language model). Use the exact `model_type` as the directory name: +value of the top-level `model_type` field in the model's `config.json`. +For multimodal models, use this top-level value, not the inner language +model's `text_config.model_type`: sub-config types differ from the +checkpoint's and change across transformers versions. Use the exact +`model_type` as the directory name: ```text modelopt_recipes/model_type/ diff --git a/plugins/modelopt/skills/ptq/references/checkpoint-validation.md b/plugins/modelopt/skills/ptq/references/checkpoint-validation.md index 92ea1f9fc59..a60cb2c9dc5 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 2d536af6aa3..ce9f574ef6c 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): - **transformers >= 5.0**: Unified fused experts (`gate_up_proj` + `down_proj` 3D tensors) → auto-detected by `register_fused_experts_on_the_fly`, handled by `_QuantFusedExperts`. Covers Mixtral, Qwen, DeepSeek, Jamba, OlMoE, etc. - **transformers < 5.0**: Sequential per-expert `nn.Linear` with `gate` + `experts` → 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 | | --- | --- | --- | @@ -164,7 +164,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 @@ -180,6 +180,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/examples/hf_ptq/test_hf_ptq_model_type.py b/tests/examples/hf_ptq/test_hf_ptq_model_type.py new file mode 100644 index 00000000000..3fa73245335 --- /dev/null +++ b/tests/examples/hf_ptq/test_hf_ptq_model_type.py @@ -0,0 +1,69 @@ +# 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. +"""hf_ptq keys its model-specific behavior on the root model's Hugging Face ``model_type``.""" + +import importlib +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +_EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" + + +@pytest.fixture +def hf_ptq(monkeypatch: pytest.MonkeyPatch): + monkeypatch.syspath_prepend(str(_EXAMPLES_DIR)) + return importlib.import_module("hf_ptq") + + +@pytest.mark.parametrize( + ("model_type", "skips_generation"), + [ + ("deepseek_v3", True), + ("deepseek_v32", True), + ("kimi_k2", True), # DeepSeek-V3 architecture under its own model_type + ("kimi_k25", True), # VLM wrapping a kimi_k2 language model + ("nemotron_h", True), + ("llama", False), + ], +) +def test_pre_quantize_generation_skip(hf_ptq, model_type: str, skips_generation: bool): + generate_calls = [] + full_model = SimpleNamespace(generate=lambda *args, **kwargs: generate_calls.append(1) or "out") + args = SimpleNamespace(specdec_offline_dataset=None, skip_generate=False) + batch = {"input_ids": torch.ones(2, 4, dtype=torch.long)} + + _, _, generated = hf_ptq.pre_quantize(args, full_model, model_type, None, [batch], False) + + assert (generated is None) == skips_generation + assert len(generate_calls) == (0 if skips_generation else 1) + + +@pytest.mark.parametrize( + ("model_type", "expected"), + [ + ("t5", True), + ("mt5", True), + ("mbart", True), + ("whisper", True), + ("t5gemma", False), + ("florence2", False), # VLM whose language model is BART; the root type decides + ("qwen3_vl", False), + ], +) +def test_is_enc_dec(hf_ptq, model_type: str, expected: bool): + assert hf_ptq.is_enc_dec(model_type) is expected diff --git a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py index 904f0128729..4a7ec7c09d7 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py +++ b/tests/examples/hf_ptq/test_hf_ptq_vision_quantization.py @@ -63,7 +63,7 @@ def test_image_calibration_model_target_follows_recipe( ) monkeypatch.setattr(hf_ptq, "get_model", lambda *args, **kwargs: full_model) - monkeypatch.setattr(hf_ptq, "get_model_type", lambda model: "qwen3_vl") + monkeypatch.setattr(hf_ptq, "hf_model_type", lambda model: "qwen3_vl") monkeypatch.setattr(hf_ptq, "is_nemotron_vl", lambda model: False) monkeypatch.setattr( hf_ptq.AutoProcessor, @@ -73,7 +73,7 @@ def test_image_calibration_model_target_follows_recipe( def extract_language_model(model): extraction_calls.append(model) - return extracted_language_model, "qwen3" + return extracted_language_model monkeypatch.setattr( hf_ptq, "extract_and_prepare_language_model_from_vl", extract_language_model diff --git a/tests/unit/torch/export/trtllm/test_trtllm_decoder_type.py b/tests/unit/torch/export/trtllm/test_trtllm_decoder_type.py new file mode 100644 index 00000000000..65289754016 --- /dev/null +++ b/tests/unit/torch/export/trtllm/test_trtllm_decoder_type.py @@ -0,0 +1,84 @@ +# 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. +"""Unit tests for :mod:`modelopt.torch.export.trtllm.decoder_type` and its deprecated aliases.""" + +import pytest +import torch.nn as nn + +import modelopt.torch.export as mte +from modelopt.torch.export import model_utils +from modelopt.torch.export.trtllm import model_config_export +from modelopt.torch.export.trtllm.decoder_type import MODEL_NAME_TO_DECODER_TYPE, get_decoder_type + + +def _named_module(name: str) -> nn.Module: + return type(name, (nn.Module,), {})() + + +@pytest.mark.parametrize( + ("class_name", "expected"), + [ + ("LlamaForCausalLM", "llama"), + ("Llama4ForConditionalGeneration", "llama4"), + ("DiffusionGemmaForCausalLM", "diffusion_gemma"), + ("Gemma3ForCausalLM", "gemma3"), + ("WhisperForConditionalGeneration", "whisper"), + ("DeepseekV3ForCausalLM", "deepseek"), + ("UnknownForCausalLM", None), + ], +) +def test_get_decoder_type(class_name: str, expected: str | None): + assert get_decoder_type(_named_module(class_name)) == expected + + +def test_get_model_type_is_deprecated(): + with pytest.warns(DeprecationWarning, match="get_model_type"): + assert mte.get_model_type(_named_module("LlamaForCausalLM")) == "llama" + + +def test_model_name_to_type_is_deprecated(): + with pytest.warns(DeprecationWarning, match="MODEL_NAME_TO_TYPE"): + assert model_utils.MODEL_NAME_TO_TYPE is MODEL_NAME_TO_DECODER_TYPE + + +def test_exporter_detects_decoder_type(monkeypatch: pytest.MonkeyPatch): + seen = {} + + def fake_dtype(model: nn.Module): + raise RuntimeError("stop after decoder_type resolution") + + monkeypatch.setattr(model_config_export, "get_dtype", fake_dtype) + monkeypatch.setattr( + model_config_export, + "get_decoder_type", + lambda model: seen.setdefault("decoder_type", get_decoder_type(model)), + ) + with pytest.raises(RuntimeError, match="stop after"): + next( + model_config_export._torch_to_tensorrt_llm_checkpoint(_named_module("GPT2LMHeadModel")) + ) + assert seen["decoder_type"] == "gpt" + + +def test_exporter_falls_back_for_undetectable_decoder_type(monkeypatch: pytest.MonkeyPatch): + def fake_dtype(model: nn.Module): + raise RuntimeError("stop after decoder_type resolution") + + monkeypatch.setattr(model_config_export, "get_dtype", fake_dtype) + with ( + pytest.warns(UserWarning, match="Unknown decoder_type for Unknown"), + pytest.raises(RuntimeError, match="stop after"), + ): + next(model_config_export._torch_to_tensorrt_llm_checkpoint(_named_module("Unknown"))) diff --git a/tests/unit/torch/models/test_model_specs.py b/tests/unit/torch/models/test_model_specs.py index 74bc6a09102..bfe580955b7 100644 --- a/tests/unit/torch/models/test_model_specs.py +++ b/tests/unit/torch/models/test_model_specs.py @@ -585,6 +585,7 @@ def test_fused_experts_shortcut_ignores_unrelated_attributes(): ("gate_proj", "up_proj"), ), ("gpt_oss", ("GptOssMLP", "GptOssMoE"), ("gate_up_proj", "down_proj"), True, None), + ("llama4", ("Llama4TextMoe",), ("gate_up_proj", "down_proj"), True, None), ("minimax", ("MiniMaxSparseMoeBlock",), ("w1", "w2", "w3"), False, ("w1", "w3")), ("mixtral", ("MixtralSparseMoeBlock",), ("w1", "w2", "w3"), False, ("w1", "w3")), ("phimoe", ("PhimoeSparseMoeBlock",), ("w1", "w2", "w3"), False, ("w1", "w3")), @@ -617,6 +618,13 @@ def test_fused_experts_shortcut_ignores_unrelated_attributes(): False, ("gate_proj", "up_proj"), ), + ( + "qwen3_vl_moe", + ("Qwen3VLMoeTextSparseMoeBlock",), + ("gate_up_proj", "down_proj"), + True, + None, + ), ] @@ -751,6 +759,7 @@ def test_iterable_experts_matches_pre_refactor_support(): "gemma4": "Gemma4ForConditionalGeneration", "gemma4_text": "Gemma4TextForCausalLM", "gpt_oss": "GptOssForCausalLM", + "llama4": "Llama4ForConditionalGeneration", "minimax": "MiniMaxForCausalLM", "mixtral": "MixtralForCausalLM", "phimoe": "PhimoeForCausalLM", @@ -759,6 +768,7 @@ def test_iterable_experts_matches_pre_refactor_support(): "qwen3_5_moe": "Qwen3_5MoeForCausalLM", "qwen3_moe": "Qwen3MoeForCausalLM", "qwen3_next": "Qwen3NextForCausalLM", + "qwen3_vl_moe": "Qwen3VLMoeForConditionalGeneration", } moe_specs = [s for s in get_specs() if _layouts(s)] assert {s.model_type for s in moe_specs} == set(root_class_names), ( diff --git a/tests/unit/torch/quantization/plugins/test_fused_experts.py b/tests/unit/torch/quantization/plugins/test_fused_experts.py index 73db66a539a..113dfc0eaa9 100644 --- a/tests/unit/torch/quantization/plugins/test_fused_experts.py +++ b/tests/unit/torch/quantization/plugins/test_fused_experts.py @@ -1529,7 +1529,7 @@ def _experts_type(): def test_registration_matches_installed_layout(self): """transformers>=5.12 experts must be left to the generic fused-experts wrapper.""" - from modelopt.torch.quantization.plugins.huggingface import _QuantQwen3VLMoeTextExperts + from modelopt.torch.models.qwen3_vl_moe.modeling_ptq import _QuantQwen3VLMoeTextExperts registered = QuantModuleRegistry.get(self._experts_type()) if _qwen3_vl_moe_is_new_layout(): diff --git a/tests/unit/torch/quantization/plugins/test_moe_linear.py b/tests/unit/torch/quantization/plugins/test_moe_linear.py index cceaacd44d5..66b52325c24 100644 --- a/tests/unit/torch/quantization/plugins/test_moe_linear.py +++ b/tests/unit/torch/quantization/plugins/test_moe_linear.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