diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 329e5d21f05..aa60942094e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,26 +1,49 @@ Changelog ========= -0.47 (2026-xx-xx) +Experimental ^^^^^^^^^^^^^^^^^ -**New Features** +- Add pruning examples for Qwen3.5-9B and Nemotron3-Nano using the `new experimental puzzletron branch `_, this branch uses `AutoModel `_ for better parallelization and efficiency. + +0.46 (2026-08-xx) +^^^^^^^^^^^^^^^^^ **Backward Breaking Changes** +- Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained. +- Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy `_ directly together with ModelOpt PTQ in ``examples/llm_ptq``. +- Remove the deprecated ``examples/llm_qad`` Megatron-LM QAD example (deprecated in 0.45). Use the `megatron_bridge QAD example `_ instead, which provides a simpler Python-based interface and better model coverage. + **Deprecations** -**Bug Fixes** +- ``examples/hf_ptq`` AutoQuantize is now driven by an **AutoQuantize recipe** (``--recipe``). The ``--auto_quantize_bits``, ``--auto_quantize_method``, ``--auto_quantize_score_size``, ``--auto_quantize_cost_model``, and ``--auto_quantize_active_moe_expert_ratio`` flags are **deprecated** but still work: they are converted into an ``AutoQuantizeConfig`` on the fly (emitting a ``DeprecationWarning``) and will be removed in a future release. Prefer a recipe under ``modelopt_recipes/general/auto_quantize/``. See ``examples/hf_ptq/README.md``. -0.46 (2026-08-17) -^^^^^^^^^^^^^^^^^ +- Renamed ``examples/llm_ptq`` to ``examples/hf_ptq`` to reflect that it covers Hugging Face LLM **and** VLM PTQ. A relative symlink ``examples/llm_ptq`` -> ``hf_ptq`` keeps existing paths and commands working; it will be removed in a future release. Please update references to the new ``examples/hf_ptq`` path. +- Consolidated ``examples/vlm_ptq`` into ``examples/hf_ptq``. Vision-language model PTQ now shares the ``hf_ptq.py`` entry point and ``scripts/huggingface_example.sh``; pass ``--vlm`` to run the TensorRT-LLM multimodal quickstart smoke test. The ``examples/vlm_ptq/scripts/huggingface_example.sh`` entry point is deprecated: it now prints a warning and forwards to the ``hf_ptq`` script with ``--vlm``, and will be removed in a future release. See `examples/hf_ptq/README.md `__. +- Dropped VILA / NVILA vision-language model support in ``examples/hf_ptq``. VILA's modeling code requires ``transformers<=4.50.0``, which conflicts with ModelOpt's minimum supported ``transformers`` version. The VILA-specific bootstrap (repo clone, ``requirements-vila.txt``) and loading paths in ``example_utils.py`` have been removed. +- Bump minimum nemo container requirement to ``nemo:26.04`` (recommended ``nemo:26.06``) for Megatron-Bridge / Megatron-LM optimization features. +- Python 3.10 support will be dropped in the next release as it is reaching EOL. **New Features** -*Quantization* - +- 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). +- Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. +- Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. - Add a fused Triton fast path for ``local_hessian`` NVFP4 weight-scale search (the Hessian-weighted FP8-E4M3 scale sweep). For each NVFP4 block it minimizes ``dwᵀ H dw`` over the 126 candidate scales using the per-cin-block local Hessian on tensor cores, replacing the per-weight Python reference sweep — roughly **34x** faster on a single 8192x4096 weight and bit-exact with the reference for fp32/fp16 weights. Used automatically during ``local_hessian`` calibration for both dense and fused-MoE expert weights; falls back to the reference sweep on CPU, when Triton is unavailable, or via ``MODELOPT_NVFP4_TRITON_SWEEP=0``. +- Add **context-parallel (CP)** and **data-parallel (DP)** support to the shared Megatron-Core inference/calibration utilities. Under CP, ``get_megatron_calibration_forward_loop`` and ``megatron_mmlu`` partition each sequence across CP ranks (zigzag load-balanced), ``megatron_prefill`` accepts a CP-partitioned ``position_ids`` and lets the CP-aware causal attention build the mask, and MMLU gathers per-rank logits back to the full sequence for last-token scoring. Under DP, calibration shards the dataset across data-parallel ranks (``DistributedSampler``; amax is max-reduced across the DP group inside ``mtq``) and ``megatron_mmlu`` shards whole batches across DP ranks and all-reduces the per-subject counts. DP is implicit (``world_size / (tp * pp * cp)``); ``examples/megatron_bridge/quantize.py`` gains a ``--cp_size`` flag. +- Add Minitron pruning support for Megatron-Core models with the following new attention and MoE variants. For these, only ``hidden_size`` is pruned (alongside the usual ``ffn_hidden_size`` / ``num_layers`` / MoE dimensions); the variant-internal dimensions noted below are not pruned: +- Add support for ONNX Q/DQ node placement for DLA via the new flag ``--target_dla``. + + - **GatedDeltaNet** (linear attention) and **gated attention** (``attention_output_gate``), such as Qwen3.5 (hybrid GatedDeltaNet + gated-attention) language models, including MoE variants — attention / linear-attention heads are not pruned. + - **Multi-Latent Attention (MLA)**, such as DeepSeek — MLA latent ranks are not pruned. + - **Latent MoE**, such as Nemotron-3-Super — ``hidden_size`` pruning resizes the latent projections while the MoE latent dim itself is not pruned. +- Add Minitron pruning support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/prune_minitron.py``. The language model is pruned while the vision tower is left intact and the full VLM is saved back; ``hidden_size`` is not pruned if it is shared with the vision projector. Pruning importance is estimated from image-text calibration (the full VLM forward over vision-conditioned activations) by default, or from a text dataset for text-only ablations. +- Add PTQ support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/quantize.py``. Only the language model is quantized (vision tower + projector left in full precision) and the full VLM is saved as a Megatron checkpoint. The calibration modality is inferred from ``--calib_dataset_name``: an image-text dataset drives the full VLM forward (vision-conditioned activations), while a text dataset runs text-only calibration of the language model. Image-text calibration shards across data-parallel ranks (context parallelism is supported only for text-only calibration). HuggingFace unified export of a quantized VLM is not yet supported. +- Add Megatron-Bridge distillation and QAD support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/distill.py``. - Add NVFP4 Four-Over-Six (4/6) weight quantization (``mtq.NVFP4_FOUR_OVER_SIX_CFG``): MSE weight calibration picks, per block, between an M=6 and an M=4 dynamic range (the choice is folded into the FP8 per-block scales), with the ``four_over_six: true`` flag normalizing those scales by 256 (vs 448) for M=4 headroom. Supported via ``mtq.quantize`` and HF / Megatron export only -- **not** ``mtq.compress``, which does not preserve the per-block M=4/M=6 choice - Add dLLM (tied-weight PTQ and HF-checkpoint export) support for diffusion-based encoder-decoder LLMs (e.g. DiffusionGemma) whose encoder/decoder stacks share parameters via HF ``_tied_weights_keys``. @@ -29,6 +52,7 @@ Changelog - The exported state_dict is also **reordered (decoder keys win instead of encoder)** so canonical-side keys per HF's ``_tied_weights_keys`` declaration win the data_ptr dedup; gated to the DiffusionGemma model class in ``_reorder_canonical_first``, no-op for every other model. - New DiffusionGemma model-specific recipe under ``modelopt_recipes/huggingface/diffusion_gemma/ptq/`` (``nvfp4_experts_only.yaml`` + its ``disabled_quantizers.yaml`` unit) adds the ``*self_conditioning*`` exclude on top of the standard default, leaving the shared ``default_disabled_quantizers`` unit clean for non-diffusion models — pattern matches the existing ``phi4mm`` / ``nemotron_vl`` model-specific recipes. - ``hf_ptq.py`` also unwraps ``ModelOutput`` dataclasses from ``.generate()`` so the preview decode works on diffusion models. Non-tied models see no behavioral change. +- Add **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet. - Add Torch-TensorRT FP8 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships a ViT-tuned FP8 PTQ recipe under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: it quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline. - Add **AutoQuantize recipe** support: ``mtq.auto_quantize`` can be driven declaratively from a YAML recipe (``RecipeType.AUTO_QUANTIZE`` / ``AutoQuantizeConfig``) specifying candidate formats, the ``effective_bits`` target, cost model (incl. ``active_moe`` and ``excluded_module_name_patterns``), scoring method, and disabled layers. Adds an ``effective_bits`` cost-model override on ``QuantizeConfig`` / ``QuantizerAttributeConfig`` (block-scale-accurate NVFP4 = 4.5 via ``configs/numerics/nvfp4``). Shipped recipes live under ``modelopt_recipes/general/auto_quantize/`` and model-specific ones under ``modelopt_recipes/huggingface//auto_quantize/``. - Add module-specific AutoQuantize search spaces through ``mtq.auto_quantize(..., module_search_spaces=...)`` and recipe-level ``auto_quantize.module_search_spaces``. Glob-matched runtime decision groups can override global candidate formats and control whether BF16/no-quant is solver-selectable with ``allow_no_quant``. Recipes can instead reuse a normal PTQ ``quantize`` config as the fixed baseline and explicitly list only genuinely searched modules; fixed and searched groups remain in one calibration, scoring, effective-bits, checkpoint, and export flow. Rules cannot partially split runtime-fused groups, fixed groups are isolated from unrelated calibration algorithms, infeasible resolved budgets fail before calibration, and checkpoint replay validates the fixed baseline, resolved groups, candidate choices, scoring boundaries, and cost weights before reusing calibration or sensitivity state. @@ -37,48 +61,6 @@ Changelog - Add ``MaxCalibConfig.skip_forward_without_activation_calib`` (opt-in, default ``False``): when enabled, max calibration skips the ``forward_loop`` if no enabled quantizer needs data-driven activation statistics — e.g. an experts-only recipe whose activation quantizers all use ``constant_amax`` / ``use_constant_amax``, or dynamic / MX (MXFP4/MXFP8) quantization (and none has a static ``bias_calibrator``). Weight calibration still runs on the weight tensors directly, so the quantized weights are unchanged; only the wasted forward is avoided. It is opt-in because the ``forward_loop`` can carry caller-side effects (notably materializing sharded parameters under DeepSpeed ZeRO-3). The advanced algorithms that always need activations (MSE, local Hessian, SmoothQuant, SVDQuant, GPTQ) call ``max_calibrate`` directly and are unaffected. The ``nvfp4_experts_only_input_scale1-kv_fp8_cast`` recipe enables it. - Add ``examples/minimax_m3/hf_ptq_mixed_mxfp8_nvfp4.py`` for streaming MiniMax-M3 export and a model-specific ``hf_ptq.py`` recipe that produces an MXFP8 language-model base with MSE-calibrated NVFP4 routed experts directly from BF16. The NVFP4 expert ``input_scale`` is fixed to 1.0. -*Speculative Decoding* - -- 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). -- Add **streaming** speculative-decoding training (EAGLE3 / DFlash): the draft trains on base-model hidden states produced on the fly by a co-located ``vllm serve`` (no disk dump), moved trainer-side over NIXL RDMA, scaling to multi-node (dedicated serve replicas + DDP trainers). New launcher examples for NVFP4 Kimi-K2.5 / K2.6 on GB200/aarch64 under ``tools/launcher/examples/moonshotai/``. -- Add **Domino** speculative-decoding training: the parallel DFlash draft backbone plus a lightweight GRU causal correction head, selected via ``dflash_architecture_config.projector_type=domino``. Trained with a base/final dual loss whose ``dflash_lambda_base_start``/``dflash_lambda_base_decay_ratio`` curriculum decays the base-loss weight 1→0. Exports in the z-lab drafter format; recipe at ``modelopt_recipes/general/speculative_decoding/domino.yaml``. Training only — the inference path is not wired up yet. - -*Megatron Framework (M-LM / M-Bridge)* - -- Add Minitron pruning support for Megatron-Core models with the following new attention and MoE variants. For these, only ``hidden_size`` is pruned (alongside the usual ``ffn_hidden_size`` / ``num_layers`` / MoE dimensions); the variant-internal dimensions noted below are not pruned: - - - **GatedDeltaNet** (linear attention) and **gated attention** (``attention_output_gate``), such as Qwen3.5 (hybrid GatedDeltaNet + gated-attention) language models, including MoE variants — attention / linear-attention heads are not pruned. - - **Multi-Latent Attention (MLA)**, such as DeepSeek — MLA latent ranks are not pruned. - - **Latent MoE**, such as Nemotron-3-Super — ``hidden_size`` pruning resizes the latent projections while the MoE latent dim itself is not pruned. -- Optimize Minitron pruning support for MoE models using the fused **grouped GEMM** experts (``TEGroupedMLP``) in addition to the existing ``SequentialMLP`` path. ``examples/megatron_bridge/prune_minitron.py`` now uses grouped GEMM by default (pass ``--no_moe_grouped_gemm`` to fall back to ``TESequentialMLP``). -- Add Minitron pruning support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/prune_minitron.py``. The language model is pruned while the vision tower is left intact and the full VLM is saved back; ``hidden_size`` is not pruned if it is shared with the vision projector. Pruning importance is estimated from image-text calibration (the full VLM forward over vision-conditioned activations) by default, or from a text dataset for text-only ablations. -- Add PTQ support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/quantize.py``. Only the language model is quantized (vision tower + projector left in full precision) and the full VLM is saved as a Megatron checkpoint. The calibration modality is inferred from ``--calib_dataset_name``: an image-text dataset drives the full VLM forward (vision-conditioned activations), while a text dataset runs text-only calibration of the language model. Image-text calibration shards across data-parallel ranks (context parallelism is supported only for text-only calibration). HuggingFace unified export of a quantized VLM is not yet supported. -- Add Megatron-Bridge distillation and Quantization-Aware Distillation (QAD) support for the language model part of vision-language models (e.g. Qwen3.5-VL, Gemma3-VL) via ``examples/megatron_bridge/distill.py``. -- Add **context-parallel (CP)** and **data-parallel (DP)** support to the shared Megatron-Core inference/calibration utilities. Under CP, ``get_megatron_calibration_forward_loop`` and ``megatron_mmlu`` partition each sequence across CP ranks (zigzag load-balanced), ``megatron_prefill`` accepts a CP-partitioned ``position_ids`` and lets the CP-aware causal attention build the mask, and MMLU gathers per-rank logits back to the full sequence for last-token scoring. Under DP, calibration shards the dataset across data-parallel ranks (``DistributedSampler``; amax is max-reduced across the DP group inside ``mtq``) and ``megatron_mmlu`` shards whole batches across DP ranks and all-reduces the per-subject counts. DP is implicit (``world_size / (tp * pp * cp)``); ``examples/megatron_bridge/quantize.py`` gains a ``--cp_size`` flag. -- Add support for retaining all Megatron-Bridge distillation checkpoints via ``distill.py --checkpoint_keep_last -1`` and exporting all or selected iterations with ``export_distilled_megatron_to_hf.py --export_iterations``. -- 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 `_. - -*Misc* - -- Add the ``day0-release`` agent skill (``.agents/skills/day0-release/``), a deterministic end-to-end driver that chains the PTQ → evaluation → comparison skills (the evaluation stage deploys the checkpoint itself) with an enforced gate after each stage and returns a publish decision (ACCEPT / REGRESSION / ANOMALOUS / INFEASIBLE). Ships three GPU-free, unit-tested gate scripts (``gate_ptq.py``, ``gate_run.py``, ``gate_compare.py``) that validate checkpoint coverage, evaluation-run completeness, and baseline-vs-candidate accuracy threshold. v1 reports and stops on regression; the recipe-search loop is deferred. -- Add support for ONNX Q/DQ node placement for DLA via the new flag ``--target_dla``. -- (Experimental) Add pruning examples for Qwen3.5-9B and Nemotron3-Nano using the `new experimental puzzletron branch `_, this branch uses `AutoModel `_ for better parallelization and efficiency. - -**Backward Breaking Changes** - -- Remove the ``examples/diffusers/eval`` image-quality evaluation example (ImageReward / CLIP-IQA / CLIP metrics) and its references in ``examples/diffusers/README.md``. The example was deprecated in 0.45 and is no longer maintained. -- Remove the deprecated ``examples/llm_autodeploy`` example (deprecated in 0.45). Use TensorRT-LLM's `AutoDeploy `_ directly together with ModelOpt PTQ in ``examples/llm_ptq``. -- Remove the deprecated ``examples/llm_qad`` Megatron-LM QAD example (deprecated in 0.45). Use the `megatron_bridge QAD example `_ instead, which provides a simpler Python-based interface and better model coverage. - -**Deprecations** - -- ``examples/hf_ptq`` AutoQuantize is now driven by an **AutoQuantize recipe** (``--recipe``). The ``--auto_quantize_bits``, ``--auto_quantize_method``, ``--auto_quantize_score_size``, ``--auto_quantize_cost_model``, and ``--auto_quantize_active_moe_expert_ratio`` flags are **deprecated** but still work: they are converted into an ``AutoQuantizeConfig`` on the fly (emitting a ``DeprecationWarning``) and will be removed in a future release. Prefer a recipe under ``modelopt_recipes/general/auto_quantize/``. See ``examples/hf_ptq/README.md``. -- Renamed ``examples/llm_ptq`` to ``examples/hf_ptq`` to reflect that it covers Hugging Face LLM **and** VLM PTQ. A relative symlink ``examples/llm_ptq`` to ``hf_ptq`` keeps existing paths and commands working; it will be removed in a future release. Please update references to the new ``examples/hf_ptq`` path. -- Consolidated ``examples/vlm_ptq`` into ``examples/hf_ptq``. Vision-language model PTQ now shares the ``hf_ptq.py`` entry point and ``scripts/huggingface_example.sh``; pass ``--vlm`` to run the TensorRT-LLM multimodal quickstart smoke test. The ``examples/vlm_ptq/scripts/huggingface_example.sh`` entry point is deprecated: it now prints a warning and forwards to the ``hf_ptq`` script with ``--vlm``, and will be removed in a future release. See `examples/hf_ptq/README.md `__. -- Dropped VILA / NVILA vision-language model support in ``examples/hf_ptq``. VILA's modeling code requires ``transformers<=4.50.0``, which conflicts with ModelOpt's minimum supported ``transformers`` version. The VILA-specific bootstrap (repo clone, ``requirements-vila.txt``) and loading paths in ``example_utils.py`` have been removed. -- Bump minimum nemo container requirement to ``nemo:26.04`` (recommended ``nemo:26.06``) for Megatron-Bridge / Megatron-LM optimization features. -- Python 3.10 support will be dropped in the next release as it is reaching EOL. - **Bug Fixes** - Fix ``ShapeInferenceError`` during ONNX INT8 + FP16 quantization (``--high_precision_dtype fp16``) of weakly-typed models (e.g. TensorFlow exports) that carry stale rank-0 ``graph.output`` shapes or ops such as ``TopK`` that ONNX's static shape inference cannot resolve. ``clear_stale_value_info`` now reconciles stale output shapes via symbolic shape inference (keeping every output's shape field populated), and AutoCast runs ONNX shape inference in strict mode and falls back to schema-based standalone type inference when it fails, so unresolved ops no longer leave tensors untyped. diff --git a/examples/torch_onnx/README.md b/examples/torch_onnx/README.md index 06ad11164dd..f479bab3ae1 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,72 @@ python ../onnx_ptq/evaluate.py \ --model_name= ``` +## HF Embedding and Reranking Models + +> **Experimental:** Accuracy has not yet been validated for this example. + +`hf_embedding_quant_to_onnx.py` quantizes an HF text-embedding or reranking +model (bidirectional Llama encoders such as +[nvidia/llama-nemotron-embed-1b-v2](https://huggingface.co/nvidia/llama-nemotron-embed-1b-v2) +and +[nvidia/llama-nemotron-rerank-1b-v2](https://huggingface.co/nvidia/llama-nemotron-rerank-1b-v2)) +with a PTQ recipe and exports it to ONNX. Embedding models are exported with +mean pooling and L2 normalization on top of the encoder; reranking +(sequence-classification) models are exported to their relevance logits. Both +graphs take `input_ids` and `attention_mask` with dynamic batch/sequence axes. + +The default recipe +(`modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml`) +quantizes weights and activations to NVFP4 and additionally quantizes the +projection-Linear outputs. Without output-side quantization, quantized GEMMs +emit FP16 activations, so FP8/FP4 engines can use as much or more activation +memory than an unquantized FP16 engine; quantizing the projection outputs keeps +inter-layer activations in the low-precision format. An FP8 twin of the recipe +(`fp8_output_quant_proj.yaml`, pass it via `--recipe`) applies the same idea to +the FP8 preset. With TensorRT 10.16 on RTX PRO 6000 Blackwell (strongly-typed +engines, 5 dynamic-shape profiles up to 32x512), engine activation memory: + +| Model | FP16 | `fp8` preset | fp8 recipe | `nvfp4` preset | nvfp4 recipe | +|-------|-----:|-------------:|-----------:|---------------:|-------------:| +| llama-nemotron-embed-1b-v2 | 1040 MiB | 1392 MiB | 1096 MiB | 1040 MiB | 516 MiB | +| llama-nemotron-rerank-1b-v2 | 1040 MiB | 1392 MiB | 1096 MiB | 520 MiB | 331 MiB | + +### Usage + +```bash +python hf_embedding_quant_to_onnx.py \ + --model_path=nvidia/llama-nemotron-embed-1b-v2 \ + --trust_remote_code \ + --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 \ + --trust_remote_code \ + --onnx_save_path=llama_nemotron_rerank_nvfp4.onnx +``` + +### Building a TensorRT engine with trtexec + +NVFP4 requires a Blackwell GPU (SM100+) and TensorRT 10.11 or later. Build a +strongly-typed engine with dynamic shapes (add optimization profiles matching +your serving batch sizes and sequence lengths): + +```bash +trtexec --onnx=llama_nemotron_embed_nvfp4.onnx \ + --stronglyTyped \ + --saveEngine=llama_nemotron_embed_nvfp4.plan \ + --minShapes=input_ids:1x2,attention_mask:1x2 \ + --optShapes=input_ids:32x128,attention_mask:32x128 \ + --maxShapes=input_ids:32x512,attention_mask:32x512 +``` + +The exported `.onnx` references a sibling weights file (`.onnx_data`); +keep the two files in the same directory when building. To inspect the chosen +kernels and per-profile activation memory, add +`--profilingVerbosity=detailed --exportLayerInfo=.json --verbose`. + ## LLM Quantization and Export with TensorRT-Edge-LLM [TensorRT-Edge-LLM](https://github.com/NVIDIA/TensorRT-Edge-LLM) provides a complete pipeline for quantizing LLMs and VLMs using NVIDIA ModelOpt and exporting them to optimized ONNX for deployment on edge platforms such as NVIDIA Jetson and DRIVE. @@ -283,7 +354,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..923c05c17ac --- /dev/null +++ b/examples/torch_onnx/hf_embedding_quant_to_onnx.py @@ -0,0 +1,328 @@ +# 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. +"""Quantize an HF embedding or reranking model and export it to ONNX. + +The default recipe quantizes weights, inputs, and projection outputs to NVFP4 so +TensorRT can keep inter-layer activations in FP4. Embedding models export pooled, +normalized embeddings; reranking models export relevance logits. + +This example is experimental because its accuracy has not been validated. +""" + +import argparse +import os + +import numpy as np +import torch +import torch.nn.functional as F +from onnx import numpy_helper +from torch.onnx import register_custom_op_symbolic, symbolic_helper +from transformers import AutoConfig, AutoModel, AutoModelForSequenceClassification, AutoTokenizer +from transformers.integrations import sdpa_attention + +import modelopt.torch._deploy.utils.torch_onnx as torch_onnx_utils +import modelopt.torch.quantization as mtq +from modelopt.recipe import load_recipe +from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata +from modelopt.torch.quantization.export_onnx import scaled_dot_product_attention + +__all__ = [ + "EmbeddingModel", + "RerankModel", + "install_static_extent_fix", + "main", + "register_bidirectional_sdpa", +] + +DEFAULT_RECIPE = "huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj" + +# TODO: Add an accuracy evaluation pipeline for the embedding and reranking models. +CALIBRATION_TEXTS = [ + ("What is the capital of France?", "Paris is the capital and most populous city of France."), + ("How do vaccines work?", "Vaccines train the immune system to recognize pathogens."), + ("Who wrote Hamlet?", "Hamlet is a tragedy written by William Shakespeare around 1600."), + ("What causes rain?", "Rain forms when water vapor condenses into droplets that fall."), + ( + "What is photosynthesis?", + "Plants convert light energy into chemical energy in chloroplasts.", + ), + ("How fast is light?", "Light travels at approximately 299,792 kilometers per second."), + ("What is machine learning?", "Machine learning builds models that learn patterns from data."), + ("Where is the Great Barrier Reef?", "The reef lies off the coast of Queensland, Australia."), + ("What is inflation?", "Inflation is the rate at which prices for goods and services rise."), + ("How do airplanes fly?", "Wings generate lift as air flows faster over their curved tops."), + ("What is DNA?", "DNA carries genetic instructions for development and functioning."), + ("Who painted the Mona Lisa?", "Leonardo da Vinci painted the Mona Lisa in the early 1500s."), + ("What is quantum computing?", "Quantum computers use qubits that exist in superposition."), + ("Why is the sky blue?", "Rayleigh scattering disperses shorter blue wavelengths of sunlight."), + ("What is the tallest mountain?", "Mount Everest rises 8,849 meters above sea level."), + ("How do batteries store energy?", "Batteries store energy in chemical form as electricity."), +] + + +class EmbeddingModel(torch.nn.Module): + """Bidirectional encoder + mean pooling + L2 normalization.""" + + def __init__(self, base): + super().__init__() + self.base = base + + def forward(self, input_ids, attention_mask): + hidden = self.base(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state + mask = attention_mask.unsqueeze(-1).to(hidden.dtype) + pooled = (hidden * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-6) + return F.normalize(pooled, p=2, dim=1) + + +class RerankModel(torch.nn.Module): + """Sequence-classification reranker exported to its relevance logits.""" + + def __init__(self, base): + super().__init__() + self.base = base + + def forward(self, input_ids, attention_mask): + return self.base(input_ids=input_ids, attention_mask=attention_mask).logits + + +def register_bidirectional_sdpa(): + """Export SDPA as a plain MatMul->Add->Softmax->MatMul pattern. + + torch's default sdpa symbolic emits IsNaN/Where guards for the dynamic + is_causal flag traced by transformers, which prevents TensorRT attention + fusion. This model is bidirectional, so is_causal is pinned to False. + """ + + 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. + 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. + """ + 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 positive_int(value): + """Parse a strictly positive integer argument.""" + value = int(value) + if value <= 0: + raise argparse.ArgumentTypeError(f"expected a positive integer, got {value}") + return value + + +def main(): + """Run recipe-driven quantization and ONNX export.""" + 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=positive_int, + default=64, + help="Number of calibration samples.", + ) + parser.add_argument( + "--batch_size", + type=positive_int, + default=8, + help="Batch size for calibration.", + ) + parser.add_argument( + "--trust_remote_code", + action="store_true", + help="Allow model repositories to execute their custom Python code.", + ) + 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=args.trust_remote_code + ) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=args.trust_remote_code) + 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=args.trust_remote_code, dtype=torch.float32 + ) + model = RerankModel(base).to(device).eval() + else: + base = AutoModel.from_pretrained( + args.model_path, trust_remote_code=args.trust_remote_code, dtype=torch.float32 + ) + model = EmbeddingModel(base).to(device).eval() + + pairs = (CALIBRATION_TEXTS * (args.calibration_data_size // len(CALIBRATION_TEXTS) + 1))[ + : args.calibration_data_size + ] + + def forward_loop(m): + for i in range(0, len(pairs), args.batch_size): + chunk = pairs[i : i + args.batch_size] + if is_reranker: + # Rerankers score (query, passage) pairs. + batch = tokenizer( + [q for q, _ in chunk], + [p for _, p in chunk], + padding=True, + truncation=True, + max_length=512, + return_tensors="pt", + ).to(device) + else: + batch = tokenizer( + [f"question:{q} \n \n passage:{p}" for q, p in chunk], + padding=True, + truncation=True, + max_length=512, + return_tensors="pt", + ).to(device) + with torch.no_grad(): + m(batch["input_ids"], batch["attention_mask"]) + + recipe = load_recipe(args.recipe) + quant_cfg = recipe.quantize.model_dump(exclude_unset=True) + mtq.quantize(model, quant_cfg, forward_loop=forward_loop) + mtq.print_quant_summary(model) + + # Unequal lengths so padding produces a real attention mask: transformers + # otherwise drops the all-ones mask and switches sdpa to native GQA, which + # the export symbolic does not support. + example = tokenizer( + ["example query one", "example query two"], + ["example passage one", "a considerably longer example passage two with extra words"], + padding=True, + return_tensors="pt", + ).to(device) + + register_bidirectional_sdpa() + install_static_extent_fix(base.config.hidden_size) + + model_name = os.path.basename(args.onnx_save_path).replace(".onnx", "") + onnx_bytes, _ = get_onnx_bytes_and_metadata( + model=model, + dummy_input={ + "input_ids": example["input_ids"], + "attention_mask": example["attention_mask"], + }, + model_name=model_name, + weights_dtype="fp16", + dynamic_axes={ + "input_ids": {0: "batch", 1: "seq"}, + "attention_mask": {0: "batch", 1: "seq"}, + }, + ) + OnnxBytes.from_bytes(onnx_bytes).write_to_disk( + os.path.dirname(args.onnx_save_path) or ".", clean_dir=False + ) + print(f"Quantized ONNX model is saved to {args.onnx_save_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index 97bd0b60c18..e0ffc75a294 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() + _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 not in {"int8", "int4_awq"} for fmt in args.auto_quantization_formats) ) if uses_fp8_conv_input: _disable_low_channel_conv_input_quantizers(quantized_model) diff --git a/modelopt/torch/quantization/export_onnx.py b/modelopt/torch/quantization/export_onnx.py index 7b42dce5782..e5778c3c96b 100644 --- a/modelopt/torch/quantization/export_onnx.py +++ b/modelopt/torch/quantization/export_onnx.py @@ -281,7 +281,7 @@ def scaled_dot_product_attention( if hasattr(torch.onnx, "_type_utils"): from torch.onnx._type_utils import JitScalarType else: # torch >= 2.9 - from torch.onnx._internal.torchscript_exporter import JitScalarType + from torch.onnx._internal.torchscript_exporter._type_utils import JitScalarType if hasattr(torch.onnx, "symbolic_opset14"): from torch.onnx.symbolic_opset14 import _attention_scale, _causal_attention_mask @@ -643,18 +643,19 @@ def configure_linear_module_onnx_quantizers(model): For modules with block quantization (NVFP4/MXFP8): - Weight quantizers use "static" export (TRT_FP4QDQ for NVFP4, DQ-only for MXFP8) - - Input/activation quantizers use "dynamic" export (TRT_FP4DynamicQuantize, etc.) + - Input/output activation quantizers use "dynamic" export (TRT_FP4DynamicQuantize, etc.) - This must be set for ALL modules with block quantization, not just nn.Linear, - because models like ResNet have non-Linear modules (e.g., MaxPool2d) with NVFP4/MXFP8 - input quantizers that would otherwise default to the static path and produce - TRT_FP4QDQ nodes on activations (which the NVFP4 exporter cannot handle). + This must be set for ALL quantizers on activations, not just nn.Linear inputs, + because non-Linear modules (e.g., ResNet's MaxPool2d) and enabled output + quantizers would otherwise default to the static path and produce TRT_FP4QDQ + nodes on activations (which the NVFP4 exporter cannot handle). """ sentinel = object() originals: list[tuple] = [] for _, module in model.named_modules(): for attr_name, new_value in ( ("input_quantizer", "dynamic"), + ("output_quantizer", "dynamic"), ("weight_quantizer", "static"), ): quantizer = getattr(module, attr_name, None) diff --git a/modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml b/modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml new file mode 100644 index 00000000000..15665a727e3 --- /dev/null +++ b/modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml @@ -0,0 +1,43 @@ +# 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 + - quantizer_name: '*score*' + enable: false diff --git a/modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml b/modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml new file mode 100644 index 00000000000..1edf942f2f2 --- /dev/null +++ b/modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml @@ -0,0 +1,45 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Llama-Nemotron embedding PTQ recipe: NVFP4 W4A4 plus dynamic NVFP4 output +# quantizers on the projection Linears (q/k/v/o_proj, gate/up/down_proj) so +# TensorRT keeps inter-layer activations in FP4 instead of FP16. [5726458] +# +# Output quantizers must stay scoped to GEMM outputs: a DynamicQuantize on +# non-GEMM outputs (embeddings, pooling) fails to compile in TensorRT. + +imports: + base_disable_all: configs/ptq/units/base_disable_all + w4a4_nvfp4_nvfp4: configs/ptq/units/w4a4_nvfp4_nvfp4 + default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers + nvfp4: configs/numerics/nvfp4 + +metadata: + recipe_type: ptq + description: 'Llama-Nemotron embedding PTQ recipe (nvfp4 + projection output quantizers): same numerics as the general nvfp4 preset, plus dynamic NVFP4 + output quantizers on the projection Linears so TensorRT engines carry FP4 activations between layers.' +quantize: + algorithm: max + quant_cfg: + - $import: base_disable_all + - $import: w4a4_nvfp4_nvfp4 + - $import: default_disabled_quantizers + - quantizer_name: '*_proj.output_quantizer' + cfg: + $import: nvfp4 + # Reranker (sequence-classification) score head: keep in high precision like + # lm_head; its [1, hidden] weight also cannot be packed by the NVFP4 exporter. + - quantizer_name: '*score*' + enable: false diff --git a/modelopt_recipes/ptq.md b/modelopt_recipes/ptq.md index 255544ffe1e..22d6fbd5ff5 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 | `minimax_m3_vl`, `qwen3_5`, `qwen3_5_moe`, `vit` | +| **Architecture-aware `quant_cfg`** | Per-sub-module format choices a single wildcard scheme can't express | `minimax_m3_vl`, `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-*`, `models/nvidia/Mistral-Medium-3.5-128B-NVFP4` | @@ -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` — `minimax_m3_vl`, `qwen3_5`, `qwen3_5_moe`, `vit` +### Architecture-aware `quant_cfg` — `minimax_m3_vl`, `qwen3_5`, `qwen3_5_moe`, `vit`, `nemotron_llama` **`minimax_m3_vl/ptq/mxfp8_nvfp4_experts`** applies MXFP8 to the language-model linear layers and MSE-calibrated NVFP4 to routed experts, with expert @@ -276,6 +276,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, like `lm_head`. In the NVFP4 recipe, its `[1, hidden]` +weight cannot be packed by the NVFP4 exporter. *Why special:* the +general recipes never enable output quantizers, and the pattern must stay scoped +to GEMM outputs — a `DynamicQuantize` on non-GEMM outputs (embedding lookup, +pooling) fails to compile in TensorRT. + A lighter case: **`step3p5/Step3.5-Flash/ptq/nvfp4-mlp-only`** is close to `general/ptq/nvfp4_mlp_only` (NVFP4 on MoE/MLP weights+inputs, FP8 KV) but pinned to one released checkpoint and carrying instance-specific disables diff --git a/tests/_test_utils/torch/transformers_models.py b/tests/_test_utils/torch/transformers_models.py index 90a5f266546..fdbcece1c86 100644 --- a/tests/_test_utils/torch/transformers_models.py +++ b/tests/_test_utils/torch/transformers_models.py @@ -34,6 +34,7 @@ Gemma3Config, GptOssConfig, LlamaConfig, + LlamaForSequenceClassification, NemotronConfig, PreTrainedModel, Qwen3Config, @@ -595,6 +596,35 @@ def create_tiny_llama_dir( ) +def get_tiny_llama_seq_cls(**config_kwargs) -> PreTrainedModel: + set_seed(SEED) + kwargs = { + "dtype": torch.bfloat16, + "hidden_size": 32, + "intermediate_size": 32, + "num_hidden_layers": 2, + "num_attention_heads": 16, + "num_key_value_heads": 2, + "max_position_embeddings": 32, + "vocab_size": 32, + "num_labels": 1, + "pad_token_id": 0, + } + kwargs.update(config_kwargs) + return LlamaForSequenceClassification(LlamaConfig(**kwargs)) + + +def create_tiny_llama_seq_cls_dir( + tmp_path: Path | str, with_tokenizer: bool = False, **config_kwargs +) -> Path: + return _create_tiny_llm_dir( + Path(tmp_path) / "tiny_llama_seq_cls", + get_tiny_llama_seq_cls, + with_tokenizer=with_tokenizer, + **config_kwargs, + ) + + ##### T5 ##### def get_tiny_t5(**config_kwargs) -> PreTrainedModel: set_seed(SEED) diff --git a/tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py b/tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py new file mode 100644 index 00000000000..1394d6a2400 --- /dev/null +++ b/tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py @@ -0,0 +1,76 @@ +# 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 onnx +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", "recipe", "expected_op"), + [ + ( + "embedding", + "huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj", + "TRT_FP4DynamicQuantize", + ), + ( + "reranking", + "huggingface/nemotron_llama/ptq/fp8_output_quant_proj", + "QuantizeLinear", + ), + ], +) +def test_hf_embedding_quant_to_onnx(tmp_path, model_kind, recipe, expected_op): + 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), + recipe=recipe, + 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() + op_types = {node.op_type for node in onnx.load(onnx_save_path).graph.node} + assert expected_op in op_types + if expected_op == "QuantizeLinear": + assert "TRT_FP4DynamicQuantize" not in op_types 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..fe99cae9bc3 100644 --- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py +++ b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py @@ -14,6 +14,7 @@ # limitations under the License. +import onnx import pytest from _test_utils.examples.run_command import extend_cmd_parts, run_example_command @@ -46,3 +47,31 @@ 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(tmp_path): + timm_model_name, model_kwargs = _MODELS["vit_tiny"] + onnx_save_path = tmp_path / "vit_tiny.recipe.onnx" + recipe_path = tmp_path / "disable_all.yaml" + recipe_path.write_text("quant_cfg:\n - quantizer_name: '*'\n enable: false\n") + + cmd_parts = extend_cmd_parts( + ["python", "torch_quant_to_onnx.py"], + timm_model_name=timm_model_name, + model_kwargs=model_kwargs, + quantize_mode="int8", + recipe=str(recipe_path), + onnx_save_path=str(onnx_save_path), + calibration_data_size="1", + num_score_steps="1", + ) + cmd_parts.append("--no_pretrained") + run_example_command(cmd_parts, "torch_onnx") + + quantize_ops = { + "QuantizeLinear", + "DequantizeLinear", + "TRT_FP4DynamicQuantize", + "TRT_FP8QuantizeLinear", + } + assert not quantize_ops & {node.op_type for node in onnx.load(onnx_save_path).graph.node}