Files
Wei-Ming Chen 1d392999b4 [OMNIML-5570] 2/2 Compose GEMM and KV-cache AutoQuant workflows (#2273)
### What does this PR do?

Type of change: new feature.

Follow-up to merged #2272. Adds composition of existing GEMM
quantization with KV-cache AutoQuantize:

- fixed FP8 GEMM PTQ followed by mixed-KV AutoQuantize;
- gradient-based NVFP4/FP8 GEMM AutoQuantize followed by independent
mixed-KV AutoQuantize;
- an optional `kv_auto_quantize` recipe stage with independent method,
constraints, candidates, and checkpoint path;
- ordered `hf_ptq.py` orchestration that keeps selected
weight/activation QDQ active while its calibration state remains frozen
during KV candidate calibration;
- fail-closed validation when a preceding stage leaves actual K/V
quantizers enabled; and
- unified export of a uniform-weight or mixed-weight checkpoint together
with the selected per-layer KV map.

The KV search still uses the public `mtq.auto_quantize(...,
constraints={"cost_model": "kv_cache", ...})` API from #2272. On a
converted model, the API preserves existing non-KV quantizers and
requires K/V to be disabled before search. Fresh-model behavior is
unchanged and starts from a deny-all quantizer baseline.

#### Why a follow-up field instead of a generic stage list?

This PR deliberately supports the two composition forms required by
`hf_ptq.py` without replacing the stable recipe schema. Existing recipes
already express a fixed `quantize` baseline plus one primary
`auto_quantize` search. A generic ordered `stages` list would require a
broader recipe/API migration, indexed checkpoint semantics, and
compatibility rules for arbitrary stage sequences. There is not yet a
demonstrated third search stage that justifies that surface-area change.

The two searches are not combined inside `mtq.auto_quantize`: each
invocation owns one search domain, constraint model, scoring method, and
resumable checkpoint. Their ordering and independent checkpoint paths
are orchestration concerns, while candidate calibration, scoring,
selection, and state application remain in the shared public API. A
general stage pipeline can be considered separately if more than this
one optional KV follow-up is needed.

Both solvers and scoring protocols are unchanged. The KV checkpoint
compatibility signature additionally fingerprints the preceding
quantizer configuration and calibrated state. Unsupported uniform-weight
plus mixed-KV exports record `kv_cache_deployment_supported: false` in
both ModelOpt and converted HF metadata.

### Usage

Fixed FP8 GEMM PTQ followed by KV AutoQuantize:

```bash
python examples/hf_ptq/hf_ptq.py \
  --pyt_ckpt_path Qwen/Qwen3-8B \
  --recipe general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits \
  --kv_auto_quantize_checkpoint /path/to/kv_autoquant.pth \
  --export_path /path/to/qwen3-8b-fp8-and-mixed-kv
```

Weight AutoQuantize followed by KV AutoQuantize:

```bash
python examples/hf_ptq/hf_ptq.py \
  --pyt_ckpt_path Qwen/Qwen3-8B \
  --recipe general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits \
  --auto_quantize_checkpoint /path/to/weight_autoquant.pth \
  --kv_auto_quantize_checkpoint /path/to/kv_autoquant.pth \
  --export_path /path/to/qwen3-8b-autoquant-and-mixed-kv
```

KV checkpoint resume requires identical preceding non-K/V quantizer
configuration and calibrated state. If rerunning the preceding stage
changes that state, use a new KV checkpoint path to recompute
sensitivities; configuration identity alone is insufficient to reuse the
scores safely.

### Testing

- Latest changed-area validation: 126 tests passed across `hf_ptq.py`
orchestration, KV checkpoint compatibility, export metadata, and HF
configuration conversion.
- A broader local run had 604 passes, one skip, and six failures: two
socket-binding failures under the sandbox and four local Transformers
API incompatibilities. This is not a full-suite pass.
- The fixed-PTQ→KV recipe executes end to end on a tiny offline Qwen
fixture.
- Public API coverage verifies that composed KV search preserves
preceding weight quantization and rejects enabled K/V state.
- Changed-file pre-commit hooks passed; the isolated recipe validator
also passed after dependency bootstrap.

### Before your PR is "*Ready for review*"

- Is this change backward compatible?: ✅
- If you copied code from any other sources or added a new PIP
dependency, did you follow guidance in `CONTRIBUTING.md`: N/A
- Did you write any new necessary tests?: ✅
- Did you update Changelog?: ✅ (0.48.0 composition feature and KV
checkpoint flag deprecation)
- Did you get Claude approval on this PR?: ❌

### Additional Information

- This follow-up targets `main`, which contains merged #2272.
- `--auto_quantize_checkpoint` and `--kv_auto_quantize_checkpoint` are
intentionally separate because KV sensitivities depend on the preceding
GEMM state.
- Uniform-weight plus mixed-KV exports are for artifact inspection until
the runtime's uniform-weight ModelOpt configuration consumes
`kv_cache_quantized_layers`. Export emits an actionable warning and
records `kv_cache_deployment_supported: false`; this marker does not
itself add runtime support.


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added staged post-training quantization workflows for weights and KV
caches, including dedicated KV-cache checkpoints.
- Added FP8/NVFP4 recipes with configurable bit constraints and scoring.
  - KV-cache quantization now supports pre-quantized models.

- **Bug Fixes**
- Mixed weight and KV-cache quantization now exports with a warning
instead of failing.
- Improved validation and checkpoint compatibility for staged
configurations.
- Added safeguards for configurations without enabled weight quantizers.

- **Documentation**
- Clarified staged KV-cache workflows, checkpoint options, configuration
behavior, and unsupported deployment combinations.
- Documented deprecated legacy quantization options and their
replacement behavior.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: weimingc <17592131+meenchen@users.noreply.github.com>
2026-09-29 00:15:40 +00:00

50 KiB
Executable File

Post-training quantization (PTQ)

Quantization is an effective model optimization technique that compresses your models. Quantization with Model Optimizer can compress model size by 2x-4x, speeding up inference while preserving model quality.

Model Optimizer enables highly performant quantization formats including NVFP4, FP8, INT8, INT4 and supports advanced algorithms such as SmoothQuant, AWQ, SVDQuant, and Double Quantization with easy-to-use Python APIs.

This section focuses on Post-training quantization, a technique that reduces model precision after training to improve inference efficiency without requiring retraining.

Section Description Link Docs
Pre-Requisites Required & optional packages to use this technique [Link]
Getting Started Learn how to optimize your models using PTQ to reduce precision and improve inference efficiency [Link] [docs]
Support Matrix View the support matrix to see quantization compatibility and feature availability across different models [Link]
AutoQuantize Automatically chooses layers/precisions for mixed precision quantization to enhanced inference performance and accuracy tradeoffs [Link] [docs]
Real Quant Real Quant compresses model weights in a low-precision format to reduce memory requirements of quantization. [Link]
Framework Scripts Example scripts demonstrating quantization techniques for optimizing Hugging Face / Megatron-Bridge / Megatron-LM models [Link]
Evaluate Accuracy Evaluate your model's accuracy! [Link]
Exporting Checkpoints Export to Hugging Face Unified Checkpoint and deploy on TRT-LLM/vLLM/SGLang [Link] [docs]
Pre-Quantized Checkpoints Ready to deploy Hugging Face pre-quantized checkpoints [Link]
Tracking runs with MLflow Record a PTQ run on an MLflow server so it can be reproduced from its entry alone [Link]
Resources Extra links to relevant resources [Link]

Pre-Requisites

Docker

For Hugging Face models, please use the TensorRT-LLM docker image (e.g., nvcr.io/nvidia/tensorrt-llm/release:1.2.0). Visit our installation docs for more information.

Also follow the installation steps below to upgrade to the latest version of Model Optimizer and install example-specific dependencies.

Local Installation

For Hugging Face models, install Model Optimizer with hf dependencies using pip from PyPI and install the requirements for the example:

pip install -U nvidia-modelopt[hf]
pip install -r requirements.txt

For TensorRT-LLM deployment, please use the TensorRT-LLM docker image or follow their installation docs. Similarly, for vLLM or SGLang deployment, please use their installation docs.

Getting Started

1. Quantize (Post Training Quantization)

With the simple API below, you can very easily use Model Optimizer to quantize your model. Model Optimizer achieves this by converting the precision of your model to the desired precision, and then using a small dataset (typically 128-512 samples) to calibrate the quantization scaling factors. The accuracy of PTQ is typically robust across different choices of calibration data, by default Model Optimizer uses a mix of cnn_dailymail and nemotron-post-training-dataset-v2. Users can try other datasets by easily modifying the calib_set.

import modelopt.torch.quantization as mtq

# Setup the model
model = AutoModelForCausalLM.from_pretrained("...")

# Simplified example set up a calibration data loader with the desired calib_size
calib_set = get_dataloader(num_samples=calib_size)

# Prepare the calibration set and define a forward loop
def forward_loop(model):
    for batch in calib_set:
        model(batch)

# PTQ with in-place replacement to quantized modules
model = mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop)

For higher NVFP4 PTQ accuracy, we recommend using mtq.NVFP4_MLP_ONLY_CFG, mtq.NVFP4_EXPERTS_ONLY_CFG, or mtq.NVFP4_OMLP_ONLY_CFG instead of mtq.NVFP4_DEFAULT_CFG. NVFP4_MLP_ONLY_CFG applies NVFP4 quantization to MLP (and MoE) layers, leaving attention layers unquantized. NVFP4_EXPERTS_ONLY_CFG quantizes only expert layers (*mlp.experts* and *block_sparse_moe*), useful for MoE models where dense MLP and attention stay in higher precision. NVFP4_OMLP_ONLY_CFG additionally quantizes the o_proj layer. All preserve accuracy in the sensitive attention QKV projections while still providing significant compression.

2. Export Quantized Model

Once your model is quantized, you can now export that model to a checkpoint for easy deployment.
We provide two APIs to export the quantized model:

  • Unified Hugging Face checkpoints, which can be deployed on TensorRT-LLM (Pytorch and C++ backends), vLLM and SGLang.
  • (Legacy) TensorRT-LLM checkpoints, a format that works with TensorRT-LLM C++ backend only.

Unified Hugging Face Checkpoints

from modelopt.torch.export import export_hf_checkpoint

with torch.inference_mode():
    export_hf_checkpoint(
        model,  # The quantized model.
        export_dir,  # The directory where the exported files will be stored.
    )

Please reference our framework scripts and our docs for more details.

Support Matrix

Hugging Face Supported Models

Model fp8 int8_smoothquant int4_awq w4a8_awq_beta1 nvfp45
LLAMA 3.x ✅ ❌ ✅ ✅3 ✅
LLAMA 4 6 ✅ ❌ ❌ ❌ ✅
Mixtral ✅ ❌ ✅2 ❌ ✅
Phi-3,4 ✅ ✅ ✅ ✅3 -
Phi-3.5 MOE ✅ ❌ ❌ ❌ -
Llama-Nemotron Super ✅ ❌ ❌ ❌ ✅
Llama-Nemotron Ultra ✅ ❌ ❌ ❌ ❌
Gemma 3 ✅2 - ✅ - -
QWen 2, 2.5 4 ✅ ✅ ✅ ✅ ✅
QWen3, 3.5 MOE, Next 6 ✅ - - - ✅
QwQ ✅ - - - ✅
DeepSeek V3, R1, V3.1, V3.27 - - - - ✅
Kimi K315 - - - - ✅
GLM-4.78 ✅ - - - ✅
Kimi K2 - - - - ✅
MiniMax M2.1 - - - - ✅
GPT-OSS10 - - - - ✅
T5 ✅ ✅ ✅ ✅ -
Whisper9 ✅ ❌ ❌ ❌ -
Nemotron-3 ✅ ❌ ❌ ❌ ✅
Llava (VLM)11 ✅ ✅12 ✅ ✅ -
Qwen2, 2.5-VL (VLM)11 ✅ ✅12 ✅ ✅ ✅
Qwen3-VL, Qwen3.5 (VLM)11,14 ✅ - - - -
Gemma 3 (VLM)11 ✅ - - - -
Nemotron VL (VLM)11,13 ✅ - - - ✅

This is a subset of the models supported. For the full list please check the TensorRT-LLM support matrix

1.The w4a8_awq_beta is an experimental quantization scheme that may result in a higher accuracy penalty.
2.For some models, there is only support for exporting quantized checkpoints.
3.W4A8_AWQ is only available on some models but not all
4.For some models, KV cache quantization may result in a higher accuracy penalty.
5.A selective set of the popular models are internally tested. The actual model support list may be longer. NVFP4 inference requires Blackwell GPUs and TensorRT-LLM v1.2 or later
6.Some models currently support export to HF format only.
7.PTQ for DeepSeek
8.GLM-4.7 has MTP (Multi-Token Prediction) layers that are automatically loaded and excluded from quantization.
9.Running Whisper model with transformers>=5.0 requires torchcodec and other system packages (e.g. ffmpeg).
10.GPT-OSS ships with native MXFP4 weights; NVFP4 export is produced via the closed-form --cast_mxfp4_to_nvfp4 cast (see MXFP4 → NVFP4 cast).
11.Vision-language model (VLM): by default, only the language model is quantized while the vision encoder is kept in high precision. Pass --vlm to the shell script (see VLM quantization).
12.For VLMs, int8_smoothquant only supports TensorRT-LLM checkpoint export and is not compatible with the TensorRT-LLM torch backend.
13.Nemotron VL automatically calibrates with image-text pairs; see VLM calibration with image-text pairs.
14.Qwen3-VL and dense Qwen3.5 VLM checkpoints support opt-in FP8 vision encoder quantization through model-specific recipes. Vision Linear layers, including primary and deepstack merger Linears where present, are quantized; patch embedding and vision-attention BMMs remain in high precision. MoE variants are not validated by these recipes.
15.Kimi K3 uses the calibration-free streaming converter because its routed experts are released as packed MXFP4 tensors; it does not use the in-memory hf_ptq.py flow.

The accuracy loss after PTQ may vary depending on the actual model and the quantization method. Different models may have different accuracy loss and usually the accuracy loss is more significant when the base model is small. If the accuracy after PTQ is not meeting the requirement, please try either modifying hf_ptq.py and disabling the KV cache quantization or using the QAT instead. For NVFP4 quantization specifically, we recommend nvfp4_mlp_only, nvfp4_experts_only, or nvfp4_omlp_only to achieve higher accuracy by restricting quantization to the MLP/expert layers (and optionally the o_proj layer) while keeping the attention QKV projections unquantized.

You can also create your own custom config using this guide.

Vision-language models (VLMs) are listed in the support matrix above (rows marked (VLM)). PTQ for VLMs is handled by the same hf_ptq.py entry point and shell script as LLMs. By default, the language model is quantized while the vision branch remains in high precision. Dense Qwen3-VL and dense Qwen3.5 VLM checkpoints additionally support the opt-in FP8 recipes documented under VLM quantization. For detailed TensorRT-LLM torch backend multimodal support, please refer to this doc.

Framework Scripts

Hugging Face Example Script

For LLM models like Llama-3:

# Install model specific pip dependencies if needed

export HF_PATH=<the downloaded LLaMA checkpoint from the Hugging Face hub, or simply the model card>
scripts/huggingface_example.sh --model $HF_PATH --quant <QFORMAT> --tp [1|2|4|8]

QFORMAT accepts any preset basename under modelopt_recipes/configs/ptq/presets/model/ — e.g. fp8, fp8_per_channel_per_token, fp8_2d_blockwise_weight_only, int8, int8_smoothquant, int8_weight_only, int4_awq, w4a8_awq_beta, nvfp4, nvfp4_awq_lite, nvfp4_w4a4_weight_mse_fp8_sweep, nvfp4_mlp_only, nvfp4_experts_only, nvfp4_omlp_only, nvfp4_svdquant, nvfp4_w4a4_weight_local_hessian, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8, mxfp8.

By default trust_remote_code is set to false. Please turn it on if model calibration and eval requires it using --trust_remote_code.

If the Huggingface model calibration fails on a multi-GPU system due to mismatched tensor placement, please try setting CUDA_VISIBLE_DEVICES to a smaller number.

FP8 calibration over a large model with limited GPU memory is not recommended but possible with the accelerate package. Please tune the device_map setting in example_utils.py if needed for model loading and the calibration process can be slow.

Huggingface models trained with modelopt.torch.speculative can be used as regular Huggingface models in PTQ. Note: there is a known issue with Huggingface models loaded across multiple GPUs for inference (i.e., "Expected all tensors to be on the same device, but found at least two devices..."). When encountered this error in PTQ of speculative decoding models, try reducing the number of GPUs used.

Calibration by default uses left padding_side for the Huggingface tokenizer as it usually leads to lower accuracy loss. The exported tokenizer files restores the default padding_side.

If a GPU OOM error occurs during model quantization despite sufficient memory, setting the --use_seq_device_map flag can help. This enforces sequential device mapping, distributing the model across GPUs and utilizing up to 80% of each GPU's memory.

You can add --low_memory_mode to the command to lower the memory requirements of the PTQ process. With this mode, the script will compress model weights to low precision before calibration. This mode is only supported for FP8 and NVFP4 with max calibration.

Recipe-based Quantization

Instead of specifying --qformat and --kv_cache_qformat separately, you can use a recipe — a declarative YAML file that bundles the full quantization configuration. Recipes are loaded via --recipe and take precedence over --qformat.

# Using a built-in recipe name (without .yaml suffix)
python hf_ptq.py \
  --pyt_ckpt_path <huggingface_model_card> \
  --recipe general/ptq/nvfp4_default-kv_fp8_cast \
  --export_path <quantized_ckpt_path>

# Using a custom recipe YAML file path
python hf_ptq.py \
  --pyt_ckpt_path <huggingface_model_card> \
  --recipe /path/to/my_ptq.yaml \
  --export_path <quantized_ckpt_path>

Built-in recipes are located in modelopt_recipes/general/ptq/ for model-agnostic recipes and in modelopt_recipes/model_type/<model_type>/ptq/ for recipes tuned to a specific Hugging Face model_type (see modelopt_recipes/model_type/README.md). You can also provide a path to your own custom YAML recipe file or directory. See the recipe documentation for details on the YAML schema and available recipes.

When --recipe is specified, --qformat is ignored. KV cache handling depends on the recipe type: a PTQ recipe bakes KV cache into its config and ignores --kv_cache_qformat; an AutoQuantize recipe falls back to --kv_cache_qformat unless it sets an explicit kv_cache field or kv_auto_quantize follow-up.

KV Cache Quantization

KV cache quantization reduces memory usage during inference by quantizing the key-value cache. This is controlled via the --kv_cache_qformat flag (default: fp8_cast).

# FP8 KV cache with cast (no calibration needed, fast)
python hf_ptq.py --pyt_ckpt_path <model> --qformat fp8 --kv_cache_qformat fp8_cast --export_path <path>

# NVFP4 KV cache with data-driven calibration
python hf_ptq.py --pyt_ckpt_path <model> --qformat nvfp4 --kv_cache_qformat nvfp4 --export_path <path>

# Disable KV cache quantization
python hf_ptq.py --pyt_ckpt_path <model> --qformat fp8 --kv_cache_qformat none --export_path <path>

Via the shell script:

scripts/huggingface_example.sh --model $HF_PATH --quant fp8 --kv_cache_quant nvfp4

Available KV cache formats:

Format Description
fp8_cast (default) FP8 KV cache without data-driven calibration (amax set to FP8 range)
fp8 FP8 KV cache with data-driven calibration
fp8_affine FP8 KV cache with affine quantization
nvfp4_cast NVFP4 KV cache without data-driven calibration
nvfp4 NVFP4 KV cache with data-driven calibration
nvfp4_affine NVFP4 KV cache with affine quantization
nvfp4_rotate NVFP4 KV cache with rotation
none Disable KV cache quantization

Formats ending in _cast (fp8_cast, nvfp4_cast) are fast — they set the amax to the format's full range without data-driven calibration. Other formats use data-driven calibration for potentially better accuracy.

MXFP4 → NVFP4 cast (for GPT-OSS)

GPT-OSS checkpoints (openai/gpt-oss-20b, openai/gpt-oss-120b) ship with native MXFP4 weights (*_blocks + *_scales in the checkpoint, quantization_config.quant_method == "mxfp4"). Passing --cast_mxfp4_to_nvfp4 tells hf_ptq.py to read the source MXFP4 scales and produce a closed-form, bit-exact NVFP4 weight export — no GEMM-level recalibration of the weights needed.

python hf_ptq.py \
  --pyt_ckpt_path openai/gpt-oss-20b \
  --qformat nvfp4_mlp_only \
  --cast_mxfp4_to_nvfp4 \
  --export_path <quantized_ckpt_path>

The cast pins each NVFP4 block's scale_2 = 2^(k_max - 8) and _amax = 6 * 2^k_j, both derived from the source MXFP4 E8M0 scales. For blocks whose k_j lands in E4M3's representable window (k_max - k_j ≤ 17), NVFP4 dequant matches MXFP4 dequant bit-for-bit; out-of-range blocks fall back to a data-derived per-block amax.

--cast_mxfp4_to_nvfp4 requires an NVFP4-family --qformat (e.g. nvfp4_mlp_only, nvfp4_experts_only, nvfp4) and is incompatible with AutoQuantize recipes (multi-format search).

Deepseek R1

PTQ for DeepSeek shows how to quantize the DeepSeek model with FP4 and export to TensorRT-LLM.

Kimi K3

PTQ for Kimi K3 shows how to cast the source MXFP4 routed experts to NVFP4 and quantize KDA/MLA attention weights to 128x128 block FP8 without loading the full model or running calibration.

VLM quantization

Vision-language models are quantized through the same script. Add --vlm so the script runs the TensorRT-LLM multimodal quickstart as the deploy smoke test instead of the text-only one:

scripts/huggingface_example.sh --model <Hugging Face model card or checkpoint> --quant fp8 --vlm

Supported --quant values for VLMs are fp8, nvfp4, int8_smoothquant, int4_awq, and w4a8_awq_beta (see the (VLM) rows in the Support Matrix).

By default, hf_ptq.py applies --qformat only to the language model. Model-specific recipes add FP8 quantization of the vision branch for validated Qwen3-VL and dense Qwen3.5 checkpoints. Use the recipe directory matching the checkpoint's model_type: qwen3_vl or qwen3_5.

# Vision encoder only: FP8 vision Linears and merger, BF16 LLM and KV cache.
python hf_ptq.py \
  --pyt_ckpt_path <Qwen3-VL-or-Qwen3.5-checkpoint> \
  --recipe model_type/qwen3_vl/ptq/fp8_vision-kv_none \
  --calib_with_images \
  --calib_size 512 \
  --skip_generate \
  --export_path <quantized-checkpoint>

# Joint vision encoder + language model FP8 with FP8 KV-cache cast.
python hf_ptq.py \
  --pyt_ckpt_path <Qwen3-VL-or-Qwen3.5-checkpoint> \
  --recipe model_type/qwen3_vl/ptq/fp8_vision_lm-kv_fp8_cast \
  --calib_with_images \
  --calib_size 512 \
  --skip_generate \
  --export_path <quantized-checkpoint>

fp8_vision-kv_none starts from all quantizers disabled and enables FP8 only for nn.Linear weights and inputs under the visual branch, including the primary merger and any deepstack mergers. The language model, KV cache, patch embedding, and vision-attention QK/softmax/AV operations stay in high precision. fp8_vision_lm-kv_fp8_cast applies the standard FP8 model recipe to both model branches and enables FP8 KV-cache cast, while keeping patch embedding and vision-attention operands in high precision. Other precision combinations can be expressed by composing the same recipe units in a custom recipe; no model-specific Python path is required.

The exported checkpoint requires an inference runtime that supports quantized vision encoder Linears. Runtime-specific vision-attention quantization is separate from this ModelOpt checkpoint. Use the direct hf_ptq.py commands above for these recipes; the generic multimodal quickstart is not a serving validation for a vision-quantized checkpoint. Both examples pass --skip_generate because the script's text-only preview does not forward the image tensors used during calibration. For a Qwen3.5 checkpoint, replace qwen3_vl with qwen3_5 in the recipe path.

VLM calibration with image-text pairs (e.g., Nemotron VL)

For vision-language models, calibration quality can likely improve by using image-text pairs instead of text-only data, especially on visual understanding tasks:

python hf_ptq.py \
  --pyt_ckpt_path <huggingface_model_card> \
  --qformat nvfp4 \
  --export_path <quantized_ckpt_path> \
  --trust_remote_code \
  --calib_with_images \
  --calib_size 512

The same flag is exposed by the shell script:

scripts/huggingface_example.sh --model <model> --quant nvfp4 --vlm --calib_with_images --trust_remote_code

With --calib_with_images, calibration batches always pass through the complete VLM so image features reach the component selected by the preset or recipe. This also applies to the default language-model-only path: its quantizers are exercised by the complete multimodal forward.

Note: when --calib_with_images is set, --calib_size must be a single value, and the calibration dataset is nvidia/nemotron_vlm_dataset_v2. This functionality is currently in beta and has been tested on nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16.

Megatron-Bridge Example Script

Please refer to examples/megatron_bridge/README.md for example scripts for PTQ / QAD with Megatron-Bridge which is generally more performant than the Hugging Face scripts.

Megatron-LM Example Script

Megatron-LM framework PTQ and TensorRT-LLM deployment examples are maintained in the Megatron-LM GitHub repo. Please refer to the examples here.

AutoQuantize

AutoQuantize (mtq.auto_quantize) is a PTQ algorithm which quantizes a model by searching for the best quantization format per-layer while meeting performance constraints specified by the user. AutoQuantize streamlines the trade-off of model accuracy and performance.

AutoQuantize uses an effective-bits target (effective_bits) as the performance constraint (for both weight-only and weight & activation quantization) — the effective number of bits for the quantized model.

You may specify an effective_bits target such as 5.4 for mixed precision quantization using NVFP4_DEFAULT_CFG & FP8_DEFAULT_CFG. AutoQuantize will automatically quantize highly sensitive layers in FP8_DEFAULT_CFG while keeping less sensitive layers in NVFP4_DEFAULT_CFG (and even skip quantization for any extremely sensitive layers) so that the the final mixed precision quantized model has an effective quantized bits of 5.4. This model would give a better accuracy than the model quantized with vanilla NVFP4_DEFAULT_CFG configuration since the more aggressive NVFP4_DEFAULT_CFG quantization was not applied for the highly sensitive layers.

Here is an example usage for AutoQuantize algorithm (Please see auto_quantize API for more details):


    import modelopt.torch.quantization as mtq

    # Define the model & calibration dataloader
    model = ...
    calib_dataloader = ...

    # Define forward_step function.
    # forward_step should take the model and data as input and return the output
    def forward_step(model, data):
        output =  model(data)
        return output

    # Define loss function which takes the model output and data as input and returns the loss
    def loss_func(output, data):
        loss = ...
        return loss


    # Perform AutoQuantize
    model, search_state_dict = mtq.auto_quantize(
        model,
        constraints = {"effective_bits": 5.4},
        # supported quantization formats are listed in `modelopt.torch.quantization.config.choices`
        quantization_formats = ["NVFP4_DEFAULT_CFG", "FP8_DEFAULT_CFG"]
        data_loader = calib_dataloader,
        forward_step=forward_step,
        loss_func=loss_func,
        ...
        )

AutoQuantize for Hugging Face models

AutoQuantize can be performed for Huggingface LLM models like Qwen / Nemotron as shown below:

AutoQuantize is driven by an AutoQuantize recipe passed with --recipe. The recipe defines the candidate formats, optional fixed PTQ baseline, effective_bits target, cost model, scoring method, search-disabled layers, and cost-excluded layers — see AutoQuantizeConfig. Shipped recipes live in modelopt_recipes/general/auto_quantize/; model-specific recipes (carrying architecture-specific disabled layers — e.g. VL vision towers) live under modelopt_recipes/model_type/<model>/auto_quantize/.

Script

export HF_PATH=<the downloaded checkpoint from the Hugging Face hub, or simply the model card>
# --recipe selects an AutoQuantize recipe; the recipe defines the candidate formats and the
# effective-bits target (here NVFP4 + FP8 at 5.4 effective bits).
scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/nvfp4_fp8_at_5p4bits --calib_batch_size 4

The recipe quantizes the less accuracy-sensitive layers with the more aggressive format (e.g. NVFP4) and keeps the more sensitive ones at higher precision (or unquantized), so the model meets the recipe's effective_bits target. To author your own, copy a shipped recipe and adjust candidate_formats, constraints.effective_bits, auto_quantize_method (gradient / kl_div), score_size, module_search_spaces (optional per-module candidate overrides), disabled_layers (excluded from the search), and cost_excluded_layers (kept out of the bit-budget accounting — e.g. VL vision towers). Recipes can splice a shared base disabled_layers set via $import (see modelopt_recipes/configs/auto_quantize/units/base_disabled_layers).

AutoQuantize recipes support two mutually exclusive search-space styles:

  1. Set top-level auto_quantize.candidate_formats to search every unmatched quantizable module, with optional module_search_spaces overrides.
  2. Set a normal top-level quantize config as the fixed PTQ baseline, omit top-level candidate_formats, and use module_search_spaces to list only the modules AutoQuantize should search. The fixed and searched modules still run through one integrated calibration, scoring, cost, and export flow.

For example, this keeps unmatched modules at the normal W4A16 NVFP4 PTQ setting while searching only attention between W4A16 and FP8:

imports:
  w4a16_nvfp4: configs/ptq/presets/model/w4a16_nvfp4
  fp8: configs/ptq/presets/model/fp8

quantize:
  $import: w4a16_nvfp4

auto_quantize:
  constraints:
    effective_bits: 6.0
  module_search_spaces:
    - module_name_patterns: ["*self_attn*", "*linear_attn*"]
      candidate_formats:
        - $import: w4a16_nvfp4
        - $import: fp8
      allow_no_quant: false

bf16 (no quantization) is an implicit per-layer choice for the top-level candidate_formats, so a single format (e.g. [fp8]) gives a {fp8, bf16} per-layer search. A module_search_spaces rule can set allow_no_quant: false to exclude bf16 from the solver choices for matching modules. Use the top-level quantize baseline, rather than a one-candidate search rule, for modules that are fixed and not actually searched.

The fixed baseline may also reuse a model-specific PTQ configuration. For example, the Qwen3.6 MoE AutoQuantize recipe imports the same model-specific quant_cfg used by model_type/qwen3_5_moe/ptq/w4a16_nvfp4-fp8_attn-kv_fp8_cast, reproduces that recipe's quantize section, and lists only shared experts, attention, and lm_head under module_search_spaces. A loader test asserts that the inherited fixed baseline remains equal to the original PTQ recipe while leaving the original recipe unchanged.

For models without backprop support (e.g. Llama-4), use the kl_div scoring method — see the shipped general/auto_quantize/nvfp4_fp8_kl_div_at_5p4bits recipe.

Weight AutoQuantize recipes still apply KV cache as a uniform post-step and fall back to --kv_cache_qformat (default fp8_cast) unless they set an explicit kv_cache field.

To optimize GEMM and KV cache in one invocation, compose ordered stages in the same recipe. A fixed quantize block followed by a KV-domain auto_quantize first calibrates the GEMM weight/activation configuration, then searches K/V while the existing GEMM QDQ remains enabled with calibration frozen. See general/auto_quantize/fp8_ptq_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits.

A weight-domain auto_quantize can instead add a kv_auto_quantize follow-up with its own method, constraints, candidates, score size, and disabled layers. This supports, for example, a gradient-based GEMM search followed by a KL-divergence KV search; see general/auto_quantize/nvfp4_fp8_gradient_then_kv_fp8_nvfp4_cast_kl_div_at_5p4bits. When the follow-up is present, the recipe owns KV configuration and suppresses the CLI's uniform --kv_cache_qformat fallback. Use --auto_quantize_checkpoint for the weight search and --kv_auto_quantize_checkpoint for the KV search.

KV-cache AutoQuantize recipes use the same mtq.auto_quantize API and set constraints.cost_model: kv_cache with an effective_bits target. Their candidate_formats are complete K/V cache configs whose config-level effective_bits includes packed scale overhead. The width-weighted budget covers eligible layers; disabled_layers are preserved and excluded. BF16 is used only as the isolated-KL reference, not as a solver choice. The shipped recipe searches FP8-cast K/V (8.0 bits/scalar) and NVFP4-cast K/V (4.5 bits/scalar) at 5.4 bits/scalar. It intentionally excludes FP8-K/NVFP4-V because the companion vLLM implementation does not support that asymmetric per-layer format:

python hf_ptq.py \
  --pyt_ckpt_path Qwen/Qwen3.8-27B \
  --recipe general/auto_quantize/kv_fp8_nvfp4_cast_kl_div_at_5p4bits \
  --kv_auto_quantize_checkpoint /path/to/kv_autoquant.pth \
  --export_path /path/to/qwen3.8-27b-mixed-kv

Each candidate uses an explicit constant scale, avoiding an additional calibration pass while keeping persistent K/V scales in the unified HF checkpoint. Unified export records the selected formats in kv_cache_quantized_layers. mtq.auto_quantize returns the sensitivity scores and selected recipe in its search state; --kv_auto_quantize_checkpoint stores that resumable state, including the candidate quantizer tensors needed for replay.

KV sensitivity scoring runs one reference forward plus one forward per eligible-layer candidate for every scoring step. Choose the recipe's auto_quantize.score_size with the number of eligible layers and candidates in mind. Peak scoring memory also grows with the number of selected tokens times the model vocabulary because reference and candidate log probabilities are computed in FP32.

Note

Layer-wise KV checkpoints require the companion vLLM mixed-KV metadata consumer or a later vLLM release containing it. The repository's currently pinned vLLM 0.26.0 does not consume kv_cache_quantized_layers, so these checkpoints are export-only in that stock environment. The companion consumer also does not yet apply the layer map when uniform FP8/NVFP4 weights are present; export warns for that composed combination and records kv_cache_deployment_supported: false. Do not deploy either unsupported case. Full FP8 K/V and full NVFP4 K/V use existing vLLM kernels once the relevant metadata path is available.

--auto_quantize_checkpoint saves/restores weight-search state. Every KV-domain search uses --kv_auto_quantize_checkpoint; a KV-primary recipe temporarily accepts the former flag as a deprecated fallback. Composed recipes therefore keep the weight and KV search states separate. A KV checkpoint is compatible only with the exact preceding non-K/V quantizer configuration and calibrated state; if that state cannot be reproduced exactly, use a new checkpoint path to recompute the KV sensitivities:

scripts/huggingface_example.sh --model $HF_PATH --recipe general/auto_quantize/nvfp4_fp8_at_5p4bits \
  --auto_quantize_checkpoint /path/to/auto_quantize.pth --calib_batch_size 4

The example scripts above also have an additional flag --tasks, where the actual tasks run in the script can be customized. The allowed tasks are quant,mmlu,lm_eval,livecodebench,simple_eval specified in the script parser. The tasks combo can be specified with a comma-separated task list. Some tasks like mmlu can take a long time to run. To run lm_eval tasks, please also specify the --lm_eval_tasks flag with comma separated lm_eval tasks here.

If GPU out-of-memory error is reported running the scripts, please try editing the scripts and reducing the max batch size to save GPU memory.

NOTE: AutoQuantize requires backpropagation of the model. Models without backpropagation support (e.g., Llama-4) will not work with AutoQuantize when using the gradient method. The kl_div method does not require backpropagation.

Real Quant

When working with large language models, memory constraints can be a significant challenge. ModelOpt provides a workflow for initializing HF models with compressed weights across multiple GPUs to dramatically reduce memory usage. Check --low_memory_mode option in hf_ptq.py for more details.

import modelopt.torch.quantization as mtq
from modelopt.torch.quantization.plugins import init_quantized_weights
from transformers import AutoModelForCausalLM, AutoConfig

# Step 1: Initialize the model with compressed weights
with init_quantized_weights(mtq.NVFP4_DEFAULT_CFG):
    model = AutoModelForCausalLM.from_pretrained(ckpt_path)

# Step 2: Calibrate the model
mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop)

Multi-Node Post-Training Quantization with FSDP2

ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the --use_fsdp2 flag on the standard hf_ptq.py entry point.

KV-cache AutoQuantize recipes are not supported with --use_fsdp2 and are rejected before model loading. Distributed KV sensitivity scoring, selection, and checkpoint writes must be synchronized before this combination can be enabled safely. Existing weight AutoQuantize recipes retain their previous experimental warning with FSDP2.

Usage

Slurm orchestrates launching the job on every node for you, so this is the easiest way to run a multi-node PTQ. A ready-to-run example that quantizes Nemotron-3-Super to NVFP4 is provided in slurm/multinode_fsdp2_ptq.slurm. Edit the CONFIG block (container image, model path, export path, recipe) and submit:

sbatch --nodes=2 slurm/multinode_fsdp2_ptq.slurm

Manual (run on each node)

Without Slurm, start torchrun on every node yourself:

torchrun \
    --nnodes=<num_nodes> --node_rank=<current_node_rank> \
    --master_addr=<node0_ip_addr> --master_port=<port> \
    --nproc_per_node=<num_gpus_per_node> \
    hf_ptq.py \
    --pyt_ckpt_path <path_to_model> \
    --recipe general/ptq/nvfp4_default-kv_fp8_cast \
    --batch_size <calib_batch_size> \
    --calib_size <num_calib_samples> \
    --export_path <export_path> \
    --use_fsdp2

See Recipe-based Quantization for the recipe format and built-in recipe names. The exported checkpoint can be deployed using TensorRT-LLM/ vLLM/ SGLang. For more details refer to the deployment section of this document.

Performance Note: FSDP2 is designed for training workloads and may result in longer calibration and export times. For faster calibration, maximize the batch size based on available GPU memory and choose the right number of GPUs to avoid unnecessary communication.

Evaluate Accuracy

TensorRT-LLM Validation

A list of accuracy validation benchmarks are provided in the llm_eval directory. Right now MMLU is supported in this example by specifying the --tasks flag running the scripts mentioned above.

The benchmark_suite.py script is used as a fast performance benchmark. For details, please refer to the TensorRT-LLM documentation

This example also covers the lm_evaluation_harness, MMLU and the human eval accuracy benchmarks, whose details can be found here. The supported lm_eval evaluation tasks are listed here

Exporting Checkpoints

Model Optimizer supports provide two paths to export the quantized model:

  • Unified Hugging Face checkpoints, which can be deployed on TensorRT-LLM (Pytorch and C++ backends), vLLM and SGLang.
  • (Legacy) TensorRT-LLM checkpoints, a format that works with TensorRT-LLM C++ backend only.

The unified checkpoint1 format design reflects two key characteristics: 1. The layer structures and tensor names remain aligned with the original Hugging Face checkpoint, and 2. The same checkpoint can be deployed across multiple inference frameworks without modification. A unified checkpoint can be exported using the following commands:

1.Unified checkpoint export currently does not support sparsity. Speculative decoding is only supported in unified checkpoint export. For legacy deployment, exported unified checkpoint then needs a TensorRT-LLM checkpoint converter (e.g., this) to convert and build the TensorRT engine(s) for deployment. Alternatively, call TensorRT-LLM LLM-API to deploy the unified checkpoints e.g., check examples here.

API

from modelopt.torch.export import export_hf_checkpoint

with torch.inference_mode():
    export_hf_checkpoint(
        model,  # The quantized model.
        export_dir,  # The directory where the exported files will be stored.
    )

Quantize and Export

python hf_ptq.py --pyt_ckpt_path <huggingface_model_card> --qformat fp8 --export_path <quantized_ckpt_path> --trust_remote_code

For exporting fake-quantized models for vLLM serving (e.g., for research or kernels not yet supported in real-quant), use the --vllm_fakequant_export flag. See vllm_serve/README.md for details.

Hugging Face framework Script

Alternatively, the framework script huggingface_example.sh also supports quantize and export:

scripts/huggingface_example.sh --model <huggingface_model_card> --quant fp8

Deployment


TRT-LLM

from tensorrt_llm import LLM

llm_fp8 = LLM(model="<the exported model path>")
print(llm_fp8.generate(["What's the age of the earth? "]))

vLLM

from vllm import LLM

llm_fp8 = LLM(model="<the exported model path>", quantization="modelopt")
print(llm_fp8.generate(["What's the age of the earth? "]))

SGLang

import sglang as sgl

llm_fp8 = sgl.Engine(model_path="<the exported model path>", quantization="modelopt")
print(llm_fp8.generate(["What's the age of the earth? "]))

Unified HF Checkpoint Deployment Model Support Matrix

The deployment support matrix — which model families and quantization formats are covered on TRT-LLM, vLLM, and SGLang, including vision-language models, speculative decoding drafters, and diffusion models — lives in the documentation so there is a single copy to keep current:

Unified HF Checkpoint → Model Support Matrix

Each entry there is drawn from tests/examples/hf_ptq/test_deploy.py, which loads the exported checkpoint in each framework and generates from short text prompts. That file is also the place to look for the exact checkpoint, tensor-parallel size, and minimum SM version behind each entry.

Note: those cases are marked release and run out-of-band — no workflow currently passes --run-release — and each is a load-and-generate smoke check on the text path. Read the legend in the docs before treating an entry as verified support.

Note: the matrix records what modelopt validates, not the full set of what will run. vLLM, SGLang, and TRT-LLM load unified HF checkpoints generically, so unlisted models frequently deploy without any modelopt change — check the serving framework's own model support list and try it.

(Legacy) TensorRT-LLM Checkpoints

The user can specify the inference time TP and PP size and the export API will organize the weights to fit the target GPUs.

from modelopt.torch.export.trtllm import export_tensorrt_llm_checkpoint

with torch.inference_mode():
    export_tensorrt_llm_checkpoint(
        model,  # The quantized model.
        decoder_type,  # The type of the model, e.g gpt, gptj, or llama.
        dtype,  # The exported weights data type.
        export_dir,  # The directory where the exported files will be stored.
        inference_tensor_parallel,  # The number of GPUs used in the inference time tensor parallel.
        inference_pipeline_parallel,  # The number of GPUs used in the inference time pipeline parallel.
        use_nfs_workspace,  # If exporting in a multi-node setup, please specify a shared directory like NFS for cross-node communication.
    )

Build the TensorRT-LLM engines

After the TensorRT-LLM checkpoint export, you can use the trtllm-build build command to build the engines from the exported checkpoints. Please check the TensorRT-LLM Build API documentation for reference.

Pre-Quantized Checkpoints

Tracking runs with MLflow

Set MLflow's own MLFLOW_TRACKING_URI, or pass --mlflow <tracking-uri>, to record a PTQ run on an MLflow server so it can be reproduced later from its MLflow entry alone:

python hf_ptq.py \
  --pyt_ckpt_path <huggingface_model_card> \
  --recipe general/ptq/nvfp4_default-kv_fp8_cast \
  --export_path <quantized_ckpt_path> \
  --mlflow https://<your-mlflow-server>/

The run is opened before the model loads, so a bad URI or a missing token fails within seconds rather than after a full calibration.

Uploaded artifacts
Artifact Contents
command.txt The full invocation, copy-pasteable, with credentials masked
version.txt The ModelOpt version that ran
experiment.json The experiment name, run id and run URL — the same file written into --export_path
recipe/resolved_recipe.yaml The --recipe with its $imports expanded, so it stands alone
logs/hf_ptq.log The run's Python stdout/stderr, including the traceback if it crashed
summary/quant_summary.txt The per-quantizer summary (unless --no-verbose)
summary/moe.html Per-expert calibration token counts, when the run produces them

Every command-line argument is also logged as a searchable param, alongside user / hostname / modelopt_version / git_sha tags. A run that fails is still recorded, with status FAILED and its log attached.

A tracked run also drops .experiment.json into --export_path, so a checkpoint found on disk names the run that produced it:

cat <quantized_ckpt_path>/.experiment.json
# {"tracking_uri": ..., "experiment_name": ..., "experiment_id": ..., "run_id": ...,
#  "run_name": ..., "run_url": ...}

The local file is written only once the export itself completes, so a run that fails earlier leaves whatever checkpoint is already in --export_path — and its pointer — untouched. The experiment.json artifact is uploaded for every run that opened, so a failed run stays traceable from the server. An export that is not tracked removes any pointer it would otherwise inherit, from a reused --export_path or from a tracked source checkpoint.

Other flags:

  • --mlflow_experiment — defaults to $USER/hf_ptq/<checkpoint basename>-<recipe name>, falling back to --qformat when no --recipe is used.
  • --mlflow_run_name — defaults to the UTC start time, YYYYmmdd-HHMMSS.
  • $MLFLOW_TRACKING_URI enables tracking on its own; --mlflow overrides it. A URI taken from the environment is best-effort — if the client is missing or the server is unreachable the run warns and continues untracked, since the variable is often exported for other tooling. An explicit --mlflow fails loudly instead.

Authentication uses MLflow's own environment variables (MLFLOW_TRACKING_TOKEN, or MLFLOW_TRACKING_USERNAME / MLFLOW_TRACKING_PASSWORD).

The tracking itself lives in modelopt.torch.utils.mlflow (MlflowRunLogger), so other example scripts can record runs the same way; hf_ptq.py only supplies the params and artifacts specific to PTQ.

Note: only the main rank uploads, so --use_fsdp2 runs produce a single run. The log captures Python output; output written directly by native libraries (NCCL, CUDA) goes to the terminal only. On SLURM, keep the job's own .out file for those.

Resources

Technical Resources

There are many quantization schemes supported in the example scripts:

  1. The FP8 format is available on the Hopper and Ada GPUs with CUDA compute capability greater than or equal to 8.9.

  2. The INT8 SmoothQuant, developed by MIT HAN Lab and NVIDIA, is designed to reduce both the GPU memory footprint and inference latency of LLM inference.

  3. The INT4 AWQ is an INT4 weight only quantization and calibration method. INT4 AWQ is particularly effective for low batch inference where inference latency is dominated by weight loading time rather than the computation time itself. For low batch inference, INT4 AWQ could give lower latency than FP8/INT8 and lower accuracy degradation than INT8.

  4. The W4A8 AWQ is an extension of the INT4 AWQ quantization that it also uses FP8 for activation for more speed up and acceleration.

  5. The NVFP4 is one of the new FP4 formats supported by NVIDIA Blackwell GPU and demonstrates good accuracy compared with other 4-bit alternatives. NVFP4 can be applied to both model weights as well as activations, providing the potential for both a significant increase in math throughput and reductions in memory footprint and memory bandwidth usage compared to the FP8 data format on Blackwell. For higher accuracy with NVFP4 PTQ, we recommend nvfp4_mlp_only, nvfp4_experts_only, or nvfp4_omlp_only. nvfp4_mlp_only restricts NVFP4 quantization to MLP (and MoE) layers only, leaving attention layers in higher precision. nvfp4_experts_only quantizes only expert layers (*mlp.experts* and *block_sparse_moe*), ideal for MoE models. nvfp4_omlp_only extends MLP-only by also quantizing the o_proj layer, providing a middle ground between full NVFP4 and MLP-only quantization.