diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index b8e4b6204c..103cf19025 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/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 ----------------------------- diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index 0c2680d20a..dfbe9e372c 100644 --- a/py/torch_tensorrt/_compile.py +++ b/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:")]`` 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 + ` 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 @@ -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 "_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/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index ab4d3e3d49..dacbea140a 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/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) diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index 0e249959b9..ada72c58b4 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/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,173 @@ 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() + + +# --- 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") + from executorch.exir import EdgeCompileConfig + + import torch_tensorrt._compile as tc + + 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() diff --git a/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py b/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py new file mode 100644 index 0000000000..64233f45cd --- /dev/null +++ b/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py @@ -0,0 +1,227 @@ +import importlib.util +import os +import shutil + +import pytest + +executorch = pytest.importorskip("executorch.exir") + +import torch # noqa: E402 + + +def _cuda_backend_available() -> bool: + # find_spec() imports the parent packages of a dotted path and raises + # ModuleNotFoundError if a parent is absent (it only returns None when the + # parent exists but the leaf module doesn't). Guard against that so collection + # still succeeds where executorch.backends.cuda isn't packaged. + try: + spec = importlib.util.find_spec("executorch.backends.cuda.cuda_partitioner") + except ModuleNotFoundError: + return False + return spec is not None + + +CUDA_BACKEND_AVAILABLE = _cuda_backend_available() + + +def _nvcc_available() -> bool: + # The AOTInductor compile in the CUDA backend needs nvcc. It's on PATH in a + # standard CUDA install, but some environments ship the toolkit via CUDA_HOME + # and the compile resolves nvcc from there rather than PATH -- accept either. + if shutil.which("nvcc") is not None: + return True + try: + from torch.utils.cpp_extension import CUDA_HOME + except Exception: + CUDA_HOME = None + for home in (CUDA_HOME, os.environ.get("CUDA_HOME"), os.environ.get("CUDA_PATH")): + if home and os.path.exists(os.path.join(home, "bin", "nvcc")): + return True + return False + + +def _skip_reason(): + # Building a coalesced TRT + CUDA .pte needs a GPU, ExecuTorch's CUDA backend, + # and a CUDA toolkit (nvcc/ptxas) for the AOTInductor compile run during export. + if not torch.cuda.is_available(): + return "CUDA GPU required for the CUDA backend" + if not CUDA_BACKEND_AVAILABLE: + return "executorch.backends.cuda not installed" + if not _nvcc_available(): + return "nvcc required for the CUDA (AOTInductor) backend compile" + return None + + +@pytest.fixture(autouse=True) +def requires_cuda_backend(): + # Evaluate the GPU/backend/nvcc gate at TEST RUNTIME, not at import. A + # module-level pytest.mark.skipif is resolved during collection; on remote-GPU + # runners collection happens off the GPU host, so the skip would be baked in + # and the runner would never attach a GPU for the (pre-skipped) test. + reason = _skip_reason() + if reason is not None: + pytest.skip(reason) + + +def _cuda_partitioner(): + """A ``CudaPartitioner`` catch-all so ops TensorRT does not take route to the + ExecuTorch CUDA (AOTInductor) backend. This is the supported, flag-free way to + build a coalesced TensorRT + CUDA ``.pte`` via ``save(partitioners=[...])``.""" + from executorch.backends.cuda.cuda_backend import CudaBackend + from executorch.backends.cuda.cuda_partitioner import CudaPartitioner + + return CudaPartitioner([CudaBackend.generate_method_name_compile_spec("forward")]) + + +def _delegate_ids(pte_path): + """Backend ids of every delegate in the serialized program, in order.""" + from executorch.exir._serialize._program import deserialize_pte_binary + + program = deserialize_pte_binary(pte_path.read_bytes()).program + return [ + delegate.id for plan in program.execution_plan for delegate in plan.delegates + ] + + +# NOTE: these tests are COMPOSITION-ONLY. They assert the serialized .pte carries +# both a TensorRTBackend and a CudaBackend delegate (and, below, that external +# CUDA weights are persisted as a .ptd). They do NOT load or run the program: a +# coalesced ATen-mode .pte cannot be loaded yet because memory-planned CUDA +# buffers get a CPU data pointer, so Method::init fails the CUDA backend's device +# check (tensor_parser_aten hardcodes CPU). A load-run-allclose test should be +# added once that runtime fix lands (follow-up). + + +def test_erfinv_routes_to_cuda_backend(tmp_path): + """A genuinely TRT-unsupported op (erfinv) is routed to the CUDA backend via a + CudaPartitioner catch-all -- no torch_executed_ops pin needed.""" + import torch_tensorrt + + # tanh keeps values in erfinv's (-1, 1) domain. + class ErfinvModel(torch.nn.Module): + def forward(self, x): + return torch.cos(torch.erfinv(torch.tanh(x))) + + model = ErfinvModel().eval().to("cuda") + inputs = (torch.randn(64, 64, device="cuda"),) + exported = torch.export.export(model, inputs) + trt_gm = torch_tensorrt.dynamo.compile( + exported, + inputs=list(inputs), + min_block_size=1, + truncate_double=True, + ) + + out = tmp_path / "erfinv_cuda.pte" + torch_tensorrt.save( + trt_gm, + str(out), + output_format="executorch", + retrace=False, + arg_inputs=list(inputs), + partitioners=[_cuda_partitioner()], + ) + + delegate_ids = _delegate_ids(out) + assert ( + "CudaBackend" in delegate_ids + ), f"erfinv did not route to the CUDA backend; delegates={delegate_ids}" + assert ( + "TensorRTBackend" in delegate_ids + ), f"tanh/cos were not delegated to TensorRT; delegates={delegate_ids}" + + +def test_weighted_partition_persists_external_data(tmp_path): + """A CUDA partition that carries weights must have its external data persisted + next to the .pte. + + The CUDA (AOTInductor) backend emits weights as external named data + (save_data_externally). If save() only writes the .pte and not the .ptd data + file, the blob is lost and the program cannot load. Here mm(x, w) carries a + weight and is pinned out of TensorRT via torch_executed_ops, so it lands on the + CudaPartitioner and its weight becomes external CUDA data; tanh/cos stay on + TensorRT. We assert a .ptd is written alongside the .pte.""" + import torch_tensorrt + + class WeightedModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.w = torch.nn.Parameter(torch.randn(64, 64)) + + def forward(self, x): + x = torch.mm(x, self.w) + x = torch.tanh(x) + return torch.cos(x) + + model = WeightedModel().eval().to("cuda") + inputs = (torch.randn(64, 64, device="cuda"),) + exported = torch.export.export(model, inputs) + # Pin mm out of TensorRT so its weight lands on the CudaPartitioner catch-all + # and is emitted as external CUDA data; tanh/cos still go to TensorRT. + trt_gm = torch_tensorrt.dynamo.compile( + exported, + inputs=list(inputs), + min_block_size=1, + truncate_double=True, + torch_executed_ops={"torch.ops.aten.mm.default"}, + ) + + out = tmp_path / "weighted_cuda.pte" + torch_tensorrt.save( + trt_gm, + str(out), + output_format="executorch", + retrace=False, + arg_inputs=list(inputs), + partitioners=[_cuda_partitioner()], + ) + + delegate_ids = _delegate_ids(out) + assert ( + "CudaBackend" in delegate_ids + ), f"the pinned mm did not route to the CUDA backend; delegates={delegate_ids}" + # Regression guard for dropped external CUDA weights: the .ptd data file must + # be written next to the .pte. + ptd_files = list(tmp_path.glob("*.ptd")) + assert ptd_files, f"no external .ptd data file was written next to {out}" + assert all( + p.stat().st_size > 0 for p in ptd_files + ), f"external .ptd data file is empty: {ptd_files}" + + +def test_trt_only_writes_no_ptd(tmp_path): + """A TRT-only program (no CudaPartitioner) carries no external data, so no + .ptd is written -- exercises the real write_to_file / _tensor_data path.""" + import torch_tensorrt + + class Model(torch.nn.Module): + def forward(self, x): + return torch.cos(torch.tanh(x)) + + model = Model().eval().to("cuda") + inputs = (torch.randn(64, 64, device="cuda"),) + exported = torch.export.export(model, inputs) + trt_gm = torch_tensorrt.dynamo.compile( + exported, + inputs=list(inputs), + min_block_size=1, + truncate_double=True, + ) + + out = tmp_path / "trt_only.pte" + torch_tensorrt.save( + trt_gm, + str(out), + output_format="executorch", + retrace=False, + arg_inputs=list(inputs), + ) + + assert out.exists() + delegate_ids = _delegate_ids(out) + assert ( + "CudaBackend" not in delegate_ids + ), f"unexpected CUDA delegate for a TRT-only model; delegates={delegate_ids}" + assert not list( + tmp_path.glob("*.ptd") + ), "TRT-only program must not write an external .ptd" diff --git a/tests/py/dynamo/executorch/test_edge_cases.py b/tests/py/dynamo/executorch/test_edge_cases.py index 631453c451..935b67604c 100644 --- a/tests/py/dynamo/executorch/test_edge_cases.py +++ b/tests/py/dynamo/executorch/test_edge_cases.py @@ -1,10 +1,13 @@ +import os from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch from torch_tensorrt._compile import ( _count_executorch_engine_nodes, _validate_executorch_engine_info, + _write_external_tensor_data, ) from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( REQUIRES_OUTPUT_ALLOCATOR_IDX, @@ -47,3 +50,39 @@ def test_count_executorch_engine_nodes_handles_execute_and_placeholder(): ) assert _count_executorch_engine_nodes(exp_program) == 2 + + +@pytest.mark.unit +def test_write_external_tensor_data_writes_when_present(tmp_path): + # A program with external named data (e.g. a CudaPartitioner delegate's + # weights) must have its .ptd written into the .pte's directory. + prog = SimpleNamespace( + _tensor_data={"forward": b"weights"}, + write_tensor_data_to_file=MagicMock(), + ) + pte = tmp_path / "model.pte" + _write_external_tensor_data(prog, str(pte)) + prog.write_tensor_data_to_file.assert_called_once_with( + os.path.dirname(os.path.abspath(str(pte))) + ) + + +@pytest.mark.unit +def test_write_external_tensor_data_noop_when_empty(tmp_path): + # TRT-only programs have empty _tensor_data (falsy) -> no .ptd written. + prog = SimpleNamespace( + _tensor_data={}, + write_tensor_data_to_file=MagicMock(), + ) + _write_external_tensor_data(prog, str(tmp_path / "model.pte")) + prog.write_tensor_data_to_file.assert_not_called() + + +@pytest.mark.unit +def test_write_external_tensor_data_fails_loud_without_attr(tmp_path): + # _tensor_data always exists on a real ExecutorchProgram; it is accessed + # directly (no getattr default) so a future rename fails loudly instead of + # silently skipping the .ptd write and reintroducing the null-weights crash. + prog = SimpleNamespace(write_tensor_data_to_file=MagicMock()) + with pytest.raises(AttributeError): + _write_external_tensor_data(prog, str(tmp_path / "model.pte"))