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(