[5726458] Add NVFP4 projection-output-quantizer recipe and HF embedding ONNX export example#1981
[5726458] Add NVFP4 projection-output-quantizer recipe and HF embedding ONNX export example#1981ajrasane wants to merge 7 commits into
Conversation
|
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. |
📝 WalkthroughWalkthroughAdds 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. ChangesQuantize-to-ONNX workflow
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: 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
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
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.
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 liftPrevent
--recipefrom disagreeing with--quantize_mode.A command such as
--quantize_mode=fp8 --recipe=nvfp4loads NVFP4 numerics but injects FP8 overrides here; Lines 613-624 also continue selecting export workarounds fromquantize_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 winDocument the new public test helpers.
get_tiny_llama_seq_clsandcreate_tiny_llama_seq_cls_dirare 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
📒 Files selected for processing (10)
CHANGELOG.rstexamples/torch_onnx/README.mdexamples/torch_onnx/hf_embedding_quant_to_onnx.pyexamples/torch_onnx/torch_quant_to_onnx.pymodelopt/torch/quantization/export_onnx.pymodelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yamlmodelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yamltests/_test_utils/torch/transformers_models.pytests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.pytests/examples/torch_onnx/test_torch_quant_to_onnx.py
…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>
05db974 to
58d9422
Compare
There was a problem hiding this comment.
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.
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
📒 Files selected for processing (10)
CHANGELOG.rstexamples/torch_onnx/README.mdexamples/torch_onnx/hf_embedding_quant_to_onnx.pyexamples/torch_onnx/torch_quant_to_onnx.pymodelopt/torch/quantization/export_onnx.pymodelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yamlmodelopt_recipes/huggingface/nemotron_llama/ptq/nvfp4_output_quant_proj.yamltests/_test_utils/torch/transformers_models.pytests/examples/torch_onnx/test_hf_embedding_quant_to_onnx.pytests/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
…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
left a comment
There was a problem hiding this comment.
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_quantizersnow typesoutput_quantizeras"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 intests/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_formatsvalues 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--recipeflag, 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_fixpermanently replacestorch_onnx.quantize_weights, andregister_bidirectional_sdpaoverwritessdpa_attention.use_gqa_in_sdpaglobally. 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.
|
There was a problem hiding this comment.
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.
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
📒 Files selected for processing (2)
modelopt/torch/quantization/export_onnx.pymodelopt_recipes/ptq.md
🚧 Files skipped from review as they are similar to previous changes (1)
- modelopt/torch/quantization/export_onnx.py
Codecov Report✅ All modified and coverable lines are covered by tests. 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
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
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? |
|
@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>
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:
fp8presetnvfp4presetEngine 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 generalnvfp4preset plus dynamic NVFP4 output quantizers scoped to the projection Linears (*_proj.output_quantizer). Scoping matters: aDynamicQuantizeon non-GEMM outputs (embedding lookup, pooling) fails to compile in TensorRT. The sequence-classificationscorehead is kept unquantized: final heads stay in high precision likelm_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 forReshape → TRT_FP4DynamicQuantize).modelopt/torch/quantization/export_onnx.py—configure_linear_module_onnx_quantizersnow types output quantizers as"dynamic"(they previously fell to the static path, which the NVFP4 weight exporter rejects on activations), and the sdpa symbolic'sJitScalarTypeimport is fixed for torch >= 2.11.examples/torch_onnx/torch_quant_to_onnx.py— replaces themtq.*_CFGmodule-constant table with YAML recipe loading and adds a--recipeflag (preset basename orQuantizeConfigYAML path);--auto_quantization_formatsvalues 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 andLlamaForSequenceClassification, sized to the NVFP4 block size).Usage
Testing
test_hf_embedding_quant_to_onnx.pyruns both model kinds (embedding + reranking) through quantize → ONNX export on tiny random-weight checkpoints;test_torch_onnx_recipe_flagcovers the--recipeflag.TRT_FP4DynamicQuantize; FP8 graphs carry the matching static Q/DQ placement).ICudaEngine.get_device_memory_size_for_profile_v2(table above) and verified kernel selection (FP8 recipe →e4m3e4m3_e4m3FP8-out GEMMs; NVFP4 → block-scaled GEMMs).tools/precommit/check_modelopt_recipes.py;pre-commitgreen 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.).torch_quant_to_onnx.py's--auto_quantization_formatsvalues 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).CONTRIBUTING.md: N/Atest_hf_embedding_quant_to_onnxfor both model kinds,test_torch_onnx_recipe_flag)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
Summary by CodeRabbit
New Features
--recipeflag and refined auto-quantization format handling.Bug Fixes
Documentation
Tests