Skip to content

Commit 847d3fe

Browse files
authored
Fix ONNX save crash when ModelProto.ByteSize() raises EncodeError (microsoft#2355)
## Summary This PR fixes a crash in ONNX model saving for large/complex models where `model.ByteSize()` can raise a protobuf `EncodeError` (`Failed to serialize proto`). ## Problem In `model_proto_to_file`, we probe model size with `model.ByteSize()` when `save_as_external_data=False`. For some models, that call fails before save, causing the pass to crash. ## Root Cause `ByteSize()` can trigger protobuf serialization internals and fail with: `google.protobuf.message.EncodeError: Failed to serialize proto`. ## Fix - Wrap `model.ByteSize()` in `try/except`. - If size probing fails, log a warning and force `save_as_external_data=True`. - Preserve existing large-model fallback behavior (`<= 0` or `>= onnx.checker.MAXIMUM_PROTOBUF`). ## Validation - Reproduced failure before fix in `mq:matmulnbitstoqdq` flow. - Re-ran with fix and workflow completed without the previous `ByteSize()` crash. ## Impact - Improves robustness for large ONNX exports. - No behavior change for normal models where `ByteSize()` works.
1 parent 43f81fb commit 847d3fe

1 file changed

Lines changed: 19 additions & 9 deletions

File tree

olive/passes/onnx/common.py

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -130,15 +130,25 @@ def model_proto_to_file(
130130
output_dir = output_path.parent
131131
output_dir.mkdir(parents=True, exist_ok=True)
132132

133-
model_size = model.ByteSize()
134-
# model size for large models might be negative (overflow?) on Windows
135-
# see https://github.com/onnx/onnx/issues/5861
136-
if not save_as_external_data and (model_size <= 0 or model_size >= onnx.checker.MAXIMUM_PROTOBUF):
137-
save_as_external_data = True
138-
logger.debug(
139-
"Model is too large to save as a single file but 'save_as_external_data' is False. Saving tensors as"
140-
" external data, regardless."
141-
)
133+
# model size probing may fail for very large models/external data. Only probe when needed.
134+
if not save_as_external_data:
135+
try:
136+
model_size = model.ByteSize()
137+
except Exception as e:
138+
logger.warning(
139+
"Failed to compute model size with ByteSize (%s). Saving tensors as external data.",
140+
e,
141+
)
142+
save_as_external_data = True
143+
else:
144+
# model size for large models might be negative (overflow?) on Windows
145+
# see https://github.com/onnx/onnx/issues/5861
146+
if model_size <= 0 or model_size >= onnx.checker.MAXIMUM_PROTOBUF:
147+
save_as_external_data = True
148+
logger.debug(
149+
"Model is too large to save as a single file but 'save_as_external_data' is False. Saving"
150+
" tensors as external data, regardless."
151+
)
142152

143153
if not save_as_external_data:
144154
# Add olive version to metadata

0 commit comments

Comments
 (0)