[6403893][ONNX][AutoCast] Fix false-success TRT parse for large ModelProto

[galagam]

What does this PR do?

Type of change: ? Bugfix

get_custom_layers silently returned 0 layers/tensors for in-memory ModelProtos at/above the protobuf 2GiB limit. Route such models through a temporary external-data file and parse_from_file(), matching the string-path behavior. Verified on the 8.12GB VLA trunk (28 layers/5468 tensors) and a synthetic 2.2GB model. Adds a regression test forcing the file-backed path.

Usage

# Add a code snippet demonstrating how to use this

Testing

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?: ❌
• Did you get Claude approval on this PR?: ✅

Additional Information

Summary by CodeRabbit

• Bug Fixes
  • Improved ONNX parsing for large models by automatically switching to a more reliable file-backed path when needed.
  • Kept parsing results consistent between in-memory and file-based handling, reducing failures on bigger or complex models.
• Tests
  • Added regression coverage to verify that both parsing paths return the same layer and tensor information.

[galagam]

/claude review

[coderabbitai]

Caution

Review failed

An error occurred during the review process. Please try again later.

📝 Walkthrough

Walkthrough

Adds a helper to determine when an in-memory ONNX model exceeds the protobuf size limit, and updates get_custom_layers() to route such models through a temporary external-data file before TensorRT parsing. A regression test verifies parsing results match across file-path, in-memory, and forced file-backed modes.

Changes

File-backed ONNX parsing support

Layer / File(s)	Summary
ByteSize check helper modelopt/onnx/trt_utils.py	Adds copy import and _requires_file_backed_parse() to check ModelProto.ByteSize() against onnx.checker.MAXIMUM_PROTOBUF, defaulting to True on computation failure with debug logging.
Three-way parsing flow in get_custom_layers modelopt/onnx/trt_utils.py	Updates get_custom_layers() to use parse_from_file for string paths, deep-copy + onnx.save external-data + parse_from_file for large in-memory models, and direct parse(SerializeToString()) otherwise, preserving existing error aggregation.
Regression test for parity tests/gpu/onnx/quantization/test_plugin.py	Imports get_custom_layers and adds test_get_custom_layers_file_backed_matches_in_memory, comparing results across file-path, in-memory, and monkeypatched forced file-backed parsing.

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

Sequence Diagram(s)

sequenceDiagram
  participant Caller
  participant get_custom_layers
  participant _requires_file_backed_parse
  participant OnnxSaveTempFile
  participant TrtParser

  Caller->>get_custom_layers: onnx_path (str or ModelProto)
  alt onnx_path is string
    get_custom_layers->>TrtParser: parse_from_file(onnx_path)
  else onnx_path is ModelProto
    get_custom_layers->>_requires_file_backed_parse: check ByteSize()
    alt requires file-backed parse
      _requires_file_backed_parse-->>get_custom_layers: True
      get_custom_layers->>OnnxSaveTempFile: deepcopy + save external data
      get_custom_layers->>TrtParser: parse_from_file(temp_path)
    else
      _requires_file_backed_parse-->>get_custom_layers: False
      get_custom_layers->>TrtParser: parse(SerializeToString())
    end
  end
  TrtParser-->>get_custom_layers: parsed layers or errors

Loading

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

Warning

Tools execution failed with the following error:

Failed to run tools: 13 INTERNAL: Received RST_STREAM with code 2 (Internal server error)

---

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

[claude]

[SUGGESTION] copy.deepcopy(onnx_path) transiently doubles peak RAM for exactly the models this path targets — an 8GB ModelProto needs ~16GB before onnx.save streams weights to disk. The deep-copy is necessary for correctness (onnx.save(save_as_external_data=True) mutates the proto in place, stripping raw_data from the caller's model), so this isn't a bug — but on the 8GB VLA trunk case cited in the PR description it may be tight on constrained hosts.

If it ever bites, an alternative that avoids the copy is to externalize in place and restore afterward (e.g. onnx.save(onnx_path, ...) then onnx.load_external_data_for_model(onnx_path, tmpdir) to re-hydrate the caller's proto), or wrap in try/finally to guarantee re-hydration on parse failure. Non-blocking; the current approach is the safest correctness-wise.

[galagam]

The peak memory usage is at ReferenceRunner where we define an auxiliary graph with every intermediate tensor being a graph output. The 2x memory usage here is not a bottleneck.

[gcunhase]

Agree!

[claude]

Claude review passed — no blocking issues found. LGTM

Scope: Reviewed the 2 files in this PR's authoritative changed-file list (modelopt/onnx/trt_utils.py, tests/gpu/onnx/quantization/test_plugin.py). The two-dot git diff origin/main HEAD surfaced many unrelated files due to the shallow-checkout divergence between the local origin/main tip and the true merge base — those are not part of this PR and were excluded.

Findings: CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 1

Assessment — low risk. Focused, well-targeted bugfix. get_custom_layers previously silently returned an empty network (0 layers/0 tensors) for in-memory ModelProtos at/above the protobuf 2 GiB limit because TensorRT's parser.parse(bytes) fails quietly while SerializeToString() succeeds. The fix routes over-limit (or un-sizeable) in-memory models through a temporary external-data file + parse_from_file(), matching the existing string-path branch and ModelOpt's file-backed policy elsewhere in this module.

Verified:

• onnx.checker.MAXIMUM_PROTOBUF is a valid constant, already used with the same >= MAXIMUM_PROTOBUF idiom at trt_utils.py:382.
• The ByteSize()-overflow except branch conservatively routes to the file path, consistent with load_onnx_model's handling.
• copy.deepcopy correctly shields the caller's model from in-place externalization by onnx.save; both in-memory callers pass a model the caller still owns.
• The regression test forces the file-backed path via monkeypatch and cross-checks results against both the in-memory fast path and the string-path baseline. Backward compatible — the small-model fast path is unchanged.

The single SUGGESTION concerns transient peak-RAM doubling from the deep-copy on multi-GB models — non-blocking, and the current approach is the correctness-safe choice.

[codecov]

Codecov Report

❌ Patch coverage is 52.63158% with 9 lines in your changes missing coverage. Please review.
✅ Project coverage is 61.35%. Comparing base (75b5803) to head (b52dc3f).
⚠️ Report is 5 commits behind head on main.

Files with missing lines	Patch %	Lines
modelopt/onnx/trt_utils.py	52.63%	9 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1928      +/-   ##
==========================================
+ Coverage   61.21%   61.35%   +0.14%     
==========================================
  Files         515      515              
  Lines       57245    58006     +761     
==========================================
+ Hits        35043    35592     +549     
- Misses      22202    22414     +212     

Flag	Coverage Δ
examples	32.59% <52.63%> (-0.28%)	⬇️
gpu	20.51% <0.00%> (-0.01%)	⬇️
unit	54.94% <21.05%> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[gcunhase]

Do you think the comment from Claude review should be a concern?

The single SUGGESTION concerns transient peak-RAM doubling from the deep-copy on multi-GB models.

Otherwise, LGTM!

[galagam]

Do you think the comment from Claude review should be a concern?

The single SUGGESTION concerns transient peak-RAM doubling from the deep-copy on multi-GB models.

Otherwise, LGTM!

@gcunhase Thanks for holding me accountable... Replied in the original thread.
#1928 (comment)