[6421642][ONNX][Quantization] Fix NVFP4 exporter node ordering#1985
[6421642][ONNX][Quantization] Fix NVFP4 exporter node ordering#1985gcunhase wants to merge 5 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. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughNVFP4 ONNX post-processing now reuses matching input casts and deterministically topologically sorts graph nodes after FP4QDQ replacement and initializer cleanup. CPU tests cover cast reuse, FP4QDQ removal, and ONNX model validation. 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
gcunhase
changed the title
Share the protobuf-only ONNX graph node sorter from modelopt.onnx.utils and document why it avoids GraphSurgeon dtype round-tripping for enum dtypes such as BF16. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
gcunhase
changed the title
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 (1)
modelopt/onnx/utils.py (1)
69-78: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUse a heap for the ready queue.
pop(0)plus sorting on every iteration can make wide graphsO(V² log V). Aheapqmin-heap preserves the current stable index ordering while reducing queue maintenance toO((V+E) log V).🤖 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/utils.py` around lines 69 - 78, Replace the ready-list queue in the topological sorting logic with a heapq min-heap, updating initialization, extraction, and insertion accordingly. Preserve ascending node-index ordering and the existing dependency processing in the loop, while removing repeated pop(0) and ready.sort() operations.
🤖 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/utils.py`:
- Around line 59-67: Update topologically_sort_graph_nodes to account for
tensors captured by nested If/Loop subgraphs when building dependencies, either
by recursively collecting those outer-scope references or by explicitly
rejecting graphs containing unsupported control-flow nodes. Preserve direct
input dependency handling and add a regression test covering a producer and
control-flow consumer connected through an implicit capture.
---
Nitpick comments:
In `@modelopt/onnx/utils.py`:
- Around line 69-78: Replace the ready-list queue in the topological sorting
logic with a heapq min-heap, updating initialization, extraction, and insertion
accordingly. Preserve ascending node-index ordering and the existing dependency
processing in the loop, while removing repeated pop(0) and ready.sort()
operations.
🪄 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: 5d326a8d-4f4d-4e77-bc32-522926e15de9
📒 Files selected for processing (2)
modelopt/onnx/export/nvfp4_exporter.pymodelopt/onnx/utils.py
| for consumer_index, node in enumerate(nodes): | ||
| for input_name in node.input: | ||
| producer_index = producer_by_tensor.get(input_name) | ||
| if producer_index is None: | ||
| continue | ||
| if producer_index == consumer_index: | ||
| raise ValueError(f"Node {node.name!r} consumes its own output {input_name!r}.") | ||
| dependencies[consumer_index].add(producer_index) | ||
| dependents[producer_index].add(consumer_index) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
git ls-files modelopt/onnx/utils.py modelopt/onnx
echo "--- utils.py outline ---"
ast-grep outline modelopt/onnx/utils.py --view expanded || true
echo "--- relevant lines ---"
wc -l modelopt/onnx/utils.py
sed -n '1,220p' modelopt/onnx/utils.py
echo "--- references to nested subgraphs / control-flow ---"
rg -n "subgraph|If\\(|Loop\\(|scan\\(|GraphProto|node.input|dependencies|producer_by_tensor" modelopt/onnx -SRepository: NVIDIA/Model-Optimizer
Length of output: 50378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "--- walk_subgraphs_recursive ---"
sed -n '1,220p' modelopt/onnx/autocast/utils.py
echo "--- topological sort call sites ---"
rg -n "topologically_sort_graph_nodes\\(" -S modelopt tests || true
echo "--- tests mentioning subgraphs / control-flow ---"
rg -n "subgraph|If\\(|Loop\\(|walk_subgraphs_recursive|topologically_sort_graph_nodes" tests modelopt -g '*test*' -S || trueRepository: NVIDIA/Model-Optimizer
Length of output: 13016
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "--- nvfp4_exporter call site context ---"
sed -n '360,470p' modelopt/onnx/export/nvfp4_exporter.py
echo "--- subgraph-related utilities in onnx/utils.py ---"
rg -n "walk_subgraphs_recursive|subgraph|AttributeProto\\.GRAPH|AttributeProto\\.GRAPHS" modelopt/onnx/utils.py modelopt/onnx/export/nvfp4_exporter.py -S || true
echo "--- tests for topological sorting or nvfp4 exporter ---"
rg -n "topologically_sort_graph_nodes|nvfp4_exporter|topological sort" tests -S || trueRepository: NVIDIA/Model-Optimizer
Length of output: 6565
Account for captured outer-scope tensors in graph sorting. topologically_sort_graph_nodes only tracks direct node.input edges, so an If/Loop node can be moved ahead of a producer that its subgraph reads implicitly. Either recurse through subgraphs or reject control-flow graphs here, and add a regression test.
🤖 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/utils.py` around lines 59 - 67, Update
topologically_sort_graph_nodes to account for tensors captured by nested If/Loop
subgraphs when building dependencies, either by recursively collecting those
outer-scope references or by explicitly rejecting graphs containing unsupported
control-flow nodes. Preserve direct input dependency handling and add a
regression test covering a producer and control-flow consumer connected through
an implicit capture.
Cache Cast outputs per input and precision dtype so multiple NVFP4-lowered MatMul consumers reuse a single activation Cast instead of producing duplicate tensor names. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
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 tests/unit/torch/quantization/test_onnx_export_cpu.py::test_nvfp4_shared_activation_reuses_cast -q.ruff format modelopt/onnx/export/nvfp4_exporter.py modelopt/onnx/utils.py tests/unit/torch/quantization/test_onnx_export_cpu.pyandruff check modelopt/onnx/export/nvfp4_exporter.py modelopt/onnx/utils.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
TRT_FP4QDQnodes after conversion, ONNX validity checks, and cast reuse behavior.