From 7c3b53591678b5da31bca7fd3405fb652b6aa6b9 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 15 Nov 2026 04:08:45 -0800 Subject: [PATCH] fix(executorch): move compile_and_save's engines into the program instead of copying them `compile_and_save` writes each engine to a temporary file as soon as it is built. The ExecuTorch save then copies that file into the `.pte` (or the `.ptd`), and the file is deleted only when the call returns. So while the program is written, every engine is on disk twice. For a model whose engine is tens of gigabytes, the copy needs that much more free disk, and an export can fail on a device that has room for the result. The backend now hands each of these engine files to ExecuTorch as data that owns the file (`FileBackedData(path, cleanup=True)`); it is deleted after the save anyway. Inside `compile_and_save`, the program is then written with ExecuTorch's `save(file_path, consume=True)` and the `.ptd` with `write_tensor_data_to_file(..., consume=True)`, so ExecuTorch moves each engine into the program: it copies it back to front and truncates the file as it goes, and each engine is on disk once. ExecuTorch does this only when the temporary directory, where the engine files are, is on the same file system as the program; otherwise it copies. Its `save` takes only a `.pte` path, so a program saved under another name is written as before, with a copy. With `generate_etrecord`, the save copies the engines instead. The ETRecord is written after the program and reads the engines, so a move would leave it nothing to read. The program is still written first, as before, so it is kept when the ETRecord can't be written. The `save` docstring now says that `compile_and_save` can write an ETRecord. This needs the `consume` argument of `save` and `write_tensor_data_to_file` from pytorch/executorch#25099. It lands after that change, together with an ExecuTorch pin that has it, which also raises the minimum ExecuTorch version. Only the move passes `consume`, so every other save calls ExecuTorch exactly as before and keeps working with releases that don't have it. Test plan: - `tests/py/dynamo/executorch/test_api.py`, CPU, with the lowering stubbed: - New `test_compile_and_save_moves_the_engines`: inside `compile_and_save`'s offload block, the program and the `.ptd` are written with `consume=True`. It fails without this change. - New `test_compile_and_save_copies_the_engines_into_a_path_without_pte`: the same block with a `.bin` path writes both without `consume`. - New `test_compile_and_save_keeps_the_program_when_the_etrecord_fails`: in the same block with an ETRecord, the program and the `.ptd` are written first, without `consume`, and the `.pte` is kept when the ETRecord save raises. - New `test_save_executorch_keeps_the_program_when_the_etrecord_fails`: the same for a plain save. - New `test_save_executorch_does_not_consume`: a plain save writes without `consume`. The stubbed ExecuTorch raises TypeError for any `consume` other than `True`, as releases without it do, so this test fails if a plain save passes `consume=False`. - `tests/py/dynamo/executorch/test_compile_and_save.py` on an A100: - The streaming test now also checks that each engine reaches ExecuTorch as data that owns its file. - New `test_engines_move_into_the_program_without_a_second_copy` checks that ExecuTorch copied no engine and that both are in the `.pte`. - New `test_an_etrecord_makes_the_save_copy_the_engines`: with an ETRecord, the engine files keep their bytes, and both the `.pte` and the ETRecord are written. --- .../runtime_performance/saving_models.rst | 2 +- py/torch_tensorrt/_compile.py | 35 ++++- .../dynamo/runtime/_offloaded_engines.py | 1 + py/torch_tensorrt/executorch/_export.py | 1 + py/torch_tensorrt/executorch/backend.py | 4 +- tests/py/dynamo/executorch/test_api.py | 128 +++++++++++++++++- .../executorch/test_compile_and_save.py | 71 +++++++++- 7 files changed, 231 insertions(+), 11 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index f4676e03fe..bba749e1db 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -281,7 +281,7 @@ warning and are ignored. With this format, an unknown keyword raises a ``TypeErr :ref:`Zero-copy KV cache `. * ``generate_etrecord``: also write an ETRecord, a debug record for the ExecuTorch Developer Tools ``Inspector``, as ``_etrecord.bin``. It does not work for a program with TensorRT - engines: writing the record fails with a ``SerializeError``. + engines: writing the record fails with a ``SerializeError``. ``compile_and_save`` can write one. :ref:`executorch_deployment` explains the run time side of ``use_cuda_graphs``, ``weight_streaming_budget_per_engine`` and ``external_engine_data``. diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index bd2d57e6b6..8454507cae 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -77,6 +77,7 @@ from torch_tensorrt.dynamo._tracer import trace as dynamo_trace from torch_tensorrt.dynamo.runtime import TorchTensorRTModule from torch_tensorrt.dynamo.runtime._offloaded_engines import ( + defer_engine_setup_to_load, offload_engines_to_disk, ) from torch_tensorrt.dynamo.utils import get_torch_inputs @@ -873,6 +874,7 @@ def save( ``_etrecord.bin`` next to the ``.pte`` for use with the ExecuTorch Inspector. It does not work for a program with TensorRT engines: writing the record raises a ``SerializeError``. + ``compile_and_save`` can write one. ``weight_streaming_budget_per_engine=`` takes an ``Optional[int]`` number of bytes of engine weights that may stay resident in GPU memory, with the rest streamed from host memory. It applies to @@ -1550,6 +1552,11 @@ def compile_and_save( but each engine is written to a file as soon as it is built, so no built engine is held on the GPU or in host memory while the next ones build. With ``output_format="executorch"`` no engine is ever loaded: the save streams each file into the ``.pte`` (or the ``.ptd`` with ``external_engine_data``). + Without ``generate_etrecord``, the save moves each file there instead, so no engine is on disk twice, + when ``file_path`` ends with ``.pte`` and the temporary directory (``TMPDIR``) is on the same file + system. With + ``generate_etrecord``, the save copies the engines and writes the ETRecord after the program, so the + program is kept if the ETRecord can't be written. The other formats save each engine from a loaded one, so their engines are loaded once all of them are built, and their save needs as much memory as after ``compile``. @@ -1600,7 +1607,9 @@ def compile_and_save( save(trt_gm, os.fspath(file_path), output_format=output_format, **save_kwargs) -def _write_external_tensor_data(executorch_program: Any, file_path: str) -> None: +def _write_external_tensor_data( + executorch_program: Any, file_path: str, consume: bool = False +) -> None: """Write an ExecuTorch program's external named tensor data (``.ptd``) next to the ``.pte``. The CUDA (AOTInductor) backend emits its weights as external named data @@ -1616,7 +1625,11 @@ def _write_external_tensor_data(executorch_program: Any, file_path: str) -> None # Direct attribute access (no getattr default) so a future rename fails loud. if executorch_program._tensor_data: out_dir = os.path.dirname(os.path.abspath(file_path)) - executorch_program.write_tensor_data_to_file(out_dir) + # ExecuTorch releases before the move take no consume argument. + if consume: + executorch_program.write_tensor_data_to_file(out_dir, consume=True) + else: + executorch_program.write_tensor_data_to_file(out_dir) # write_tensor_data_to_file adds .ptd unless the name already ends with it. paths = [ os.path.join(out_dir, name if name.endswith(".ptd") else f"{name}.ptd") @@ -1752,9 +1765,21 @@ def _lower_and_write_executorch( # the same call is the caller's to make, because finalization happens on # ExecuTorch's own manager and out of this library's sight. check_zero_copy_kv(executorch_program) - with open(file_path, "wb") as f: - executorch_program.write_to_file(f) - _write_external_tensor_data(executorch_program, file_path) + # Inside compile_and_save the engine files are deleted after this save, so it moves them into + # the program. The ETRecord reads the engines after the program is written, so with one the save + # copies them, which also keeps the program when the ETRecord can't be saved. ExecuTorch moves + # data only in its own save, which takes a .pte path. + consume = ( + defer_engine_setup_to_load() + and not generate_etrecord + and file_path.endswith(".pte") + ) + if consume: + executorch_program.save(file_path, consume=True) + else: + with open(file_path, "wb") as f: + executorch_program.write_to_file(f) + _write_external_tensor_data(executorch_program, file_path, consume) if generate_etrecord: # Follows ExecuTorch's example convention (e.g. examples/cuda/scripts/export.py): # persist the ETRecord as "_etrecord.bin" next to the .pte. diff --git a/py/torch_tensorrt/dynamo/runtime/_offloaded_engines.py b/py/torch_tensorrt/dynamo/runtime/_offloaded_engines.py index 629c996bff..3f23d89ae7 100644 --- a/py/torch_tensorrt/dynamo/runtime/_offloaded_engines.py +++ b/py/torch_tensorrt/dynamo/runtime/_offloaded_engines.py @@ -8,6 +8,7 @@ engines built so far are then held neither on the GPU nor in host memory while the next ones build. ``compile_module`` leaves them in their files: an ExecuTorch save streams each file into the program, and ``compile_and_save`` loads them before any other save. The files are deleted when the block ends. +A file the save moved stays mapped, empty, until then, so nothing may read its engine after the save. """ import contextlib diff --git a/py/torch_tensorrt/executorch/_export.py b/py/torch_tensorrt/executorch/_export.py index 9a6235d9f2..e4e70b9718 100644 --- a/py/torch_tensorrt/executorch/_export.py +++ b/py/torch_tensorrt/executorch/_export.py @@ -599,6 +599,7 @@ def export( generate_etrecord (bool): Ask ExecuTorch for an ETRecord for later debugging. This copies the whole program, engines included. Saving the record does not work for a program with TensorRT engines: it raises a ``SerializeError``. + ``torch_tensorrt.compile_and_save`` can write one. weight_streaming_budget_per_engine (Optional[int]): Bytes of engine weights that may stay resident in GPU memory, with the rest streamed from host memory. It applies to **each** TensorRT engine separately, not as a total for the diff --git a/py/torch_tensorrt/executorch/backend.py b/py/torch_tensorrt/executorch/backend.py index ef061a7211..beaec1dd08 100644 --- a/py/torch_tensorrt/executorch/backend.py +++ b/py/torch_tensorrt/executorch/backend.py @@ -699,7 +699,9 @@ def preprocess( engine_data: Any if engine_file is not None: - engine_data = FileBackedData(engine_file) + # compile_and_save deletes the file after the save, so the data can own it, which lets + # a consuming write move the engine into the program instead of copying it. + engine_data = FileBackedData(engine_file, cleanup=True) engine_data_sha256 = engine_data.sha256() else: assert isinstance(engine_bytes, bytes) diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index 61f2338339..511d01d7c1 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/tests/py/dynamo/executorch/test_api.py @@ -2824,7 +2824,7 @@ def test_lift_keeps_parameter_for_any_dtype(dtype, requires_grad): # next to the .pte. -def _patch_executorch_lowering(monkeypatch, captured): +def _patch_executorch_lowering(monkeypatch, captured, log=None, etrecord_error=None): """Spy the seams _save_as_executorch delegates to after #4440, without running a real TRT lowering. Records the kwargs forwarded into torch_tensorrt.executorch.export() (captured["export_kwargs"]); lets the real @@ -2832,20 +2832,48 @@ def _patch_executorch_lowering(monkeypatch, captured): ExecuTorch lowering, so the kwargs export() finally hands that lowering -- including the get_edge_compile_config() it substitutes for an omitted compile_config -- land in `captured`; and records backend_config from - to_executorch(). Fills `captured`; returns nothing.""" + to_executorch(). Fills `captured`; returns nothing. With `log`, records the order of + the ETRecord, program and .ptd writes, and whether each write consumes; + `etrecord_error` makes the ETRecord save raise.""" import executorch.exir as exir import torch_tensorrt._compile as tc import torch_tensorrt.executorch as tte class _FakeETRecord: def save(self, path): + if log is not None: + log.append("etrecord") + if etrecord_error is not None: + raise etrecord_error with open(path, "wb") as fh: fh.write(b"etrecord") + def _consumes(kwargs): + # ExecuTorch releases before the move reject consume, so only a move may pass it. + if kwargs not in ({}, {"consume": True}): + raise TypeError(f"unexpected keyword arguments {kwargs}") + return bool(kwargs) + class _FakeExec: + _tensor_data = {"model": b""} + + def save(self, path, **kwargs): + consume = _consumes(kwargs) + if log is not None: + log.append(("program", consume)) + with open(path, "wb") as fh: + fh.write(b"") + def write_to_file(self, f): + if log is not None: + log.append(("program", False)) f.write(b"") + def write_tensor_data_to_file(self, outdir, **kwargs): + consume = _consumes(kwargs) + if log is not None: + log.append(("ptd", consume)) + def get_etrecord(self): return _FakeETRecord() @@ -2875,7 +2903,6 @@ def _spy_export(source, **kw): return real_export(source, **kw) monkeypatch.setattr(tte, "export", _spy_export) - monkeypatch.setattr(tc, "_write_external_tensor_data", lambda prog, path: None) # ENABLED_FEATURES is an immutable namedtuple; swap the whole module attribute. monkeypatch.setattr( tc, "ENABLED_FEATURES", types.SimpleNamespace(torch_tensorrt_runtime=True) @@ -2933,6 +2960,101 @@ def test_save_executorch_forwards_lowering_kwargs(monkeypatch, tmp_path): assert (tmp_path / "model_etrecord.bin").exists() +@pytest.mark.unit +def test_save_executorch_keeps_the_program_when_the_etrecord_fails( + monkeypatch, tmp_path +): + pytest.importorskip("executorch.exir") + import torch_tensorrt._compile as tc + + log = [] + _patch_executorch_lowering( + monkeypatch, {}, log=log, etrecord_error=RuntimeError("cannot serialize") + ) + out = tmp_path / "model.pte" + ep = torch.export.export(_AddOne(), (torch.randn(2, 2),)) + with pytest.raises(RuntimeError, match="cannot serialize"): + tc._save_as_executorch(ep, str(out), generate_etrecord=True) + + assert log == [("program", False), ("ptd", False), "etrecord"] + assert out.exists() + + +@pytest.mark.unit +def test_save_executorch_does_not_consume(monkeypatch, tmp_path): + pytest.importorskip("executorch.exir") + import torch_tensorrt._compile as tc + + log = [] + _patch_executorch_lowering(monkeypatch, {}, log=log) + ep = torch.export.export(_AddOne(), (torch.randn(2, 2),)) + tc._save_as_executorch(ep, str(tmp_path / "model.pte")) + + assert log == [("program", False), ("ptd", False)] + + +@pytest.mark.unit +def test_compile_and_save_moves_the_engines(monkeypatch, tmp_path): + pytest.importorskip("executorch.exir") + import torch_tensorrt._compile as tc + from torch_tensorrt.dynamo.runtime._offloaded_engines import ( + offload_engines_to_disk, + ) + + log = [] + _patch_executorch_lowering(monkeypatch, {}, log=log) + ep = torch.export.export(_AddOne(), (torch.randn(2, 2),)) + # The block compile_and_save saves in, where the engine files are deleted after the save. + with offload_engines_to_disk(): + tc._save_as_executorch(ep, str(tmp_path / "model.pte")) + + assert log == [("program", True), ("ptd", True)] + + +@pytest.mark.unit +def test_compile_and_save_copies_the_engines_into_a_path_without_pte( + monkeypatch, tmp_path +): + pytest.importorskip("executorch.exir") + import torch_tensorrt._compile as tc + from torch_tensorrt.dynamo.runtime._offloaded_engines import ( + offload_engines_to_disk, + ) + + log = [] + _patch_executorch_lowering(monkeypatch, {}, log=log) + ep = torch.export.export(_AddOne(), (torch.randn(2, 2),)) + # ExecuTorch moves data only in its save, which needs a .pte path. + with offload_engines_to_disk(): + tc._save_as_executorch(ep, str(tmp_path / "model.bin")) + + assert log == [("program", False), ("ptd", False)] + + +@pytest.mark.unit +def test_compile_and_save_keeps_the_program_when_the_etrecord_fails( + monkeypatch, tmp_path +): + pytest.importorskip("executorch.exir") + import torch_tensorrt._compile as tc + from torch_tensorrt.dynamo.runtime._offloaded_engines import ( + offload_engines_to_disk, + ) + + log = [] + _patch_executorch_lowering( + monkeypatch, {}, log=log, etrecord_error=OSError("No space left on device") + ) + out = tmp_path / "model.pte" + ep = torch.export.export(_AddOne(), (torch.randn(2, 2),)) + with offload_engines_to_disk(): + with pytest.raises(OSError, match="No space left"): + tc._save_as_executorch(ep, str(out), generate_etrecord=True) + + assert log == [("program", False), ("ptd", False), "etrecord"] + assert out.exists() + + @pytest.mark.unit def test_save_executorch_defaults_when_lowering_kwargs_omitted(monkeypatch, tmp_path): pytest.importorskip("executorch.exir") diff --git a/tests/py/dynamo/executorch/test_compile_and_save.py b/tests/py/dynamo/executorch/test_compile_and_save.py index 8b497644d0..cc8e299533 100644 --- a/tests/py/dynamo/executorch/test_compile_and_save.py +++ b/tests/py/dynamo/executorch/test_compile_and_save.py @@ -31,7 +31,7 @@ def forward(self, x): @pytest.mark.unit def test_engines_stream_from_their_files_into_the_program(tmp_path, monkeypatch): monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) - setups, streamed = [], {} + setups, streamed, owned = [], {}, [] real_setup = TorchTensorRTModule.setup_engine real_add = NamedDataStore.add_named_data @@ -43,6 +43,8 @@ def recording_add(self, key, data, *args, **kwargs): if isinstance(data, FileBackedData): with open(data._path, "rb") as f: streamed[data._path] = f.read() + # Data that owns its file is what a consuming save may move. + owned.append(data._finalizer is not None) else: streamed[key] = data return real_add(self, key, data, *args, **kwargs) @@ -65,6 +67,7 @@ def recording_add(self, key, data, *args, **kwargs): # No engine was ever set up, and each one reached the program as its own file, from the cache directory. assert setups == [] assert len(streamed) == 2 + assert owned and all(owned) offloaded = str(_cache_root() / "offloaded_engines") assert all( path.endswith(".engine") and os.path.dirname(path) == offloaded @@ -74,3 +77,69 @@ def recording_add(self, key, data, *args, **kwargs): assert all(pte.find(engine) >= 0 for engine in streamed.values()) # The files are gone once the program is saved. assert not glob.glob(str(tmp_path / "**" / "*.engine"), recursive=True) + + +@pytest.mark.unit +def test_engines_move_into_the_program_without_a_second_copy(tmp_path, monkeypatch): + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + engines, copied = {}, [] + real_add = NamedDataStore.add_named_data + real_write = FileBackedData.write_to_file + + def recording_add(self, key, data, *args, **kwargs): + if isinstance(data, FileBackedData): + with open(data._path, "rb") as f: + engines[data._path] = f.read() + return real_add(self, key, data, *args, **kwargs) + + def recording_write(self, outfile): + copied.append(self._path) + real_write(self, outfile) + + monkeypatch.setattr(NamedDataStore, "add_named_data", recording_add) + monkeypatch.setattr(FileBackedData, "write_to_file", recording_write) + torch.manual_seed(0) + x = torch.randn(4, 64, device="cuda") + pte_path = tmp_path / "model.pte" + torch_tensorrt.compile_and_save( + torch.export.export(_TwoEngines().eval().cuda(), (x,)), + pte_path, + output_format="executorch", + arg_inputs=[x], + min_block_size=1, + torch_executed_ops={torch.ops.aten.sin.default}, + ) + # ExecuTorch copies data it can't move through write_to_file, so no call there means each + # engine was moved, and the program holds them. + assert len(engines) == 2 and copied == [] + pte = pte_path.read_bytes() + assert all(pte.find(engine) >= 0 for engine in engines.values()) + + +@pytest.mark.unit +def test_an_etrecord_makes_the_save_copy_the_engines(tmp_path, monkeypatch): + monkeypatch.setattr(tempfile, "tempdir", str(tmp_path)) + left = [] + real_write = FileBackedData.write_to_file + + def recording_write(self, outfile): + real_write(self, outfile) + left.append(os.path.getsize(self._path)) + + monkeypatch.setattr(FileBackedData, "write_to_file", recording_write) + torch.manual_seed(0) + x = torch.randn(4, 64, device="cuda") + pte_path = tmp_path / "model.pte" + torch_tensorrt.compile_and_save( + torch.export.export(_TwoEngines().eval().cuda(), (x,)), + pte_path, + output_format="executorch", + arg_inputs=[x], + min_block_size=1, + torch_executed_ops={torch.ops.aten.sin.default}, + save_options={"generate_etrecord": True}, + ) + assert pte_path.stat().st_size > 0 + assert (tmp_path / "model_etrecord.bin").stat().st_size > 0 + # The ETRecord reads the engines after the program is written, so they stayed in their files. + assert left and all(size > 0 for size in left)