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)