fix(executorch): copy the engine in bulk instead of byte by byte - #4473
Open
shoumikhin wants to merge 1 commit into
Open
fix(executorch): copy the engine in bulk instead of byte by byte#4473shoumikhin wants to merge 1 commit into
shoumikhin wants to merge 1 commit into
Conversation
Serializing an engine to bytes went through bytes(tensor.untyped_storage()). That looks like a bulk copy but iterates the storage one element at a time in Python, at roughly two seconds per megabyte: 16.8 MB bytes(untyped_storage()) 36.27 s 16.8 MB memoryview(... .numpy()) 0.01 s For a model whose graph produces many engines, that dominates the whole export. A program with 8.8 GB of engines spends about five hours here, and one with 70 GB spends closer to two days. Reading the storage is also wrong for a tensor that is a view. contiguous() is a no-op for a row slice, so the storage carries neighbouring bytes and the serialized engine comes out longer than the tensor: tensor [0, 1, 2, 3] bytes(untyped_storage()) -> [0, 1, 2, 3, 4, 5, 6, 7] memoryview(numpy()) -> [0, 1, 2, 3] Copy through a memoryview over the tensor's own buffer. This keeps the property the previous code was after, no extra full-size buffer for a large engine, while copying in one shot and respecting the tensor's bounds.
shoumikhin
force-pushed
the
trt-engine-serialization-perf
branch
from
August 9, 2026 23:37
21ac6db to
4063483
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The problem
Serializing a TensorRT engine to bytes goes through
bytes(tensor.untyped_storage()).That reads like a bulk copy, but it iterates the storage one element at a time in Python,
at roughly two seconds per megabyte:
For a model whose graph produces many engines, this step dominates the entire export. A
program with 8.8 GB of engines spends about five hours here. One with 70 GB spends closer
to two days. It is easy to mistake for a slow engine build, because the export simply sits
there making no visible progress.
It is also wrong for a view
contiguous()is a no-op for a tensor that is already contiguous, including a row slice ofa larger tensor. In that case the storage holds neighbouring bytes, so the serialized
engine comes out longer than the tensor it came from:
Today's engines happen to own their whole storage, so this does not bite in practice yet.
It is a trap for anyone who later produces the engine buffer as a slice.
The fix
Copy through a
memoryviewover the tensor's own buffer:The comment on the previous code explains what it was avoiding:
.numpy().tobytes()allocates a second full-size buffer, which roughly doubles peak memory for a multi-gigabyte
engine.
memoryviewkeeps that property, since it is a view rather than a copy, whiledoing the copy in one shot and respecting the tensor's bounds.
Testing
Verified byte-for-byte equality against the previous path for the shapes this code sees:
The third row is the bug above: the old path returned the whole storage rather than the
tensor.
Also exported a multi-method model that produces ten engines totalling 8.8 GB. Before the
change the serialization step ran for over an hour and a half without finishing; after it,
the same step completes in minutes and produces a program that loads and runs.