[6425069][ONNX][Autocast] Fix autocast metadata propagation#1983
[6425069][ONNX][Autocast] Fix autocast metadata propagation#1983gcunhase wants to merge 8 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe changes preserve public graph I/O metadata during precision conversion, add standard-operator shape inference for custom-op propagation, and synchronize element-type metadata when folding Constant-to-Cast patterns. Tests cover scalar and ranked shapes, rank changes, and strict ONNX validation. ChangesONNX autocast updates
Estimated code review effort: 3 (Moderate) | ~25 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant PrecisionConverter
participant StandardOps
participant FakePlugin
PrecisionConverter->>StandardOps: infer missing standard-op output shapes
StandardOps-->>PrecisionConverter: return ranked or scalar tensor shape
PrecisionConverter->>FakePlugin: propagate shaped tensors through custom-op nodes
FakePlugin-->>PrecisionConverter: provide custom-op output shape
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1983 +/- ##
==========================================
- Coverage 77.80% 77.67% -0.13%
==========================================
Files 518 518
Lines 58269 58386 +117
==========================================
+ Hits 45334 45350 +16
- Misses 12935 13036 +101
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:
|
57492c5 to
8768ce3
Compare
8768ce3 to
21d0069
Compare
|
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. |
Update ONNX autocast cleanup to keep folded Constant metadata in sync and infer GatherND output shapes in custom-op mode without running global ONNX shape inference. This preserves valid ONNX metadata after redundant Cast folding and avoids incorrect GatherND rank propagation when TensorRT plugin custom ops are present. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
028620c to
ea8f814
Compare
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.
Small, focused bug fix (+126/-1, 3 files) for ONNX autocast metadata propagation. Two changes, both correct and well-tested:
-
_sync_value_info_elem_type(utils.py) — after folding a Constant→Cast, propagates the new dtype into matching value_info/input/output entries, recursing into subgraphs. Theand tensor_type.elem_typeguard sensibly skips undefined-type entries. The regression test folds a real Constant→Cast and verifies both the value_info dtype and thatonnx.shape_inferencepasses withstrict_mode=True, check_type=True. -
_infer_gathernd_op_shape(precisionconverter.py) — in custom-op mode, computes GatherND output shape directly instead of copying input-0's shape. The formulaindices_shape[:-1] + data_shape[batch_dims + index_rank:]matches the ONNX GatherND spec; verified against the test (data [1,4,2], indices [1,2,2], batch_dims=1 → [1,2]). The old input-0 fallback is preserved viaeliffor non-GatherND nodes.
Nested helper functions are consistent with the existing style in this module. No new function-local imports, no licensing changes (no header edits/new files/license files), no injection attempts in PR content. CHANGELOG not updated, but this is an internal metadata-propagation bug fix.
Minor (non-blocking): if shape := _infer_gathernd_op_shape(node): treats an empty-list result (rank-0 GatherND output) as falsy and would fall through to the input-0 fallback — an unlikely edge case, and no worse than prior behavior.
![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: 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 300-317: Update the GatherND shape handling near the conversion
logic to check whether _infer_gathernd_op_shape(node) returns None rather than
relying on truthiness, preserving a valid empty list for scalar outputs. Add a
regression test covering scalar GatherND output shape inference and confirm it
does not fall back to the input shape.
In `@modelopt/onnx/utils.py`:
- Around line 1537-1540: Update the value-info synchronization loop to match the
requested tensor name only when value_info.type.HasField("tensor_type"), then
assign tensor_type.elem_type unconditionally so TensorProto.UNDEFINED (0)
declarations are updated. Add a regression covering an undefined tensor
declaration produced during Constant-to-Cast folding.
🪄 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: fa5f317b-1c4e-428c-a011-a13092348708
📒 Files selected for processing (3)
modelopt/onnx/autocast/precisionconverter.pymodelopt/onnx/utils.pytests/unit/onnx/autocast/test_precisionconverter.py
Treat an empty inferred GatherND shape as a valid scalar shape instead of falling through to the generic input-shape fallback in custom-op propagation. Add a regression test for the rank-0 GatherND output case. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Use a walrus expression with an explicit None check so scalar GatherND shapes remain valid while keeping the custom-op propagation branch concise. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Synchronize matching ONNX tensor declarations even when the previous element type is TensorProto.UNDEFINED, and cover that Constant Cast folding case in the regression test. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
modelopt/onnx/autocast/precisionconverter.py (1)
428-435: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winDo not use the data shape as a fallback for
GatherND.When
_infer_gathernd_op_shapereturnsNone, Line [430] falls through to the generic input-0 fallback and can assigndata.shapeto theGatherNDoutput. That is generally incorrect becauseGatherNDderives its output shape from both inputs and can produce a scalar. (onnx.ai)Keep the
GatherNDbranch separate: leave its shape unset when inference is unavailable, and apply the input-0 fallback only to actual custom-op outputs.🤖 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 428 - 435, Update the GatherND shape-handling branch around _infer_gathernd_op_shape so that an unsuccessful inference leaves out.shape unset, without falling through to inp’s constant or generic shape fallbacks. Restrict the input-0 fallback logic to actual custom-op outputs handled outside the GatherND branch.
🤖 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.
Outside diff comments:
In `@modelopt/onnx/autocast/precisionconverter.py`:
- Around line 428-435: Update the GatherND shape-handling branch around
_infer_gathernd_op_shape so that an unsuccessful inference leaves out.shape
unset, without falling through to inp’s constant or generic shape fallbacks.
Restrict the input-0 fallback logic to actual custom-op outputs handled outside
the GatherND branch.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f1d3f0ea-3339-4664-bef6-24d9dfa813df
📒 Files selected for processing (1)
modelopt/onnx/autocast/precisionconverter.py
|
/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
Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0
Focused, well-tested bug fix for ONNX autocast metadata propagation (+168/−1, 3 files). Reviewed all three changed files across correctness, export compatibility, and backward compatibility.
What I verified:
_infer_gathernd_op_shape(precisionconverter.py): the output-shape formulaindices_shape[:-1] + data_shape[batch_dims + index_rank:]matches the ONNX GatherND spec. Checked against both tests — standard[1,4,2]+[1,2,2],batch_dims=1 →[1,2], and scalar[4]+[1]→[]. Guards for dynamic (non-int) dims and out-of-rangesuffix_startare sound.- The prior truthiness concern (rank-0 GatherND result treated as falsy) raised on earlier commits is already resolved at HEAD — the code now uses
is not None, andtest_custom_op_mode_preserves_scalar_gathernd_shapecovers the scalar case explicitly. _sync_value_info_elem_type(utils.py): correctly propagates the folded Constant→Cast dtype into value_info/input/output entries and recurses into GRAPH/GRAPHS subgraph attributes. TheHasField("tensor_type")guard with an unconditionalelem_typeassignment meansUNDEFINEDdeclarations are updated — the parametrized test covers bothFLOATandUNDEFINED, andinfer_shapes(strict_mode=True, check_type=True)validates the result.- No public API/signature changes; both new helpers are private and module-internal. Backward compatible.
Note (non-blocking, already tracked by CodeRabbit): when _infer_gathernd_op_shape returns None for a real GatherND (shapes unknown), it falls through to the input-0 shape fallback, which is technically wrong for GatherND. This is pre-existing behavior — the PR strictly improves the common case and introduces no regression — so it doesn't block approval.
Risk: Low. Internal metadata-propagation fix with good regression coverage. CHANGELOG intentionally not updated (internal bug fix).
Preserve full public I/O metadata when keep_io_types=True and avoid applying the input-0 shape fallback to ordinary standard operators in custom-op mode. Infer shapes for standard Gather, Unsqueeze, Shape, and GatherND paths so strict ONNX checking passes without downstream metadata restoration. 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 (1)
modelopt/onnx/autocast/precisionconverter.py (1)
523-526: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate constant-shape extraction logic; reuse
_get_const_values.This mirrors exactly what
_get_const_values(Line 306-311) already does, just accessing.shapeinstead of the values. Keeping two implementations of the same "is this a Constant / embedded Constant node" check risks drift if one is updated without the other.♻️ Proposed refactor
- elif isinstance(inp, gs.Constant): - out.shape = inp.values.shape - elif inp.inputs and inp.inputs[0].op == "Constant": - out.shape = inp.inputs[0].attrs["value"].values.shape + elif (const_values := _get_const_values(inp)) is not None: + out.shape = const_values.shape🤖 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 523 - 526, Replace the duplicated Constant and embedded Constant-node checks in the output-shape handling with the existing _get_const_values helper, then assign out.shape from the returned values’ shape. Preserve the current behavior for both gs.Constant inputs and inputs whose first node is a Constant.
🤖 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 520-529: Update the output-shape guard in the surrounding
conversion logic to distinguish an unknown shape (None) from a known scalar
shape ([]). Only run _infer_standard_op_shape and the subsequent fallback chain
when out.shape is None, preserving an already-inferred empty shape and
preventing the custom-op fallback from overwriting scalar outputs.
---
Nitpick comments:
In `@modelopt/onnx/autocast/precisionconverter.py`:
- Around line 523-526: Replace the duplicated Constant and embedded
Constant-node checks in the output-shape handling with the existing
_get_const_values helper, then assign out.shape from the returned values’ shape.
Preserve the current behavior for both gs.Constant inputs and inputs whose first
node is a Constant.
🪄 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: a4c66aa3-d41d-4817-b8f7-d77dcd603612
📒 Files selected for processing (2)
modelopt/onnx/autocast/precisionconverter.pytests/unit/onnx/autocast/test_precisionconverter.py
Treat only None as a missing tensor shape so known scalar shapes are not clobbered by custom-op fallback propagation. Add a regression for a custom op with scalar output metadata. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
Apply formatting requested by pre-commit for the scalar-shape propagation follow-up without changing behavior. Signed-off-by: Gwena Cunha <4861122+gcunhase@users.noreply.github.com>
2 participants
What does this PR do?
Type of change: Bug fix
Constant -> Castpatterns, so the publishedvalue_infodtype matches the converted Constant payload.GatherNDoutput shapes directly when autocast is running in custom-op mode, avoiding a generic input-0 shape copy for this shape-changing standard ONNX operator.GatherNDshape propagation.Usage
Testing
0.44.0and with main ToTcba8a5c62a1a54fe89fb69bfd484ea0a653c633abefore the fix.GatherNDoutput shape after the fix.pytest tests/unit/onnx/autocast/test_precisionconverter.py::test_folded_constant_cast_updates_value_info_type tests/unit/onnx/autocast/test_precisionconverter.py::test_custom_op_mode_uses_schema_shape_for_standard_gathernd -qpytest tests/unit/onnx/autocast/test_precisionconverter.py -qBefore 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
GatherND,Gather,Unsqueeze, andShape), with correct scalar (rank-0) handling and stricter input-shape propagation.CastafterConstant, element-type metadata is now updated consistently across the graph and nested subgraphs.Constant -> Cast -> Identityfolding and forGatherND/rank-change shape propagation.