[6385267][ONNX][Autocast] Preserve graph-output Cast producers by gcunhase · Pull Request #1943 · NVIDIA/Model-Optimizer

gcunhase · 2026-07-08T17:06:15Z

What does this PR do?

Type of change: Bug fix

Updates modelopt/onnx/autocast/precisionconverter.py so pre-existing floating-point Cast nodes are preserved when they are valid producers of graph outputs.
The cleanup now keeps a graph-output Cast only when the Cast target type matches the declared graph output element type. This avoids dangling graph outputs while preserving existing keep_io_types=True behavior for mismatched output Casts that should be removed and repaired.
Adds regression coverage in tests/unit/onnx/autocast/test_precisionconverter.py for a valid ONNX model whose graph output is produced by a consecutive Cast chain.

Usage

$ python -m modelopt.onnx.autocast --onnx=model.onnx

Testing

Reproduced the failure locally before the fix: conversion raised Exception("Sanity Check Failed") after ONNX validation reported that graph output t2 was not produced by any node.
Verified the same repro after the fix: conversion completed and reported Converted 2/2 nodes (100.00%) to fp16.
Ran pytest tests/unit/onnx/autocast/test_precisionconverter.py::test_convert_preserves_cast_chain_graph_output -q

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.).

Is this change backward compatible?: ✅
If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: N/A
Did you write any new necessary tests?: ✅
Did you update Changelog?: N/A

Summary by CodeRabbit

Bug Fixes
- Improved mixed-precision conversion to preserve cast nodes only when their output aligns with the graph output type, avoiding unintended removal of valid output casts.
Tests
- Added a unit test covering cast-chain conversion to confirm the resulting model passes validation and retains the expected cast behavior at the graph output.

Preserve pre-existing floating Cast nodes when they are valid producers of graph outputs so autocast cleanup does not leave dangling graph outputs. Signed-off-by: Gwenaelle Cunha Sergio <gcunhasergio@nvidia.com>

copy-pr-bot · 2026-07-08T17:06:27Z

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.

coderabbitai · 2026-07-08T17:06:29Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ca284580-d06d-435d-a57c-21d267420095

📥 Commits

Reviewing files that changed from the base of the PR and between 16cb877 and 9b113e3.

📒 Files selected for processing (1)

tests/unit/onnx/autocast/test_precisionconverter.py

🚧 Files skipped from review as they are similar to previous changes (1)

tests/unit/onnx/autocast/test_precisionconverter.py

📝 Walkthrough

Walkthrough

Modifies the condition for preserving pre-existing Cast nodes in PrecisionConverter._remove_preexisting_casts, requiring the cast output to match a network output and the cast target type to match that output’s element type. Adds a unit test that converts a Cast-chain model to mixed precision and checks the graph output remains produced by a single Cast node.

Changes

Cast Preservation Fix

Layer / File(s)	Summary
Cast preservation condition update `modelopt/onnx/autocast/precisionconverter.py`	Changes preservation logic so a pre-existing Cast is kept only when its output tensor is a network output and `cast_to_type` matches that output’s element type.
Test coverage for cast chain preservation `tests/unit/onnx/autocast/test_precisionconverter.py`	Adds import of `convert_to_mixed_precision` and a new test that builds a DOUBLE-to-FLOAT Cast chain model, converts it to fp16, and asserts exactly one Cast node produces the graph output.

Estimated code review effort: 2 (Simple) | ~10 minutes

🚥 Pre-merge checks | ✅ 6

✅ Passed checks (6 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly summarizes the main change: preserving ONNX autocast Cast nodes that produce graph outputs.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns	✅ Passed	Only ONNX autocast logic and a regression test changed; neither file adds disallowed torch/numpy loads, trust_remote_code=True, eval/exec, # nosec, or new deps.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

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.

👉 Steps to fix this

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 91-119: The test currently uses a DOUBLE→DOUBLE→FLOAT cast chain,
which won’t hit the preexisting-cast removal path; update the graph in
test_convert_preserves_cast_chain_graph_output to use an FP32→FP16→FP32 chain
instead. Adjust the input/output tensor types and the two Cast nodes in this
test so the converted_model path exercises _remove_preexisting_casts and
verifies the output-preservation behavior directly.

🪄 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: 54296094-8937-434b-b86e-19fe095c922d

📥 Commits

Reviewing files that changed from the base of the PR and between b012e4c and 16cb877.

📒 Files selected for processing (2)

modelopt/onnx/autocast/precisionconverter.py
tests/unit/onnx/autocast/test_precisionconverter.py

cjluo-nv

Bot review — DM the bot to share feedback.

Small, focused bug fix (+37/-3, 2 files) in the ONNX AutoCast PrecisionConverter._remove_preexisting_casts.

What it does: Previously a graph-output Cast was preserved only when both its input and output were graph outputs (the output→cast→output pattern). A valid cast chain like in → Cast → t1 → Cast → t2(output) had its final cast removed and bypassed, leaving t2 dangling and failing the sanity check. The new condition keeps a Cast when its output is a graph output AND the cast target type matches the declared output element type — the correct, broader predicate that still preserves the old chain case.

Correctness: Verified the logic. The change also broadens preservation to casts whose input isn't a graph output (e.g. test_multiple_output_node_casted_to_output, test_casted_input_to_output_model) but those tests assert only check_model, which stays valid whether the cast is kept or removed. The type-match guard prevents keeping mismatched output casts that keep_io_types=True should still repair.

Tests: New test_convert_preserves_cast_chain_graph_output reproduces the reported failure end-to-end through convert_to_mixed_precision (exercises the default keep_io_types=False path). I confirmed via the GraphSanitizer that the DOUBLE casts are converted to FP32 before _remove_preexisting_casts, so the test genuinely hits the fixed branch. It asserts the output is produced by exactly one Cast node and that the model validates.

No licensing changes (existing Apache-licensed file, headers untouched), no new abstractions/subsystems, no prompt-injection in the untrusted PR content. Minor nit (non-blocking): unlike neighboring tests, the new test isn't parametrized over keep_io_types/bf16.

Use an FP32 to FP16 to FP32 Cast chain so the test exercises pre-existing Cast removal directly. Signed-off-by: Gwenaelle Cunha Sergio <gcunhasergio@nvidia.com>

Fix ONNX autocast output Cast preservation

16cb877

Preserve pre-existing floating Cast nodes when they are valid producers of graph outputs so autocast cleanup does not leave dangling graph outputs. Signed-off-by: Gwenaelle Cunha Sergio <gcunhasergio@nvidia.com>

gcunhase marked this pull request as ready for review July 8, 2026 17:14

gcunhase requested review from a team as code owners July 8, 2026 17:14

gcunhase requested a review from galagam July 8, 2026 17:14

coderabbitai Bot reviewed Jul 8, 2026

View reviewed changes

Comment thread tests/unit/onnx/autocast/test_precisionconverter.py

cjluo-nv approved these changes Jul 8, 2026

View reviewed changes

Update autocast output Cast regression test

9b113e3

Use an FP32 to FP16 to FP32 Cast chain so the test exercises pre-existing Cast removal directly. Signed-off-by: Gwenaelle Cunha Sergio <gcunhasergio@nvidia.com>

gcunhase requested a review from ajrasane July 8, 2026 18:01

coderabbitai Bot approved these changes Jul 8, 2026

View reviewed changes

Merge branch 'main' into dev/agent/6385267_preserve_output_cast

7a86aff

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[6385267][ONNX][Autocast] Preserve graph-output Cast producers#1943

[6385267][ONNX][Autocast] Preserve graph-output Cast producers#1943
gcunhase wants to merge 3 commits into
NVIDIA:mainfrom
gcunhase:dev/agent/6385267_preserve_output_cast

gcunhase commented Jul 8, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

copy-pr-bot Bot commented Jul 8, 2026

Uh oh!

coderabbitai Bot commented Jul 8, 2026 •

edited

Loading

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Uh oh!

Uh oh!

cjluo-nv left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

gcunhase commented Jul 8, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Usage

Testing

Before your PR is "Ready for review"

Summary by CodeRabbit

Uh oh!

copy-pr-bot Bot commented Jul 8, 2026

Uh oh!

coderabbitai Bot commented Jul 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cjluo-nv left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

gcunhase commented Jul 8, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Jul 8, 2026 •

edited

Loading