[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 6 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
Caution Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional.
❌ Failed checks (1 error, 1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
| from torch.onnx import register_custom_op_symbolic, symbolic_helper | ||
|
|
||
| from modelopt.torch.quantization.export_onnx import scaled_dot_product_attention |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Move these imports to module scope or document why they must remain local.
The imports are neither accompanied by an optional-dependency, circular-dependency, nor heavy-import justification.
As per coding guidelines, local imports require a brief explanatory comment.
Also applies to: 136-136, 148-151
🤖 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/hf_embedding_quant_to_onnx.py` around lines 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.
Source: Coding guidelines
| parser.add_argument( | ||
| "--calibration_data_size", | ||
| type=int, | ||
| default=64, | ||
| help="Number of calibration samples.", | ||
| ) | ||
| parser.add_argument( | ||
| "--batch_size", | ||
| type=int, | ||
| default=8, | ||
| help="Batch size for calibration.", | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reject non-positive calibration and batch sizes at argument parsing.
--batch_size=0 crashes range() at Line 251, while non-positive calibration sizes produce no calibration samples. Use type=positive_int or validate immediately after parsing.
As per coding guidelines and path instructions, external CLI input must be validated at the interface boundary.
🤖 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/hf_embedding_quant_to_onnx.py` around lines 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.
Sources: Coding guidelines, Path instructions
| tokenizer = AutoTokenizer.from_pretrained(args.model_path, trust_remote_code=True) | ||
| if tokenizer.pad_token is None: | ||
| tokenizer.pad_token = tokenizer.eos_token | ||
| config = AutoConfig.from_pretrained(args.model_path, trust_remote_code=True) | ||
| is_reranker = any("ForSequenceClassification" in a for a in (config.architectures or [])) | ||
| if is_reranker: | ||
| base = AutoModelForSequenceClassification.from_pretrained( | ||
| args.model_path, trust_remote_code=True, dtype=torch.float32 | ||
| ) | ||
| model = RerankModel(base).to(device).eval() | ||
| else: | ||
| base = AutoModel.from_pretrained( | ||
| args.model_path, trust_remote_code=True, dtype=torch.float32 | ||
| ) |
There was a problem hiding this comment.
🔒 Security & Privacy | 🔴 Critical | ⚡ Quick win
Do not hardcode trust_remote_code=True.
--model_path accepts arbitrary Hub IDs or local paths, so these calls execute Python supplied by an untrusted model repository. Expose an opt-in CLI flag defaulting to False and pass it to all four loaders.
As per coding guidelines and path instructions, trust_remote_code must be caller-configurable and default to false.
🧰 Tools
🪛 ast-grep (0.44.1)
[warning] 232-232: Loading a HuggingFace model, tokenizer, dataset or pipeline with trust_remote_code=True executes arbitrary code from the remote repository. Only enable it for fully trusted sources; prefer artifacts that do not require it.
Context: trust_remote_code=True
Note: [CWE-94] Improper Control of Generation of Code ('Code Injection').
(transformers-trust-remote-code-python)
[warning] 236-236: Loading a HuggingFace model, tokenizer, dataset or pipeline with trust_remote_code=True executes arbitrary code from the remote repository. Only enable it for fully trusted sources; prefer artifacts that do not require it.
Context: trust_remote_code=True
Note: [CWE-94] Improper Control of Generation of Code ('Code Injection').
(transformers-trust-remote-code-python)
[warning] 241-241: Loading a HuggingFace model, tokenizer, dataset or pipeline with trust_remote_code=True executes arbitrary code from the remote repository. Only enable it for fully trusted sources; prefer artifacts that do not require it.
Context: trust_remote_code=True
Note: [CWE-94] Improper Control of Generation of Code ('Code Injection').
(transformers-trust-remote-code-python)
🤖 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/hf_embedding_quant_to_onnx.py` around lines 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.
Sources: Coding guidelines, Path instructions, Linters/SAST tools
| uses_fp8_conv_input = args.quantize_mode in ("fp8", "mxfp8", "nvfp4") or ( | ||
| args.quantize_mode == "auto" | ||
| and any(fmt != "INT8_DEFAULT_CFG" for fmt in args.auto_quantization_formats) | ||
| and any(fmt != "int8" for fmt in args.auto_quantization_formats) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Do not classify int4_awq as using FP8 Conv inputs.
Auto mode overrides int4_awq Conv2d quantizers to INT8, but this predicate still disables the low-channel input quantizer. Exclude both INT8-based formats.
Proposed fix
- and any(fmt != "int8" for fmt in args.auto_quantization_formats)
+ and any(fmt not in {"int8", "int4_awq"} for fmt in args.auto_quantization_formats)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| and any(fmt != "int8" for fmt in args.auto_quantization_formats) | |
| and any(fmt not in {"int8", "int4_awq"} for fmt in args.auto_quantization_formats) |
🤖 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` at 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.
| quant_cfg: | ||
| - $import: base_disable_all | ||
| - $import: w8a8_fp8_fp8 | ||
| - $import: default_disabled_quantizers | ||
| - quantizer_name: '*_proj.output_quantizer' | ||
| cfg: | ||
| $import: fp8 |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Exclude the reranking score head from FP8 quantization.
This recipe claims reranking support, but the base W8A8 configuration remains active on score. Add the same explicit exclusion used by the NVFP4 recipe to preserve the unsupported/high-precision head.
Proposed fix
- quantizer_name: '*_proj.output_quantizer'
cfg:
$import: fp8
+ - quantizer_name: '*score*'
+ enable: false📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| quant_cfg: | |
| - $import: base_disable_all | |
| - $import: w8a8_fp8_fp8 | |
| - $import: default_disabled_quantizers | |
| - quantizer_name: '*_proj.output_quantizer' | |
| cfg: | |
| $import: fp8 | |
| 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 |
🤖 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 `@modelopt_recipes/huggingface/nemotron_llama/ptq/fp8_output_quant_proj.yaml`
around lines 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.
| cmd_parts = extend_cmd_parts( | ||
| ["python", "hf_embedding_quant_to_onnx.py"], | ||
| model_path=str(model_dir), | ||
| onnx_save_path=str(onnx_save_path), | ||
| calibration_data_size="2", | ||
| batch_size="2", | ||
| ) | ||
| run_example_command(cmd_parts, "torch_onnx") | ||
|
|
||
| assert onnx_save_path.exists() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Make the test distinguish recipe loading from the default path.
quantize_mode="nvfp4" and recipe="configs/ptq/presets/model/nvfp4" select the same configuration. If --recipe were accepted but ignored, this test would still pass. Use a behaviorally distinct recipe or assert a recipe-specific exported characteristic.
As per path instructions, tests must exercise the behavior they claim to validate.
🤖 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/examples/torch_onnx/test_hf_embedding_quant_to_onnx.py` around lines 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.
Source: Path instructions
| def test_torch_onnx_recipe_flag(): | ||
| timm_model_name, model_kwargs = _MODELS["vit_tiny"] | ||
|
|
||
| cmd_parts = extend_cmd_parts( | ||
| ["python", "torch_quant_to_onnx.py"], | ||
| timm_model_name=timm_model_name, | ||
| model_kwargs=model_kwargs, | ||
| quantize_mode="nvfp4", | ||
| recipe="configs/ptq/presets/model/nvfp4", | ||
| onnx_save_path="vit_tiny.recipe.onnx", | ||
| calibration_data_size="1", | ||
| num_score_steps="1", | ||
| ) | ||
| cmd_parts.extend(["--no_pretrained", "--trt_build"]) | ||
| run_example_command(cmd_parts, "torch_onnx") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Make the recipe-flag test verify recipe application.
The explicit recipe matches the nvfp4 mode’s default recipe, so this test would pass even if the CLI ignored the recipe contents. Use a distinct recipe or assert a recipe-specific result in the exported model.
As per path instructions, tests must exercise the behavior they claim to validate.
🤖 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/examples/torch_onnx/test_torch_quant_to_onnx.py` around lines 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.
Source: Path instructions
…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
force-pushed
the
ajrasane/nvbug_5726458_nvfp4_output_quant
branch
from
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
| import argparse | ||
| import os | ||
|
|
||
| import torch | ||
| import torch.nn.functional as F | ||
| from transformers import AutoConfig, AutoModel, AutoModelForSequenceClassification, AutoTokenizer | ||
|
|
||
| import modelopt.torch.quantization as mtq | ||
| from modelopt.recipe import load_recipe | ||
| from modelopt.torch._deploy.utils import OnnxBytes, get_onnx_bytes_and_metadata | ||
|
|
||
| """ | ||
| Quantize an HF embedding or reranking model (bidirectional Llama encoder, e.g. | ||
| nvidia/llama-nemotron-embed-1b-v2 or nvidia/llama-nemotron-rerank-1b-v2) with a | ||
| PTQ recipe and export it to ONNX for TensorRT deployment. | ||
|
|
||
| The default recipe quantizes weights and activations to NVFP4 and additionally | ||
| quantizes the projection-Linear outputs so the TensorRT engine keeps | ||
| inter-layer activations in FP4 (see modelopt_recipes/huggingface/nemotron_llama/). | ||
| Embedding models are exported with mean pooling and L2 normalization on top of | ||
| the encoder; reranking (sequence-classification) models are exported to their | ||
| relevance logits. Both graphs take (input_ids, attention_mask) with dynamic | ||
| batch/sequence axes. | ||
| """ |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Make the description the actual module docstring and declare the public API.
Because the string starts after imports, module.__doc__ remains unset and the expression is dead code. Move it above imports, add __all__, and document main().
Proposed fix
+"""Quantize an HF embedding or reranking model and export it to ONNX."""
+
import argparse
import os
...
-"""
-Quantize an HF embedding or reranking model ...
-"""
+__all__ = [
+ "EmbeddingModel",
+ "RerankModel",
+ "install_static_extent_fix",
+ "main",
+ "register_bidirectional_sdpa",
+]
...
def main():
+ """Run the quantization and ONNX export CLI."""As per coding guidelines, define __all__, document higher-level APIs, and remove dead code.
Also applies to: 196-199
🤖 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/hf_embedding_quant_to_onnx.py` around lines 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.
Source: Coding guidelines
| def load_quant_config(recipe: str) -> dict: | ||
| """Load a quantization config from a recipe YAML. | ||
|
|
||
| ``recipe`` is either a preset basename under | ||
| ``modelopt_recipes/configs/ptq/presets/model/`` (e.g. ``nvfp4``) or a path to a | ||
| ``QuantizeConfig`` YAML file (filesystem or built-in recipe library). | ||
| """ | ||
| if "/" not in recipe and not recipe.endswith((".yml", ".yaml")): | ||
| recipe = f"{MODEL_QUANT_PRESET_DIR}/{recipe}" | ||
| return load_config(recipe, schema_type=QuantizeConfig).model_dump(exclude_unset=True) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the relevant files and inspect the surrounding code.
git ls-files | rg '^(examples/torch_onnx/torch_quant_to_onnx\.py|.*QuantizeConfig.*|.*config.*\.py|.*schema.*\.py)$' || true
printf '\n--- torch_quant_to_onnx.py (around load_quant_config and usages) ---\n'
sed -n '1,220p' examples/torch_onnx/torch_quant_to_onnx.py
printf '\n--- search for quant_cfg usages in the example ---\n'
rg -n '"quant_cfg"|quant_cfg' examples/torch_onnx/torch_quant_to_onnx.py
printf '\n--- locate QuantizeConfig definition ---\n'
rg -n 'class QuantizeConfig|QuantizeConfig\(' -S .
printf '\n--- locate load_config implementation ---\n'
rg -n 'def load_config\(' -S .
printf '\n--- inspect likely schema/default definitions if found ---\n'
python3 - <<'PY'
from pathlib import Path
import re
paths = [
Path('examples/torch_onnx/torch_quant_to_onnx.py'),
]
for p in paths:
if p.exists():
print(f"\nFILE: {p}")
text = p.read_text()
for pat in ['load_quant_config', 'quant_cfg', 'exclude_unset=True', 'exclude_unset=False']:
if pat in text:
print(f" contains: {pat}")
PYRepository: NVIDIA/Model-Optimizer
Length of output: 18134
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Narrow the search to the repository for the schema model and loader.
rg -n "class QuantizeConfig|QuantizeConfig\\b|def load_config\\(" . -SRepository: NVIDIA/Model-Optimizer
Length of output: 31184
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '\n--- QuantizeConfig definition around quant_cfg ---\n'
sed -n '1496,1605p' modelopt/torch/quantization/config.py
printf '\n--- helper that also uses exclude_unset=True ---\n'
sed -n '1568,1588p' modelopt/torch/quantization/config.py
printf '\n--- override logic around lines 120-160 ---\n'
sed -n '120,170p' examples/torch_onnx/torch_quant_to_onnx.py
printf '\n--- override logic around lines 380-410 ---\n'
sed -n '380,410p' examples/torch_onnx/torch_quant_to_onnx.py
printf '\n--- inspect load_config contract in modelopt/recipe/loader.py ---\n'
sed -n '200,255p' modelopt/recipe/loader.py
printf '\n--- inspect load_config implementation signature/behavior ---\n'
sed -n '619,760p' modelopt/torch/opt/config_loader.pyRepository: NVIDIA/Model-Optimizer
Length of output: 13256
Preserve quant_cfg when dumping QuantizeConfig. exclude_unset=True drops the schema default for valid recipes that omit quant_cfg, so later config["quant_cfg"].extend(...) calls in both get_quant_config() and the auto-quant format path can hit KeyError. Dump with defaults or merge the schema defaults before applying overrides.
🤖 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 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.
…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
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
| engine activation memory on these models. The sequence-classification `score` | ||
| head stays unquantized (final heads stay high precision like `lm_head`, and its | ||
| `[1, hidden]` weight cannot be packed by the NVFP4 exporter). *Why special:* the |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Make the score-head rationale recipe-specific.
The [1, hidden] packing limitation applies to the NVFP4 exporter, not the FP8 recipe. Clarify that both recipes leave score unquantized, while only NVFP4 requires this additional exporter-specific explanation.
Proposed wording
-The sequence-classification `score` head stays unquantized (final heads stay high precision like `lm_head`, and its `[1, hidden]` weight cannot be packed by the NVFP4 exporter).
+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.📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| engine activation memory on these models. The sequence-classification `score` | |
| head stays unquantized (final heads stay high precision like `lm_head`, and its | |
| `[1, hidden]` weight cannot be packed by the NVFP4 exporter). *Why special:* the | |
| 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 |
🧰 Tools
🪛 LanguageTool
[uncategorized] ~280-~280: If this is a compound adjective that modifies the following noun, use a hyphen.
Context: ...ead stays unquantized (final heads stay high precision like lm_head, and its [1, hidden] w...
(EN_COMPOUND_ADJECTIVE_INTERNAL)
🤖 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 `@modelopt_recipes/ptq.md` around lines 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.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1981 +/- ##
==========================================
- Coverage 77.44% 75.25% -2.20%
==========================================
Files 522 525 +3
Lines 58452 61716 +3264
==========================================
+ Hits 45271 46445 +1174
- Misses 13181 15271 +2090
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:
|
2 participants
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 example hardcodes
trust_remote_code=Truebecause the target model family requires remote modeling code; flagged here per the security guidelines. 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.Bug Fixes
Documentation
Tests