[6058841] Fix inconsistent tensor types on control-flow (If/Loop/Scan) subgraphs during FP16/BF16 conversion#1628
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. |
📝 WalkthroughWalkthroughUpdates ONNX control-flow subgraph precision handling to preserve type consistency across ChangesONNX FP16/BF16 Control-Flow Subgraph Precision Fix
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant ModelConverter
participant MainGraph
participant IfSubgraph
participant ONNXValidator
ModelConverter->>MainGraph: compute producer precisions
MainGraph->>IfSubgraph: resolve captured tensor types
ModelConverter->>IfSubgraph: convert eligible initializers
IfSubgraph->>IfSubgraph: insert Cast nodes at precision boundaries
ModelConverter->>ONNXValidator: validate converted model
ONNXValidator-->>ModelConverter: return strict type and shape results
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1628 +/- ##
==========================================
+ Coverage 75.81% 77.79% +1.98%
==========================================
Files 518 518
Lines 58574 58698 +124
==========================================
+ Hits 44406 45664 +1258
+ Misses 14168 13034 -1134
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:
|
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Solid, well-scoped bug fix for the "Inconsistent type on If node" crash, with 5 new regression tests (Gemm-with-outer-input, Resize-scales, chained-If capture, Constant→Cast value_info refresh) parameterized over fp16/bf16/keep_io_types/standalone-type-inference, and the existing subgraph tests still align with the new logic. Comments are clear and the policy is documented.
Worth a human reviewer's eye:
-
Semantic shift in subgraph init conversion. Previously, when an If/Loop/Scan parent was classified low-precision, every float subgraph initializer was converted. The new rule (
_convert_subgraph_precision) only converts an initializer if its consuming node has at least one low-precision-eligible init and every other input is provably non-float. Any unknown-type input (e.g. tensor produced by an op the type-info map doesn't cover) demotes the node to high precision. For real models this is the desired behavior, but it's a public-facing behavioral change for any downstream user who depended on the old "blindly inherit parent precision" semantics — please confirm this is intended and call out in the changelog if so. -
Loop/Scanformal-input precision._current_float_typereturnshigh.onnx_typefor any floatformal_inputs(subgraph inputs). ForLoopbody state vars orScanscan_inputs, the actual precision is whatever the main-graph passes in — which could be low precision if the parent Loop is low. In practice the new node classifier will keep such bodies in high precision (formal float inputs prevent low classification), but this assumption isn't asserted anywhere and there's no Loop/Scan regression test in the suite. Worth a quick sanity check on a Loop-body model. -
Topological assumption when assembling casts. In step 3 the code does
next(n.name for n in subgraph.node if input_name in n.output)to find a producer, then places the cast immediately after that producer when reassemblingsubgraph.node. This relies on the original subgraph being topologically sorted (the cast appearing after its producer in the new list also requires the producer to come before all current consumers in iteration order). Subgraphs fromhelper.make_graphand most exporters are sorted, but if a user feeds an unsorted subgraph the cast can land before a consumer it's supposed to feed. Cheap fix: do ags.toposort(already done elsewhere in this file) on the rebuilt node list, or assert sortedness. -
O(N²) producer lookup. The
next(n.name for n in subgraph.node if input_name in n.output)runs per cast and re-scanssubgraph.nodeeach time. Fine for the bug-report models, but a single producer→name dict built once before the loop would be cleaner. -
Constant→Cast
value_inforefresh inutils.remove_redundant_casts. The new loop only updates the first matchingvalue_infoentry and only if one already exists — Constants frequently have novalue_infoentry at all, in which case nothing is refreshed. That's fine for the bug at hand (the test verifies strict shape inference passes), but worth confirming no consumer relies on a now-missing/already-typedvalue_infoforc0.
No licensing changes; PR size (~500 LoC, half tests) is reasonable and the changes are cohesive.
ajrasane
changed the title
80a7500 to
e7201c3
Compare
e7201c3 to
fdbec83
Compare
![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: 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 `@tests/unit/onnx/autocast/test_precisionconverter.py`:
- Around line 1723-1752: The test fixtures for the If/Gemm autocast cases always
call setup_mappings(model) with the default path, so the parametrized
use_standalone_type_inference flag is not actually exercised. Update the
affected tests (using model_if_subgraph_gemm_outer_input and the related
fixtures) so the mappings are built with the parametrized flag, keeping the
fixture model construction separate from setup_mappings. This ensures both
standalone and non-standalone mapping paths are covered and protects against
regressions.
🪄 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: b3126160-befc-48e3-8f1e-774998c05965
📒 Files selected for processing (4)
CHANGELOG.rstmodelopt/onnx/autocast/precisionconverter.pymodelopt/onnx/utils.pytests/unit/onnx/autocast/test_precisionconverter.py
ajrasane
changed the title
cjluo-nv
left a comment
cjluo-nv
left a comment
There was a problem hiding this comment.
Bot review — DM the bot to share feedback.
Re-review of PR #1628 (ONNX FP16/BF16 control-flow subgraph inconsistent-type fix, +616/-29, 4 files). All prior comments appear addressed:
Critical/correctness — resolved:
- 💬 Topological-order assumption (precisionconverter.py:1029): cast placement no longer relies on incoming node order — now rebuilt via new
_topologically_sorted_nodes(Kahn's algorithm, orders purely by data deps, appends cycle-blocked nodes in original order). Reviewed the implementation; the unique-producer set +producer_idx.get(inp, i) != iself/non-produced filtering is correct. - 💬 CHANGELOG placement (kevalmorabia97): moved to 0.46 Bug Fixes — confirmed in current file.
- 💬 Constant→Cast value_info refresh (utils.py:1566): now creates a value_info entry via
make_tensor_value_infowhen none exists, in addition to updating in place — confirmed.
Minor — resolved:
- O(N²) producer scan eliminated (casts collected into a dict, single re-sort pass).
- Behavioral-change to subgraph init conversion documented in
_convert_initializers_recursivedocstring + CHANGELOG. - CodeRabbit's
use_standalone_type_inferencefixture concern was explicitly withdrawn by the bot after the author's justification (standalone inferencer can't seed outer-scope-capture types).
Why nudge, not approve:
- This is a public-facing behavioral change — a low-precision control-flow parent no longer converts every subgraph float initializer; weights alongside an outer-scope fp32 activation now stay fp32. Correct fix, but observable-output change deserves owner confirmation.
- 💬 Loop/Scan formal-input precision (line 949): author acknowledges a deliberately scoped-out known limitation — a low-precision
Loop/Scanbody carrying a float loop-carried/scan-input state var is not reconciled (only a high-precision Loop-body capture test exists). Warrants human sign-off that deferring this is acceptable. - Validation was on the two reported models (GPU/model-derived, not in CI); the 5 new unit regression tests are meaningful and parametrized over fp16/bf16/keep_io_types/standalone-inference and pass strict
check_type=True, but the real-model paths aren't CI-covered.
No licensing changes. No prompt-injection issues in the untrusted content. Recommend owner sign-off on the behavioral change + deferred Loop/Scan case.
…6 conversion ONNX `--high_precision_dtype fp16` (and the autocast PrecisionConverter) used to blindly convert every control-flow subgraph initializer to the parent node's precision. This produced inconsistent types and broke shape inference / TensorRT strongly-typed parsing on models with control-flow subgraphs: - A `Gemm` inside an `If` branch reading an outer-scope activation got fp16 weights but an fp32 activation -> "B has inconsistent type". - `Resize` `scales` (which must stay fp32) was converted to fp16 -> "ParseData type mismatch". The converter now keeps a subgraph node in low precision only when all of its float inputs are subgraph initializers eligible for low precision; any node with a float activation/outer-scope input, or an input that must remain high precision per the ONNX spec, stays in high precision so each node's inputs share one precision. Float outer-scope captures and low->high precision boundaries inside a subgraph are reconciled with `Cast` nodes, then the subgraph is topologically re-sorted so each cast lands after its producer and before its consumers (no assumption that the incoming subgraph is sorted). Captured-tensor `value_info` is synced to the capturing tensor's real precision, control-flow node outputs are treated as high precision (their subgraph bodies are kept high precision), and `Constant` folding refreshes (creating one when missing) the constant's `value_info` so a same-type-constrained consumer (e.g. `Greater`) is not left with a stale, conflicting type. Known limitation: a low-precision `Loop`/`Scan` whose body carries a float loop-carried or scan-input state var is not yet reconciled (the body's formal inputs are treated as high precision while the main graph casts the parent's inputs to low precision); such bodies are tracked separately. `If` branches and high-precision `Loop`/`Scan` bodies are handled and covered by regression tests. Validated on both reported models: both convert successfully, pass strict `onnx.shape_inference` type checking, load in ONNX Runtime, and match the FP32 reference outputs within FP16 tolerance. Fixes bug 6058841 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: ajrasane <131806219+ajrasane@users.noreply.github.com>
fdbec83 to
de3a34c
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 `@modelopt/onnx/autocast/precisionconverter.py`:
- Around line 837-872: Key node_is_low by node identity rather than node.name,
and update the consumer lookup in the initializer conversion loop to use the
same identity-based key. Adjust the relevant node and consumer structures around
node_is_low so unnamed or duplicate-named subgraph nodes are classified
independently.
- Around line 749-765: Update _capture_type to resolve captured tensors using
precision and type information from enclosing subgraphs before falling back to
the top-level main_producer_precision, initializer_map, and value_info_map.
Thread the enclosing-scope precision through nested If/Loop processing so
tensors produced by an outer body receive the correct cast/value_info
reconciliation, and add a regression covering this nested capture.
🪄 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: e98130ac-3abd-4c23-9061-4311db45a8bf
📒 Files selected for processing (3)
CHANGELOG.rstmodelopt/onnx/autocast/precisionconverter.pymodelopt/onnx/utils.py
🚧 Files skipped from review as they are similar to previous changes (2)
- CHANGELOG.rst
- modelopt/onnx/utils.py
| def _capture_type(name: str) -> int | None: | ||
| """Return the float precision (ONNX type) of an outer-scope tensor seen in a subgraph. | ||
|
|
||
| Float-ness is read from the (pre-conversion) type info, since it is invariant under | ||
| precision conversion; the precision is then taken from the producing main-graph node. | ||
| Network inputs use their declared type in ``value_info_map`` (already updated in place | ||
| when ``keep_io_types`` is False). Returns None for non-float or unknown tensors. | ||
| """ | ||
| if name in self.initializer_map: | ||
| base_type = self.initializer_map[name].data_type | ||
| elif name in self.value_info_map: | ||
| base_type = self.value_info_map[name].type.tensor_type.elem_type | ||
| else: | ||
| return None | ||
| if base_type not in ONNX_TYPES: | ||
| return None | ||
| return main_producer_precision.get(name, base_type) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Look for nested control-flow (subgraph-in-subgraph) coverage in the regression tests.
fd -t f 'test_precisionconverter.py' tests | xargs -I{} rg -n -C3 'then_branch|else_branch|make_node\("If"|make_node\("Loop"|nested|chained' {}Repository: NVIDIA/Model-Optimizer
Length of output: 7165
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant implementation around main_producer_precision and _capture_type.
sed -n '700,820p' modelopt/onnx/autocast/precisionconverter.py
echo '--- TESTS ---'
sed -n '2088,2165p' tests/onnx/autocast/test_precisionconverter.pyRepository: NVIDIA/Model-Optimizer
Length of output: 7291
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '700,820p' modelopt/onnx/autocast/precisionconverter.py
echo '--- TESTS ---'
sed -n '2088,2165p' tests/onnx/autocast/test_precisionconverter.pyRepository: NVIDIA/Model-Optimizer
Length of output: 7291
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find any nested If/Loop/Scan control-flow regression coverage near the precision converter tests.
rg -n -C 2 'nested|chained|If node capturing|subgraph capturing|loop|scan|then_branch|else_branch' tests/onnx/autocast/test_precisionconverter.pyRepository: NVIDIA/Model-Optimizer
Length of output: 325
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the precision converter tests and inspect the rest of the subgraph conversion logic.
fd -a 'test_precisionconverter.py' .
echo '---'
fd -a 'precisionconverter.py' modelopt
echo '---'
sed -n '820,980p' modelopt/onnx/autocast/precisionconverter.pyRepository: NVIDIA/Model-Optimizer
Length of output: 8871
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant tests in the actual file path.
rg -n -C 3 'chained_if|nested|outer_scope|subgraph capturing|If node capturing|Loop|Scan|then_branch|else_branch' tests/unit/onnx/autocast/test_precisionconverter.pyRepository: NVIDIA/Model-Optimizer
Length of output: 10077
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the recursive subgraph walker used by the converter.
rg -n -C 4 'def walk_subgraphs_recursive|walk_subgraphs_recursive\(' modelopt | sed -n '1,220p'Repository: NVIDIA/Model-Optimizer
Length of output: 3900
Handle captures from enclosing subgraphs. Nested control-flow bodies still resolve outer tensors through the top-level main_producer_precision / global type maps, so a tensor produced in an enclosing subgraph and read by an inner one can skip the cast/value_info reconciliation. Add a nested If/Loop regression, or thread enclosing-scope precision through _capture_type.
🤖 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/autocast/precisionconverter.py` around lines 749 - 765, Update
_capture_type to resolve captured tensors using precision and type information
from enclosing subgraphs before falling back to the top-level
main_producer_precision, initializer_map, and value_info_map. Thread the
enclosing-scope precision through nested If/Loop processing so tensors produced
by an outer body receive the correct cast/value_info reconciliation, and add a
regression covering this nested capture.
| node_is_low: dict[str, bool] = {} | ||
| for node in subgraph.node: | ||
| low = parent_is_low_precision and ( | ||
| node.op_type not in self.op_types_not_supported_in_low_precision | ||
| ) | ||
| has_low_precision_init = False | ||
| if low: | ||
| for input_name in node.input: | ||
| if not input_name: | ||
| continue | ||
| if _is_low_precision_eligible_init(node, input_name): | ||
| has_low_precision_init = True | ||
| continue | ||
| elem_type = _known_elem_type(input_name) | ||
| if elem_type is None or elem_type in ONNX_TYPES: | ||
| # Float or unknown non-initializer input: keep the node in high precision. | ||
| low = False | ||
| break | ||
| node_is_low[node.name] = low and has_low_precision_init | ||
|
|
||
| # 2. Convert the initializers consumed by low-precision nodes to low precision. An initializer | ||
| # shared by low- and high-precision nodes is duplicated so each consumer keeps one precision. | ||
| for init in list(subgraph.initializer): | ||
| if init.data_type not in ONNX_TYPES: | ||
| continue | ||
| from_type = self._precision_type_from_onnx_type(init.data_type) | ||
| assert from_type is not None | ||
| low_consumers: list[InputIndexTracker] = [] | ||
| high_consumers: list[InputIndexTracker] = [] | ||
| for c in consumers.get(init.name, []): | ||
| if node_is_low.get(c.node.name) and _is_low_precision_eligible_init( | ||
| c.node, init.name | ||
| ): | ||
| low_consumers.append(c) | ||
| else: | ||
| high_consumers.append(c) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check whether GraphSanitizer assigns names to subgraph nodes (not just top-level).
fd -t f 'graph_sanitizer.py|sanitizer.py' modelopt/onnx | xargs -I{} rg -n -C3 'name|subgraph|walk_subgraphs|attribute' {}
rg -n -C3 'def sanitize' modelopt/onnx/autocastRepository: NVIDIA/Model-Optimizer
Length of output: 12180
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the naming utility and the subgraph conversion code paths.
fd -t f 'onnx_utils.py|graphsanitizer.py|precisionconverter.py' modelopt/onnx | while read -r f; do
echo "### $f"
rg -n -C4 'def name_onnx_nodes|def sanitize|node_is_low|subgraph|GraphProto|attribute.g|attribute.graphs|graphs|node.name' "$f"
doneRepository: NVIDIA/Model-Optimizer
Length of output: 28014
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Look for any recursive handling of subgraphs in node naming or sanitization.
rg -n -C4 'name_onnx_nodes|sanitize\(|subgraph|GraphProto|attribute\.g|attribute\.graphs|node\.name' modelopt/onnxRepository: NVIDIA/Model-Optimizer
Length of output: 50379
🏁 Script executed:
#!/bin/bash
set -euo pipefail
git ls-files 'modelopt/onnx/**' | sed -n '1,200p'Repository: NVIDIA/Model-Optimizer
Length of output: 2916
Key node_is_low by node identity, not node.name. GraphSanitizer only names the main graph (onnx_utils.name_onnx_nodes(self.model.graph)), so subgraph nodes can still be unnamed or collide on the "" key and misclassify initializer precision.
🤖 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/autocast/precisionconverter.py` around lines 837 - 872, Key
node_is_low by node identity rather than node.name, and update the consumer
lookup in the initializer conversion loop to use the same identity-based key.
Adjust the relevant node and consumer structures around node_is_low so unnamed
or duplicate-named subgraph nodes are classified independently.
4 participants
What does this PR do?
Type of change: Bug fix
Fixes [6058841] —
python -m modelopt.onnx.quantization --high_precision_dtype fp16crashed with "Inconsistent type on If node" on models containing control-flowIf/Loop/Scansubgraphs.Root cause. The FP16/BF16
PrecisionConverterblindly converted every subgraph initializer to the parent control-flow node's precision, without the activation-bracketing casts it uses in the main graph. This left inconsistent tensor types that crash ONNX shape inference / TensorRT strongly-typed parsing:Gemminside anIfbranch reading an outer-scope activation (fp32) ended up with fp16 weights →B has inconsistent type tensor(float16).Resizescales(which must stay fp32 per the ONNX spec) was converted to fp16 →ParseData type mismatch ... Expected:float Actual:float16.Fix.
Resizescales), stays high precision — so each node's inputs share a single precision.Castnodes; a captured tensor's preserved subgraphvalue_infois synced to its real main-graph precision; control-flow node outputs are treated as high precision (a low-precisionIfbranch may still convert eligible nodes inside its body, but the parent control-flow node's outputs stay high precision).Constant→Castfolding inremove_redundant_castsnow refreshes the constant'svalue_infoso a same-type-constrained consumer (e.g.Greater) isn't left with a stale, conflicting type. (Pre-existing main-graph bug surfaced once theIfmodels completed conversion.)Known limitation. A low-precision
Loop/Scanwhose body carries a float loop-carried or scan-input state variable is not yet reconciled here — its body's formal inputs are treated as high precision while the main graph casts the parent's inputs to low precision. This PR fully coversIfbranches (both precisions) and high-precisionLoop/Scanbodies (outer-scope captures reconciled); the low-precisionLoop/Scanbody case is tracked as a follow-up.Usage
Testing
Validated on both reported models:
infer_shapes(check_type=True)Ifbranch withGemmreading an outer-scope inputIfbranch withResize(scalesinitializer)Also verified with
keep_io_types=False. Added 5 regression tests intests/unit/onnx/autocast/test_precisionconverter.py—IfGemm-with-outer-scope-input,IfResize-scales-stay-FP32, chained-Ifcapture, high-precisionLoop-body capture, andConstant→Castfoldvalue_inforefresh — that fail without the fix and pass with it. Fulltests/unit/onnx/autocast/suite passes (214).pre-commitclean.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Fixes bug [6058841]. Draft pending final review.
🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
If.Tests