[6721556] Fix mixed dtypes in NVFP4 FP16 ONNX export

[ajrasane]

What does this PR do?

Type of change: Bug fix

FP16 export of NVFP4-quantized ONNX graphs could invoke TensorRT parsing while elementwise operations still had mixed FP32 and FP16 inputs. TensorRT rejected the first mixed-type Add before AutoCast could normalize the graph.

This change:

• Routes FP16 NVFP4 exports through the existing Q/DQ-aware FP16 conversion path before TensorRT custom-node discovery.
• Normalizes the ONNX opset after NVFP4 weight post-processing so FP4 DequantizeLinear nodes use their required opset-23 schema.
• Keeps NVFP4 MatMul and Gemm compute internally in FP16 or BF16, casts an optional Gemm bias to that compute type, and restores original FP32 graph boundaries with explicit output casts.
• Preserves pure native-BF16 NVFP4 export as a weight-focused no-op.
• Rejects unsupported FP32-to-BF16 and BF16-to-FP16 NVFP4 conversions with clear errors.
• Adds real CPU NVFP4 graph coverage for opset validity and consistent MatMul, Gemm, and elementwise types, plus focused routing and conversion-guard coverage.
• Updates the changelog.

Usage

cd examples/torch_onnx
python torch_quant_to_onnx.py \
  --timm_model_name swin_tiny_patch4_window7_224 \
  --model_kwargs '{"depths": [1, 1, 1, 1]}' \
  --qformat nvfp4 \
  --onnx_save_path swin_tiny.nvfp4.onnx \
  --calibration_data_size 1 \
  --num_score_steps 1 \
  --no_pretrained \
  --trt_build

Testing

• Added regression-first coverage that exposed the opset-20 FP4 graph, the premature BF16 AutoCast/TensorRT path, mixed FP32/FP16 MatMul output boundaries, and a Gemm with a mismatched FP32 bias. Before the final fix, full ONNX checking failed with C has inconsistent type tensor(float).

• Ran the focused current-head CPU coverage with no GPU visible:

  CUDA_VISIBLE_DEVICES="" python -m pytest -o addopts= -q \
    tests/unit/torch/quantization/test_onnx_export_cpu.py::test_nvfp4_exported_onnx_is_topologically_sorted \
    tests/unit/torch/quantization/test_onnx_export_cpu.py::test_nvfp4_deploy_export_has_consistent_elementwise_types \
    tests/unit/torch/quantization/test_onnx_export_cpu.py::test_nvfp4_deploy_export_has_consistent_gemm_types \
    tests/unit/torch/quantization/test_onnx_export_cpu.py::test_nvfp4_shared_activation_reuses_cast \
    tests/unit/torch/quantization/test_onnx_export_cpu.py::test_nvfp4_packed_weights_match_eager \
    tests/unit/torch/deploy/utils/test_torch_onnx_utils.py::test_nvfp4_export_rejects_unsupported_dtype_conversion \
    tests/unit/torch/deploy/utils/test_torch_onnx_utils.py::test_nvfp4_export_selects_precision_converter \
    tests/unit/onnx/quantization/test_qdq_utils.py::TestFP4QDQTo2DQ::test_fp4qdq_conversion

  Result: 20 passed in 1.03 seconds (19.93 seconds including container startup).

• The real CPU regressions cover FP32 preservation, FP32-to-FP16 conversion, and native-BF16 no-op behavior across MatMul plus explicit bias Add and biased Gemm topologies. They output opset-23 graphs, retain dynamic FP4 activation quantization, pass full ONNX checking and strict shape inference, and verify the expected internal compute and graph-boundary types.

• Ran changed-file hooks after merging current main. All applicable hooks passed, including Ruff, mypy, RST, license, and Bandit checks in 7.81 seconds.

• The Blackwell TensorRT engine build was not rerun after the final exporter change; current-head hardware acceptance remains pending.

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?: ❌ No. NVFP4 export now raises ValueError for weights_dtype="bf16" with FP32 source parameters and for weights_dtype="fp16" with BF16 source parameters, rather than attempting unsupported graph-wide conversions. Use FP32-to-FP16 export or native-BF16 export instead.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
• Did you write any new necessary tests?: ✅
• Did you update Changelog?: ✅
• Did you get Claude approval on this PR?: ❌

Additional Information

• BF16 handling is intentionally limited to a pure native-BF16 NVFP4 model with weights_dtype="bf16"; graph-wide quantized FP32-to-BF16 conversion remains unsupported and fails early.
• The default weights_dtype="fp32" path keeps NVFP4 weighted operations in their supported low-precision compute type and restores FP32 at the surrounding graph boundary.
• The public no-pretrained reproducer is an export and engine-build acceptance case, not an accuracy oracle.

Summary by CodeRabbit

• Bug Fixes
  • Improved mixed-precision ONNX export for NVFP4 and FP8 models.
  • Added clearer validation when converting BF16 model parameters to FP16, helping prevent unsupported export configurations.
  • Ensured export uses the appropriate precision conversion path for FP16 and BF16 outputs.

🤖 Generated by Codex (AI agent).

[copy-pr-bot]

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]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 19262493-a0c1-4b57-b61f-3044d7d7e4f9

📥 Commits

Reviewing files that changed from the base of the PR and between 1ffc34f and 2b8c579.

📒 Files selected for processing (3)

• CHANGELOG.rst
• modelopt/onnx/export/nvfp4_exporter.py
• tests/unit/torch/quantization/test_onnx_export_cpu.py

🚧 Files skipped from review as they are similar to previous changes (1)

• CHANGELOG.rst

Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review.

---

📝 Walkthrough

Walkthrough

The change updates NVFP4 and FP8 ONNX precision handling. It adds FP4 graph opset normalization, preserves compatible BF16 paths, aligns casted output types, and adds regression tests for conversion errors and mixed-precision graphs.

Changes

ONNX precision conversion and NVFP4 export

Layer / File(s)	Summary
Conversion validation and export handling modelopt/torch/_deploy/utils/torch_onnx.py	FP4 and FP8 graphs support BF16 no-op handling, reject unsupported conversions, and normalize FP4 graphs to the required opset.
NVFP4 graph casting modelopt/onnx/export/nvfp4_exporter.py	The exporter casts Gemm inputs to the selected precision and restores original output types when required.
NVFP4 regression coverage tests/unit/torch/deploy/utils/test_torch_onnx_utils.py, tests/unit/torch/quantization/test_onnx_export_cpu.py, CHANGELOG.rst	Tests cover conversion errors, converter selection, BF16 preservation, ONNX validity, and consistent elementwise and Gemm dtypes. The changelog records the ONNX parsing fix.

Priority: ➖ Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Suggested reviewers: gcunhase

Merge Risk: ⚪ Minimal · up to 2b8c5

The NVFP4 precision conversion and graph-casting changes preserve dtype boundaries and are covered by CPU export regression tests.

🚥 Pre-merge checks | ✅ 5 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 11.11% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 18 functions across 4 files. (1 skipped: …	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 security anti-pattern is introduced. The PR changes only two modelopt Python files, two test files, and CHANGELOG.rst; no examples or dependency manifests changed. Added-line scans found no unsafe …
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: fixing mixed data types in NVFP4 FP16 ONNX export.

Full details: Docstring Coverage

Explanation

Docstring coverage is 11.11% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 18 functions across 4 files. (1 skipped: 1 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches 💡 1

📝 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-6721556-swin-mixed-dtype-add

---

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

[github-actions]

PR Preview Action v1.8.1
🚀 View preview at https://NVIDIA.github.io/Model-Optimizer/pr-preview/pr-2415/
Built to branch gh-pages at 2026-09-12 04:56 UTC. Preview will be ready when the GitHub Pages deployment is complete.

[codecov]

Codecov Report

❌ Patch coverage is 97.22222% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 78.94%. Comparing base (51de53e) to head (2b8c579).

Files with missing lines	Patch %	Lines
modelopt/onnx/export/nvfp4_exporter.py	95.83%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2415      +/-   ##
==========================================
+ Coverage   71.39%   78.94%   +7.54%     
==========================================
  Files         590      590              
  Lines       64646    64675      +29     
==========================================
+ Hits        46154    51056    +4902     
+ Misses      18492    13619    -4873     

Flag	Coverage Δ
examples-diffusers	20.87% <8.33%> (-0.01%)	⬇️
examples-gpt-oss	13.38% <0.00%> (-0.02%)	⬇️
examples-hf_ptq	22.44% <0.00%> (-0.06%)	⬇️
examples-llm_distill	13.45% <0.00%> (-0.02%)	⬇️
examples-llm_eval	17.36% <0.00%> (-0.02%)	⬇️
examples-llm_qat	17.69% <0.00%> (-0.03%)	⬇️
examples-llm_sparsity	15.92% <0.00%> (-0.02%)	⬇️
examples-megatron_bridge	26.26% <0.00%> (-0.15%)	⬇️
examples-specdec_bench	13.14% <0.00%> (-0.02%)	⬇️
examples-speculative_decoding	17.77% <0.00%> (-0.09%)	⬇️
examples-torch_onnx	21.55% <86.11%> (-0.31%)	⬇️
examples-torch_trt	15.20% <0.00%> (-0.02%)	⬇️
gpu	58.33% <80.55%> (+25.93%)	⬆️
regression	15.14% <0.00%> (+0.27%)	⬆️
unit	57.89% <97.22%> (+0.12%)	⬆️

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.

[cjluo-nv]

Bot review (claude-opus-5) — DM the bot to share feedback.

Nudge: the branch change itself is sound, but the new tests mock out both converters so nothing verifies the actual mixed-dtype fix, and the FP16 NVFP4 path now skips AutoCast's opset upgrade.

Needs action:

• Confirm the opset regression in modelopt/torch/_deploy/utils/torch_onnx.py: convert_to_f16 bumped NVFP4 graphs to opset 23 (get_min_opset_for_precisions), the convert_float_to_float16 branch does not, so default exports now ship opset-20 metadata over opset-23 FP4 DequantizeLinear — fix it here or say why it's acceptable.
• Add an end-to-end check that the exported FP16 NVFP4 graph has no mixed float-type elementwise inputs; tests/unit/torch/quantization/test_onnx_export_cpu.py::test_nvfp4_exported_onnx_is_topologically_sorted shows a CPU-feasible NVFP4 export pattern.
• Explain whether weights_dtype="bf16" + NVFP4 still hits the same TensorRT parse failure, since it keeps routing through convert_to_f16 → GraphSanitizer.find_custom_nodes.

No action needed:

• The added weights_dtype == "fp16" guard is a no-op for BF16 (the earlier ValueError already covers FP8/other quantizers), so the BF16 path is genuinely preserved.
• PR body notes the change is Codex-authored without Claude approval, so a human pass is warranted regardless.

[ajrasane]

@cjluo-nv Addressed the three review items in 1ffc34f89:

• Opset: NVFP4 Q/DQ precision requirements are now derived after weight post-processing, and the graph is upgraded to opset 23 before precision conversion without invoking TensorRT custom-node discovery.
• End-to-end coverage: added a real CPU NVFP4 MatMul→Add export regression. It retains dynamic FP4 activation quantization, passes full ONNX checking and strict shape inference, and asserts that every elementwise node has consistent floating input types.
• BF16: pure native-BF16 NVFP4 export is now an explicit weight-focused no-op, so it does not enter AutoCast/TensorRT pre-parsing. FP32-to-BF16 and BF16-to-FP16 NVFP4 conversions fail early. The exporter also keeps MatMul/Gemm output annotations consistent with its inserted BF16/FP16 casts.

Validation: 128 affected CPU tests passed; the public Swin NVFP4 reproducer exported an opset-23 graph with zero mixed floating-type elementwise operations and completed a strongly typed TensorRT 11.1 engine build on Blackwell.

🤖 Generated by Codex (AI agent).

[cjluo-nv]

Bot review (claude-opus-5) — DM the bot to share feedback.

Nudge: all three prior concerns are fixed with real code and real CPU coverage, but the PR quietly turns two previously-working NVFP4 conversions into hard errors and the agent-authored change still wants a human pass.

Needs action:

• Document the new hard failures in CHANGELOG.rst: weights_dtype="bf16" with FP32 params + NVFP4, and weights_dtype="fp16" with BF16 params + NVFP4, previously ran through convert_to_f16 and now raise ValueError in modelopt/torch/_deploy/utils/torch_onnx.py. The PR body still marks this backward compatible.
• Confirm the weights_dtype="fp32" NVFP4 path: the new if uses_fp4: opset block runs there too and nvfp4_exporter.post_process now stamps MatMul outputs as FP16, so that graph can still carry mixed FP16/FP32 elementwise inputs — say whether it is out of scope or add it to test_nvfp4_deploy_export_has_consistent_elementwise_types.
• Have a human review the Codex-authored commits before merge.

No action needed:

• ✔️ Resolved since the last review: the opset-23 upgrade before precision conversion, the real CPU NVFP4 export test asserting no mixed-float elementwise inputs, and the BF16 no-op routing.
• Refactor of test_nvfp4_exported_onnx_is_topologically_sorted into _make_cpu_nvfp4_model keeps coverage intact.

[ajrasane]

Addressed the current-head review items in c97ae4a5d and merged the latest main in 2b8c579cb:

• Expanded the changelog and corrected the PR compatibility declaration to document the new FP32-to-BF16 and BF16-to-FP16 ValueError behavior.
• Fixed the accepted weights_dtype="fp32" path: NVFP4 MatMul/Gemm operations stay in their supported low-precision compute type, an optional Gemm bias is cast to that type, and an explicit output cast restores the FP32 graph boundary.
• Added real CPU NVFP4 coverage for both MatMul plus bias Add and biased Gemm, parameterized across FP32 preservation, FP32-to-FP16 conversion, and native-BF16 no-op behavior. The graphs retain dynamic FP4 activation quantization, use opset 23, and pass full ONNX checking and strict shape inference.
• Reran 20 focused CPU cases after the base merge; all passed. Changed-file hooks also passed.

The PR is now conflict-free. Current-head Blackwell/TensorRT acceptance has not been rerun, and human approval remains external.

🤖 Generated by Codex (AI agent).

[cjluo-nv]

Bot review (claude-opus-5) — DM the bot to share feedback.

Nudge: the two prior code concerns are resolved with real CPU coverage, but hardware acceptance is stale on current head and one BF16 branch in the exporter is unexplained.

Needs action:

• Rerun or explicitly waive the Blackwell/TensorRT engine build for the current head; the PR body says it was not repeated after the final exporter change.
• Explain (ideally as a code comment) why the BF16 branch in modelopt/onnx/export/nvfp4_exporter.py::_cast_input_dtypes re-stamps the node output value_info to BF16 instead of inserting a boundary Cast like the FP16 path — on a graph with any BF16 initializer this silently retypes outputs whose consumers are still FP32.
• Have a human review the Codex-authored commits and the declared backward-incompatible ValueError behavior before merge.

No action needed:

• ✔️ Resolved since the last review: the CHANGELOG.rst entry for the new FP32→BF16 / BF16→FP16 hard failures, and the weights_dtype="fp32" mixed-type path (Gemm bias cast + output boundary cast, covered by the new parameterized elementwise/gemm CPU tests).
• The output cast is skipped when a node output has no value_info (e.g. after a failed shape inference), which just restores pre-PR behavior rather than regressing.