From 9d885670b475627e9f08df56817465fd2695b1e0 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 16 Jul 2026 00:41:40 +0000 Subject: [PATCH 1/6] Load quantization configs from YAML recipes in torch_quant_to_onnx example Replace the mtq.*_CFG module-constant table with recipe YAML loading and add a --recipe flag to select a preset basename or a QuantizeConfig YAML path. Auto-quantize format names switch to preset basenames accordingly. Co-Authored-By: Claude Fable 5 Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- examples/torch_onnx/torch_quant_to_onnx.py | 83 ++++++++++++------- .../torch_onnx/test_torch_quant_to_onnx.py | 17 ++++ 2 files changed, 72 insertions(+), 28 deletions(-) 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/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") From e4fcdaf914379ecfe1c53a2f6c126822823f0fe1 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 16 Jul 2026 00:42:03 +0000 Subject: [PATCH 2/6] [5726458] Add NVFP4 projection-output-quantizer recipe and HF embedding ONNX export example Without output-side quantization, TensorRT's NVFP4 GEMMs emit FP16 activations, so FP4 engines can use more activation memory than FP8 engines whose output-side Q/DQ enables FP8-out kernels. Quantizing the projection outputs keeps inter-layer activations in FP4, roughly halving engine activation memory on the Llama-Nemotron embedding model. - modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml: nvfp4 preset plus dynamic NVFP4 output quantizers scoped to the projection Linears (DynamicQuantize on non-GEMM outputs fails to compile in TensorRT). - examples/torch_onnx/hf_embedding_quant_to_onnx.py: end-to-end recipe-driven quantize -> ONNX export for bidirectional Llama embedding encoders, with the export shims needed for a TensorRT-fusable graph (bidirectional sdpa symbolic, static blocked-axis extents for DynamicQuantize). - configure_linear_module_onnx_quantizers now types output quantizers as dynamic (the static path fails on activations in the NVFP4 exporter) and the sdpa symbolic's JitScalarType import is fixed for torch >= 2.11. - README section covering the example and trtexec engine-build steps. Co-Authored-By: Claude Fable 5 Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 1 + examples/torch_onnx/README.md | 54 +++- .../torch_onnx/hf_embedding_quant_to_onnx.py | 279 ++++++++++++++++++ modelopt/torch/quantization/export_onnx.py | 6 +- .../ptq/nvfp4_output_quant_proj.yaml | 41 +++ 5 files changed, 378 insertions(+), 3 deletions(-) create mode 100644 examples/torch_onnx/hf_embedding_quant_to_onnx.py create mode 100644 modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df3cdefdb5e..b83c6ec0258 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,7 @@ Changelog **New Features** +- Add an NVFP4 PTQ recipe with projection-output quantizers for Llama-Nemotron embedding models (``modelopt_recipes/huggingface/nemotron_llama/``) and an end-to-end HF embedding 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..7bed690e514 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,53 @@ python ../onnx_ptq/evaluate.py \ --model_name= ``` +## HF Embedding Models + +`hf_embedding_quant_to_onnx.py` quantizes an HF text-embedding model +(a bidirectional Llama encoder such as +[nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2)) +with a PTQ recipe and exports it to ONNX. The exported graph wraps the encoder +with mean pooling and L2 normalization and takes `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, TensorRT's NVFP4 +GEMMs emit FP16 activations, so FP4 engines can use more activation memory than +FP8 engines whose output-side Q/DQ enables FP8-out kernels; quantizing the +projection outputs keeps inter-layer activations in FP4 and roughly halves the +engine's activation memory relative to the plain `nvfp4` preset on this model. + +### 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 +``` + +### 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 +335,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..786ef4fb063 --- /dev/null +++ b/examples/torch_onnx/hf_embedding_quant_to_onnx.py @@ -0,0 +1,279 @@ +# 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 AutoModel, 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 model (bidirectional Llama encoder, e.g. +nvidia/llama-nemotron-embed-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/). +The exported graph wraps the encoder with mean pooling and L2 normalization, +taking (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) + + +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 + 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 + ] + texts = [f"question:{q} \n \n passage:{p}" for q, p in pairs] + + def forward_loop(m): + for i in range(0, len(texts), args.batch_size): + batch = tokenizer( + texts[i : i + args.batch_size], + 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/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index 7b42dce5782..17398317010 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,7 +643,8 @@ 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.); the static path would fail on activations in the NVFP4 exporter. 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 @@ -655,6 +656,7 @@ def configure_linear_module_onnx_quantizers(model): 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/nvfp4_output_quant_proj.yaml b/modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml new file mode 100644 index 00000000000..fed09272a9c --- /dev/null +++ b/modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_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 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 From 34f1f461f33302f25c2e7b0a1ae2e920ea2cc4d2 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 16 Jul 2026 01:38:16 +0000 Subject: [PATCH 3/6] [5726458] Support the reranking model in the recipe and embedding example - Recipe: keep the sequence-classification score head unquantized; its [1, hidden] weight cannot be packed by the NVFP4 exporter and final heads stay in high precision like lm_head. No-op for the embedding model. - Example: auto-detect ForSequenceClassification architectures, load the reranker head, pair-encode calibration inputs, and export relevance logits. - README: cover both models and add TensorRT 10.16 activation-memory results. Co-Authored-By: Claude Fable 5 Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 2 +- examples/torch_onnx/README.md | 34 +++++++--- .../torch_onnx/hf_embedding_quant_to_onnx.py | 68 ++++++++++++++----- .../ptq/nvfp4_output_quant_proj.yaml | 4 ++ 4 files changed, 80 insertions(+), 28 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index b83c6ec0258..078c7c8032a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,7 +19,7 @@ Changelog **New Features** -- Add an NVFP4 PTQ recipe with projection-output quantizers for Llama-Nemotron embedding models (``modelopt_recipes/huggingface/nemotron_llama/``) and an end-to-end HF embedding 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 an NVFP4 PTQ recipe 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 7bed690e514..2f4ea6ee58b 100644 --- a/examples/torch_onnx/README.md +++ b/examples/torch_onnx/README.md @@ -98,14 +98,17 @@ python ../onnx_ptq/evaluate.py \ --model_name= ``` -## HF Embedding Models - -`hf_embedding_quant_to_onnx.py` quantizes an HF text-embedding model -(a bidirectional Llama encoder such as -[nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2)) -with a PTQ recipe and exports it to ONNX. The exported graph wraps the encoder -with mean pooling and L2 normalization and takes `input_ids` and -`attention_mask` with dynamic batch/sequence axes. +## 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`) @@ -113,8 +116,14 @@ quantizes weights and activations to NVFP4 and additionally quantizes the projection-Linear outputs. Without output-side quantization, TensorRT's NVFP4 GEMMs emit FP16 activations, so FP4 engines can use more activation memory than FP8 engines whose output-side Q/DQ enables FP8-out kernels; quantizing the -projection outputs keeps inter-layer activations in FP4 and roughly halves the -engine's activation memory relative to the plain `nvfp4` preset on this model. +projection outputs keeps inter-layer activations in FP4. With TensorRT 10.16 on +RTX PRO 6000 Blackwell (strongly-typed engines, 5 dynamic-shape profiles up to +32x512), engine activation memory vs the plain `nvfp4` preset: + +| Model | `nvfp4` preset | This recipe | +|-------|---------------:|------------:| +| llama-nemotron-embed-1b-v2 | 1040 MiB | 516 MiB | +| llama-nemotron-rerank-1b-v2 | 520 MiB | 331 MiB | ### Usage @@ -123,6 +132,11 @@ 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 diff --git a/examples/torch_onnx/hf_embedding_quant_to_onnx.py b/examples/torch_onnx/hf_embedding_quant_to_onnx.py index 786ef4fb063..852cb446410 100644 --- a/examples/torch_onnx/hf_embedding_quant_to_onnx.py +++ b/examples/torch_onnx/hf_embedding_quant_to_onnx.py @@ -18,22 +18,24 @@ import torch import torch.nn.functional as F -from transformers import AutoModel, AutoTokenizer +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 model (bidirectional Llama encoder, e.g. -nvidia/llama-nemotron-embed-1b-v2) with a PTQ recipe and export it to ONNX for -TensorRT deployment. +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/). -The exported graph wraps the encoder with mean pooling and L2 normalization, -taking (input_ids, attention_mask) with dynamic batch/sequence axes. +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" @@ -75,6 +77,17 @@ def forward(self, input_ids, attention_mask): 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. @@ -217,23 +230,44 @@ def main(): tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - base = AutoModel.from_pretrained(args.model_path, trust_remote_code=True, dtype=torch.float32) - model = EmbeddingModel(base).to(device).eval() + 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 ] - texts = [f"question:{q} \n \n passage:{p}" for q, p in pairs] def forward_loop(m): - for i in range(0, len(texts), args.batch_size): - batch = tokenizer( - texts[i : i + args.batch_size], - padding=True, - truncation=True, - max_length=512, - return_tensors="pt", - ).to(device) + 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"]) 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 index fed09272a9c..1edf942f2f2 100644 --- a/modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml +++ b/modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml @@ -39,3 +39,7 @@ quantize: - 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 From c7b6a5867abacb024d66e0d3e6174bd596f6759a Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 16 Jul 2026 02:02:11 +0000 Subject: [PATCH 4/6] [5726458] Add end-to-end example test for HF embedding and reranking export Tiny plain-Llama stand-ins (a causal encoder and a LlamaForSequenceClassification with the score head the recipe excludes) run hf_embedding_quant_to_onnx.py through quantize -> ONNX export for both model kinds. Projection output extents are sized to the NVFP4 block size (16). Co-Authored-By: Claude Fable 5 Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- .../_test_utils/torch/transformers_models.py | 30 ++++++++++ .../test_hf_embedding_quant_to_onnx.py | 56 +++++++++++++++++++ 2 files changed, 86 insertions(+) create mode 100644 tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py 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() From 58d942200708a0c34c832e07f11a5f14c27a233e Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 16 Jul 2026 03:01:33 +0000 Subject: [PATCH 5/6] [5726458] Add FP8 projection-output-quantizer recipe Applies the same output-side quantization to the FP8 preset: FP8-output GEMM kernels replace the FP16-out + shadow-copy pattern, cutting engine activation memory 21% vs the fp8 preset (1392 -> 1096 MiB on both Llama-Nemotron models, TensorRT 10.16, RTX PRO 6000). README gains the full FP16/FP8/NVFP4 table. Co-Authored-By: Claude Fable 5 Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- CHANGELOG.rst | 2 +- examples/torch_onnx/README.md | 23 ++++++----- .../ptq/fp8_output_quant_proj.yaml | 41 +++++++++++++++++++ 3 files changed, 54 insertions(+), 12 deletions(-) create mode 100644 modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 078c7c8032a..ce1f493a28a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,7 +19,7 @@ Changelog **New Features** -- Add an NVFP4 PTQ recipe 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 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 2f4ea6ee58b..3e7a4d4eaca 100644 --- a/examples/torch_onnx/README.md +++ b/examples/torch_onnx/README.md @@ -113,17 +113,18 @@ 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, TensorRT's NVFP4 -GEMMs emit FP16 activations, so FP4 engines can use more activation memory than -FP8 engines whose output-side Q/DQ enables FP8-out kernels; quantizing the -projection outputs keeps inter-layer activations in FP4. With TensorRT 10.16 on -RTX PRO 6000 Blackwell (strongly-typed engines, 5 dynamic-shape profiles up to -32x512), engine activation memory vs the plain `nvfp4` preset: - -| Model | `nvfp4` preset | This recipe | -|-------|---------------:|------------:| -| llama-nemotron-embed-1b-v2 | 1040 MiB | 516 MiB | -| llama-nemotron-rerank-1b-v2 | 520 MiB | 331 MiB | +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 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 From 53e68d4cd0d0f246c31a1db1e802f4464d5b5134 Mon Sep 17 00:00:00 2001 From: ajrasane <131806219+ajrasane@users.noreply.github.com> Date: Thu, 16 Jul 2026 17:04:36 +0000 Subject: [PATCH 6/6] [5726458] Fix CI: document nemotron_llama recipes in ptq.md, fix docstring RST test_every_model_specific_ptq_dir_is_mentioned requires every huggingface/ model recipe dir to appear in modelopt_recipes/ptq.md; add nemotron_llama to the architecture-aware kind with its delta. Restore single-line bullets in the configure_linear_module_onnx_quantizers docstring (docutils rejects multi-line bullet continuations in that list) and fold the output-quantizer rationale into the following paragraph. Co-Authored-By: Claude Fable 5 Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com> --- modelopt/torch/quantization/export_onnx.py | 11 +++++------ modelopt_recipes/ptq.md | 16 ++++++++++++++-- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index 17398317010..e5778c3c96b 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -643,13 +643,12 @@ 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/output activation quantizers use "dynamic" export (TRT_FP4DynamicQuantize, - etc.); the static path would fail on activations in the NVFP4 exporter. + - 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] = [] 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