executorch: persist external .ptd weights and preserve lifted constant dtype/device

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
@@ -749,6 +750,13 @@ def save(
                 ``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.
     """
     if isinstance(module, CudaGraphsTorchTensorRTModule):
         module = module.compiled_module
@@ -1327,6 +1335,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.
 
@@ -1395,6 +1426,7 @@ def _save_as_executorch(exp_program: Any, file_path: str, **kwargs: Any) -> None
     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)
 
 
 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)

tests/py/dynamo/executorch/test_api.py

@@ -7,8 +7,8 @@
 import pytest
 import torch
 from torch._library.fake_class_registry import FakeScriptObject
-
-from torch_tensorrt.dynamo._exporter import _resolve_lifted_custom_obj
+from torch._subclasses.fake_tensor import FakeTensor
+from torch_tensorrt.dynamo._exporter import _resolve_lifted_custom_obj, lift
 
 
 @pytest.mark.unit
@@ -266,3 +266,70 @@ def fake_info(ep, node):
     assert d0 == b"cuda:0"
     assert d1 == b"cuda:1"
     assert d0 != d1
+
+
+# --- lift() preserves constant dtype/device ----------------------------------
+# lift() replaces get_attr constants with placeholders and copies a fake meta
+# "val". It must fakify the source constant through the graph's own fake_mode
+# (fake_mode.from_tensor), preserving dtype, device and stride; otherwise a
+# non-fp32 or non-CPU lifted constant silently gets fp32/cpu meta.
+
+
+def _traced_gm_with_parameter(dtype, device):
+    """A symbolically-traced GraphModule with one get_attr parameter (`c`) of the
+    given dtype/device, plus a stub graph_signature lift() can mutate."""
+    from torch._subclasses.fake_tensor import FakeTensorMode
+
+    class M(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.c = torch.nn.Parameter(
+                torch.zeros(3, 3, dtype=dtype, device=device), requires_grad=False
+            )
+
+        def forward(self, x):
+            return x + self.c
+
+    gm = torch.fx.symbolic_trace(M())
+    # lift() calls detect_fake_mode over the placeholder "val" metas, so the user
+    # input needs a FakeTensor meta.
+    fake_mode = FakeTensorMode()
+    with fake_mode:
+        fake_x = torch.empty(3, 3)
+    for node in gm.graph.nodes:
+        if node.op == "placeholder":
+            node.meta["val"] = fake_x
+    sig = types.SimpleNamespace(input_specs=[], output_specs=[], user_inputs=["x"])
+    return gm, sig
+
+
+def _lifted_constant_meta(gm, sig):
+    lifted_gm, _, _, _ = lift(gm, sig)
+    lifted = [
+        n for n in lifted_gm.graph.nodes if n.op == "placeholder" and n.name != "x"
+    ]
+    assert len(lifted) == 1, f"expected 1 lifted constant, got {len(lifted)}"
+    return lifted[0].meta["val"]
+
+
+@pytest.mark.unit
+@pytest.mark.parametrize(
+    "dtype, device",
+    [
+        (torch.bfloat16, "cpu"),
+        (torch.float32, "cuda"),
+    ],
+)
+def test_lift_preserves_constant_dtype_device(dtype, device):
+    # Runtime gate (not a module-level skipif, which resolves at collection time
+    # and is fragile on remote-GPU runners): skip the CUDA case only when no GPU.
+    if device == "cuda" and not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    gm, sig = _traced_gm_with_parameter(dtype, device)
+    val = _lifted_constant_meta(gm, sig)
+    assert isinstance(val, FakeTensor)
+    assert val.dtype == dtype
+    assert val.device.type == 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()