feat(executorch): forward ExecuTorch lowering kwargs through save(output_format="executorch")

docsrc/user_guide/runtime_performance/saving_models.rst

@@ -22,6 +22,7 @@ specifying the `output_format` flag. Here are the options `output_format` will a
 * `exported_program` : This is the default. We perform transformations on the graphmodule first and use `torch.export.save` to save the module.
 * `torchscript` : We trace the graphmodule via `torch.jit.trace` and save it via `torch.jit.save`.
 * `PT2 Format` : This is a next generation runtime for PyTorch models, allowing them to run in Python and in C++
+* `executorch` : We lower the graphmodule to an ExecuTorch ``.pte`` program, delegating the TensorRT engines to the ExecuTorch backend. Linux-only; requires the ``executorch`` package.
 
 a) ExportedProgram
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -219,6 +220,64 @@ For dynamic shapes, C++ deployment, and a full comparison with the ExportedProgr
 see :ref:`aot_inductor`.
 
 
+.. _executorch_save:
+
+c) ExecuTorch (.pte)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The ``executorch`` output format lowers the compiled module to an ExecuTorch
+``.pte`` program, delegating the TensorRT engines to the Torch-TensorRT ExecuTorch
+backend. It requires the ``executorch`` package (``pip install
+"torch_tensorrt[executorch]"``) and is Linux-only.
+
+.. code-block:: python
+
+    import torch
+    import torch_tensorrt
+
+    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,
+    )
+
+**Coalesced TensorRT + CUDA .pte**
+
+To run the ops TensorRT does not take on ExecuTorch's CUDA (AOTInductor) backend
+instead of leaving them non-delegated, pass a ``CudaPartitioner`` via
+``partitioners=``. It is appended after the TensorRT partitioner, so TensorRT
+claims what it can and the ``CudaPartitioner`` picks up the rest as a catch-all:
+
+.. code-block:: python
+
+    from executorch.backends.cuda.cuda_backend import CudaBackend
+    from executorch.backends.cuda.cuda_partitioner import CudaPartitioner
+
+    torch_tensorrt.save(
+        trt_gm, "trt.pte", output_format="executorch",
+        retrace=False, arg_inputs=inputs,
+        partitioners=[
+            CudaPartitioner(
+                [CudaBackend.generate_method_name_compile_spec("forward")]
+            )
+        ],
+    )
+
+This needs ExecuTorch's CUDA backend and a CUDA toolkit (nvcc/ptxas) at export
+time, and produces a ``.pte`` that requires a CUDA runtime at load. Any external
+CUDA weights are written as ``.ptd`` data file(s) next to the ``.pte``; the runtime
+must be pointed at those data files to load them.
+
+.. warning::
+
+    The CUDA backend names its external weight blob per-device (e.g.
+    ``aoti_cuda_blob.ptd``), not per-model, so saving two different coalesced
+    ``.pte`` into the same directory overwrites the blob and the first ``.pte``
+    will fail to load. Save each coalesced model into its own directory.
+
+
 Saving torch.compile models
 -----------------------------
 

py/torch_tensorrt/_compile.py

@@ -4,6 +4,7 @@
 import importlib
 import inspect
 import logging
+import os
 import platform
 import warnings
 from enum import Enum
@@ -742,13 +743,33 @@ 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:<i>")]`` to override the default target device (``cuda:0``).
                 ``backend_config=`` takes an ``Optional[ExecutorchBackendConfig]``
                 and is forwarded to ``to_executorch(config=...)`` to customize
                 ExecuTorch lowering (e.g. memory planning or device placement).
+                ``partitioners=`` appends caller partitioners after the TensorRT
+                partitioner, so ops TensorRT does not take can be routed to another
+                delegate -- e.g. a ``CudaPartitioner`` for a coalesced TensorRT +
+                CUDA ``.pte``. See :ref:`the ExecuTorch save guide
+                <executorch_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
+                ``<model>_etrecord.bin`` next to the ``.pte`` for use with the
+                ExecuTorch Inspector.
     """
     if isinstance(module, CudaGraphsTorchTensorRTModule):
         module = module.compiled_module
@@ -776,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(
@@ -976,6 +1001,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(
@@ -1042,6 +1071,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(
@@ -1129,6 +1162,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(
@@ -1327,6 +1364,29 @@ def _replace_execute_engine_for_executorch(exp_program: Any) -> Any:
     return exp_program
 
 
+def _write_external_tensor_data(executorch_program: Any, file_path: str) -> 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
+    (``save_data_externally``), which ExecuTorch serializes only via
+    ``write_tensor_data_to_file`` -- ``ExecutorchProgram.write_to_file`` does not
+    persist it. So a partition that carries external weights (e.g. a
+    ``CudaPartitioner`` delegate) would lose its blob and the ``.pte`` could not
+    load. This writes the ``.ptd`` data file(s) into the ``.pte``'s directory when
+    the program has any; the runtime must then be given those data file(s) (e.g.
+    via the ExecuTorch ``Module`` data-files argument) to load the weights.
+    """
+    # 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)
+        logger.info(
+            "Wrote external delegate weights (.ptd) to %s; point the runtime's "
+            "data-files at this directory to load them.",
+            out_dir,
+        )
+
+
 def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None:
     """Save an ExportedProgram (with TensorRT execute_engine nodes) as an ExecuTorch .pte file.
 
@@ -1387,14 +1447,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 = bool(kwargs.get("generate_etrecord"))
     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 "<pte_base>_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:

py/torch_tensorrt/dynamo/_exporter.py

@@ -1,13 +1,12 @@
 import base64
 import copy
 import operator
-from typing import Any, Dict, Optional, Sequence, Tuple, cast
+from typing import Any, Dict, Optional, Sequence, Tuple
 
 import torch
 from torch._export.non_strict_utils import make_constraints
 from torch._guards import detect_fake_mode
 from torch._library.fake_class_registry import FakeScriptObject
-from torch._subclasses.fake_tensor import FakeTensor
 from torch.export import ExportedProgram, ExportGraphSignature
 from torch.export._trace import _combine_args
 from torch.export.exported_program import (
@@ -202,12 +201,11 @@ def lift(
                 # Copy the node meta into this new placeholder node
                 const_placeholder_node.meta = node.meta
                 if isinstance(lift_val, torch.Tensor):
-                    const_placeholder_node.meta["val"] = cast(
-                        FakeTensor,
-                        torch.empty_strided(
-                            tuple(lift_val.shape),
-                            tuple([1] * len(lift_val.shape)),
-                        ),
+                    # Fakify the source constant through the graph's own fake_mode
+                    # so the placeholder meta preserves its dtype, device and
+                    # stride without allocating a real tensor on lift_val.device.
+                    const_placeholder_node.meta["val"] = fake_mode.from_tensor(
+                        lift_val, static_shapes=True
                     )
 
                 node.replace_all_uses_with(const_placeholder_node)