Skip to content

[5726458] Add NVFP4 projection-output-quantizer recipe and HF embedding ONNX export example#1981

Open
ajrasane wants to merge 7 commits into
mainfrom
ajrasane/nvbug_5726458_nvfp4_output_quant
Open

[5726458] Add NVFP4 projection-output-quantizer recipe and HF embedding ONNX export example#1981
ajrasane wants to merge 7 commits into
mainfrom
ajrasane/nvbug_5726458_nvfp4_output_quant

Conversation

@ajrasane

@ajrasane ajrasane commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Type of change: new example

Without output-side quantization, TensorRT's quantized GEMMs emit FP16 activations: every quantized GEMM input adds a low-precision copy on top of the FP16 tensors instead of replacing them, so FP8/FP4 engines can use as much or more activation memory than an unquantized FP16 engine ([5726458]). Quantizing the projection-Linear outputs makes the engine carry inter-layer activations in the low-precision format. This PR ships that as recipes for the Llama-Nemotron embedding/reranking family (NVFP4 and FP8 variants) plus an end-to-end example.

Measured on RTX PRO 6000 Blackwell with TensorRT 10.16 (strongly-typed engines, 5 dynamic-shape profiles up to 32x512), activation memory per profile:

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

Engine sizes (dominated by weights): FP16 ≈ 2374 MiB, FP8 ≈ 1453 MiB, NVFP4 ≈ 1050–1075 MiB. The presets only shrink weights — their activation memory matches (or exceeds, for FP8) the FP16 engine because every quantized GEMM still emits FP16; the output-quantizer recipes are what reduce activation memory (fp8: −21% vs its preset; nvfp4: −50% vs its preset and 2x below FP16).

  • modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml — the general nvfp4 preset plus dynamic NVFP4 output quantizers scoped to the projection Linears (*_proj.output_quantizer). Scoping matters: a DynamicQuantize on non-GEMM outputs (embedding lookup, pooling) fails to compile in TensorRT. The sequence-classification score head is kept unquantized: final heads stay in high precision like lm_head, and its [1, hidden] weight cannot be packed by the NVFP4 exporter.
  • modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml — the FP8 twin: per-tensor FP8 output quantizers on the projection Linears, switching the engine from FP16-out GEMMs (e4m3f16..._bias_f16) to FP8-out GEMMs (e4m3e4m3_e4m3).
  • examples/torch_onnx/hf_embedding_quant_to_onnx.py — minimal end-to-end recipe-driven quantize → ONNX export for HF bidirectional Llama embedding and reranking encoders (auto-detected from the model architecture; embedding models export mean-pooled L2-normalized embeddings, rerankers export relevance logits), with the export shims needed for a TensorRT-fusable graph (bidirectional sdpa symbolic with single-precision attention constants; static blocked-axis extents for Reshape → TRT_FP4DynamicQuantize).
  • modelopt/torch/quantization/export_onnx.pyconfigure_linear_module_onnx_quantizers now types output quantizers as "dynamic" (they previously fell to the static path, which the NVFP4 weight exporter rejects on activations), and the sdpa symbolic's JitScalarType import is fixed for torch >= 2.11.
  • examples/torch_onnx/torch_quant_to_onnx.py — replaces the mtq.*_CFG module-constant table with YAML recipe loading and adds a --recipe flag (preset basename or QuantizeConfig YAML path); --auto_quantization_formats values switch from config-constant names to preset basenames.
  • tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py — end-to-end example test running both model kinds through quantize → export with tiny random-weight stand-ins (plain Llama encoder and LlamaForSequenceClassification, sized to the NVFP4 block size).
  • README section for the new example, including the results table and trtexec engine-build steps; CHANGELOG entry.

Usage

# Embedding model (default recipe: nvfp4 + projection output quantizers)
python examples/torch_onnx/hf_embedding_quant_to_onnx.py \
    --model_path=nvidia/llama-nemotron-embed-1b-v2 \
    --trust_remote_code \
    --onnx_save_path=llama_nemotron_embed_nvfp4.onnx

# Reranking model (auto-detected), FP8 variant of the recipe
python examples/torch_onnx/hf_embedding_quant_to_onnx.py \
    --model_path=nvidia/llama-nemotron-rerank-1b-v2 \
    --trust_remote_code \
    --recipe=huggingface/nemotron_llama/ptq/fp8_output_quant_proj \
    --onnx_save_path=llama_nemotron_rerank_fp8.onnx

# Build a strongly-typed TensorRT engine (Blackwell, TensorRT >= 10.11)
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

Testing

  • New end-to-end example test test_hf_embedding_quant_to_onnx.py runs both model kinds (embedding + reranking) through quantize → ONNX export on tiny random-weight checkpoints; test_torch_onnx_recipe_flag covers the --recipe flag.
  • Ran the example end-to-end on GPU for both real models and both recipes; each graph parses with TensorRT strongly-typed mode (NVFP4 graphs carry 112 input-side + 112 output-side TRT_FP4DynamicQuantize; FP8 graphs carry the matching static Q/DQ placement).
  • Built strongly-typed 5-profile engines on RTX PRO 6000 Blackwell with TensorRT 10.16 for FP16 and both presets/recipes on both models; compared per-profile activation memory via ICudaEngine.get_device_memory_size_for_profile_v2 (table above) and verified kernel selection (FP8 recipe → e4m3e4m3_e4m3 FP8-out GEMMs; NVFP4 → block-scaled GEMMs).
  • Recipes pass tools/precommit/check_modelopt_recipes.py; pre-commit green on all changed files.

Before your PR is "Ready for review"

Make sure you read and follow Contributor guidelines and your commits are signed (git commit -s -S).

Make sure you read and follow the Security Best Practices (e.g. avoiding hardcoded trust_remote_code=True, torch.load(..., weights_only=False), pickle, etc.).

  • Is this change backward compatible?: ❌ — torch_quant_to_onnx.py's --auto_quantization_formats values are renamed from config-constant names (e.g. NVFP4_AWQ_LITE_CFG) to preset basenames (e.g. nvfp4_awq_lite); the loaded configs are identical. Core APIs are backward compatible (the output-quantizer export typing is additive).
  • 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?: ✅ (test_hf_embedding_quant_to_onnx for both model kinds, test_torch_onnx_recipe_flag)
  • Did you update Changelog?: ✅

Additional Information

The target model family requires remote modeling code, so the usage examples opt in explicitly with --trust_remote_code. Accuracy parity (embedding quality / reranking scores) of the output-quantized recipes has not been evaluated yet and should be validated before recommending them as defaults.

🤖 Generated with Claude Code

🤖 Generated by Claude (AI agent).

Summary by CodeRabbit

  • New Features

    • Added NVFP4 and FP8 projection-output PTQ recipes for Llama-Nemotron embedding/reranking.
    • Added a Hugging Face “quantize-to-ONNX” workflow for TensorRT export.
    • Enhanced Torch→ONNX quantization to load YAML recipes via a --recipe flag and refined auto-quantization format handling.
  • Bug Fixes

    • Improved ONNX export for scaled dot-product attention and TensorRT DynamicQuantize compatibility (including output-quantizer behavior).
  • Documentation

    • Updated Torch→ONNX example documentation and PTQ recipe guidance.
  • Tests

    • Added/expanded ONNX graph assertions for recipe and Hugging Face flows.

🤖 Generated by Codex (AI agent).

@copy-pr-bot

copy-pr-bot Bot commented Jul 16, 2026

Copy link
Copy Markdown

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

@coderabbitai

coderabbitai Bot commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

Adds Nemotron FP8/NVFP4 PTQ recipes, a Hugging Face embedding/reranking quantize-to-ONNX example, YAML recipe selection for Torch→ONNX quantization, ONNX export fixes, documentation, and end-to-end tests.

Changes

Quantize-to-ONNX workflow

Layer / File(s) Summary
Projection-output recipes and ONNX quantizer export
modelopt_recipes/huggingface/nemotron_llama/ptq/*, modelopt/torch/quantization/export_onnx.py, modelopt_recipes/ptq.md
Adds FP8 and NVFP4 projection-output recipes, documents the Nemotron-specific quantization scheme, and enables dynamic input/output activation quantizers for block-quantized ONNX exports.
Hugging Face embedding and reranking export
examples/torch_onnx/hf_embedding_quant_to_onnx.py, tests/_test_utils/torch/transformers_models.py, tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py
Adds embedding and reranking wrappers, calibration and PTQ orchestration, SDPA/export fixes, static FP4 reshape correction, tiny sequence-classification fixtures, and ONNX graph tests.
YAML recipe loading and CLI integration
examples/torch_onnx/torch_quant_to_onnx.py, examples/torch_onnx/README.md, tests/examples/torch_onnx/test_torch_quant_to_onnx.py, CHANGELOG.rst
Replaces preset constant lookup with recipe loading, adds --recipe, updates auto-format names, documents both workflows, and tests explicit recipe selection.

Estimated code review effort: 4 (Complex) | ~45 minutes

Suggested reviewers: meenchen, vishalpandya1990

Sequence Diagram(s)

sequenceDiagram
  participant CLI as hf_embedding_quant_to_onnx.py
  participant Model as EmbeddingModel/RerankModel
  participant PTQ as mtq.quantize
  participant Export as ONNX exporter
  participant File as Quantized ONNX file
  CLI->>Model: Prepare wrapped model and calibration inputs
  CLI->>PTQ: Apply selected PTQ recipe
  CLI->>Export: Register SDPA and static-extent fixes
  Export->>File: Write dynamic-shape quantized ONNX output
Loading
🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 21.05% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (5 passed)
Check name Status Explanation
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns ✅ Passed No prohibited patterns found in changed modelopt/examples Python files; trust_remote_code is CLI-gated default False, and no torch.load/numpy.load/nosec/eval/exec abuse appears.
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly matches the main change: a new NVFP4 projection-output recipe and an HF embedding ONNX export example.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch ajrasane/nvbug_5726458_nvfp4_output_quant

Comment @coderabbitai help to get the list of available commands.

@ajrasane
ajrasane marked this pull request as ready for review July 16, 2026 03:42
@ajrasane
ajrasane requested review from a team as code owners July 16, 2026 03:42

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Warning

CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.

Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.

👉 Steps to fix this

Actionable comments posted: 8

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
examples/torch_onnx/torch_quant_to_onnx.py (1)

125-150: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift

Prevent --recipe from disagreeing with --quantize_mode.

A command such as --quantize_mode=fp8 --recipe=nvfp4 loads NVFP4 numerics but injects FP8 overrides here; Lines 613-624 also continue selecting export workarounds from quantize_mode. Validate recipe compatibility or derive these behaviors from the loaded configuration.

Also applies to: 458-468, 598-598

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/torch_onnx/torch_quant_to_onnx.py` around lines 125 - 150, Update
get_quant_config and the related export-workaround selection to prevent recipe
and quantize_mode from diverging: validate that a supplied recipe matches the
requested mode before applying mode-specific overrides, or derive all
override/workaround behavior from the loaded configuration. Ensure mismatched
combinations such as fp8 with an nvfp4 recipe are rejected or consistently
handled.
🧹 Nitpick comments (1)
tests/_test_utils/torch/transformers_models.py (1)

599-625: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Document the new public test helpers.

get_tiny_llama_seq_cls and create_tiny_llama_seq_cls_dir are public module-level helpers but have no docstrings. Add concise documentation covering the sequence-classification output and tokenizer/model directory behavior.

As per coding guidelines, public and higher-level APIs must be documented with docstrings.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/_test_utils/torch/transformers_models.py` around lines 599 - 625, Add
concise docstrings to the public helpers get_tiny_llama_seq_cls and
create_tiny_llama_seq_cls_dir. Document that the first creates a tiny Llama
sequence-classification model and that the second creates its tokenizer/model
directory, including the with_tokenizer behavior.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/torch_onnx/hf_embedding_quant_to_onnx.py`:
- Around line 215-226: Validate both --calibration_data_size and --batch_size as
strictly positive integers during argument parsing, using the existing
positive_int helper if available or an equivalent argparse type. Update the
argument definitions in the parser setup so invalid values are rejected before
the calibration workflow runs.
- Around line 230-243: Add an opt-in CLI flag for remote code trust, defaulting
to false, and use its value for trust_remote_code in
AutoTokenizer.from_pretrained, AutoConfig.from_pretrained,
AutoModelForSequenceClassification.from_pretrained, and
AutoModel.from_pretrained. Remove the hardcoded True values while preserving the
existing model-loading behavior when the flag is explicitly enabled.
- Around line 98-100: Move the imports for register_custom_op_symbolic,
symbolic_helper, and scaled_dot_product_attention from their local locations
into module scope in hf_embedding_quant_to_onnx.py. Apply the same change to the
additional local imports referenced at lines 136 and 148-151, preserving
behavior and avoiding duplicate imports.

In `@examples/torch_onnx/README.md`:
- Around line 113-127: Qualify the shared activation-memory reduction claim to
match the measurements: update examples/torch_onnx/README.md lines 113-127 to
say NVFP4 reduces memory by up to roughly half, or state the per-model
reductions consistently, and update CHANGELOG.rst line 22 to remove the blanket
“roughly halving” wording using the same qualified claim.

In `@examples/torch_onnx/torch_quant_to_onnx.py`:
- Line 624: Update the auto-quantization format predicate near the Conv2d input
quantizer logic to exclude both INT8-based formats, “int8” and “int4_awq”.
Ensure int4_awq follows the same path as int8 and does not disable the
low-channel input quantizer or get classified as using FP8 Conv inputs.

In `@modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml`:
- Around line 35-41: Update the quant_cfg in fp8_output_quant_proj.yaml to
explicitly disable or exclude the reranking score head from FP8 quantization,
matching the existing exclusion pattern used by the NVFP4 recipe. Preserve the
current FP8 output projection configuration for all other modules.

In `@tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py`:
- Around line 47-56: Update the test invoking hf_embedding_quant_to_onnx so it
validates recipe loading rather than only successful export: use a recipe
configuration behaviorally different from the default path, or assert an
exported characteristic uniquely determined by the supplied recipe. Keep the
existing export existence assertion while ensuring the test would fail if
--recipe were accepted but ignored.

In `@tests/examples/torch_onnx/test_torch_quant_to_onnx.py`:
- Around line 51-65: Update test_torch_onnx_recipe_flag so it validates that the
explicit recipe is actually applied rather than merely matching the nvfp4
default. Use a distinct recipe configuration or inspect the exported model for a
recipe-specific result, and retain the existing command execution flow.

---

Outside diff comments:
In `@examples/torch_onnx/torch_quant_to_onnx.py`:
- Around line 125-150: Update get_quant_config and the related export-workaround
selection to prevent recipe and quantize_mode from diverging: validate that a
supplied recipe matches the requested mode before applying mode-specific
overrides, or derive all override/workaround behavior from the loaded
configuration. Ensure mismatched combinations such as fp8 with an nvfp4 recipe
are rejected or consistently handled.

---

Nitpick comments:
In `@tests/_test_utils/torch/transformers_models.py`:
- Around line 599-625: Add concise docstrings to the public helpers
get_tiny_llama_seq_cls and create_tiny_llama_seq_cls_dir. Document that the
first creates a tiny Llama sequence-classification model and that the second
creates its tokenizer/model directory, including the with_tokenizer behavior.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 0707ca97-7db7-477c-9a84-ca7c1e7afb9a

📥 Commits

Reviewing files that changed from the base of the PR and between 95ee9c4 and 05db974.

📒 Files selected for processing (10)
  • CHANGELOG.rst
  • examples/torch_onnx/README.md
  • examples/torch_onnx/hf_embedding_quant_to_onnx.py
  • examples/torch_onnx/torch_quant_to_onnx.py
  • modelopt/torch/quantization/export_onnx.py
  • modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml
  • modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml
  • tests/_test_utils/torch/transformers_models.py
  • tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py
  • tests/examples/torch_onnx/test_torch_quant_to_onnx.py

Comment thread examples/torch_onnx/hf_embedding_quant_to_onnx.py Outdated
Comment thread examples/torch_onnx/hf_embedding_quant_to_onnx.py
Comment thread examples/torch_onnx/hf_embedding_quant_to_onnx.py Outdated
Comment thread examples/torch_onnx/README.md
Comment thread examples/torch_onnx/torch_quant_to_onnx.py Outdated
Comment thread tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py
Comment thread tests/examples/torch_onnx/test_torch_quant_to_onnx.py Outdated
ajrasane and others added 5 commits July 16, 2026 04:42
…ample

Replace the mtq.*_CFG module-constant table with recipe YAML loading and add
a --recipe flag to select a preset basename or a QuantizeConfig YAML path.
Auto-quantize format names switch to preset basenames accordingly.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
…ng ONNX export example

Without output-side quantization, TensorRT's NVFP4 GEMMs emit FP16
activations, so FP4 engines can use more activation memory than FP8 engines
whose output-side Q/DQ enables FP8-out kernels. Quantizing the projection
outputs keeps inter-layer activations in FP4, roughly halving engine
activation memory on the Llama-Nemotron embedding model.

- modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml:
  nvfp4 preset plus dynamic NVFP4 output quantizers scoped to the projection
  Linears (DynamicQuantize on non-GEMM outputs fails to compile in TensorRT).
- examples/torch_onnx/hf_embedding_quant_to_onnx.py: end-to-end recipe-driven
  quantize -> ONNX export for bidirectional Llama embedding encoders, with the
  export shims needed for a TensorRT-fusable graph (bidirectional sdpa
  symbolic, static blocked-axis extents for DynamicQuantize).
- configure_linear_module_onnx_quantizers now types output quantizers as
  dynamic (the static path fails on activations in the NVFP4 exporter) and the
  sdpa symbolic's JitScalarType import is fixed for torch >= 2.11.
- README section covering the example and trtexec engine-build steps.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
…mple

- Recipe: keep the sequence-classification score head unquantized; its
  [1, hidden] weight cannot be packed by the NVFP4 exporter and final heads
  stay in high precision like lm_head. No-op for the embedding model.
- Example: auto-detect ForSequenceClassification architectures, load the
  reranker head, pair-encode calibration inputs, and export relevance logits.
- README: cover both models and add TensorRT 10.16 activation-memory results.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
…export

Tiny plain-Llama stand-ins (a causal encoder and a LlamaForSequenceClassification
with the score head the recipe excludes) run hf_embedding_quant_to_onnx.py through
quantize -> ONNX export for both model kinds. Projection output extents are sized
to the NVFP4 block size (16).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
Applies the same output-side quantization to the FP8 preset: FP8-output GEMM
kernels replace the FP16-out + shadow-copy pattern, cutting engine activation
memory 21% vs the fp8 preset (1392 -> 1096 MiB on both Llama-Nemotron models,
TensorRT 10.16, RTX PRO 6000). README gains the full FP16/FP8/NVFP4 table.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
@ajrasane
ajrasane force-pushed the ajrasane/nvbug_5726458_nvfp4_output_quant branch from 05db974 to 58d9422 Compare July 16, 2026 04:43

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Warning

CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.

Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.

👉 Steps to fix this

Actionable comments posted: 2

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/torch_onnx/hf_embedding_quant_to_onnx.py`:
- Around line 16-39: Move the descriptive string to the top of the module so it
becomes the actual module docstring, before all imports. Add an __all__
declaration exposing the module’s intended public API, and add a docstring to
the main() entry point describing its purpose and usage.

In `@examples/torch_onnx/torch_quant_to_onnx.py`:
- Around line 61-70: The load_quant_config function currently omits the
QuantizeConfig default quant_cfg when serializing. Update its model_dump call to
preserve schema defaults, ensuring config["quant_cfg"] is present for subsequent
extensions in get_quant_config and the auto-quant format path while retaining
explicitly configured values.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f3d7693c-6c50-4568-a81e-1114670f9b96

📥 Commits

Reviewing files that changed from the base of the PR and between 05db974 and 58d9422.

📒 Files selected for processing (10)
  • CHANGELOG.rst
  • examples/torch_onnx/README.md
  • examples/torch_onnx/hf_embedding_quant_to_onnx.py
  • examples/torch_onnx/torch_quant_to_onnx.py
  • modelopt/torch/quantization/export_onnx.py
  • modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml
  • modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml
  • tests/_test_utils/torch/transformers_models.py
  • tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py
  • tests/examples/torch_onnx/test_torch_quant_to_onnx.py
🚧 Files skipped from review as they are similar to previous changes (7)
  • modelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yaml
  • tests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py
  • tests/_test_utils/torch/transformers_models.py
  • examples/torch_onnx/README.md
  • modelopt/torch/quantization/export_onnx.py
  • modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml
  • CHANGELOG.rst

Comment thread examples/torch_onnx/hf_embedding_quant_to_onnx.py Outdated
Comment thread examples/torch_onnx/torch_quant_to_onnx.py Outdated
…tring RST

test_every_model_specific_ptq_dir_is_mentioned requires every huggingface/
model recipe dir to appear in modelopt_recipes/ptq.md; add nemotron_llama to
the architecture-aware kind with its delta. Restore single-line bullets in the
configure_linear_module_onnx_quantizers docstring (docutils rejects multi-line
bullet continuations in that list) and fold the output-quantizer rationale into
the following paragraph.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>

@cjluo-nv cjluo-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Bot review (bedrock-claude-opus-4-8) — DM the bot to share feedback.

New example + two Llama-Nemotron PTQ recipes (NVFP4/FP8 projection-output quantizers) + a small core export change. +630/-31, 10 files. Overall solid and cohesive; recipes correctly reuse the existing $import/units recipe framework and load_recipe/load_config (no new subsystem — design review satisfied), and all new source files carry the canonical multi-line NVIDIA LICENSE_HEADER, so licensing is clean. The install_static_extent_fix monkeypatch correctly targets the module-global torch_onnx.quantize_weights, which get_onnx_bytes_and_metadata resolves at call time, so the patch is effective. Tests are included for both example flows and the --recipe flag, and the tiny-model dims are chosen for NVFP4 block-size divisibility.

Why nudge (owner sign-off) rather than approve:

  • Core-library behavior change: configure_linear_module_onnx_quantizers now types output_quantizer as "dynamic". It's additive and scoped to block-quantized quantizers, but the only coverage is the new GPU-only example test — no focused unit test in tests/unit. Worth a maintainer's eye given it affects all block-quant ONNX exports, not just this example.
  • Backward-incompatible CLI change: torch_quant_to_onnx.py's --auto_quantization_formats values are renamed from config-constant names (NVFP4_AWQ_LITE_CFG) to preset basenames (nvfp4_awq_lite). The PR body flags this, but the CHANGELOG entry only calls out the new --recipe flag, not the value rename.
  • Author explicitly notes accuracy parity of the output-quantized recipes has not been evaluated yet and should be validated before recommending as defaults — this warrants human judgment before merge.
  • Example-only global side effects (never restored): install_static_extent_fix permanently replaces torch_onnx.quantize_weights, and register_bidirectional_sdpa overwrites sdpa_attention.use_gqa_in_sdpa globally. Acceptable for a single-shot example script but worth an owner glance.

No prompt-injection content in the untrusted blocks. Minor nit (non-blocking): several imports in hf_embedding_quant_to_onnx.py are function-local without a stated reason.

@github-actions

Copy link
Copy Markdown
Contributor
PR Preview Action v1.8.1

QR code for preview link

🚀 View preview at
https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-1981/

Built to branch gh-pages at 2026-07-16 17:08 UTC.
Preview will be ready when the GitHub Pages deployment is complete.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Warning

CodeRabbit couldn't request changes on this pull request because it doesn't have sufficient GitHub permissions.

Please grant CodeRabbit Pull requests: Read and write permission and re-run the review.

👉 Steps to fix this

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@modelopt_recipes/ptq.md`:
- Around line 279-281: Update the sequence-classification score-head explanation
in the recipe documentation so it states that both FP8 and NVFP4 recipes leave
score unquantized, while limiting the [1, hidden] weight packing limitation to
the NVFP4 exporter. Keep the existing high-precision final-head rationale and
make the exporter-specific distinction explicit.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 4f9641c2-77de-4c0b-bc9a-95fb7ef5223c

📥 Commits

Reviewing files that changed from the base of the PR and between 58d9422 and 53e68d4.

📒 Files selected for processing (2)
  • modelopt/torch/quantization/export_onnx.py
  • modelopt_recipes/ptq.md
🚧 Files skipped from review as they are similar to previous changes (1)
  • modelopt/torch/quantization/export_onnx.py

Comment thread modelopt_recipes/ptq.md Outdated
@codecov

codecov Bot commented Jul 16, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 66.93%. Comparing base (95ee9c4) to head (ed7db4d).
⚠️ Report is 17 commits behind head on main.

Additional details and impacted files
@@             Coverage Diff             @@
##             main    #1981       +/-   ##
===========================================
- Coverage   77.44%   66.93%   -10.52%     
===========================================
  Files         522      525        +3     
  Lines       58452    61716     +3264     
===========================================
- Hits        45271    41310     -3961     
- Misses      13181    20406     +7225     
Flag Coverage Δ
examples 43.40% <100.00%> (-0.04%) ⬇️
gpu 32.04% <0.00%> (-26.11%) ⬇️
regression 15.07% <0.00%> (+0.07%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@vishalpandya1990

Copy link
Copy Markdown
Contributor

I see example created for embedding model - is this relevant for LLMs too? i.e. Does the problem of high activation VRAM is seen with NVFP4 LLMs too?

@ajrasane

Copy link
Copy Markdown
Contributor Author

@vishalpandya1990, yes we will run into this issue for other LLMs too if we are using the default recipes with TensorRT deployment.

Make remote code execution opt-in and validate positive CLI sizes. Preserve
recipe defaults and align the FP8 score-head handling with NVFP4. Strengthen
the example tests so the two recipes and custom recipe loading are distinct.

Co-Authored-By: OpenAI Codex <noreply@openai.com>
Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants