From f5c6d8008447809945a588d8572cbc9ef1361aad Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Fri, 24 Jul 2026 16:22:23 -0700 Subject: [PATCH] feat(executorch): forward ExecuTorch lowering kwargs through save(output_format="executorch") torch_tensorrt.save already forwards partitioners= and backend_config= to the ExecuTorch lowering. Forward the remaining to_edge_transform_and_lower kwargs the same way -- pop them in save(), thread them through the three _save_as_executorch call sites, and pass them to to_edge_transform_and_lower(...): - constant_methods: bake constant methods (e.g. LLM-runner metadata) into the .pte - transform_passes: run caller-supplied edge transform passes - compile_config: override the EdgeCompileConfig. Defaults to _check_ir_validity=False (the TRT execute_engine placeholder graph fails edge IR validation); a caller-supplied config is forwarded verbatim so an explicit _check_ir_validity is honored. - generate_etrecord: generate + persist an ETRecord as "_etrecord.bin" next to the .pte, following ExecuTorch's example convention (e.g. examples/cuda/scripts/export.py) These kwargs only apply to output_format="executorch"; passing any of them with a different output_format now logs an ignored-kwarg warning, matching partitioners=, compile_specs= and backend_config=. Tests: - forward each kwarg through _save_as_executorch, covering an omitted compile_config (defaults to _check_ir_validity=False) vs a caller-supplied one (forwarded verbatim), and the ETRecord sidecar write. - forward the kwargs through the public torch_tensorrt.save() across all three dispatch branches: an ExportedProgram, a GraphModule with retrace=True, and a GraphModule with retrace=False. - a real lowering with generate_etrecord=True whose sidecar parses back into an Inspector-consumable ETRecord via executorch.devtools. Docs: document the lowering options, the compile_config default, the _etrecord.bin sidecar, and an ETRecord parse/Inspector example in the ExecuTorch saving guide. --- .../runtime_performance/saving_models.rst | 47 +++ py/torch_tensorrt/_compile.py | 71 ++++- tests/py/dynamo/executorch/test_api.py | 269 ++++++++++++++++++ 3 files changed, 385 insertions(+), 2 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 103cf19025..76f00485b3 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -277,6 +277,53 @@ must be pointed at those data files to load them. ``.pte`` into the same directory overwrites the blob and the first ``.pte`` will fail to load. Save each coalesced model into its own directory. +**ExecuTorch lowering options** + +When ``output_format="executorch"``, ``torch_tensorrt.save`` forwards the following +keyword arguments to ExecuTorch's ``to_edge_transform_and_lower(...)``. They are +only consulted for the ``executorch`` format; passing them with any other +``output_format`` logs a warning and is otherwise ignored. + +* ``constant_methods`` — a ``dict`` of extra constant methods to embed in the + ``.pte`` (e.g. ``{"get_max_seq_len": 2048}`` for an LLM runner). +* ``transform_passes`` — additional edge-dialect transform passes to run before + lowering. +* ``compile_config`` — an ``EdgeCompileConfig``. When omitted, Torch-TensorRT + supplies a default with ``_check_ir_validity=False`` (the TensorRT + ``execute_engine`` placeholder graph does not pass edge-IR validation). A + caller-supplied config is forwarded **verbatim**, so if you pass your own and + your graph carries TensorRT engines, set ``_check_ir_validity=False`` explicitly. +* ``backend_config`` — an ``ExecutorchBackendConfig`` forwarded to + ``to_executorch(...)``. +* ``generate_etrecord`` — a ``bool`` (default ``False``). When ``True``, an + `ETRecord `_ is written + next to the ``.pte`` as ``_etrecord.bin`` (e.g. ``trt.pte`` → + ``trt_etrecord.bin``) for use with the ExecuTorch Developer Tools ``Inspector``. + +.. code-block:: python + + from executorch.exir import EdgeCompileConfig + + 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, + ) + +The ETRecord sidecar can be parsed back and paired with a runtime ETDump in the +Developer Tools ``Inspector``: + +.. code-block:: python + + from executorch.devtools import Inspector + from executorch.devtools.etrecord import parse_etrecord + + etrecord = parse_etrecord("trt_etrecord.bin") + inspector = Inspector(etdump_path="etdump.etdp", etrecord=etrecord) + inspector.print_data_tabular() + Saving torch.compile models ----------------------------- diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index a4d1feaabf..e2220e8703 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -743,7 +743,9 @@ def save( - If both dynamic_shapes and Input objects are provided, the explicit dynamic_shapes parameter takes precedence. kwargs: Additional format-specific kwargs. ``partitioners=``, - ``compile_specs=``, and ``backend_config=`` are only used with + ``compile_specs=``, ``backend_config=``, ``constant_methods=``, + ``transform_passes=``, ``compile_config=`` and + ``generate_etrecord=`` are only used with ``output_format="executorch"``; otherwise they are ignored with a warning. Pass ``compile_specs=[CompileSpec("target_device", b"cuda:")]`` to override the default target device (``cuda:0``). @@ -757,6 +759,17 @@ def save( ` for the ``CudaPartitioner`` recipe, its export-time requirements (CUDA backend + nvcc), and the external ``.ptd`` weight caveats. + ``constant_methods=``, ``transform_passes=`` and + ``compile_config=`` are forwarded to + ``to_edge_transform_and_lower(...)``. When ``compile_config`` is + omitted it defaults to ``_check_ir_validity=False`` (the TRT engine + graph fails edge IR validation); a caller-supplied + ``compile_config`` is forwarded verbatim, so set + ``_check_ir_validity=False`` on it explicitly when the graph carries + TRT engines. + ``generate_etrecord=True`` writes an ETRecord to + ``_etrecord.bin`` next to the ``.pte`` for use with the + ExecuTorch Inspector. """ if isinstance(module, CudaGraphsTorchTensorRTModule): module = module.compiled_module @@ -784,6 +797,10 @@ def save( executorch_partitioners = kwargs.pop("partitioners", None) executorch_compile_specs = kwargs.pop("compile_specs", None) executorch_backend_config = kwargs.pop("backend_config", None) + executorch_constant_methods = kwargs.pop("constant_methods", None) + executorch_transform_passes = kwargs.pop("transform_passes", None) + executorch_compile_config = kwargs.pop("compile_config", None) + executorch_generate_etrecord = kwargs.pop("generate_etrecord", False) if output_format not in accepted_formats: raise ValueError( @@ -915,6 +932,26 @@ def _extract_tensor(obj: Any) -> Any: "backend_config= is only used with output_format='executorch' and will " f"be ignored for output_format='{output_format}'." ) + if executorch_constant_methods and output_format != "executorch": + logger.warning( + "constant_methods= is only used with output_format='executorch' and will " + f"be ignored for output_format='{output_format}'." + ) + if executorch_transform_passes and output_format != "executorch": + logger.warning( + "transform_passes= is only used with output_format='executorch' and will " + f"be ignored for output_format='{output_format}'." + ) + if executorch_compile_config and output_format != "executorch": + logger.warning( + "compile_config= is only used with output_format='executorch' and will " + f"be ignored for output_format='{output_format}'." + ) + if executorch_generate_etrecord and output_format != "executorch": + logger.warning( + "generate_etrecord= is only used with output_format='executorch' and will " + f"be ignored for output_format='{output_format}'." + ) if output_format == "aot_inductor" and platform.system() != "Linux": raise ValueError( f"The AOT Inductor format is only supported on Linux, {platform.system()} is not a supported platform for this format" @@ -984,6 +1021,10 @@ def _extract_tensor(obj: Any) -> Any: partitioners=executorch_partitioners, compile_specs=executorch_compile_specs, backend_config=executorch_backend_config, + constant_methods=executorch_constant_methods, + transform_passes=executorch_transform_passes, + compile_config=executorch_compile_config, + generate_etrecord=executorch_generate_etrecord, ) else: raise RuntimeError( @@ -1050,6 +1091,10 @@ def _extract_tensor(obj: Any) -> Any: partitioners=executorch_partitioners, compile_specs=executorch_compile_specs, backend_config=executorch_backend_config, + constant_methods=executorch_constant_methods, + transform_passes=executorch_transform_passes, + compile_config=executorch_compile_config, + generate_etrecord=executorch_generate_etrecord, ) else: raise RuntimeError( @@ -1137,6 +1182,10 @@ def _extract_tensor(obj: Any) -> Any: partitioners=executorch_partitioners, compile_specs=executorch_compile_specs, backend_config=executorch_backend_config, + constant_methods=executorch_constant_methods, + transform_passes=executorch_transform_passes, + compile_config=executorch_compile_config, + generate_etrecord=executorch_generate_etrecord, ) else: raise RuntimeError( @@ -1418,15 +1467,33 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None # would fail trying to cast the CustomObjArgument placeholder to a real C++ Engine object. exp_program = _replace_execute_engine_for_executorch(exp_program) + # Default to _check_ir_validity=False (get_edge_compile_config): the TRT + # execute_engine placeholder graph fails edge IR validation. A caller who passes + # their own compile_config owns it -- it is forwarded verbatim and we do NOT + # override _check_ir_validity (EdgeCompileConfig defaults it to True, so a caller + # whose graph carries TRT engines should set _check_ir_validity=False explicitly). + compile_config = kwargs.get("compile_config") + if compile_config is None: + compile_config = get_edge_compile_config() + + generate_etrecord = kwargs.get("generate_etrecord", False) edge_program = to_edge_transform_and_lower( exp_program, + transform_passes=kwargs.get("transform_passes"), partitioner=partitioners, - compile_config=get_edge_compile_config(), + compile_config=compile_config, + constant_methods=kwargs.get("constant_methods"), + generate_etrecord=generate_etrecord, ) executorch_program = edge_program.to_executorch(config=kwargs.get("backend_config")) with open(file_path, "wb") as f: executorch_program.write_to_file(f) _write_external_tensor_data(executorch_program, file_path) + 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. + etrecord_path = os.path.splitext(file_path)[0] + "_etrecord.bin" + executorch_program.get_etrecord().save(etrecord_path) def _normalize_engine_constants_to_python(exp_program: "ExportedProgram") -> None: diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index 5d88822716..a4f5b97b10 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/tests/py/dynamo/executorch/test_api.py @@ -333,3 +333,272 @@ def test_lift_preserves_constant_dtype_device(dtype, device): # The source constant is a contiguous 3x3, so from_tensor must carry its real # (3, 1) stride onto the meta, not the old all-ones synthetic stride. assert val.stride() == torch.zeros(3, 3).stride() + + +# --- save(output_format="executorch") forwards ExecuTorch lowering kwargs ------- +# to_edge_transform_and_lower accepts transform_passes / constant_methods / +# compile_config / generate_etrecord; save() should forward each. compile_config +# defaults to _check_ir_validity=False (the TRT engine placeholder graph fails edge +# IR validation) when omitted, but a caller-supplied config is forwarded verbatim. +# generate_etrecord persists a "_etrecord.bin" next to the .pte. + + +def _patch_executorch_lowering(monkeypatch, captured): + """Stub the ExecuTorch lowering + TRT-specific pre/post steps in _save_as_executorch + so the test exercises only kwarg forwarding. Returns nothing; fills `captured`.""" + import executorch.exir as exir + import torch_tensorrt._compile as tc + + class _FakeETRecord: + def save(self, path): + with open(path, "wb") as fh: + fh.write(b"etrecord") + + class _FakeExec: + def write_to_file(self, f): + f.write(b"") + + def get_etrecord(self): + return _FakeETRecord() + + class _FakeEdge: + def to_executorch(self, config=None): + captured["backend_config"] = config + return _FakeExec() + + def _fake_lower(exp_program, **kw): + captured.update(kw) + return _FakeEdge() + + monkeypatch.setattr(exir, "to_edge_transform_and_lower", _fake_lower) + monkeypatch.setattr(tc, "_count_executorch_engine_nodes", lambda ep: 0) + monkeypatch.setattr(tc, "_replace_execute_engine_for_executorch", lambda ep: ep) + 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) + ) + + +@pytest.mark.unit +def test_save_executorch_forwards_lowering_kwargs(monkeypatch, tmp_path): + pytest.importorskip("executorch.exir") + import torch_tensorrt._compile as tc + from executorch.exir import EdgeCompileConfig + + captured = {} + _patch_executorch_lowering(monkeypatch, captured) + + sentinel_passes = [object()] + sentinel_methods = {"get_max_seq_len": 128} + caller_cfg = EdgeCompileConfig(_check_ir_validity=True) + out = str(tmp_path / "model.pte") + + tc._save_as_executorch( + object(), + out, + partitioners=[], + compile_specs=[], + backend_config=None, + constant_methods=sentinel_methods, + transform_passes=sentinel_passes, + compile_config=caller_cfg, + generate_etrecord=True, + ) + + assert captured["transform_passes"] is sentinel_passes + assert captured["constant_methods"] is sentinel_methods + assert captured["generate_etrecord"] is True + # A caller-supplied compile_config is forwarded verbatim (explicit override + # respected, not overridden even though it sets _check_ir_validity=True). + assert captured["compile_config"] is caller_cfg + assert captured["compile_config"]._check_ir_validity is True + # ETRecord persisted next to the .pte per ET's "_etrecord.bin" convention. + assert (tmp_path / "model_etrecord.bin").exists() + + +@pytest.mark.unit +def test_save_executorch_defaults_when_lowering_kwargs_omitted(monkeypatch, tmp_path): + pytest.importorskip("executorch.exir") + import torch_tensorrt._compile as tc + + captured = {} + _patch_executorch_lowering(monkeypatch, captured) + + out = str(tmp_path / "model.pte") + tc._save_as_executorch(object(), out) + + # Falls back to get_edge_compile_config() (also _check_ir_validity=False). + assert captured["compile_config"]._check_ir_validity is False + assert captured["transform_passes"] is None + assert captured["generate_etrecord"] is False + # No etrecord written when generate_etrecord is falsy. + assert not (tmp_path / "model_etrecord.bin").exists() + + +# --- the same lowering kwargs flow through the *public* torch_tensorrt.save() ----- +# save() pops the ExecuTorch-only kwargs and forwards them to _save_as_executorch +# from three dispatch branches: an ExportedProgram input, a GraphModule with +# retrace=True (re-exported here), and a GraphModule with retrace=False (routed +# through the dynamo exporter). Each must extract the options and forward them. + + +class _AddOne(torch.nn.Module): + def forward(self, x): + return x + 1 + + +def _stub_save_as_executorch(monkeypatch): + """Capture the (module, file_path, kwargs) that save() forwards to + _save_as_executorch without running the real lowering.""" + import torch_tensorrt._compile as tc + + calls = [] + + def _fake(module, file_path, **kw): + calls.append({"module": module, "file_path": file_path, "kwargs": kw}) + + monkeypatch.setattr(tc, "_save_as_executorch", _fake) + return calls + + +def _assert_lowering_kwargs_forwarded(kw, methods, passes, cfg): + assert kw["constant_methods"] is methods + assert kw["transform_passes"] is passes + assert kw["compile_config"] is cfg + assert kw["generate_etrecord"] is True + assert kw["partitioners"] == [] + assert kw["compile_specs"] == [] + assert kw["backend_config"] is None + + +@pytest.mark.unit +def test_public_save_forwards_lowering_kwargs_exported_program(monkeypatch, tmp_path): + pytest.importorskip("executorch.exir") + import torch_tensorrt + from executorch.exir import EdgeCompileConfig + + calls = _stub_save_as_executorch(monkeypatch) + methods = {"get_max_seq_len": 128} + passes = [object()] + cfg = EdgeCompileConfig(_check_ir_validity=False) + out = str(tmp_path / "ep.pte") + + ep = torch.export.export(_AddOne(), (torch.randn(2, 2),)) + torch_tensorrt.save( + ep, + out, + output_format="executorch", + partitioners=[], + compile_specs=[], + backend_config=None, + constant_methods=methods, + transform_passes=passes, + compile_config=cfg, + generate_etrecord=True, + ) + + assert len(calls) == 1 + assert calls[0]["module"] is ep + _assert_lowering_kwargs_forwarded(calls[0]["kwargs"], methods, passes, cfg) + + +@pytest.mark.unit +def test_public_save_forwards_lowering_kwargs_graphmodule_retrace( + monkeypatch, tmp_path +): + pytest.importorskip("executorch.exir") + import torch_tensorrt + from executorch.exir import EdgeCompileConfig + + calls = _stub_save_as_executorch(monkeypatch) + methods = {"get_max_seq_len": 128} + passes = [object()] + cfg = EdgeCompileConfig(_check_ir_validity=False) + out = str(tmp_path / "gm_retrace.pte") + + gm = torch.export.export(_AddOne(), (torch.randn(2, 2),)).module() + torch_tensorrt.save( + gm, + out, + output_format="executorch", + retrace=True, + arg_inputs=(torch.randn(2, 2),), + partitioners=[], + compile_specs=[], + backend_config=None, + constant_methods=methods, + transform_passes=passes, + compile_config=cfg, + generate_etrecord=True, + ) + + assert len(calls) == 1 + _assert_lowering_kwargs_forwarded(calls[0]["kwargs"], methods, passes, cfg) + + +@pytest.mark.unit +def test_public_save_forwards_lowering_kwargs_graphmodule_no_retrace( + monkeypatch, tmp_path +): + pytest.importorskip("executorch.exir") + import torch_tensorrt + import torch_tensorrt.dynamo._exporter as _exporter + from executorch.exir import EdgeCompileConfig + + calls = _stub_save_as_executorch(monkeypatch) + # retrace=False routes through the dynamo exporter (TRT-specific graph surgery); + # stub it so the test isolates save()'s option extraction + forwarding. + sentinel_ep = object() + monkeypatch.setattr(_exporter, "export", lambda *a, **k: sentinel_ep) + + methods = {"get_max_seq_len": 128} + passes = [object()] + cfg = EdgeCompileConfig(_check_ir_validity=False) + out = str(tmp_path / "gm_no_retrace.pte") + + gm = torch.export.export(_AddOne(), (torch.randn(2, 2),)).module() + torch_tensorrt.save( + gm, + out, + output_format="executorch", + retrace=False, + partitioners=[], + compile_specs=[], + backend_config=None, + constant_methods=methods, + transform_passes=passes, + compile_config=cfg, + generate_etrecord=True, + ) + + assert len(calls) == 1 + assert calls[0]["module"] is sentinel_ep + _assert_lowering_kwargs_forwarded(calls[0]["kwargs"], methods, passes, cfg) + + +@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 + devtools can parse back into an Inspector-consumable ETRecord.""" + pytest.importorskip("executorch.exir") + pytest.importorskip("executorch.devtools") + import torch_tensorrt + from executorch.devtools.etrecord import parse_etrecord + from torch_tensorrt._features import ENABLED_FEATURES + + if not ENABLED_FEATURES.torch_tensorrt_runtime: + pytest.skip("output_format='executorch' requires the torch_tensorrt runtime") + + ep = torch.export.export(_AddOne(), (torch.randn(4, 4),)) + out = str(tmp_path / "tiny.pte") + torch_tensorrt.save(ep, out, output_format="executorch", generate_etrecord=True) + + etrecord_path = tmp_path / "tiny_etrecord.bin" + assert etrecord_path.exists() + + record = parse_etrecord(str(etrecord_path)) + assert record is not None + # The parsed record carries the edge-dialect program the Inspector correlates + # runtime events against. + assert getattr(record, "edge_dialect_program", None) is not None