[6421642][ONNX][Quantization] Fix NVFP4 exporter node ordering#1985
[6421642][ONNX][Quantization] Fix NVFP4 exporter node ordering#1985gcunhase wants to merge 3 commits into
Conversation
Ensure NVFP4 ONNX post-processing returns graphs with producers before consumers after TRT_FP4QDQ lowering and Cast insertion. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
|
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. |
📝 WalkthroughWalkthroughNVFP4 ONNX post-processing now deterministically topologically sorts graph nodes after FP4QDQ replacement and initializer cleanup. A CPU test exports, converts, and validates an NVFP4 ONNX model. ChangesNVFP4 ONNX Export
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
303dda2 to
226f0f3
Compare
Keep the NVFP4 ONNX exporter regression with the existing CPU Torch ONNX export coverage instead of a standalone test file. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1985 +/- ##
==========================================
+ Coverage 76.50% 77.14% +0.63%
==========================================
Files 525 525
Lines 59257 59429 +172
==========================================
+ Hits 45337 45847 +510
+ Misses 13920 13582 -338
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:
|
226f0f3 to
d3f7199
Compare
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review (gpt-5.6-sol) — DM the bot to share feedback.
The fix is straightforward and consistent with existing exporter patterns: GraphSurgeon topologically sorts the graph after the NVFP4 post-processing code appends producer nodes, and the new CPU regression test exercises the real Torch→ONNX→NVFP4 conversion path and validates the result with ONNX checker. Imports are correctly placed, the existing ONNX optional dependency already includes GraphSurgeon, and no licensing changes are introduced.
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude review passed — no blocking issues found. LGTM
Scope: The PR's authoritative changed-file list is 2 files — modelopt/onnx/export/nvfp4_exporter.py and tests/unit/torch/quantization/test_onnx_export_cpu.py. (Local base drift also surfaced trt_utils.py/test_plugin.py hunks, but those are not part of this PR and were excluded.) Both PR files were reviewed in full.
Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0
Assessment (low risk):
- The fix appends a GraphSurgeon
import_onnx → toposort() → export_onnxround-trip at the end ofNVFP4QuantExporter.post_process(), correctly reordering the newly-insertedDequantizeLinear/Castproducer nodes ahead of theTranspose/MatMulconsumers they feed. This matches the established pattern used by the sibling FP8/INT4/INT8 exporters andgraph_utils, andonnx-graphsurgeon>=0.6.1is already a declared dependency (no new import risk). - The sole caller (
_deploy/utils/torch_onnx.py:477) reassignsprocess_model()'s return, so returning a fresh ModelProto fromgs.export_onnx()is safe — nothing relies on in-place identity. - The new CPU regression test drives the real Torch→ONNX→NVFP4 path, asserts
TRT_FP4QDQis present pre-conversion and gone post-conversion, and validates withonnx.checker.check_model(which itself enforces topological ordering) — a real guard for this bug.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
60e28c2 to
3ccb0bd
Compare
3ccb0bd to
7c448e1
Compare
Use a protobuf-only graph node topological sort so NVFP4 post-processing preserves BF16 ONNX metadata while still returning producer-before-consumer node order. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
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
🧹 Nitpick comments (2)
modelopt/onnx/export/nvfp4_exporter.py (2)
467-467: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert the ordering contract directly in the regression test.
The supplied test checks FP4QDQ removal and
onnx.checker.check_model, but does not explicitly verify that every in-graph producer precedes its consumer. Add that edge-order assertion, ideally with a branched graph, so this new sorting call cannot silently regress.🤖 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/onnx/export/nvfp4_exporter.py` at line 467, Extend the regression test covering the export flow around _topologically_sort_graph_nodes to assert that every in-graph producer node appears before its consumer, preferably using a branched graph. Keep the existing FP4QDQ-removal and onnx.checker.check_model checks, and make the new assertion validate the sorted node order directly.
62-71: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUse a heap for the ready queue.
pop(0)shifts the remaining list, andready.sort()runs on every iteration. For wide ONNX graphs this can degrade toward quadratic/logarithmic behavior and make export unnecessarily slow. Useheapqkeyed by the original node index to preserve stable ordering withO((V+E) log V)behavior.Proposed change
+import heapq + - ready = [node_index for node_index, dependency in enumerate(dependencies) if not dependency] + ready = [node_index for node_index, dependency in enumerate(dependencies) if not dependency] + heapq.heapify(ready) sorted_nodes: list[onnx.NodeProto] = [] while ready: - node_index = ready.pop(0) + node_index = heapq.heappop(ready) sorted_nodes.append(nodes[node_index]) for dependent_index in sorted(dependents[node_index]): dependencies[dependent_index].remove(node_index) if not dependencies[dependent_index]: - ready.append(dependent_index) - ready.sort() + heapq.heappush(ready, dependent_index)🤖 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/onnx/export/nvfp4_exporter.py` around lines 62 - 71, Replace the list-based ready queue in the topological sorting flow with a heapq min-heap keyed by original node indices. Heapify the initial ready indices, use heap push/pop operations when nodes become ready, and remove the per-iteration sorting while preserving stable ascending node ordering and the existing dependency processing.
🤖 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/onnx/export/nvfp4_exporter.py`:
- Line 467: Update _cast_input_dtypes() to cache and reuse Cast outputs keyed by
(input_name, precision_dtype), ensuring shared activations produce only one Cast
node and tensor. Preserve the existing <input>_f16 naming where applicable, and
ensure _topologically_sort_graph_nodes() receives a graph without duplicate
producers.
---
Nitpick comments:
In `@modelopt/onnx/export/nvfp4_exporter.py`:
- Line 467: Extend the regression test covering the export flow around
_topologically_sort_graph_nodes to assert that every in-graph producer node
appears before its consumer, preferably using a branched graph. Keep the
existing FP4QDQ-removal and onnx.checker.check_model checks, and make the new
assertion validate the sorted node order directly.
- Around line 62-71: Replace the list-based ready queue in the topological
sorting flow with a heapq min-heap keyed by original node indices. Heapify the
initial ready indices, use heap push/pop operations when nodes become ready, and
remove the per-iteration sorting while preserving stable ascending node ordering
and the existing dependency processing.
🪄 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: 5894027d-c666-4213-b109-316bdf9515e4
📒 Files selected for processing (1)
modelopt/onnx/export/nvfp4_exporter.py
| graph.initializer.extend(new_initializers) | ||
| logger.info(f"Removed {len(initializers_to_delete)} initializers") | ||
|
|
||
| _topologically_sort_graph_nodes(graph) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant file and nearby symbols first.
ast-grep outline modelopt/onnx/export/nvfp4_exporter.py --view expanded
echo
echo "---- lines around _cast_input_dtypes ----"
sed -n '360,430p' modelopt/onnx/export/nvfp4_exporter.py
echo
echo "---- lines around the added topological sort call ----"
sed -n '440,490p' modelopt/onnx/export/nvfp4_exporter.py
echo
echo "---- search for Duplicate producer / cast insertion helpers ----"
rg -n "Duplicate producer|_cast_input_dtypes|topologically_sort|Cast" modelopt/onnx/export -SRepository: NVIDIA/Model-Optimizer
Length of output: 8139
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "---- _topologically_sort_graph_nodes ----"
sed -n '1,80p' modelopt/onnx/export/nvfp4_exporter.py
echo
echo "---- _replace_fp4qdq_with_2dq ----"
sed -n '113,235p' modelopt/onnx/export/nvfp4_exporter.py
echo
echo "---- consumer map helper ----"
rg -n "def get_tensor_consumer_nodes|get_tensor_consumer_nodes\(" -S modelopt | head -n 20Repository: NVIDIA/Model-Optimizer
Length of output: 8802
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "---- graph consumer helper ----"
sed -n '260,340p' modelopt/onnx/quantization/graph_utils.py
echo
echo "---- surrounding export logic in nvfp4_exporter ----"
sed -n '330,470p' modelopt/onnx/export/nvfp4_exporter.py
echo
echo "---- search for any reuse/dedup of cast outputs in exporter code ----"
rg -n "_f16\"|cast_output_name|make_node\\(\"Cast\"|Cast nodes for each input" modelopt/onnx/export modelopt/onnx/quantization -SRepository: NVIDIA/Model-Optimizer
Length of output: 10110
Prevent duplicate Cast outputs for shared activations. _cast_input_dtypes() always names the cast tensor <input>_f16, so two MatMul nodes that consume the same activation will insert duplicate Cast producers and make _topologically_sort_graph_nodes() fail. Cache/reuse the Cast per (input_name, precision_dtype) or make the output name unique per consumer.
🤖 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/onnx/export/nvfp4_exporter.py` at line 467, Update
_cast_input_dtypes() to cache and reuse Cast outputs keyed by (input_name,
precision_dtype), ensuring shared activations produce only one Cast node and
tensor. Preserve the existing <input>_f16 naming where applicable, and ensure
_topologically_sort_graph_nodes() receives a graph without duplicate producers.
2 participants
What does this PR do?
Type of change: Bug fix
NVFP4QuantExporter.post_process()returns.DequantizeLinearandCastproducer nodes appear before theTranspose/MatMulnodes that consume them without round-tripping BF16 metadata through ONNX GraphSurgeon.Usage
Testing
mainbefore this patch.producer_after_consumer_edges=0,onnx.checker.check_model(converted)passes, and the repro exits withrc=0.python -m pytest tests/unit/torch/quantization/test_onnx_export_cpu.py::test_nvfp4_exported_onnx_is_topologically_sorted -q.ruff format modelopt/onnx/export/nvfp4_exporter.py tests/unit/torch/quantization/test_onnx_export_cpu.pyandruff check modelopt/onnx/export/nvfp4_exporter.py tests/unit/torch/quantization/test_onnx_export_cpu.py.tests/gpu/torch/quantization/test_nvfp4_onnx_export.py::test_simple_linear[BFloat16]by avoiding an ONNX GraphSurgeon import/export round-trip in the fix path.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.).CONTRIBUTING.md: N/ASummary by CodeRabbit
Bug Fixes
Tests