diff --git a/cpp/src/torch_tensorrt/executorch/README.md b/cpp/src/torch_tensorrt/executorch/README.md index da49329d6f..a764f2236a 100644 --- a/cpp/src/torch_tensorrt/executorch/README.md +++ b/cpp/src/torch_tensorrt/executorch/README.md @@ -430,7 +430,6 @@ torch_tensorrt.save( trt_module, "model.pte", output_format="executorch", - arg_inputs=inputs, use_cuda_graphs=True, ) ``` diff --git a/docsrc/user_guide/runtime_performance/executorch.rst b/docsrc/user_guide/runtime_performance/executorch.rst index 8af7f32071..5eb8737a36 100644 --- a/docsrc/user_guide/runtime_performance/executorch.rst +++ b/docsrc/user_guide/runtime_performance/executorch.rst @@ -80,11 +80,9 @@ To pick up a change before it ships, or to report a problem against ``main``: Compile and Save ----------------- -The usual ``ir="dynamo"`` path, with two extra arguments to ``torch_tensorrt.save``: - -* ``output_format="executorch"`` selects the ``.pte`` writer. -* ``retrace=False`` keeps the compiled graph as it is. Re-exporting would drop the TensorRT - engines before ExecuTorch's partitioner sees them. +The usual ``ir="dynamo"`` path, with one extra argument to ``torch_tensorrt.save``: +``output_format="executorch"`` selects the ``.pte`` writer. For this format ``save`` keeps +the compiled graph as it is rather than re-exporting it, so it needs no example inputs. .. code-block:: python @@ -101,13 +99,7 @@ The usual ``ir="dynamo"`` path, with two extra arguments to ``torch_tensorrt.sav min_block_size=1, ) - torch_tensorrt.save( - trt_gm, - "model.pte", - output_format="executorch", - arg_inputs=example_inputs, - retrace=False, - ) + torch_tensorrt.save(trt_gm, "model.pte", output_format="executorch") :ref:`executorch_save` covers the rest: dynamic shapes, several methods in one file, a zero-copy KV cache, and the two-step ``torch_tensorrt.executorch.export()`` path for @@ -129,8 +121,6 @@ CUDA backend compiles them instead, so nothing leaves the GPU: trt_gm, "coalesced.pte", output_format="executorch", - arg_inputs=example_inputs, - retrace=False, partitioners=[ CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")]) ], @@ -164,8 +154,6 @@ copies are pure overhead: trt_gm, "device_resident.pte", output_format="executorch", - arg_inputs=example_inputs, - retrace=False, backend_config=ExecutorchBackendConfig( propagate_device_config=PropagateDeviceConfig( skip_h2d_for_method_inputs=True, @@ -397,8 +385,6 @@ launch many short kernels, where CPU launch work dominates each call. trt_gm, "model.pte", output_format="executorch", - arg_inputs=example_inputs, - retrace=False, use_cuda_graphs=True, ) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 04547bd3cb..19e326cdf6 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -275,10 +275,7 @@ Default: ``save()`` model = MyModel().eval().cuda() inputs = [torch.randn((1, 3, 224, 224)).cuda()] trt_gm = torch_tensorrt.compile(model, ir="dynamo", arg_inputs=inputs) - torch_tensorrt.save( - trt_gm, "trt.pte", output_format="executorch", - retrace=False, arg_inputs=inputs, - ) + torch_tensorrt.save(trt_gm, "trt.pte", output_format="executorch") ``save`` writes both the ``.pte`` and any external ``.ptd`` tensor-data files. @@ -479,7 +476,6 @@ one silently would break a runner built before this feature. torch_tensorrt.save( trt_gm, "decode.pte", output_format="executorch", - arg_inputs=inputs, retrace=False, zero_copy_kv=True, ) @@ -539,7 +535,6 @@ claims what it can and the ``CudaPartitioner`` picks up the rest as a catch-all: torch_tensorrt.save( trt_gm, "trt.pte", output_format="executorch", - retrace=False, arg_inputs=inputs, partitioners=[ CudaPartitioner( [CudaBackend.generate_method_name_compile_spec("forward")] @@ -629,7 +624,6 @@ and ``zero_copy_kv`` is read by ``save`` on both sides of that boundary. torch_tensorrt.save( trt_gm, "trt.pte", output_format="executorch", - retrace=False, arg_inputs=inputs, constant_methods={"get_max_seq_len": 2048}, compile_config=EdgeCompileConfig(_check_ir_validity=False), generate_etrecord=True, diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index a7d2f6a0d9..11250afab2 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -739,7 +739,7 @@ def save( inputs: Optional[Sequence[torch.Tensor | Input]] = None, arg_inputs: Optional[Sequence[torch.Tensor | Input]] = None, kwarg_inputs: Optional[Dict[str, Any]] = None, - retrace: bool = True, + retrace: Optional[bool] = None, use_legacy_exporter: Optional[bool] = None, pickle_protocol: int = 2, dynamic_shapes: Optional[Dict[str, Any]] = None, @@ -755,7 +755,10 @@ def save( arg_inputs (Tuple[Union[torch.Tensor, torch_tensorrt.Input], ...]): Same as inputs. Alias for better understanding with kwarg_inputs. kwarg_inputs (dict[str, Union[torch.Tensor, torch_tensorrt.Input]]): Optional, kwarg inputs to the module forward function. output_format (str): Format to save the model. Options include exported_program | torchscript | aot_inductor | executorch. - retrace (bool): When the module type is a fx.GraphModule, this option re-exports the graph using torch.export.export(strict=False) to save it. + retrace (Optional[bool]): When the module type is a fx.GraphModule, this option re-exports the graph using torch.export.export(strict=False) to save it. + + Defaults to True, except for ``output_format="executorch"``, where it defaults to False: + the compiled graph is lowered as it is, so no example inputs are needed. For TRT-compiled modules with dynamic shapes, both retrace=True and retrace=False are supported: @@ -879,6 +882,10 @@ def save( """ if isinstance(module, CudaGraphsTorchTensorRTModule): module = module.compiled_module + if retrace is None: + # Match torch_tensorrt.executorch.export, which lowers a compiled GraphModule as it + # is, so writing a .pte needs no example inputs. + retrace = output_format != "executorch" module_type = _parse_module_type(module) accepted_formats = {"exported_program", "torchscript", "aot_inductor", "executorch"} if arg_inputs is not None and not all( diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index e4fe24f000..b5123e6f85 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/tests/py/dynamo/executorch/test_api.py @@ -3053,6 +3053,47 @@ def test_public_save_forwards_lowering_kwargs_graphmodule_no_retrace( _assert_lowering_kwargs_forwarded(calls[0]["kwargs"], methods, passes, cfg) +@pytest.mark.unit +@_needs_torch_tensorrt +def test_public_save_graphmodule_to_executorch_skips_retrace_by_default( + monkeypatch, tmp_path +): + pytest.importorskip("executorch.exir") + import torch_tensorrt + import torch_tensorrt.dynamo._exporter as _exporter + + calls = _stub_save_as_executorch(monkeypatch) + stub_ep = torch.export.export(_AddOne(), (torch.randn(2, 2),)) + exporter_calls = [] + + def _fake_export(*args, **kwargs): + exporter_calls.append(kwargs) + return stub_ep + + monkeypatch.setattr(_exporter, "export", _fake_export) + + # Neither retrace nor inputs: a retrace would refuse to run without inputs. + gm = torch.export.export(_AddOne(), (torch.randn(2, 2),)).module() + torch_tensorrt.save(gm, str(tmp_path / "gm.pte"), output_format="executorch") + + assert len(exporter_calls) == 1 + assert exporter_calls[0]["use_legacy_exporter"] is True + assert len(calls) == 1 + assert calls[0]["module"] is stub_ep + + +@pytest.mark.unit +@_needs_torch_tensorrt +def test_public_save_graphmodule_to_exported_program_still_retraces_by_default( + tmp_path, +): + import torch_tensorrt + + gm = torch.export.export(_AddOne(), (torch.randn(2, 2),)).module() + with pytest.raises(ValueError, match="retrace is True"): + torch_tensorrt.save(gm, str(tmp_path / "gm.ep")) + + @pytest.mark.unit def test_save_executorch_real_etrecord_is_inspector_consumable(tmp_path): """A real lowering with generate_etrecord=True writes a sidecar that ExecuTorch's