[6425069][ONNX][Autocast] Fix autocast metadata propagation#1983
[6425069][ONNX][Autocast] Fix autocast metadata propagation#1983gcunhase wants to merge 4 commits into
Conversation
|
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 (2)
📝 WalkthroughWalkthroughThe changes add GatherND-specific shape inference for custom-op propagation and synchronize element-type metadata when redundant Constant-to-Cast patterns are folded. Tests cover standard and scalar GatherND shapes plus strict ONNX inference after cast removal. ChangesONNX autocast updates
Estimated code review effort: 3 (Moderate) | ~20 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant FakeTensorRTPlugin
participant PrecisionConverter
participant GatherND
FakeTensorRTPlugin->>PrecisionConverter: provide custom-op tensor shapes
PrecisionConverter->>GatherND: infer output shape from data and indices
GatherND-->>PrecisionConverter: return shaped output
🚥 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.38% 77.26% -0.12%
==========================================
Files 524 524
Lines 59263 59298 +35
==========================================
- Hits 45860 45818 -42
- Misses 13403 13480 +77
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).
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
GatherNDoutput shape inference during precision conversion, including custom-op scenarios and scalar (rank-0) outputs.CastafterConstant,elem_typemetadata is now updated consistently across the main graph and any nested subgraphs.Constant -> Cast -> Identityfolding, ensuringvalue_infoelement-type updates and passing strict ONNX shape/type inference.GatherNDshape propagation tests (standard schema behavior and scalar rank preservation).