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

[gcunhase]

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.

[copy-pr-bot]

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]

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.