From 40634836a71c68a5b149e93ceac1349f7019a42b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 9 Aug 2026 14:53:14 -0700 Subject: [PATCH] fix(executorch): copy the engine in bulk instead of byte by byte 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. --- py/torch_tensorrt/executorch/backend.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index b73d50eea2..2da1f73d03 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -282,16 +282,16 @@ def preprocess( _validate_engine_info(engine_info) serialized_engine = engine_info[ENGINE_IDX] if isinstance(serialized_engine, torch.Tensor): - # Single copy out of the underlying storage. The prior - # `.numpy().tobytes()` path allocated a fresh bytes buffer - # on top of the numpy view, which for a >2 GB engine - # roughly doubled peak memory at this step. `.cpu()` and - # `.contiguous()` are no-ops when already host-side and - # contiguous (the common case for the uint8 buffer this - # backend produces). - engine_info[ENGINE_IDX] = bytes( - serialized_engine.cpu().contiguous().untyped_storage() - ) + # A single copy out of the tensor's own memory. `bytes(storage)` looks + # equivalent but iterates the storage element by element in Python, + # which costs about two seconds per megabyte and turns serializing a + # multi-gigabyte set of engines into hours. `memoryview` hands the + # buffer to the copy in one shot without the intermediate numpy view + # that would double peak memory for a large engine. `.cpu()` and + # `.contiguous()` are no-ops for the host-side uint8 buffer this + # backend normally produces. + engine_bytes = serialized_engine.cpu().contiguous().view(torch.uint8) + engine_info[ENGINE_IDX] = bytes(memoryview(engine_bytes.numpy())) elif not isinstance(serialized_engine, (bytes, bytearray)): engine_info[ENGINE_IDX] = bytes(serialized_engine) input_names = _reorder_input_names_for_executorch(