Skip to content

[5532019][AutoCast] Update GraphSanitizer to convert Double to Float32 #364

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
Sep 25, 2025
+332 −0
Merged

[5532019][AutoCast] Update GraphSanitizer to convert Double to Float32 #364

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
1b85355
[AutoCast] restore outputs if keep_io_types is set
ajrasane Sep 24, 2025
1e4d697
Remove check for keep_io_types
ajrasane Sep 24, 2025
01822ec
Simplify restore_original_io_types
ajrasane Sep 25, 2025
36fc519
Update GraphSanitizer to convert Double to Float32
ajrasane Sep 25, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions modelopt/onnx/autocast/precisionconverter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,9 @@ def convert(
# Remove redundant casts
self._cleanup()

if self.keep_io_types:
self._restore_original_io_types()

self._sanity_check()

return self.model
Expand All @@ -210,6 +213,32 @@ def _ensure_types_are_defined(self):
if vi.type.tensor_type.elem_type == onnx.TensorProto.UNDEFINED:
vi.type.tensor_type.elem_type = self.low_precision_type.onnx_type

def _restore_original_io_types(self):
"""Restore original I/O types."""
# Restore input types
for input_tensor in self.model.graph.input:
if input_tensor.name in self.original_network_io:
original_type = self.original_network_io[input_tensor.name]
if input_tensor.type.tensor_type.elem_type != original_type:
input_tensor.type.tensor_type.elem_type = original_type
# Update value_info_map if tensor exists there
if input_tensor.name in self.value_info_map:
self.value_info_map[
input_tensor.name
].type.tensor_type.elem_type = original_type

# Restore output types
for output_tensor in self.model.graph.output:
if output_tensor.name in self.original_network_io:
original_type = self.original_network_io[output_tensor.name]
if output_tensor.type.tensor_type.elem_type != original_type:
output_tensor.type.tensor_type.elem_type = original_type
# Update value_info_map if tensor exists there
if output_tensor.name in self.value_info_map:
self.value_info_map[
output_tensor.name
].type.tensor_type.elem_type = original_type

def _propagate_types_shapes_custom_ops(self, model):
"""Propagate types and shapes after insertion of 'Cast' nodes or other graph modifications."""
logger.info("Propagating tensor shapes and types in model with custom ops.")
Expand Down
Loading