feat(executorch): do not retrace by default when saving a .pte

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,
 )
 ```

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,
     )
 

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,

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(

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