diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df3cdefdb5e..ce1f493a28a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,7 @@ Changelog **New Features** +- Add NVFP4 and FP8 PTQ recipes with projection-output quantizers for Llama-Nemotron embedding and reranking models (``modelopt_recipes/huggingface/nemotron_llama/``) and an end-to-end HF embedding/reranking quantize-to-ONNX example (``examples/torch_onnx/hf_embedding_quant_to_onnx.py``). Quantizing the projection-Linear outputs keeps TensorRT inter-layer activations in FP4, roughly halving engine activation memory versus the plain ``nvfp4`` preset. NVFP4/MXFP8 output quantizers now export through the dynamic quantize path. ``examples/torch_onnx/torch_quant_to_onnx.py`` also gains a ``--recipe`` flag to load quantization configs from YAML recipes instead of the removed ``mtq.*_CFG`` module-constant table. - Add the ``prepare_megatron_data_blend`` utility to prepare weighted Megatron data blends from YAML configs, including optional token-budgeted subsets for distillation workflows. See the `Megatron data preparation guide `_. - Add Learned Scale Quantization (LSQ) and Dual-LSQ support for quantization-aware distillation, including learnable ``amax`` parameters, tied-scale and pre-scale options, focused NVFP4 recipes, and scale-only training. - Add the **D-PACE** loss objective for DFlash speculative-decoding training (`arXiv:2605.18810 `_) and make it the default (``dflash_loss_objective: dpace``). It replaces the static exponential position decay with dynamic, confidence-derived per-position weights that adapt to whichever block positions currently limit acceptance. Smoothing is controlled by ``dflash_dpace_alpha`` (default 0.5); set ``dflash_loss_objective: decay`` to restore the previous static schedule. Training-only and detached from the gradient (no architecture or inference change). diff --git a/examples/torch_onnx/README.md b/examples/torch_onnx/README.md index 06ad11164dd..3e7a4d4eaca 100644 --- a/examples/torch_onnx/README.md +++ b/examples/torch_onnx/README.md @@ -69,6 +69,11 @@ python torch_quant_to_onnx.py \ --onnx_save_path= ``` +Quantization configs are loaded from the YAML preset recipes under +`modelopt_recipes/configs/ptq/presets/model/`, selected by `--quantize_mode`. Pass +`--recipe=` to use a different +recipe (e.g. `--recipe=nvfp4_awq_lite` or `--recipe=/path/to/my_quant_cfg.yaml`). + ### Conv2d Quantization Override TensorRT only supports FP8 and INT8 for convolution operations. When quantizing models with Conv2d layers (like SwinTransformer), the script automatically applies the following overrides: @@ -93,6 +98,68 @@ python ../onnx_ptq/evaluate.py \ --model_name= ``` +## HF Embedding and Reranking Models + +`hf_embedding_quant_to_onnx.py` quantizes an HF text-embedding or reranking +model (bidirectional Llama encoders such as +[nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) +and +[nvidia/llama-nemotron-rerank-1b-v2](https://huggingface.co/nvidia/llama-nemotron-rerank-1b-v2)) +with a PTQ recipe and exports it to ONNX. Embedding models are exported with +mean pooling and L2 normalization on top of the encoder; reranking +(sequence-classification) models are exported to their relevance logits. Both +graphs take `input_ids` and `attention_mask` with dynamic batch/sequence axes. + +The default recipe +(`modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml`) +quantizes weights and activations to NVFP4 and additionally quantizes the +projection-Linear outputs. Without output-side quantization, quantized GEMMs +emit FP16 activations, so FP8/FP4 engines can use as much or more activation +memory than an unquantized FP16 engine; quantizing the projection outputs keeps +inter-layer activations in the low-precision format. An FP8 twin of the recipe +(`fp8_output_quant_proj.yaml`, pass it via `--recipe`) applies the same idea to +the FP8 preset. With TensorRT 10.16 on RTX PRO 6000 Blackwell (strongly-typed +engines, 5 dynamic-shape profiles up to 32x512), engine activation memory: + +| Model | FP16 | `fp8` preset | fp8 recipe | `nvfp4` preset | nvfp4 recipe | +|-------|-----:|-------------:|-----------:|---------------:|-------------:| +| llama-nemotron-embed-1b-v2 | 1040 MiB | 1392 MiB | 1096 MiB | 1040 MiB | 516 MiB | +| llama-nemotron-rerank-1b-v2 | 1040 MiB | 1392 MiB | 1096 MiB | 520 MiB | 331 MiB | + +### Usage + +```bash +python hf_embedding_quant_to_onnx.py \ + --model_path=nvidia/llama-nemotron-embed-1b-v2 \ + --recipe=huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj \ + --onnx_save_path=llama_nemotron_embed_nvfp4.onnx + +# Reranking variant (auto-detected from the model architecture) +python hf_embedding_quant_to_onnx.py \ + --model_path=nvidia/llama-nemotron-rerank-1b-v2 \ + --onnx_save_path=llama_nemotron_rerank_nvfp4.onnx +``` + +### Building a TensorRT engine with trtexec + +NVFP4 requires a Blackwell GPU (SM100+) and TensorRT 10.11 or later. Build a +strongly-typed engine with dynamic shapes (add optimization profiles matching +your serving batch sizes and sequence lengths): + +```bash +trtexec --onnx=llama_nemotron_embed_nvfp4.onnx \ + --stronglyTyped \ + --saveEngine=llama_nemotron_embed_nvfp4.plan \ + --minShapes=input_ids:1x2,attention_mask:1x2 \ + --optShapes=input_ids:32x128,attention_mask:32x128 \ + --maxShapes=input_ids:32x512,attention_mask:32x512 +``` + +The exported `.onnx` references a sibling weights file (`.onnx_data`); +keep the two files in the same directory when building. To inspect the chosen +kernels and per-profile activation memory, add +`--profilingVerbosity=detailed --exportLayerInfo=.json --verbose`. + ## LLM Quantization and Export with TensorRT-Edge-LLM [TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM) provides a complete pipeline for quantizing LLMs and VLMs using NVIDIA ModelOpt and exporting them to optimized ONNX for deployment on edge platforms such as NVIDIA Jetson and DRIVE. @@ -283,7 +350,7 @@ The `auto` mode enables mixed precision quantization by searching for the optima python torch_quant_to_onnx.py \ --timm_model_name=vit_base_patch16_224 \ --quantize_mode=auto \ - --auto_quantization_formats NVFP4_AWQ_LITE_CFG FP8_DEFAULT_CFG \ + --auto_quantization_formats nvfp4_awq_lite fp8 \ --effective_bits=4.8 \ --num_score_steps=128 \ --calibration_data_size=512 \ diff --git a/examples/torch_onnx/hf_embedding_quant_to_onnx.py b/examples/torch_onnx/hf_embedding_quant_to_onnx.py new file mode 100644 index 00000000000..852cb446410 --- /dev/null +++ b/examples/torch_onnx/hf_embedding_quant_to_onnx.py @@ -0,0 +1,313 @@ +# 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. + +import argparse +import os + +import torch +import torch.nn.functional as F +from transformers import AutoConfig, AutoModel, AutoModelForSequenceClassification, AutoTokenizer + +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata + +""" +Quantize an HF embedding or reranking model (bidirectional Llama encoder, e.g. +nvidia/llama-nemotron-embed-1b-v2 or nvidia/llama-nemotron-rerank-1b-v2) with a +PTQ recipe and export it to ONNX for TensorRT deployment. + +The default recipe quantizes weights and activations to NVFP4 and additionally +quantizes the projection-Linear outputs so the TensorRT engine keeps +inter-layer activations in FP4 (see modelopt_recipes/huggingface/nemotron_llama/). +Embedding models are exported with mean pooling and L2 normalization on top of +the encoder; reranking (sequence-classification) models are exported to their +relevance logits. Both graphs take (input_ids, attention_mask) with dynamic +batch/sequence axes. +""" + +DEFAULT_RECIPE = "huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj" + +CALIBRATION_TEXTS = [ + ("What is the capital of France?", "Paris is the capital and most populous city of France."), + ("How do vaccines work?", "Vaccines train the immune system to recognize pathogens."), + ("Who wrote Hamlet?", "Hamlet is a tragedy written by William Shakespeare around 1600."), + ("What causes rain?", "Rain forms when water vapor condenses into droplets that fall."), + ( + "What is photosynthesis?", + "Plants convert light energy into chemical energy in chloroplasts.", + ), + ("How fast is light?", "Light travels at approximately 299,792 kilometers per second."), + ("What is machine learning?", "Machine learning builds models that learn patterns from data."), + ("Where is the Great Barrier Reef?", "The reef lies off the coast of Queensland, Australia."), + ("What is inflation?", "Inflation is the rate at which prices for goods and services rise."), + ("How do airplanes fly?", "Wings generate lift as air flows faster over their curved tops."), + ("What is DNA?", "DNA carries genetic instructions for development and functioning."), + ("Who painted the Mona Lisa?", "Leonardo da Vinci painted the Mona Lisa in the early 1500s."), + ("What is quantum computing?", "Quantum computers use qubits that exist in superposition."), + ("Why is the sky blue?", "Rayleigh scattering disperses shorter blue wavelengths of sunlight."), + ("What is the tallest mountain?", "Mount Everest rises 8,849 meters above sea level."), + ("How do batteries store energy?", "Batteries store energy in chemical form as electricity."), +] + + +class EmbeddingModel(torch.nn.Module): + """Bidirectional encoder + mean pooling + L2 normalization.""" + + def __init__(self, base): + super().__init__() + self.base = base + + def forward(self, input_ids, attention_mask): + hidden = self.base(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state + mask = attention_mask.unsqueeze(-1).to(hidden.dtype) + pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-6) + return F.normalize(pooled, p=2, dim=1) + + +class RerankModel(torch.nn.Module): + """Sequence-classification reranker exported to its relevance logits.""" + + def __init__(self, base): + super().__init__() + self.base = base + + def forward(self, input_ids, attention_mask): + return self.base(input_ids=input_ids, attention_mask=attention_mask).logits + + +def register_bidirectional_sdpa(): + """Export SDPA as a plain MatMul->Add->Softmax->MatMul pattern. + + torch's default sdpa symbolic emits IsNaN/Where guards for the dynamic + is_causal flag traced by transformers, which prevents TensorRT attention + fusion. This model is bidirectional, so is_causal is pinned to False. + """ + from torch.onnx import register_custom_op_symbolic, symbolic_helper + + from modelopt.torch.quantization.export_onnx import scaled_dot_product_attention + + def bidirectional_sdpa( + g, + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=None, + enable_gqa=False, + ): + if not isinstance(dropout_p, float): + dropout_p = symbolic_helper._maybe_get_const(dropout_p, "f") + # Keep the attention chain in a single precision: the symbolic emits the + # scale and boolean-mask constants in f32 otherwise, upcasting softmax. + q_dtype = query.type().scalarType() + if scale is not None and not symbolic_helper._is_none(scale) and q_dtype is not None: + scale = g.op("Cast", scale, to_i=symbolic_helper.cast_pytorch_to_onnx[q_dtype]) + if ( + attn_mask is not None + and not symbolic_helper._is_none(attn_mask) + and attn_mask.type().scalarType() == "Bool" + ): + mdtype = torch.float16 if q_dtype == "Half" else torch.float32 + zero = g.op("Constant", value_t=torch.tensor([0.0], dtype=mdtype)) + neg = g.op("Constant", value_t=torch.tensor([torch.finfo(mdtype).min], dtype=mdtype)) + attn_mask = g.op("Where", attn_mask, zero, neg) + return scaled_dot_product_attention( + g, query, key, value, attn_mask, dropout_p, False, scale, False + ) + + register_custom_op_symbolic("aten::scaled_dot_product_attention", bidirectional_sdpa, 14) + + # Repeat KV heads explicitly: the sdpa symbolic does not implement enable_gqa. + from transformers.integrations import sdpa_attention + + sdpa_attention.use_gqa_in_sdpa = lambda *args, **kwargs: False + + +def install_static_extent_fix(hidden_size): + """Give Reshape -> TRT_FP4DynamicQuantize a static blocked-axis extent. + + torch traces the attention-output reshape as Reshape(target=Concat(B, S, -1)); + TensorRT's DynamicQuantize requires the blocked (last) axis extent to be known + at build time, so the trailing -1 is rewired to the hidden size. + """ + import numpy as np + from onnx import numpy_helper + + import modelopt.torch._deploy.utils.torch_onnx as torch_onnx_utils + + orig_quantize_weights = torch_onnx_utils.quantize_weights + + def fix_reshapes(onnx_graph): + producers = {out: n for n in onnx_graph.graph.node for out in n.output} + inits = {i.name: i for i in onnx_graph.graph.initializer} + fixed = 0 + for node in onnx_graph.graph.node: + if node.op_type != "TRT_FP4DynamicQuantize": + continue + src = producers.get(node.input[0]) + while src is not None and src.op_type in ("Cast", "Identity"): + src = producers.get(src.input[0]) + if src is None or src.op_type != "Reshape": + continue + shape_src = producers.get(src.input[1]) + if shape_src is None or shape_src.op_type != "Concat": + continue + last = shape_src.input[-1] + arr = None + if last in inits: + arr = numpy_helper.to_array(inits[last]) + else: + cn = producers.get(last) + if cn is not None and cn.op_type == "Constant": + arr = numpy_helper.to_array(cn.attribute[0].t) + if arr is None or arr.size != 1 or int(arr.reshape(-1)[0]) != -1: + continue + new_init = numpy_helper.from_array( + np.array([hidden_size], dtype=np.int64), + name=f"{src.name or src.output[0]}_static_dim", + ) + onnx_graph.graph.initializer.append(new_init) + shape_src.input[len(shape_src.input) - 1] = new_init.name + fixed += 1 + print(f"Static-extent fix applied to {fixed} Reshape->DynamicQuantize sites") + return onnx_graph + + def quantize_weights_and_fix(model, onnx_graph): + return fix_reshapes(orig_quantize_weights(model, onnx_graph)) + + torch_onnx_utils.quantize_weights = quantize_weights_and_fix + + +def main(): + parser = argparse.ArgumentParser( + description="Quantize an HF embedding model with a PTQ recipe and export to ONNX." + ) + parser.add_argument( + "--model_path", + default="nvidia/llama-nemotron-embed-1b-v2", + help="HF hub id or local path of the embedding model.", + ) + parser.add_argument( + "--recipe", + default=DEFAULT_RECIPE, + help="PTQ recipe: a path under modelopt_recipes/ or a recipe YAML file.", + ) + parser.add_argument( + "--onnx_save_path", + required=True, + help="The save path for the exported ONNX model.", + ) + parser.add_argument( + "--calibration_data_size", + type=int, + default=64, + help="Number of calibration samples.", + ) + parser.add_argument( + "--batch_size", + type=int, + default=8, + help="Batch size for calibration.", + ) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True) + is_reranker = any("ForSequenceClassification" in a for a in (config.architectures or [])) + if is_reranker: + base = AutoModelForSequenceClassification.from_pretrained( + args.model_path, trust_remote_code=True, dtype=torch.float32 + ) + model = RerankModel(base).to(device).eval() + else: + base = AutoModel.from_pretrained( + args.model_path, trust_remote_code=True, dtype=torch.float32 + ) + model = EmbeddingModel(base).to(device).eval() + + pairs = (CALIBRATION_TEXTS * (args.calibration_data_size // len(CALIBRATION_TEXTS) + 1))[ + : args.calibration_data_size + ] + + def forward_loop(m): + for i in range(0, len(pairs), args.batch_size): + chunk = pairs[i : i + args.batch_size] + if is_reranker: + # Rerankers score (query, passage) pairs. + batch = tokenizer( + [q for q, _ in chunk], + [p for _, p in chunk], + padding=True, + truncation=True, + max_length=512, + return_tensors="pt", + ).to(device) + else: + batch = tokenizer( + [f"question:{q} \n \n passage:{p}" for q, p in chunk], + padding=True, + truncation=True, + max_length=512, + return_tensors="pt", + ).to(device) + with torch.no_grad(): + m(batch["input_ids"], batch["attention_mask"]) + + recipe = load_recipe(args.recipe) + quant_cfg = recipe.quantize.model_dump(exclude_unset=True) + mtq.quantize(model, quant_cfg, forward_loop=forward_loop) + mtq.print_quant_summary(model) + + # Unequal lengths so padding produces a real attention mask: transformers + # otherwise drops the all-ones mask and switches sdpa to native GQA, which + # the export symbolic does not support. + example = tokenizer( + ["example query one", "example query two"], + ["example passage one", "a considerably longer example passage two with extra words"], + padding=True, + return_tensors="pt", + ).to(device) + + register_bidirectional_sdpa() + install_static_extent_fix(base.config.hidden_size) + + model_name = os.path.basename(args.onnx_save_path).replace(".onnx", "") + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model=model, + dummy_input={ + "input_ids": example["input_ids"], + "attention_mask": example["attention_mask"], + }, + model_name=model_name, + weights_dtype="fp16", + dynamic_axes={ + "input_ids": {0: "batch", 1: "seq"}, + "attention_mask": {0: "batch", 1: "seq"}, + }, + ) + OnnxBytes.from_bytes(onnx_bytes).write_to_disk( + os.path.dirname(args.onnx_save_path) or ".", clean_dir=False + ) + print(f"Quantized ONNX model is saved to {args.onnx_save_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index 97bd0b60c18..923476a6e84 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/examples/torch_onnx/torch_quant_to_onnx.py @@ -14,13 +14,13 @@ # limitations under the License. import argparse -import copy import json import re import subprocess import sys import warnings from pathlib import Path +from typing import Any # Add onnx_ptq to path for shared modules sys.path.insert(0, str(Path(__file__).parent.parent / "onnx_ptq")) @@ -34,6 +34,9 @@ from evaluation import evaluate import modelopt.torch.quantization as mtq +from modelopt.recipe import load_config +from modelopt.recipe.presets import MODEL_QUANT_PRESET_DIR +from modelopt.torch.quantization.config import QuantizeConfig """ Quantize a timm vision model and export to ONNX for TensorRT deployment. @@ -54,13 +57,18 @@ mp.set_start_method("spawn", force=True) # Needed for data loader with multiple workers -QUANT_CONFIG_DICT: dict[str, dict] = { - "fp8": mtq.FP8_DEFAULT_CFG, - "int8": mtq.INT8_DEFAULT_CFG, - "mxfp8": mtq.MXFP8_DEFAULT_CFG, - "nvfp4": mtq.NVFP4_DEFAULT_CFG, - "int4_awq": mtq.INT4_AWQ_CFG, -} + +def load_quant_config(recipe: str) -> dict: + """Load a quantization config from a recipe YAML. + + ``recipe`` is either a preset basename under + ``modelopt_recipes/configs/ptq/presets/model/`` (e.g. ``nvfp4``) or a path to a + ``QuantizeConfig`` YAML file (filesystem or built-in recipe library). + """ + if "/" not in recipe and not recipe.endswith((".yml", ".yaml")): + recipe = f"{MODEL_QUANT_PRESET_DIR}/{recipe}" + return load_config(recipe, schema_type=QuantizeConfig).model_dump(exclude_unset=True) + _FP8_CONV_OVERRIDE: list = [ { @@ -104,27 +112,28 @@ }, ] -# Auto-quantize format configs that use block quantization and need Conv2d overrides for TRT. +# Auto-quantize format presets that use block quantization and need Conv2d overrides for TRT. # TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors. _NEEDS_FP8_CONV_OVERRIDE: set[str] = { - "NVFP4_AWQ_LITE_CFG", - "NVFP4_DEFAULT_CFG", - "MXFP8_DEFAULT_CFG", + "nvfp4_awq_lite", + "nvfp4", + "mxfp8", } -_NEEDS_INT8_CONV_OVERRIDE: set[str] = {"INT4_AWQ_CFG"} +_NEEDS_INT8_CONV_OVERRIDE: set[str] = {"int4_awq"} -def get_quant_config(quantize_mode): +def get_quant_config(quantize_mode, recipe=None): """Get quantization config, overriding Conv2d for TRT compatibility. - TensorRT only supports FP8 and INT8 for Conv layers. + The config is loaded from ``recipe`` when given, else from the preset YAML + matching ``quantize_mode``. TensorRT only supports FP8 and INT8 for Conv layers. - For FP8: add MHA-aware LayerNorm output quantizer so TRT fuses shared Q/DQ into downstream attention matmuls. Softmax-output Q/DQ is inserted by the FP8 ONNX exporter's post-processing (fixed 1/448 scale, no calibration needed). - For MXFP8, NVFP4: override Conv2d to FP8 - For INT4_AWQ: override Conv2d to INT8 """ - config: dict = copy.deepcopy(QUANT_CONFIG_DICT[quantize_mode]) + config: dict = load_quant_config(recipe or quantize_mode) if quantize_mode == "fp8": config["quant_cfg"].extend(_FP8_MHA_OVERRIDE) elif quantize_mode in ("mxfp8", "nvfp4"): @@ -366,7 +375,8 @@ def auto_quantize_model( Args: model: PyTorch model to quantize data_loader: DataLoader with image-label dict batches - quantization_formats: List of quantization format config names or dicts + quantization_formats: List of quantization recipe names (preset basenames, + e.g. ``nvfp4_awq_lite``) or config dicts effective_bits: Target effective bits constraint num_calib_steps: Number of calibration steps num_score_steps: Number of scoring steps for sensitivity analysis @@ -381,10 +391,10 @@ def auto_quantize_model( # TRT DynamicQuantize requires 2D/3D input, but Conv2d operates on 4D tensors. # By including the overrides in the format configs, the auto_quantize search # correctly accounts for Conv2d being FP8/INT8 in the effective_bits budget. - format_configs = [] + format_configs: list[dict[str, Any] | str] = [] for fmt in quantization_formats: if isinstance(fmt, str): - config = copy.deepcopy(getattr(mtq, fmt)) + config = load_quant_config(fmt) if fmt in _NEEDS_FP8_CONV_OVERRIDE: config["quant_cfg"].extend(_FP8_CONV_OVERRIDE) elif fmt in _NEEDS_INT8_CONV_OVERRIDE: @@ -445,6 +455,17 @@ def main(): default="mxfp8", help="Type of quantization to apply. Default is MXFP8.", ) + parser.add_argument( + "--recipe", + type=str, + default=None, + help=( + "Quantization config recipe: a preset basename under " + "modelopt_recipes/configs/ptq/presets/model/ or a path to a QuantizeConfig " + "YAML. Defaults to the preset matching --quantize_mode. Not supported with " + "--quantize_mode=auto." + ), + ) parser.add_argument( "--onnx_save_path", required=True, @@ -480,14 +501,14 @@ def main(): "--auto_quantization_formats", nargs="+", choices=[ - "NVFP4_AWQ_LITE_CFG", - "FP8_DEFAULT_CFG", - "MXFP8_DEFAULT_CFG", - "INT8_DEFAULT_CFG", - "INT4_AWQ_CFG", + "nvfp4_awq_lite", + "fp8", + "mxfp8", + "int8", + "int4_awq", ], - default=["NVFP4_AWQ_LITE_CFG", "FP8_DEFAULT_CFG"], - help="Quantization formats to search from for auto mode (e.g., NVFP4_AWQ_LITE_CFG FP8_DEFAULT_CFG)", + default=["nvfp4_awq_lite", "fp8"], + help="Quantization preset recipes to search from for auto mode (e.g., nvfp4_awq_lite fp8)", ) parser.add_argument( "--effective_bits", @@ -520,6 +541,12 @@ def main(): args = parser.parse_args() + if args.recipe and args.quantize_mode == "auto": + parser.error( + "--recipe is not supported with --quantize_mode=auto; " + "use --auto_quantization_formats instead." + ) + # Create model and move to appropriate device device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model_kwargs = json.loads(args.model_kwargs) if args.model_kwargs else {} @@ -568,7 +595,7 @@ def main(): # Note: MXFP8 is dynamic and does not need calibration itself, but when # Conv2d layers are overridden to FP8 (for TRT compatibility), those FP8 # quantizers require calibration data. - config = get_quant_config(args.quantize_mode) + config = get_quant_config(args.quantize_mode, args.recipe) data_loader = load_calibration_data( model, @@ -594,7 +621,7 @@ def main(): # no tactic for that 3-channel Q→Conv fusion. Skip for pure INT8 (unaffected). uses_fp8_conv_input = args.quantize_mode in ("fp8", "mxfp8", "nvfp4") or ( args.quantize_mode == "auto" - and any(fmt != "INT8_DEFAULT_CFG" for fmt in args.auto_quantization_formats) + and any(fmt != "int8" for fmt in args.auto_quantization_formats) ) if uses_fp8_conv_input: _disable_low_channel_conv_input_quantizers(quantized_model) diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index 7b42dce5782..e5778c3c96b 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -281,7 +281,7 @@ def scaled_dot_product_attention( if hasattr(torch.onnx, "_type_utils"): from torch.onnx._type_utils import JitScalarType else: # torch >= 2.9 - from torch.onnx._internal.torchscript_exporter import JitScalarType + from torch.onnx._internal.torchscript_exporter._type_utils import JitScalarType if hasattr(torch.onnx, "symbolic_opset14"): from torch.onnx.symbolic_opset14 import _attention_scale, _causal_attention_mask @@ -643,18 +643,19 @@ def configure_linear_module_onnx_quantizers(model): For modules with block quantization (NVFP4/MXFP8): - Weight quantizers use "static" export (TRT_FP4QDQ for NVFP4, DQ-only for MXFP8) - - Input/activation quantizers use "dynamic" export (TRT_FP4DynamicQuantize, etc.) + - Input/output activation quantizers use "dynamic" export (TRT_FP4DynamicQuantize, etc.) - This must be set for ALL modules with block quantization, not just nn.Linear, - because models like ResNet have non-Linear modules (e.g., MaxPool2d) with NVFP4/MXFP8 - input quantizers that would otherwise default to the static path and produce - TRT_FP4QDQ nodes on activations (which the NVFP4 exporter cannot handle). + This must be set for ALL quantizers on activations, not just nn.Linear inputs, + because non-Linear modules (e.g., ResNet's MaxPool2d) and enabled output + quantizers would otherwise default to the static path and produce TRT_FP4QDQ + nodes on activations (which the NVFP4 exporter cannot handle). """ sentinel = object() originals: list[tuple] = [] for _, module in model.named_modules(): for attr_name, new_value in ( ("input_quantizer", "dynamic"), + ("output_quantizer", "dynamic"), ("weight_quantizer", "static"), ): quantizer = getattr(module, attr_name, None) diff --git a/modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml b/modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml new file mode 100644 index 00000000000..f0e9881775a --- /dev/null +++ b/modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml @@ -0,0 +1,41 @@ +# 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. + +# Llama-Nemotron embedding/reranking PTQ recipe: FP8 W8A8 plus per-tensor FP8 +# output quantizers on the projection Linears, so TensorRT picks FP8-output +# GEMM kernels and carries FP8 activations between layers. Without output-side +# quantization, FP8 GEMMs emit FP16 and each input Q/DQ adds an FP8 copy on +# top of the FP16 activations, so the engine uses MORE activation memory than +# an unquantized FP16 engine. [5726458] + +imports: + base_disable_all: configs/ptq/units/base_disable_all + w8a8_fp8_fp8: configs/ptq/units/w8a8_fp8_fp8 + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + fp8: configs/numerics/fp8 + +metadata: + recipe_type: ptq + description: 'Llama-Nemotron embedding/reranking PTQ recipe (fp8 + projection output quantizers): same numerics as the general fp8 preset, plus per-tensor + FP8 output quantizers on the projection Linears so TensorRT engines carry FP8 activations between layers.' +quantize: + algorithm: max + quant_cfg: + - $import: base_disable_all + - $import: w8a8_fp8_fp8 + - $import: default_disabled_quantizers + - quantizer_name: '*_proj.output_quantizer' + cfg: + $import: fp8 diff --git a/modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml b/modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml new file mode 100644 index 00000000000..1edf942f2f2 --- /dev/null +++ b/modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml @@ -0,0 +1,45 @@ +# 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. + +# Llama-Nemotron embedding PTQ recipe: NVFP4 W4A4 plus dynamic NVFP4 output +# quantizers on the projection Linears (q/k/v/o_proj, gate/up/down_proj) so +# TensorRT keeps inter-layer activations in FP4 instead of FP16. [5726458] +# +# Output quantizers must stay scoped to GEMM outputs: a DynamicQuantize on +# non-GEMM outputs (embeddings, pooling) fails to compile in TensorRT. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + w4a4_nvfp4_nvfp4: configs/ptq/units/w4a4_nvfp4_nvfp4 + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + +metadata: + recipe_type: ptq + description: 'Llama-Nemotron embedding PTQ recipe (nvfp4 + projection output quantizers): same numerics as the general nvfp4 preset, plus dynamic NVFP4 + output quantizers on the projection Linears so TensorRT engines carry FP4 activations between layers.' +quantize: + algorithm: max + quant_cfg: + - $import: base_disable_all + - $import: w4a4_nvfp4_nvfp4 + - $import: default_disabled_quantizers + - quantizer_name: '*_proj.output_quantizer' + cfg: + $import: nvfp4 + # Reranker (sequence-classification) score head: keep in high precision like + # lm_head; its [1, hidden] weight also cannot be packed by the NVFP4 exporter. + - quantizer_name: '*score*' + enable: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index f24931aa89d..9d83bad9d80 100644 --- a/modelopt_recipes/ptq.md +++ b/modelopt_recipes/ptq.md @@ -230,7 +230,7 @@ that baseline. The deviations come in four kinds: | Kind | What changes vs. the general recipe | Examples | |------|-------------------------------------|----------| -| **Architecture-aware `quant_cfg`** | Per-sub-module format choices a single wildcard scheme can't express | `qwen3_5`, `qwen3_5_moe`, `vit` | +| **Architecture-aware `quant_cfg`** | Per-sub-module format choices a single wildcard scheme can't express | `qwen3_5`, `qwen3_5_moe`, `vit`, `nemotron_llama` | | **Algorithm override** | Same numerics & scope, but the *calibration algorithm* is tweaked because the default breaks or regresses | `gemma`, `gemma4`, `mpt` | | **Extra exclusions** | Adds disabled-quantizer patterns so non-language branches stay full precision | `nemotron_vl`, `phi4mm`, `diffusion_gemma` | | **Checkpoint mirror** | A mixed-precision map reproducing one published checkpoint exactly | `models/nvidia/Nemotron-3-*` | @@ -239,7 +239,7 @@ The numerics and standard exclusions are still inherited from `configs/` wherever possible — the model folder captures *only* the delta. Each `/` folder carries a `README.md` spelling out that delta. -### Architecture-aware `quant_cfg` — `qwen3_5`, `qwen3_5_moe`, `vit` +### Architecture-aware `quant_cfg` — `qwen3_5`, `qwen3_5_moe`, `vit`, `nemotron_llama` `huggingface/qwen3_5/ptq/w4a16_nvfp4-fp8_attn-kv_fp8_cast` (and its MoE twin, which shares the same `quant_cfg` snippet) is a **mixed scheme no single general @@ -271,6 +271,18 @@ and disables the output quantizers. *Why special:* the general LLM recipes never quantize the attention BMM inputs; for ViT the whole attention block runs in FP8 so Torch-TRT can compile it end-to-end. +**`nemotron_llama/ptq/{nvfp4,fp8}_output_quant_proj`** covers the Llama-Nemotron +embedding/reranking encoders (e.g. `llama-nemotron-embed-1b-v2`): same numerics +as the general `nvfp4`/`fp8` presets, plus **output quantizers on the projection +Linears** (`*_proj.output_quantizer`) so TensorRT engines carry inter-layer +activations in the low-precision format instead of FP16 — roughly half the +engine activation memory on these models. The sequence-classification `score` +head stays unquantized (final heads stay high precision like `lm_head`, and its +`[1, hidden]` weight cannot be packed by the NVFP4 exporter). *Why special:* the +general recipes never enable output quantizers, and the pattern must stay scoped +to GEMM outputs — a `DynamicQuantize` on non-GEMM outputs (embedding lookup, +pooling) fails to compile in TensorRT. + A lighter case: **`step3p5/Step3.5-Flash/ptq/nvfp4-mlp-only`** is close to `general/ptq/nvfp4_mlp_only` (NVFP4 on MoE/MLP weights+inputs, FP8 KV) but pinned to one released checkpoint and carrying instance-specific disables diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index 90a5f266546..fdbcece1c86 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -34,6 +34,7 @@ Gemma3Config, GptOssConfig, LlamaConfig, + LlamaForSequenceClassification, NemotronConfig, PreTrainedModel, Qwen3Config, @@ -595,6 +596,35 @@ def create_tiny_llama_dir( ) +def get_tiny_llama_seq_cls(**config_kwargs) -> PreTrainedModel: + set_seed(SEED) + kwargs = { + "dtype": torch.bfloat16, + "hidden_size": 32, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "max_position_embeddings": 32, + "vocab_size": 32, + "num_labels": 1, + "pad_token_id": 0, + } + kwargs.update(config_kwargs) + return LlamaForSequenceClassification(LlamaConfig(**kwargs)) + + +def create_tiny_llama_seq_cls_dir( + tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs +) -> Path: + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_llama_seq_cls", + get_tiny_llama_seq_cls, + with_tokenizer=with_tokenizer, + **config_kwargs, + ) + + ##### T5 ##### def get_tiny_t5(**config_kwargs) -> PreTrainedModel: set_seed(SEED) diff --git a/tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py b/tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py new file mode 100644 index 00000000000..245312c7c84 --- /dev/null +++ b/tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py @@ -0,0 +1,56 @@ +# 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. + + +import pytest +from _test_utils.examples.run_command import extend_cmd_parts, run_example_command +from _test_utils.torch.transformers_models import ( + create_tiny_llama_dir, + create_tiny_llama_seq_cls_dir, +) + +# Tiny stand-ins for the target architectures: a plain encoder for the embedding +# path and a sequence-classification model (with a `score` head, kept unquantized +# by the recipe) for the reranking path. +_MODEL_FACTORIES = { + "embedding": create_tiny_llama_dir, + "reranking": create_tiny_llama_seq_cls_dir, +} + +# Every projection output extent must be divisible by the NVFP4 block size (16) +# for the recipe's output quantizers: kv output = 2 kv-heads * head_dim 16 = 32. +_TINY_CONFIG = { + "hidden_size": 64, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "intermediate_size": 32, +} + + +@pytest.mark.parametrize("model_kind", list(_MODEL_FACTORIES)) +def test_hf_embedding_quant_to_onnx(tmp_path, model_kind): + model_dir = _MODEL_FACTORIES[model_kind](tmp_path, with_tokenizer=True, **_TINY_CONFIG) + onnx_save_path = tmp_path / f"{model_kind}_nvfp4.onnx" + + cmd_parts = extend_cmd_parts( + ["python", "hf_embedding_quant_to_onnx.py"], + model_path=str(model_dir), + onnx_save_path=str(onnx_save_path), + calibration_data_size="2", + batch_size="2", + ) + run_example_command(cmd_parts, "torch_onnx") + + assert onnx_save_path.exists() diff --git a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py index 6d6e0d9de57..1086e50d121 100644 --- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py +++ b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py @@ -46,3 +46,20 @@ def test_torch_onnx(model_key, quantize_mode): ) cmd_parts.extend(["--no_pretrained", "--trt_build"]) run_example_command(cmd_parts, "torch_onnx") + + +def test_torch_onnx_recipe_flag(): + timm_model_name, model_kwargs = _MODELS["vit_tiny"] + + cmd_parts = extend_cmd_parts( + ["python", "torch_quant_to_onnx.py"], + timm_model_name=timm_model_name, + model_kwargs=model_kwargs, + quantize_mode="nvfp4", + recipe="configs/ptq/presets/model/nvfp4", + onnx_save_path="vit_tiny.recipe.onnx", + calibration_data_size="1", + num_score_steps="1", + ) + cmd_parts.extend(["--no_pretrained", "--trt_build"]) + run_example_command(cmd_parts, "torch_onnx")