From 93efd54d0b430fd91e1cfae352498e75f5c5826a Mon Sep 17 00:00:00 2001 From: shoumikhin Date: Tue, 28 Jul 2026 13:54:37 -0700 Subject: [PATCH 1/4] feat(executorch): expose composable Edge export API ## What this adds `torch_tensorrt.save(..., output_format="executorch")` turns a compiled model into an ExecuTorch `.pte` file in one step. That is the easiest path and it stays the recommendation. Some workflows need to stop earlier. You may want to look at the delegated Edge graph, run your own Edge transforms, send operations TensorRT cannot handle to another backend, keep separate methods such as `prefill` and `decode`, add constant methods, or pick the final ExecuTorch configuration yourself. This adds `torch_tensorrt.executorch.export()`. It returns ExecuTorch's standard `EdgeProgramManager`, which is the supported place to inspect and customize an Edge program before you call `to_executorch()`. ```python import torch_tensorrt.executorch edge = torch_tensorrt.executorch.export(trt_module) # Inspect or transform the Edge program here. print(edge.exported_program().graph) program = edge.to_executorch() with open("model.pte", "wb") as output: program.write_to_file(output) # Advanced callers must also persist external tensor data when present. program.write_tensor_data_to_file(".") ``` `save()` now calls the same implementation, so the simple and advanced paths share one lowering path instead of two. ## Which one should I use? Use `save()` when you just want a `.pte`. Use `torch_tensorrt.executorch.export()` when you need control before the `.pte` is created. ## Example: TensorRT with a fallback backend TensorRT always gets the first chance to claim its prebuilt engine nodes. Any partitioners you pass run afterward and can claim operations TensorRT does not support. ```python edge = torch_tensorrt.executorch.export( trt_module, partitioners=[cuda_partitioner], ) program = edge.to_executorch() ``` ## Example: multiple methods ```python edge = torch_tensorrt.executorch.export( { "prefill": prefill_program, "decode": decode_program, }, compile_specs={ "prefill": prefill_specs, "decode": decode_specs, }, ) ``` Method names are preserved, and each method can have its own partitioner and compile-spec pipeline. This does not make mutable state shared between methods. ## Accepted inputs - a TensorRT-compiled `torch.fx.GraphModule` - an engine-bearing `torch.export.ExportedProgram` - a mapping from method names to independent `ExportedProgram` objects A plain `torch.nn.Module` must be compiled with Torch-TensorRT first. ## Safety and memory behavior Rewriting engine calls changes the exported graph. To avoid damaging the program you passed in, export first stages its own copy of the graph structure, signatures, state containers, constants, and metadata. If rewriting or lowering fails, your original program is left alone. Tensor and TensorRT engine payloads are treated as read-only and shared with that staged copy, so a multi-gigabyte engine is not duplicated. Custom transform passes must not modify these shared payload objects. Metadata that describes symbolic shapes is a special case. It points back to the live shape environment that the exported program is guarded on, so it is shared rather than copied. Copying it would both fail (the shape environment holds fake tensors that cannot be duplicated) and detach the copy from the symbols the graph depends on. This is what keeps dynamic-shape models working. When a lifted TensorRT engine is replaced, the old placeholder, its graph-signature entry, and its constant are removed. Several calls that share one engine reuse a single materialized `uint8` payload buffer. ## Compatibility - Existing one-step `save(..., output_format="executorch")` behavior is preserved, including its rejection of non-list `partitioners` and `compile_specs`. - `"executorch"` is now included in the `output_format` type hint. It was already accepted at run time, so this only fixes the annotation. - `.pte` serialization and external `.ptd` persistence are preserved. - TensorRT remains the first partitioner. - Shared dynamic dimensions declared with `Input(shared_dims=...)` remain shared. - Zero-engine methods are allowed. A later partitioner may claim them, or portable operators may remain undelegated. ## Testing Tested on a CUDA GPU with TensorRT and ExecuTorch installed, building Torch-TensorRT from source so the C++ runtime was active. Verified by compiling a real model with TensorRT and then exporting it: - static-shape `export()` produces a loadable `.pte` - dynamic-shape `export()` produces a loadable `.pte`, matching the output of the previous one-step save path - the source `GraphModule` still holds its original engine nodes afterward - engine payloads are shared, not re-serialized - `save(..., output_format="executorch")` still writes a `.pte` - `save()` still raises `TypeError` for a non-list `partitioners` or `compile_specs` Also ran the ExecuTorch test directory (76 passed). Two GPU tests in `test_cuda_partitioner_composition.py` fail, and they fail the same way with this change reverted, so they are not caused by it. Static checks: formatting, import ordering, lint, and type checking on the changed files. --- .../runtime_performance/saving_models.rst | 31 + py/torch_tensorrt/_compile.py | 261 +----- py/torch_tensorrt/executorch/__init__.py | 3 + py/torch_tensorrt/executorch/_export_utils.py | 312 +++++++ py/torch_tensorrt/executorch/export.py | 350 +++++++ tests/py/dynamo/executorch/test_api.py | 11 + .../test_cuda_partitioner_composition.py | 111 ++- tests/py/dynamo/executorch/test_edge_cases.py | 80 +- tests/py/dynamo/executorch/test_export.py | 881 ++++++++++++++++++ 9 files changed, 1761 insertions(+), 279 deletions(-) create mode 100644 py/torch_tensorrt/executorch/_export_utils.py create mode 100644 py/torch_tensorrt/executorch/export.py create mode 100644 tests/py/dynamo/executorch/test_export.py diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index 103cf19025..d32351d293 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -234,6 +234,7 @@ backend. It requires the ``executorch`` package (``pip install import torch import torch_tensorrt + import torch_tensorrt.executorch model = MyModel().eval().cuda() inputs = [torch.randn((1, 3, 224, 224)).cuda()] @@ -243,6 +244,36 @@ backend. It requires the ``executorch`` package (``pip install retrace=False, arg_inputs=inputs, ) +``save`` writes both the ``.pte`` and any external ``.ptd`` tensor-data files. +Advanced ExecuTorch users can stop at the standard Edge program boundary to +inspect the delegated graph, add metadata, or control final memory planning: + +.. code-block:: python + + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=inputs, + retrace=False, + partitioners=extra_partitioners, + transform_passes=passes, + compile_config=edge_config, + constant_methods={"get_vocab_size": 256}, + ) + + # Inspect edge.exported_program() or apply additional Edge transforms here. + program = edge.to_executorch(config=backend_config) + with open("trt.pte", "wb") as output: + program.write_to_file(output) + program.write_tensor_data_to_file(".") + +``executorch.export`` accepts a TensorRT-compiled ``GraphModule``, an +engine-bearing ``ExportedProgram``, or a mapping of independently exported +methods. It always applies the TensorRT partitioner first, followed by caller +``partitioners`` in order, and returns ExecuTorch's native +``EdgeProgramManager``. Method mappings preserve independent entry points but +do not imply shared mutable state. Use ``save`` for the one-shot path where +Torch-TensorRT manages Edge lowering, finalization, and persistence. + **Coalesced TensorRT + CUDA .pte** To run the ops TensorRT does not take on ExecuTorch's CUDA (AOTInductor) backend diff --git a/py/torch_tensorrt/_compile.py b/py/torch_tensorrt/_compile.py index 32ee801886..cf17c700a3 100644 --- a/py/torch_tensorrt/_compile.py +++ b/py/torch_tensorrt/_compile.py @@ -657,7 +657,7 @@ def save( *, extra_files: Optional[dict[str, str]] = None, output_format: Literal[ - "exported_program", "torchscript", "aot_inductor" + "exported_program", "torchscript", "aot_inductor", "executorch" ] = "exported_program", inputs: Optional[Sequence[torch.Tensor | Input]] = None, arg_inputs: Optional[Sequence[torch.Tensor | Input]] = None, @@ -1147,197 +1147,6 @@ def _extract_tensor(obj: Any) -> Any: ) -def _get_engine_info_from_state(engine_obj: Any) -> List[Any]: - """Normalize TensorRT engine state into the serialized engine-info list.""" - state = engine_obj.__getstate__() - engine_info = state[0] if isinstance(state, tuple) else state - return list(engine_info) - - -def _validate_executorch_engine_info( - engine_info: Sequence[Any], *, node_name: str = "" -) -> None: - """Reject engine configurations unsupported by the ExecuTorch export path.""" - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( - REQUIRES_OUTPUT_ALLOCATOR_IDX, - ) - - if ( - len(engine_info) > REQUIRES_OUTPUT_ALLOCATOR_IDX - and str(engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX]) == "1" - ): - node_suffix = f" for node '{node_name}'" if node_name else "" - raise RuntimeError( - "ExecuTorch export does not support TensorRT engines that require " - "an output allocator (data-dependent output shapes)" - f"{node_suffix}." - ) - - -def _count_executorch_engine_nodes(exp_program: Any) -> int: - """Count TRT execute-engine nodes in an ExportedProgram.""" - count = 0 - for node in exp_program.graph_module.graph.nodes: - if node.op != "call_function": - continue - if node.target is torch.ops.tensorrt.execute_engine.default: - count += 1 - continue - target = node.target - if ( - hasattr(target, "_schema") - and target._schema.name == "tensorrt::no_op_placeholder_for_execute_engine" - ): - count += 1 - return count - - -def _replace_execute_engine_for_executorch(exp_program: Any) -> Any: - """Replace execute_engine nodes with no_op_placeholder_for_execute_engine. - - ExecuTorch's to_edge_transform_and_lower runs ExportPass subclasses that - dispatch through the C++ schema validator. The validator rejects the - ScriptObject engine arg (it arrives as a CustomObjArgument placeholder - rather than a real FakeScriptObject). Converting each execute_engine node - to no_op_placeholder_for_execute_engine (which carries all engine info as - plain strings) avoids the ScriptObject entirely so the passes succeed. - - The TRT engine bytes are stored as a ``torch.uint8`` buffer on the graph - module and referenced from the no_op call via a ``get_attr`` FX node. This - keeps the engine out of the FX-emitted Python source: CPython's tokenizer - cannot parse string literals larger than ~2 GB, so an inline base64 string - breaks ``gm.recompile()`` for any engine whose payload exceeds that limit. - """ - from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( - ENGINE_IDX, - SERIALIZATION_LEN, - ) - - gm = exp_program.graph_module - execute_engine_op = torch.ops.tensorrt.execute_engine.default - no_op = torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default - - nodes_to_replace = [ - n - for n in gm.graph.nodes - if n.op == "call_function" and n.target is execute_engine_op - ] - if not nodes_to_replace: - return exp_program - - for engine_idx_in_graph, node in enumerate(nodes_to_replace): - inputs_arg = node.args[0] - engine_node = node.args[1] - - # Retrieve the engine ScriptObject from the graph module or constants. - if engine_node.op == "get_attr": - engine_obj = getattr(gm, engine_node.target, None) - if engine_obj is None: - raise RuntimeError( - f"execute_engine node '{node.name}': get_attr target " - f"'{engine_node.target}' not found on graph module" - ) - elif engine_node.op == "placeholder": - from torch_tensorrt.dynamo._exporter import _resolve_lifted_custom_obj - - engine_obj = _resolve_lifted_custom_obj(exp_program, engine_node) - if engine_obj is None: - raise RuntimeError( - f"execute_engine node '{node.name}': placeholder engine " - f"'{engine_node.name}' did not resolve to a lifted " - f"custom-object constant (available: " - f"{sorted(getattr(exp_program, 'constants', {}) or {})})" - ) - else: - raise RuntimeError( - f"execute_engine node '{node.name}': unexpected engine arg op " - f"'{engine_node.op}'" - ) - - engine_info = _get_engine_info_from_state(engine_obj) - _validate_executorch_engine_info(engine_info, node_name=node.name) - # Ensure the engine bytes slot is a base64 string (no_op takes str args). - engine_bytes = engine_info[ENGINE_IDX] - if isinstance(engine_bytes, str): - # `_get_engine_info_from_state` returns the engine as a - # base64-encoded `str` when the engine arrived through the - # serialized TRT runtime round-trip path. Decode back to raw - # bytes so it can land in a uint8 buffer below. - import base64 - - engine_bytes = base64.b64decode(engine_bytes) - elif not isinstance(engine_bytes, (bytes, bytearray)): - engine_bytes = bytes(engine_bytes) - # Store engine payload as a uint8 buffer + get_attr ref. FX emits a - # name reference instead of an inline literal, sidestepping the - # tokenizer's >2 GB string-literal limit. - engine_tensor = torch.frombuffer(bytearray(engine_bytes), dtype=torch.uint8) - # Use FX's unique-attr-name helper so re-export passes (which may - # invoke this rewriter multiple times on the same `gm`) don't - # silently overwrite earlier engine buffers. - from torch.fx.experimental.const_fold import ( - get_unique_attr_name_in_module, - ) - - buffer_name = get_unique_attr_name_in_module(gm, "_trt_engine_0") - gm.register_buffer(buffer_name, engine_tensor, persistent=True) - exp_program.state_dict[buffer_name] = engine_tensor - - str_args = [ - str(x) if x is not None else "" for x in engine_info[:SERIALIZATION_LEN] - ] - # Build a FakeTensor mirror so downstream FX passes (FakeTensorProp, - # ExecuTorch lowering, export-serde) that read `node.meta["val"]` - # on the `get_attr` reference don't `KeyError`. Must reuse the - # graph's existing FakeTensorMode — creating a fresh one would - # fail downstream with "fake mode from input 0 doesn't match - # mode from input 1" the moment any pass mixes the two. - from torch._guards import detect_fake_mode - - fake_mode = detect_fake_mode( - [n.meta["val"] for n in gm.graph.nodes if "val" in n.meta] - ) - fake_engine = ( - fake_mode.from_tensor(engine_tensor) - if fake_mode is not None - else engine_tensor - ) - with gm.graph.inserting_before(node): - engine_attr_node = gm.graph.get_attr(buffer_name) - engine_attr_node.meta["val"] = fake_engine - no_op_args = ( - inputs_arg, - *str_args[:ENGINE_IDX], - engine_attr_node, - *str_args[ENGINE_IDX + 1 :], - ) - no_op_node = gm.graph.call_function(no_op, no_op_args) - no_op_node.meta["val"] = node.meta.get("val") - - node.replace_all_uses_with(no_op_node) - gm.graph.erase_node(node) - - # Erase the engine get_attr node if it is now unused. - if engine_node.op == "get_attr" and not engine_node.users: - gm.graph.erase_node(engine_node) - # Also drop the now-orphan attribute from the module so the - # original engine bytes aren't double-serialized into state_dict - # alongside the new uint8 buffer. Use FX's dotted-path helper - # so nested-target attrs (e.g. `submod.engine`) are deleted - # correctly — plain `delattr(gm, "a.b")` only works on - # top-level names. - from torch.fx.graph_module import _del_attr, _has_attr - - if _has_attr(gm, engine_node.target): - _del_attr(gm, engine_node.target) - - gm.graph.eliminate_dead_code() - gm.graph.lint() - gm.recompile() - - 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``. @@ -1362,69 +1171,35 @@ def _write_external_tensor_data(executorch_program: Any, file_path: str) -> None 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. - - Partitions the graph by torch.ops.tensorrt.no_op_placeholder_for_execute_engine - (execute_engine is pre-converted to avoid schema type errors in edge passes), - serializes each engine to the same blob format as the TRT runtime (vector of - strings), and embeds it in the .pte. Requires the ``executorch`` package and - torch_tensorrt_runtime. See https://pytorch.org/executorch/stable/getting-started-setup.html - """ + """Save an engine-bearing ExportedProgram as an ExecuTorch program.""" if not ENABLED_FEATURES.torch_tensorrt_runtime: raise RuntimeError( "output_format='executorch' requires the Torch-TensorRT runtime " "(torch_tensorrt_runtime). Reinstall torch_tensorrt with the runtime extension." ) try: - from executorch.exir import to_edge_transform_and_lower + from torch_tensorrt.executorch import export except ImportError: raise ImportError( "ExecuTorch is not installed. Install with: pip install " "\"torch_tensorrt[executorch]\" to use output_format='executorch'." ) import torch_tensorrt.dynamo.runtime.meta_ops.register_meta_ops # noqa: F401 - from torch_tensorrt.executorch import ( - TensorRTPartitioner, - get_edge_compile_config, - ) - - extra_partitioners = kwargs.get("partitioners") or [] - if not isinstance(extra_partitioners, (list, tuple)): - raise TypeError( - "partitioners must be a list or tuple when using " - "output_format='executorch'" - ) - # Forward any caller-provided compile_specs to TensorRTPartitioner so users - # can override the default target_device ("cuda:0") by passing e.g. - # `compile_specs=[CompileSpec("target_device", b"cuda:1")]` to save(). - # When omitted, TensorRTPartitioner auto-appends the cuda:0 default. - executorch_compile_specs = kwargs.get("compile_specs") or [] - if not isinstance(executorch_compile_specs, (list, tuple)): - raise TypeError( - "compile_specs must be a list or tuple when using " - "output_format='executorch'" - ) - partitioners = [ - TensorRTPartitioner(compile_specs=list(executorch_compile_specs)) - ] + list(extra_partitioners) - - engine_count = _count_executorch_engine_nodes(exp_program) - if engine_count > 1: - logger.warning( - "%d TRT engines detected. Multi-engine .pte exports can incur extra " - "delegate boundary overhead.", - engine_count, - ) - # Replace execute_engine nodes (ScriptObject arg) with no_op_placeholder_for_execute_engine - # (string args only) before to_edge_transform_and_lower runs ExportPass subclasses that - # would fail trying to cast the CustomObjArgument placeholder to a real C++ Engine object. - exp_program = _replace_execute_engine_for_executorch(exp_program) + # save() only ever lowers a single method, so keep rejecting the per-method + # mappings that torch_tensorrt.executorch.export() accepts. + for option in ("partitioners", "compile_specs"): + value = kwargs.get(option) + if value is not None and not isinstance(value, (list, tuple)): + raise TypeError( + f"{option} must be a list or tuple when using " + "output_format='executorch'" + ) - edge_program = to_edge_transform_and_lower( + edge_program = export( exp_program, - partitioner=partitioners, - compile_config=get_edge_compile_config(), + partitioners=kwargs.get("partitioners"), + compile_specs=kwargs.get("compile_specs"), ) executorch_program = edge_program.to_executorch(config=kwargs.get("backend_config")) with open(file_path, "wb") as f: @@ -1443,14 +1218,10 @@ def _normalize_engine_constants_to_python(exp_program: "ExportedProgram") -> Non import base64 from torch_tensorrt.dynamo.runtime._serialized_engine_layout import ENGINE_IDX - from torch_tensorrt.dynamo.runtime._TRTEngine import ( - EngineSerializer, - TRTEngine, - ) + from torch_tensorrt.dynamo.runtime._TRTEngine import EngineSerializer, TRTEngine for fqn, constant in list(exp_program.constants.items()): if isinstance(constant, (torch._C.ScriptObject, TRTEngine)): - state = constant.__getstate__() if len(state) == 2 and ( state[1] == "TRTEngine" diff --git a/py/torch_tensorrt/executorch/__init__.py b/py/torch_tensorrt/executorch/__init__.py index 123eee846d..44ac5824b9 100644 --- a/py/torch_tensorrt/executorch/__init__.py +++ b/py/torch_tensorrt/executorch/__init__.py @@ -25,9 +25,11 @@ def __getattr__(name: str) -> NoReturn: "get_edge_compile_config", "TensorRTPartitioner", "TensorRTBackend", + "export", ] else: from torch_tensorrt.executorch.backend import TensorRTBackend + from torch_tensorrt.executorch.export import export from torch_tensorrt.executorch.partitioner import TensorRTPartitioner def get_edge_compile_config() -> "EdgeCompileConfig": @@ -40,4 +42,5 @@ def get_edge_compile_config() -> "EdgeCompileConfig": "get_edge_compile_config", "TensorRTPartitioner", "TensorRTBackend", + "export", ] diff --git a/py/torch_tensorrt/executorch/_export_utils.py b/py/torch_tensorrt/executorch/_export_utils.py new file mode 100644 index 0000000000..ebacd4bbce --- /dev/null +++ b/py/torch_tensorrt/executorch/_export_utils.py @@ -0,0 +1,312 @@ +from __future__ import annotations + +import copy +from typing import Any, Sequence + +import torch +from torch._subclasses.fake_tensor import is_fake +from torch.export.graph_signature import InputKind + + +def _is_graph_bound_metadata(value: Any) -> bool: + """True if a metadata value is tied to the exported program's shape environment. + + Fake tensors and symbolic sizes reach back to a live ``ShapeEnv`` that owns + tensors deepcopy refuses to touch, and cloning them would also detach the copy + from the symbols the graph is guarded on. + """ + for leaf in torch.utils._pytree.tree_leaves(value): + if isinstance(leaf, (torch.SymInt, torch.SymFloat, torch.SymBool)): + return True + if isinstance(leaf, torch.Tensor) and is_fake(leaf): + return True + return False + + +def get_engine_info_from_state(engine_obj: Any) -> list[Any]: + """Normalize TensorRT engine state into the serialized engine-info list.""" + state = engine_obj.__getstate__() + engine_info = state[0] if isinstance(state, tuple) else state + return list(engine_info) + + +def validate_engine_info(engine_info: Sequence[Any], *, node_name: str = "") -> None: + """Reject engine configurations unsupported by ExecuTorch export.""" + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + REQUIRES_OUTPUT_ALLOCATOR_IDX, + ) + + if ( + len(engine_info) > REQUIRES_OUTPUT_ALLOCATOR_IDX + and str(engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX]) == "1" + ): + node_suffix = f" for node '{node_name}'" if node_name else "" + raise RuntimeError( + "ExecuTorch export does not support TensorRT engines that require " + "an output allocator (data-dependent output shapes)" + f"{node_suffix}." + ) + + +def _schema_name(node: Any) -> str: + target = node.target + return target._schema.name if hasattr(target, "_schema") else "" + + +def _resolve_engine_info(exported_program: Any, node: Any) -> list[Any]: + if node.target is not torch.ops.tensorrt.execute_engine.default: + return list(node.args[1:]) + + graph_module = exported_program.graph_module + engine_node = node.args[1] + if engine_node.op == "get_attr": + engine_obj = getattr(graph_module, engine_node.target, None) + if engine_obj is None: + raise RuntimeError( + f"execute_engine node '{node.name}': get_attr target " + f"'{engine_node.target}' not found on graph module" + ) + elif engine_node.op == "placeholder": + from torch_tensorrt.dynamo._exporter import _resolve_lifted_custom_obj + + engine_obj = _resolve_lifted_custom_obj(exported_program, engine_node) + if engine_obj is None: + raise RuntimeError( + f"execute_engine node '{node.name}': placeholder engine " + f"'{engine_node.name}' did not resolve to a lifted " + f"custom-object constant (available: " + f"{sorted(getattr(exported_program, 'constants', {}) or {})})" + ) + else: + raise RuntimeError( + f"execute_engine node '{node.name}': unexpected engine arg op " + f"'{engine_node.op}'" + ) + return get_engine_info_from_state(engine_obj) + + +def validate_engine_program(exported_program: Any) -> int: + """Validate all TensorRT engine nodes before any input program is mutated.""" + count = 0 + for node in exported_program.graph_module.graph.nodes: + if node.op != "call_function": + continue + if ( + node.target is not torch.ops.tensorrt.execute_engine.default + and _schema_name(node) != "tensorrt::no_op_placeholder_for_execute_engine" + ): + continue + validate_engine_info( + _resolve_engine_info(exported_program, node), node_name=node.name + ) + count += 1 + return count + + +def _payload_sharing_memo(exported_program: Any) -> dict[int, Any]: + memo: dict[int, Any] = {} + for value in ( + *exported_program.state_dict.values(), + *exported_program.constants.values(), + ): + memo[id(value)] = value + for module in exported_program.graph_module.modules(): + for value in ( + *module.parameters(recurse=False), + *module.buffers(recurse=False), + ): + memo[id(value)] = value + for value in module.__dict__.values(): + if isinstance(value, (torch.Tensor, torch.ScriptObject)): + memo[id(value)] = value + for module in exported_program.graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + continue + for node in module.graph.nodes: + for value in node.meta.values(): + for leaf in torch.utils._pytree.tree_leaves(value): + if isinstance(leaf, (torch.Tensor, torch.ScriptObject)): + memo[id(leaf)] = leaf + return memo + + +def _stage_graph_module( + graph_module: torch.fx.GraphModule, + payload_memo: dict[int, Any], +) -> torch.fx.GraphModule: + staged = copy.deepcopy(graph_module, payload_memo) + # FX gives each staged node a fresh meta dict but shares the values inside it, + # so nested containers need their own copy to keep the source unaffected. + for name, source_module in graph_module.named_modules(): + if not isinstance(source_module, torch.fx.GraphModule): + continue + staged_module = staged if not name else staged.get_submodule(name) + if not isinstance(staged_module, torch.fx.GraphModule): + continue + staged_nodes = {node.name: node for node in staged_module.graph.nodes} + source_nodes = {node.name: node for node in source_module.graph.nodes} + if staged_nodes.keys() != source_nodes.keys(): + raise RuntimeError( + f"Staged GraphModule {name or ''!r} changed node identities." + ) + for node_name, source_node in source_nodes.items(): + staged_nodes[node_name].meta = { + key: ( + value + if _is_graph_bound_metadata(value) + else copy.deepcopy(value, payload_memo) + ) + for key, value in source_node.meta.items() + } + return staged + + +def stage_exported_program(exported_program: Any) -> Any: + """Copy program structure while sharing immutable tensor and engine payloads.""" + graph_module = _stage_graph_module( + exported_program.graph_module, + _payload_sharing_memo(exported_program), + ) + return exported_program._update( + graph_module, + copy.deepcopy(exported_program.graph_signature), + state_dict=dict(exported_program.state_dict), + constants=dict(exported_program.constants), + ) + + +def _remove_lifted_engine_placeholder( + exported_program: Any, engine_node: torch.fx.Node +) -> None: + """Remove an unused lifted engine from the graph, signature, and constants.""" + if engine_node.op != "placeholder" or engine_node.users: + raise RuntimeError( + f"Engine placeholder {engine_node.name!r} must be unused before cleanup." + ) + + matching_specs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind == InputKind.CUSTOM_OBJ + and getattr(spec.arg, "name", None) == engine_node.name + ] + if len(matching_specs) != 1: + raise RuntimeError( + f"Engine placeholder {engine_node.name!r} must have exactly one " + f"CUSTOM_OBJ input spec; found {len(matching_specs)}." + ) + + engine_spec = matching_specs[0] + engine_target = engine_spec.target + if not isinstance(engine_target, str): + raise RuntimeError( + f"Engine placeholder {engine_node.name!r} has no constant target." + ) + + exported_program._graph_signature.input_specs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec is not engine_spec + ] + engine_node.graph.erase_node(engine_node) + + if not any( + spec.target == engine_target + for spec in exported_program.graph_signature.input_specs + ): + exported_program.constants.pop(engine_target, None) + + +def replace_execute_engine(exported_program: Any) -> Any: + """Replace execute_engine nodes with ExecuTorch-safe placeholder calls.""" + from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( + ENGINE_IDX, + SERIALIZATION_LEN, + ) + + graph_module = exported_program.graph_module + execute_engine_op = torch.ops.tensorrt.execute_engine.default + no_op = torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default + nodes_to_replace = [ + node + for node in graph_module.graph.nodes + if node.op == "call_function" and node.target is execute_engine_op + ] + if not nodes_to_replace: + return exported_program + + materialized_engines: dict[torch.fx.Node, tuple[torch.fx.Node, list[str]]] = {} + for node in nodes_to_replace: + inputs_arg = node.args[0] + engine_node = node.args[1] + materialized = materialized_engines.get(engine_node) + if materialized is None: + engine_info = _resolve_engine_info(exported_program, node) + engine_bytes = engine_info[ENGINE_IDX] + if isinstance(engine_bytes, str): + import base64 + + engine_bytes = base64.b64decode(engine_bytes) + elif not isinstance(engine_bytes, (bytes, bytearray)): + engine_bytes = bytes(engine_bytes) + engine_tensor = torch.frombuffer(bytearray(engine_bytes), dtype=torch.uint8) + + from torch.fx.experimental.const_fold import get_unique_attr_name_in_module + + buffer_name = get_unique_attr_name_in_module(graph_module, "_trt_engine_0") + graph_module.register_buffer(buffer_name, engine_tensor, persistent=True) + exported_program.state_dict[buffer_name] = engine_tensor + # The engine slot is replaced by engine_attr_node below, so skip it: + # str() on multi-gigabyte engine bytes would build a throwaway string + # roughly four times their size. + str_args = [ + ("" if index == ENGINE_IDX else str(value) if value is not None else "") + for index, value in enumerate(engine_info[:SERIALIZATION_LEN]) + ] + + from torch._guards import detect_fake_mode + + fake_mode = detect_fake_mode( + [ + graph_node.meta["val"] + for graph_node in graph_module.graph.nodes + if "val" in graph_node.meta + ] + ) + fake_engine = ( + fake_mode.from_tensor(engine_tensor) + if fake_mode is not None + else engine_tensor + ) + with graph_module.graph.inserting_before(node): + engine_attr_node = graph_module.graph.get_attr(buffer_name) + engine_attr_node.meta["val"] = fake_engine + materialized = (engine_attr_node, str_args) + materialized_engines[engine_node] = materialized + engine_attr_node, str_args = materialized + + with graph_module.graph.inserting_before(node): + no_op_args = ( + inputs_arg, + *str_args[:ENGINE_IDX], + engine_attr_node, + *str_args[ENGINE_IDX + 1 :], + ) + no_op_node = graph_module.graph.call_function(no_op, no_op_args) + no_op_node.meta["val"] = node.meta.get("val") + + node.replace_all_uses_with(no_op_node) + graph_module.graph.erase_node(node) + if engine_node.op == "get_attr" and not engine_node.users: + graph_module.graph.erase_node(engine_node) + from torch.fx.graph_module import _del_attr, _has_attr + + if _has_attr(graph_module, engine_node.target): + _del_attr(graph_module, engine_node.target) + elif engine_node.op == "placeholder" and not engine_node.users: + _remove_lifted_engine_placeholder(exported_program, engine_node) + + graph_module.graph.eliminate_dead_code() + graph_module.graph.lint() + graph_module.recompile() + return exported_program diff --git a/py/torch_tensorrt/executorch/export.py b/py/torch_tensorrt/executorch/export.py new file mode 100644 index 0000000000..66e0ae5785 --- /dev/null +++ b/py/torch_tensorrt/executorch/export.py @@ -0,0 +1,350 @@ +from __future__ import annotations + +import logging +import platform +from collections.abc import Mapping, Sequence +from inspect import signature +from typing import TYPE_CHECKING, Any + +import torch +from torch.export import ExportedProgram +from torch_tensorrt._Input import Input + +if TYPE_CHECKING: + from executorch.exir import EdgeCompileConfig, EdgeProgramManager + from executorch.exir.backend.compile_spec_schema import CompileSpec + from executorch.exir.backend.partitioner import Partitioner + +logger = logging.getLogger(__name__) + + +def _contains_dynamic_input(value: Any) -> bool: + if isinstance(value, Input): + return bool(value.shape_mode == Input._ShapeMode.DYNAMIC) + if isinstance(value, Mapping): + return any(_contains_dynamic_input(item) for item in value.values()) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return any(_contains_dynamic_input(item) for item in value) + return False + + +def _all_input_specs(value: Any) -> bool: + if isinstance(value, Input): + return True + if isinstance(value, Mapping): + return all(_all_input_specs(item) for item in value.values()) + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + return all(_all_input_specs(item) for item in value) + return False + + +def _dynamic_shapes_from_input(value: Any, dim_registry: dict[str, Any]) -> Any: + from torch_tensorrt.dynamo._tracer import get_dynamic_shapes + + if isinstance(value, Input): + return get_dynamic_shapes(value, dim_registry) + if isinstance(value, Mapping): + return { + key: _dynamic_shapes_from_input(item, dim_registry) + for key, item in value.items() + } + if isinstance(value, tuple): + return tuple(_dynamic_shapes_from_input(item, dim_registry) for item in value) + if isinstance(value, list): + return [_dynamic_shapes_from_input(item, dim_registry) for item in value] + raise TypeError(f"Unsupported input structure leaf: {type(value).__name__}") + + +def _to_torch_inputs(value: Any, device: torch.device) -> Any: + from torch_tensorrt.dynamo.utils import get_torch_tensor + + if isinstance(value, Input): + return get_torch_tensor(value, device) + if isinstance(value, torch.Tensor): + return value.to(device) + if isinstance(value, Mapping): + return {key: _to_torch_inputs(item, device) for key, item in value.items()} + if isinstance(value, tuple): + return tuple(_to_torch_inputs(item, device) for item in value) + if isinstance(value, list): + return [_to_torch_inputs(item, device) for item in value] + return value + + +def _prepare_graph_module( + module: torch.fx.GraphModule, + *, + arg_inputs: Sequence[Any] | None, + kwarg_inputs: Mapping[str, Any] | None, + dynamic_shapes: Any | None, + retrace: bool, +) -> ExportedProgram: + from torch_tensorrt.dynamo._defaults import default_device + from torch_tensorrt.dynamo._exporter import export as export_graph_module + from torch_tensorrt.dynamo._tracer import build_dim_registry + from torch_tensorrt.dynamo.utils import to_torch_device + + arguments = tuple(arg_inputs or ()) + keyword_arguments = dict(kwarg_inputs or {}) + all_specs = ( + bool(arguments or keyword_arguments) + and _all_input_specs(arguments) + and _all_input_specs(keyword_arguments) + ) + if dynamic_shapes is None and all_specs: + bound_arguments = ( + signature(module.forward).bind(*arguments, **keyword_arguments).arguments + ) + dim_registry = build_dim_registry(arguments, keyword_arguments) + dynamic_shapes = { + name: _dynamic_shapes_from_input(value, dim_registry) + for name, value in bound_arguments.items() + } + elif dynamic_shapes is None and ( + _contains_dynamic_input(arguments) or _contains_dynamic_input(keyword_arguments) + ): + raise ValueError( + "Mixed Tensor and dynamic Input arguments require explicit dynamic_shapes." + ) + + device = to_torch_device(default_device()) + torch_arguments = tuple(_to_torch_inputs(arguments, device)) + torch_keyword_arguments = _to_torch_inputs(keyword_arguments, device) + if retrace and not torch_arguments and not torch_keyword_arguments: + raise ValueError("retrace=True requires example inputs.") + + return export_graph_module( + module, + arg_inputs=torch_arguments, + kwarg_inputs=torch_keyword_arguments, + dynamic_shapes=dynamic_shapes, + use_legacy_exporter=not retrace, + ) + + +def _prepare_programs( + source: ExportedProgram | torch.fx.GraphModule | Mapping[str, ExportedProgram], + *, + arg_inputs: Sequence[Any] | None, + kwarg_inputs: Mapping[str, Any] | None, + dynamic_shapes: Any | None, + retrace: bool | None, +) -> tuple[ExportedProgram | dict[str, ExportedProgram], tuple[str, ...]]: + source_options_used = ( + arg_inputs is not None + or kwarg_inputs is not None + or dynamic_shapes is not None + or retrace is True + ) + if isinstance(source, ExportedProgram): + if source_options_used: + raise ValueError( + "Inputs, dynamic_shapes, and retrace=True are invalid for an " + "already-exported program." + ) + return source, ("forward",) + + if isinstance(source, torch.fx.GraphModule): + program = _prepare_graph_module( + source, + arg_inputs=arg_inputs, + kwarg_inputs=kwarg_inputs, + dynamic_shapes=dynamic_shapes, + retrace=False if retrace is None else retrace, + ) + return program, ("forward",) + + if isinstance(source, Mapping): + if not source: + raise ValueError("Method mapping must not be empty.") + if source_options_used: + raise ValueError( + "Method mappings accept only pre-exported programs; per-method " + "input and retrace options are not supported." + ) + programs: dict[str, ExportedProgram] = {} + program_names_by_id: dict[int, str] = {} + for name, program in source.items(): + if not isinstance(name, str) or not name: + raise ValueError("Every method name must be a non-empty string.") + if not isinstance(program, ExportedProgram): + raise TypeError( + "Method mappings must contain only ExportedProgram values." + ) + previous_name = program_names_by_id.get(id(program)) + if previous_name is not None: + raise ValueError( + "Method mapping contains the same ExportedProgram object for " + f"{previous_name!r} and {name!r}. Each method requires an " + "independent program." + ) + program_names_by_id[id(program)] = name + programs[name] = program + return programs, tuple(programs) + + raise TypeError( + "source must be a TensorRT-compiled GraphModule, an engine-bearing " + "ExportedProgram, or a mapping of method names to ExportedPrograms. " + "Compile nn.Module inputs with torch_tensorrt.compile() first." + ) + + +def _per_method_values( + value: Sequence[Any] | Mapping[str, Sequence[Any] | None] | None, + method_names: tuple[str, ...], + option_name: str, +) -> dict[str, list[Any]]: + if isinstance(value, Mapping): + unknown = set(value) - set(method_names) + if unknown: + raise ValueError( + f"{option_name} contains unknown methods: {sorted(unknown)}" + ) + normalized: dict[str, list[Any]] = {} + for name in method_names: + method_value = value.get(name) + if method_value is not None and ( + not isinstance(method_value, Sequence) + or isinstance(method_value, (str, bytes)) + ): + raise TypeError(f"{option_name}[{name!r}] must be a sequence or None.") + normalized[name] = list(method_value or ()) + return normalized + if value is not None and ( + not isinstance(value, Sequence) or isinstance(value, (str, bytes)) + ): + raise TypeError( + f"{option_name} must be a sequence, or a mapping of method name to " + "sequence." + ) + shared = list(value or ()) + return {name: list(shared) for name in method_names} + + +def export( + source: ExportedProgram | torch.fx.GraphModule | Mapping[str, ExportedProgram], + *, + inputs: Sequence[Any] | None = None, + arg_inputs: Sequence[Any] | None = None, + kwarg_inputs: Mapping[str, Any] | None = None, + dynamic_shapes: Any | None = None, + retrace: bool | None = None, + transform_passes: Any | None = None, + partitioners: ( + Sequence["Partitioner"] | Mapping[str, Sequence["Partitioner"] | None] | None + ) = None, + compile_specs: ( + Sequence["CompileSpec"] | Mapping[str, Sequence["CompileSpec"] | None] | None + ) = None, + compile_config: "EdgeCompileConfig | None" = None, + constant_methods: Mapping[str, Any] | None = None, + generate_etrecord: bool = False, +) -> "EdgeProgramManager": + """Prepare TensorRT-compiled programs for composable ExecuTorch lowering. + + TensorRT claims engine nodes first. Additional partitioners run afterward in + caller-provided order. The returned EdgeProgramManager is the standard + ExecuTorch inspection and customization boundary; call ``to_executorch()`` + on it when ready to perform final memory planning and serialization. + + Export stages independent graph and metadata containers while sharing tensor + and engine payload storage, avoiding copies of potentially multi-gigabyte + TensorRT engines. Transform passes must treat shared payload values as + immutable. Method mappings preserve independent entry points but do not imply + shared mutable state between them. + """ + from torch_tensorrt._features import ENABLED_FEATURES + + if platform.system() != "Linux": + raise ValueError( + f"The executorch format is only supported on Linux, {platform.system()} " + "is not a supported platform for this format" + ) + if not ENABLED_FEATURES.torch_tensorrt_runtime: + raise RuntimeError( + "ExecuTorch export requires the Torch-TensorRT runtime " + "(torch_tensorrt_runtime). Reinstall torch_tensorrt with the runtime extension." + ) + if inputs is not None and arg_inputs is not None: + raise ValueError("inputs and arg_inputs are mutually exclusive.") + arguments = inputs if inputs is not None else arg_inputs + + import torch_tensorrt.dynamo.runtime.meta_ops.register_meta_ops # noqa: F401 + from executorch.exir import to_edge_transform_and_lower + from torch_tensorrt.executorch import TensorRTPartitioner, get_edge_compile_config + from torch_tensorrt.executorch._export_utils import ( + replace_execute_engine, + stage_exported_program, + validate_engine_program, + ) + + programs, method_names = _prepare_programs( + source, + arg_inputs=arguments, + kwarg_inputs=kwarg_inputs, + dynamic_shapes=dynamic_shapes, + retrace=retrace, + ) + program_map = {"forward": programs} if not isinstance(programs, dict) else programs + extra_partitioners = _per_method_values(partitioners, method_names, "partitioners") + method_compile_specs = _per_method_values( + compile_specs, method_names, "compile_specs" + ) + + if constant_methods is not None: + collisions = set(constant_methods) & set(method_names) + if collisions: + raise ValueError( + f"constant_methods collide with executable methods: {sorted(collisions)}" + ) + + if isinstance(transform_passes, Mapping): + unknown = set(transform_passes) - set(method_names) + if unknown: + raise ValueError( + f"transform_passes contains unknown methods: {sorted(unknown)}" + ) + + engine_counts = { + name: validate_engine_program(program) for name, program in program_map.items() + } + # Zero-engine methods are allowed: later partitioners may claim their ops, + # or portable operators may remain undelegated. + staged_programs = { + name: stage_exported_program(program) for name, program in program_map.items() + } + rewritten: dict[str, ExportedProgram] = {} + method_partitioners: dict[str, list[Partitioner]] = {} + for name, program in staged_programs.items(): + if engine_counts[name] > 1: + logger.warning( + "%s contains %d TRT engines. Multi-engine .pte exports can incur " + "extra delegate boundary overhead.", + name, + engine_counts[name], + ) + rewritten[name] = replace_execute_engine(program) + method_partitioners[name] = [ + TensorRTPartitioner(compile_specs=method_compile_specs[name]), + *extra_partitioners[name], + ] + + edge_programs: ExportedProgram | dict[str, ExportedProgram] + partitioner_pipeline: list[Partitioner] | dict[str, list[Partitioner]] + if isinstance(programs, dict): + edge_programs = rewritten + partitioner_pipeline = method_partitioners + else: + edge_programs = rewritten["forward"] + partitioner_pipeline = method_partitioners["forward"] + + return to_edge_transform_and_lower( + edge_programs, + transform_passes=transform_passes, + partitioner=partitioner_pipeline, + constant_methods=dict(constant_methods) if constant_methods else None, + compile_config=( + compile_config if compile_config is not None else get_edge_compile_config() + ), + generate_etrecord=generate_etrecord, + ) diff --git a/tests/py/dynamo/executorch/test_api.py b/tests/py/dynamo/executorch/test_api.py index 5d88822716..3763f7951e 100644 --- a/tests/py/dynamo/executorch/test_api.py +++ b/tests/py/dynamo/executorch/test_api.py @@ -13,7 +13,12 @@ @pytest.mark.unit def test_lazy_import_error_when_executorch_missing(monkeypatch): + import torch_tensorrt + original_module = sys.modules.pop("torch_tensorrt.executorch", None) + original_attribute = getattr(torch_tensorrt, "executorch", None) + if hasattr(torch_tensorrt, "executorch"): + delattr(torch_tensorrt, "executorch") original_find_spec = importlib.util.find_spec def fake_find_spec(name, package=None): @@ -30,6 +35,10 @@ def fake_find_spec(name, package=None): sys.modules.pop("torch_tensorrt.executorch", None) if original_module is not None: sys.modules["torch_tensorrt.executorch"] = original_module + if original_attribute is not None: + torch_tensorrt.executorch = original_attribute + elif hasattr(torch_tensorrt, "executorch"): + delattr(torch_tensorrt, "executorch") @pytest.mark.unit @@ -59,6 +68,8 @@ def test_public_api_symbols_present(): assert "get_edge_compile_config" in module.__all__ assert "TensorRTPartitioner" in module.__all__ assert "TensorRTBackend" in module.__all__ + assert "export" in module.__all__ + assert "to_executorch" not in module.__all__ _REPO_ROOT = Path(__file__).resolve().parents[4] diff --git a/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py b/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py index 64233f45cd..455bbc16b2 100644 --- a/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py +++ b/tests/py/dynamo/executorch/test_cuda_partitioner_composition.py @@ -73,11 +73,15 @@ def _cuda_partitioner(): 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.""" +def _deserialize_program(pte_path): from executorch.exir._serialize._program import deserialize_pte_binary - program = deserialize_pte_binary(pte_path.read_bytes()).program + return deserialize_pte_binary(pte_path.read_bytes()).program + + +def _delegate_ids(pte_path): + """Backend ids of every delegate in the serialized program, in order.""" + program = _deserialize_program(pte_path) return [ delegate.id for plan in program.execution_plan for delegate in plan.delegates ] @@ -176,7 +180,11 @@ def forward(self, x): partitioners=[_cuda_partitioner()], ) + assert out.exists() delegate_ids = _delegate_ids(out) + assert ( + "TensorRTBackend" in delegate_ids + ), f"tanh/cos were not delegated to TensorRT; delegates={delegate_ids}" assert ( "CudaBackend" in delegate_ids ), f"the pinned mm did not route to the CUDA backend; delegates={delegate_ids}" @@ -189,6 +197,103 @@ def forward(self, x): ), f"external .ptd data file is empty: {ptd_files}" +def test_multimethod_trt_export_preserves_methods(tmp_path): + """Independent TRT programs remain separate delegated methods after lowering.""" + import torch_tensorrt + + class Prefill(torch.nn.Module): + def forward(self, x): + return torch.cos(x) + + class Decode(torch.nn.Module): + def forward(self, x): + return torch.tanh(x) + + inputs = (torch.randn(16, 16, device="cuda"),) + programs = {} + source_snapshots = {} + for name, model in ( + ("prefill", Prefill()), + ("decode", Decode()), + ): + exported = torch.export.export(model.eval().to("cuda"), inputs) + trt_gm = torch_tensorrt.dynamo.compile( + exported, + inputs=list(inputs), + min_block_size=1, + truncate_double=True, + ) + program = torch_tensorrt.dynamo._exporter.export( + trt_gm, + arg_inputs=inputs, + use_legacy_exporter=False, + ) + programs[name] = program + source_snapshots[name] = ( + program.graph_module.code, + tuple(program.state_dict), + tuple(program.constants), + tuple(node.target for node in program.graph.nodes), + ) + + edge = torch_tensorrt.executorch.export(programs) + out = tmp_path / "multimethod_trt.pte" + with out.open("wb") as output: + edge.to_executorch().write_to_file(output) + + execution_plans = _deserialize_program(out).execution_plan + assert {plan.name for plan in execution_plans} == {"prefill", "decode"} + assert all( + any(delegate.id == "TensorRTBackend" for delegate in plan.delegates) + for plan in execution_plans + ) + for name, program in programs.items(): + assert source_snapshots[name] == ( + program.graph_module.code, + tuple(program.state_dict), + tuple(program.constants), + tuple(node.target for node in program.graph.nodes), + ) + assert not any(key.startswith("_trt_engine_") for key in program.state_dict) + + +def test_graph_module_export_preserves_source(tmp_path): + import torch_tensorrt + + class Model(torch.nn.Module): + def forward(self, x): + return torch.cos(torch.tanh(x)) + + inputs = (torch.randn(16, 16, device="cuda"),) + exported = torch.export.export(Model().eval().to("cuda"), inputs) + trt_gm = torch_tensorrt.dynamo.compile( + exported, + inputs=list(inputs), + min_block_size=1, + truncate_double=True, + ) + source_code = trt_gm.code + source_state = tuple(trt_gm.state_dict()) + source_attrs = tuple(name for name, _ in trt_gm.named_modules()) + + edge = torch_tensorrt.executorch.export( + trt_gm, + arg_inputs=inputs, + retrace=False, + ) + out = tmp_path / "graph_module_trt.pte" + with out.open("wb") as output: + edge.to_executorch().write_to_file(output) + + assert "TensorRTBackend" in _delegate_ids(out) + assert trt_gm.code == source_code + assert tuple(trt_gm.state_dict()) == source_state + assert tuple(name for name, _ in trt_gm.named_modules()) == source_attrs + assert not any( + name.startswith("_trt_engine_") for name, _ in trt_gm.named_buffers() + ) + + 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.""" diff --git a/tests/py/dynamo/executorch/test_edge_cases.py b/tests/py/dynamo/executorch/test_edge_cases.py index 935b67604c..74b4a42c7c 100644 --- a/tests/py/dynamo/executorch/test_edge_cases.py +++ b/tests/py/dynamo/executorch/test_edge_cases.py @@ -3,21 +3,15 @@ from unittest.mock import MagicMock import pytest -import torch from torch_tensorrt._compile import ( - _count_executorch_engine_nodes, - _validate_executorch_engine_info, + _save_as_executorch, _write_external_tensor_data, ) from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( REQUIRES_OUTPUT_ALLOCATOR_IDX, SERIALIZATION_LEN, ) - - -class _SchemaTarget: - def __init__(self, name): - self._schema = SimpleNamespace(name=name) +from torch_tensorrt.executorch._export_utils import validate_engine_info @pytest.mark.unit @@ -26,36 +20,64 @@ def test_validate_executorch_engine_info_rejects_output_allocator(): engine_info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = "1" with pytest.raises(RuntimeError, match="output allocator"): - _validate_executorch_engine_info(engine_info, node_name="trt") + validate_engine_info(engine_info, node_name="trt") @pytest.mark.unit -def test_count_executorch_engine_nodes_handles_execute_and_placeholder(): - execute_node = SimpleNamespace( - op="call_function", - target=torch.ops.tensorrt.execute_engine.default, - ) - placeholder_node = SimpleNamespace( - op="call_function", - target=_SchemaTarget("tensorrt::no_op_placeholder_for_execute_engine"), +def test_save_as_executorch_uses_public_lowering_and_persists_data( + monkeypatch, tmp_path +): + import torch_tensorrt.executorch as executorch_api + + program = SimpleNamespace( + _tensor_data={"forward": b"weights"}, + write_to_file=MagicMock(), + write_tensor_data_to_file=MagicMock(), ) - other_node = SimpleNamespace( - op="call_function", - target=_SchemaTarget("aten::add"), + edge = SimpleNamespace(to_executorch=MagicMock(return_value=program)) + export = MagicMock(return_value=edge) + monkeypatch.setattr(executorch_api, "export", export) + + pte = tmp_path / "model.pte" + source = object() + partitioners = [object()] + compile_specs = [object()] + backend_config = object() + _save_as_executorch( + source, + str(pte), + partitioners=partitioners, + compile_specs=compile_specs, + backend_config=backend_config, ) - exp_program = SimpleNamespace( - graph_module=SimpleNamespace( - graph=SimpleNamespace(nodes=[execute_node, placeholder_node, other_node]) - ) + + export.assert_called_once_with( + source, + partitioners=partitioners, + compile_specs=compile_specs, ) + edge.to_executorch.assert_called_once_with(config=backend_config) + program.write_to_file.assert_called_once() + program.write_tensor_data_to_file.assert_called_once_with(str(tmp_path)) - assert _count_executorch_engine_nodes(exp_program) == 2 + +@pytest.mark.unit +@pytest.mark.parametrize("option", ["partitioners", "compile_specs"]) +def test_save_as_executorch_rejects_per_method_mapping(monkeypatch, tmp_path, option): + import torch_tensorrt.executorch as executorch_api + + export = MagicMock() + monkeypatch.setattr(executorch_api, "export", export) + + with pytest.raises(TypeError, match="must be a list or tuple"): + _save_as_executorch( + object(), str(tmp_path / "model.pte"), **{option: {"forward": []}} + ) + export.assert_not_called() @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(), @@ -69,7 +91,6 @@ def test_write_external_tensor_data_writes_when_present(tmp_path): @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(), @@ -80,9 +101,6 @@ def test_write_external_tensor_data_noop_when_empty(tmp_path): @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")) diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py new file mode 100644 index 0000000000..d1eae88b27 --- /dev/null +++ b/tests/py/dynamo/executorch/test_export.py @@ -0,0 +1,881 @@ +import importlib +from types import SimpleNamespace +from typing import NamedTuple +from unittest.mock import MagicMock + +import pytest + +pytest.importorskip("executorch.exir") + +import torch # noqa: E402 +import torch_tensorrt # noqa: E402 +from torch.export.graph_signature import ( # noqa: E402 + CustomObjArgument, + ExportGraphSignature, + InputKind, + InputSpec, + TensorArgument, +) +from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ( # noqa: E402 + ENGINE_IDX, + REQUIRES_OUTPUT_ALLOCATOR_IDX, + SERIALIZATION_LEN, +) + + +class FakeExportedProgram: + pass + + +class FakeTensorRTPartitioner: + def __init__(self, compile_specs): + self.compile_specs = compile_specs + + +class _EngineState: + def __init__(self, payload=b"engine-bytes", requires_output_allocator=False): + self.info = [""] * SERIALIZATION_LEN + self.info[ENGINE_IDX] = payload + self.info[REQUIRES_OUTPUT_ALLOCATOR_IDX] = str(int(requires_output_allocator)) + + def __getstate__(self): + return (self.info,) + + +def _engine_program(*, lifted, execute_count=1, requires_output_allocator=False): + graph = torch.fx.Graph() + engine = _EngineState(requires_output_allocator=requires_output_allocator) + if lifted: + engine_node = graph.placeholder("obj_engine") + root = torch.nn.Module() + engine_target = "engine_fqn" + input_specs = [ + InputSpec( + kind=InputKind.CUSTOM_OBJ, + arg=CustomObjArgument(name=engine_node.name, class_fqn=""), + target=engine_target, + ) + ] + constants = {engine_target: engine} + else: + root = torch.nn.Module() + root.engine = engine + engine_node = graph.get_attr("engine") + input_specs = [] + constants = {} + + input_node = graph.placeholder("x") + input_specs.append( + InputSpec( + kind=InputKind.USER_INPUT, + arg=TensorArgument(name=input_node.name), + target=None, + ) + ) + results = [ + graph.call_function( + torch.ops.tensorrt.execute_engine.default, + ([input_node], engine_node), + ) + for _ in range(execute_count) + ] + graph.output(results[0] if len(results) == 1 else tuple(results)) + graph_module = torch.fx.GraphModule(root, graph) + signature = ExportGraphSignature(input_specs=input_specs, output_specs=[]) + program = SimpleNamespace( + graph_module=graph_module, + graph_signature=signature, + _graph_signature=signature, + state_dict={}, + constants=constants, + ) + return program, engine_node, input_node, engine + + +def _patch_lowering(monkeypatch, engine_counts=None): + import executorch.exir + import torch_tensorrt._features as features + import torch_tensorrt.executorch as executorch_api + import torch_tensorrt.executorch._export_utils as export_utils + + monkeypatch.setattr( + features, + "ENABLED_FEATURES", + features.ENABLED_FEATURES._replace(torch_tensorrt_runtime=True), + ) + export_module = importlib.import_module("torch_tensorrt.executorch.export") + engine_counts = engine_counts or {} + lower = MagicMock(return_value=object()) + monkeypatch.setattr(executorch.exir, "to_edge_transform_and_lower", lower) + monkeypatch.setattr(executorch_api, "TensorRTPartitioner", FakeTensorRTPartitioner) + monkeypatch.setattr(executorch_api, "get_edge_compile_config", lambda: "default") + monkeypatch.setattr(export_module, "ExportedProgram", FakeExportedProgram) + monkeypatch.setattr( + export_utils, + "validate_engine_program", + lambda program: engine_counts.get(program, 1), + ) + monkeypatch.setattr(export_utils, "stage_exported_program", lambda program: program) + monkeypatch.setattr( + export_utils, + "replace_execute_engine", + lambda program: ("rewritten", program), + ) + return export_module, lower + + +@pytest.mark.unit +@pytest.mark.skipif( + not torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime, + reason="Torch-TensorRT runtime operators are not available", +) +@pytest.mark.parametrize("lifted", [False, True]) +def test_validate_and_replace_execute_engine(lifted): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + program, engine_node, input_node, _ = _engine_program(lifted=lifted) + + assert export_utils.validate_engine_program(program) == 1 + rewritten = export_utils.replace_execute_engine(program) + + assert rewritten is program + nodes = list(program.graph_module.graph.nodes) + assert not any( + node.target is torch.ops.tensorrt.execute_engine.default for node in nodes + ) + no_op_nodes = [ + node + for node in nodes + if node.target + is torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default + ] + assert len(no_op_nodes) == 1 + buffer_node = no_op_nodes[0].args[1 + ENGINE_IDX] + engine_buffer = getattr(program.graph_module, buffer_node.target) + assert engine_buffer.dtype == torch.uint8 + assert engine_buffer.device.type == "cpu" + assert bytes(engine_buffer.tolist()) == b"engine-bytes" + assert program.state_dict[buffer_node.target] is engine_buffer + assert engine_node not in nodes + assert input_node in nodes + + if lifted: + assert program.graph_signature.inputs_to_lifted_custom_objs == {} + assert program.constants == {} + assert [spec.arg.name for spec in program.graph_signature.input_specs] == [ + input_node.name + ] + else: + assert not hasattr(program.graph_module, "engine") + + +@pytest.mark.unit +@pytest.mark.skipif( + not torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime, + reason="Torch-TensorRT runtime operators are not available", +) +@pytest.mark.parametrize("lifted", [False, True]) +def test_replace_execute_engine_cleans_shared_engine_after_last_use(lifted): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + program, engine_node, input_node, _ = _engine_program( + lifted=lifted, execute_count=2 + ) + + assert export_utils.validate_engine_program(program) == 2 + export_utils.replace_execute_engine(program) + + nodes = list(program.graph_module.graph.nodes) + no_op_nodes = [ + node + for node in nodes + if node.target + is torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default + ] + assert len(no_op_nodes) == 2 + buffer_nodes = {node.args[1 + ENGINE_IDX] for node in no_op_nodes} + assert len(buffer_nodes) == 1 + assert set(program.state_dict) == {next(iter(buffer_nodes)).target} + assert engine_node not in nodes + assert input_node in nodes + if lifted: + assert program.graph_signature.inputs_to_lifted_custom_objs == {} + assert program.constants == {} + else: + assert not hasattr(program.graph_module, "engine") + + +@pytest.mark.unit +@pytest.mark.skipif( + not torch_tensorrt.ENABLED_FEATURES.torch_tensorrt_runtime, + reason="Torch-TensorRT runtime operators are not available", +) +def test_validate_engine_program_rejects_output_allocator_without_mutation(): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + program, engine_node, _, engine = _engine_program( + lifted=True, requires_output_allocator=True + ) + original_nodes = list(program.graph_module.graph.nodes) + original_specs = list(program.graph_signature.input_specs) + + with pytest.raises(RuntimeError, match="output allocator"): + export_utils.validate_engine_program(program) + + assert list(program.graph_module.graph.nodes) == original_nodes + assert program.graph_signature.input_specs == original_specs + assert program.constants == {"engine_fqn": engine} + assert engine_node.users + assert program.state_dict == {} + + +@pytest.mark.unit +def test_export_returns_edge_and_forwards_all_options(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + program = FakeExportedProgram() + extra_a = object() + extra_b = object() + compile_spec = object() + transform_pass = object() + edge_config = object() + constant_methods = {"get_vocab_size": 256} + partitioners = [extra_a, extra_b] + compile_specs = [compile_spec] + + result = export_module.export( + program, + transform_passes=[transform_pass], + partitioners=partitioners, + compile_specs=compile_specs, + compile_config=edge_config, + constant_methods=constant_methods, + generate_etrecord=True, + ) + + assert result is lower.return_value + assert lower.call_args.args == (("rewritten", program),) + assert lower.call_args.kwargs["transform_passes"] == [transform_pass] + assert lower.call_args.kwargs["compile_config"] is edge_config + assert lower.call_args.kwargs["constant_methods"] == constant_methods + assert lower.call_args.kwargs["generate_etrecord"] is True + lowered_partitioners = lower.call_args.kwargs["partitioner"] + assert isinstance(lowered_partitioners[0], FakeTensorRTPartitioner) + assert lowered_partitioners[0].compile_specs == [compile_spec] + assert lowered_partitioners[1:] == [extra_a, extra_b] + assert partitioners == [extra_a, extra_b] + assert compile_specs == [compile_spec] + + +@pytest.mark.unit +def test_export_preserves_independent_method_mapping(monkeypatch): + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + export_module, lower = _patch_lowering(monkeypatch) + prefill_extra = object() + decode_extra = object() + prefill_spec = object() + decode_spec = object() + + export_module.export( + {"prefill": prefill, "decode": decode}, + partitioners={"prefill": [prefill_extra], "decode": [decode_extra]}, + compile_specs={"prefill": [prefill_spec], "decode": [decode_spec]}, + ) + + assert lower.call_args.args == ( + {"prefill": ("rewritten", prefill), "decode": ("rewritten", decode)}, + ) + pipelines = lower.call_args.kwargs["partitioner"] + assert pipelines["prefill"][0].compile_specs == [prefill_spec] + assert pipelines["prefill"][1:] == [prefill_extra] + assert pipelines["decode"][0].compile_specs == [decode_spec] + assert pipelines["decode"][1:] == [decode_extra] + + +@pytest.mark.unit +def test_stage_exported_program_isolates_structure_and_shares_payloads(): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2)) + + def forward(self, x): + return x + self.weight + + program = torch.export.export(Model(), (torch.ones(2),)) + + class Metadata(NamedTuple): + shape: tuple[int, ...] + dtype: torch.dtype + + first_node = next(iter(program.graph.nodes)) + first_node.meta["nested"] = {"items": ["source"]} + first_node.meta["tensor_meta"] = Metadata((2,), torch.float32) + program.graph_module.meta["nested"] = {"items": ["source"]} + staged = export_utils.stage_exported_program(program) + + assert staged is not program + assert staged.graph_module is not program.graph_module + assert staged.graph is not program.graph + assert staged.state_dict is not program.state_dict + assert staged.constants is not program.constants + assert staged.state_dict["weight"] is program.state_dict["weight"] + assert isinstance(next(iter(staged.graph.nodes)).meta["tensor_meta"], Metadata) + + staged.state_dict["staged_only"] = torch.zeros(1) + staged.graph_module.meta["staged_only"] = True + staged.graph_module.meta["nested"]["items"].append("staged") + next(iter(staged.graph.nodes)).meta["nested"]["items"].append("staged") + assert "staged_only" not in program.state_dict + assert "staged_only" not in program.graph_module.meta + assert program.graph_module.meta["nested"]["items"] == ["source"] + assert first_node.meta["nested"]["items"] == ["source"] + + +@pytest.mark.unit +def test_stage_exported_program_supports_dynamic_shapes(): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.lin = torch.nn.Linear(4, 4) + + def forward(self, x): + return self.lin(x) * x.shape[0] + + batch = torch.export.Dim("batch", min=1, max=32) + program = torch.export.export( + Model(), (torch.randn(2, 4),), dynamic_shapes={"x": {0: batch}} + ) + assert program.range_constraints + + staged = export_utils.stage_exported_program(program) + + # Symbolic metadata stays bound to the original ShapeEnv, so it must be shared + # rather than copied for the staged program to remain consistent. + source_nodes = {node.name: node for node in program.graph.nodes} + symbolic_shared = 0 + for node in staged.graph.nodes: + source_node = source_nodes[node.name] + for key, source_value in source_node.meta.items(): + if not export_utils._is_graph_bound_metadata(source_value): + continue + assert node.meta[key] is source_value + symbolic_shared += 1 + assert symbolic_shared + + assert staged.graph_module is not program.graph_module + assert staged.state_dict["lin.weight"] is program.state_dict["lin.weight"] + + +@pytest.mark.unit +def test_stage_exported_program_clones_nested_graph_modules(): + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + + class Leaf(torch.nn.Module): + def forward(self, x): + return x + 1 + + root = torch.fx.Graph() + x = root.placeholder("x") + call = root.call_module("nested", (x,)) + root.output(call) + nested = torch.fx.symbolic_trace(Leaf()) + graph_module = torch.fx.GraphModule({"nested": nested}, root) + nested_node = next(iter(graph_module.nested.graph.nodes)) + nested_payload = torch.ones(1) + nested_node.meta["nested"] = { + "items": ["source"], + "payload": nested_payload, + } + payload_memo = export_utils._payload_sharing_memo( + SimpleNamespace(graph_module=graph_module, state_dict={}, constants={}) + ) + staged = export_utils._stage_graph_module(graph_module, payload_memo) + + assert staged.nested is not graph_module.nested + staged.nested.meta["staged_only"] = True + next(iter(staged.nested.graph.nodes)).meta["nested"]["items"].append("staged") + staged_nested_meta = next(iter(staged.nested.graph.nodes)).meta["nested"] + assert staged_nested_meta["payload"] is nested_payload + assert "staged_only" not in graph_module.nested.meta + assert nested_node.meta["nested"]["items"] == ["source"] + + +@pytest.mark.unit +def test_export_prepares_compiled_graph_module(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + prepared = FakeExportedProgram() + prepare = MagicMock(return_value=prepared) + monkeypatch.setattr(export_module, "_prepare_graph_module", prepare) + inputs = [torch.ones(1)] + + export_module.export(graph_module, arg_inputs=inputs, retrace=False) + + prepare.assert_called_once_with( + graph_module, + arg_inputs=inputs, + kwarg_inputs=None, + dynamic_shapes=None, + retrace=False, + ) + assert lower.call_args.args == (("rewritten", prepared),) + + +@pytest.mark.unit +@pytest.mark.parametrize("input_option", ["inputs", "arg_inputs"]) +def test_export_forwards_input_alias(monkeypatch, input_option): + export_module, _ = _patch_lowering(monkeypatch) + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + prepared = FakeExportedProgram() + prepare = MagicMock(return_value=prepared) + monkeypatch.setattr(export_module, "_prepare_graph_module", prepare) + example_inputs = [torch.ones(1)] + + export_module.export(graph_module, **{input_option: example_inputs}, retrace=False) + + prepare.assert_called_once_with( + graph_module, + arg_inputs=example_inputs, + kwarg_inputs=None, + dynamic_shapes=None, + retrace=False, + ) + + +@pytest.mark.unit +def test_export_rejects_inputs_and_arg_inputs(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + + with pytest.raises(ValueError, match="mutually exclusive"): + export_module.export( + graph_module, + inputs=[torch.ones(1)], + arg_inputs=[torch.ones(1)], + ) + + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_rejects_non_linux_platform(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + monkeypatch.setattr(export_module.platform, "system", lambda: "Windows") + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + + with pytest.raises(ValueError, match="only supported on Linux"): + export_module.export(graph_module) + + lower.assert_not_called() + + +@pytest.mark.unit +def test_prepare_graph_module_preserves_tensor_keyword_inputs(monkeypatch): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + exported = FakeExportedProgram() + export_graph_module = MagicMock(return_value=exported) + keyword_tensor = torch.ones(2, device="cuda") + monkeypatch.setattr("torch_tensorrt.dynamo._exporter.export", export_graph_module) + monkeypatch.setattr( + "torch_tensorrt.dynamo._defaults.default_device", + lambda: torch.device("cuda"), + ) + + assert ( + export_module._prepare_graph_module( + graph_module, + arg_inputs=(), + kwarg_inputs={"mask": keyword_tensor}, + dynamic_shapes=None, + retrace=True, + ) + is exported + ) + assert ( + export_graph_module.call_args.kwargs["kwarg_inputs"]["mask"] is keyword_tensor + ) + + +@pytest.mark.unit +def test_export_allows_zero_engine_program(monkeypatch): + program = FakeExportedProgram() + export_module, lower = _patch_lowering(monkeypatch, {program: 0}) + + assert export_module.export(program) is lower.return_value + lowered_partitioners = lower.call_args.kwargs["partitioner"] + assert len(lowered_partitioners) == 1 + assert isinstance(lowered_partitioners[0], FakeTensorRTPartitioner) + + +@pytest.mark.unit +def test_export_rejects_duplicate_program_identity_before_validation(monkeypatch): + program = FakeExportedProgram() + export_module, lower = _patch_lowering(monkeypatch) + validate = MagicMock(return_value=1) + monkeypatch.setattr( + "torch_tensorrt.executorch._export_utils.validate_engine_program", validate + ) + + with pytest.raises(ValueError, match="same ExportedProgram object"): + export_module.export({"prefill": program, "decode": program}) + + validate.assert_not_called() + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_validates_all_methods_before_rewriting(monkeypatch): + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + export_module, lower = _patch_lowering(monkeypatch) + validate = MagicMock(side_effect=[1, RuntimeError("decode is invalid")]) + replace = MagicMock() + monkeypatch.setattr( + "torch_tensorrt.executorch._export_utils.validate_engine_program", validate + ) + monkeypatch.setattr( + "torch_tensorrt.executorch._export_utils.replace_execute_engine", replace + ) + + with pytest.raises(RuntimeError, match="decode is invalid"): + export_module.export({"prefill": prefill, "decode": decode}) + + replace.assert_not_called() + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_rewrite_failure_leaves_sources_unchanged(monkeypatch): + import executorch.exir + import torch_tensorrt.executorch as executorch_api + import torch_tensorrt.executorch._export_utils as export_utils + + class Model(torch.nn.Module): + def forward(self, x): + return x + 1 + + prefill = torch.export.export(Model(), (torch.ones(1),)) + decode = torch.export.export(Model(), (torch.ones(1),)) + original_codes = { + "prefill": prefill.graph_module.code, + "decode": decode.graph_module.code, + } + original_state_keys = { + "prefill": set(prefill.state_dict), + "decode": set(decode.state_dict), + } + + monkeypatch.setattr(export_utils, "validate_engine_program", lambda program: 1) + monkeypatch.setattr(executorch_api, "TensorRTPartitioner", FakeTensorRTPartitioner) + monkeypatch.setattr(executorch_api, "get_edge_compile_config", lambda: "default") + lower = MagicMock() + monkeypatch.setattr(executorch.exir, "to_edge_transform_and_lower", lower) + call_count = 0 + + def fail_second_rewrite(program): + nonlocal call_count + call_count += 1 + program.state_dict["staged_only"] = torch.zeros(1) + program.graph_module.meta["staged_only"] = True + if call_count == 2: + raise RuntimeError("rewrite failed") + return program + + monkeypatch.setattr(export_utils, "replace_execute_engine", fail_second_rewrite) + + with pytest.raises(RuntimeError, match="rewrite failed"): + executorch_api.export({"prefill": prefill, "decode": decode}) + + assert prefill.graph_module.code == original_codes["prefill"] + assert decode.graph_module.code == original_codes["decode"] + assert set(prefill.state_dict) == original_state_keys["prefill"] + assert set(decode.state_dict) == original_state_keys["decode"] + assert "staged_only" not in prefill.graph_module.meta + assert "staged_only" not in decode.graph_module.meta + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_lowering_failure_leaves_source_unchanged(monkeypatch): + import executorch.exir + import torch_tensorrt.executorch as executorch_api + import torch_tensorrt.executorch._export_utils as export_utils + + class Model(torch.nn.Module): + def forward(self, x): + return x + 1 + + source = torch.export.export(Model(), (torch.ones(1),)) + original_code = source.graph_module.code + original_state_keys = set(source.state_dict) + + monkeypatch.setattr(export_utils, "validate_engine_program", lambda program: 1) + monkeypatch.setattr(executorch_api, "TensorRTPartitioner", FakeTensorRTPartitioner) + monkeypatch.setattr(executorch_api, "get_edge_compile_config", lambda: "default") + + def mutate_staged_program(program): + program.state_dict["staged_only"] = torch.zeros(1) + program.graph_module.meta["staged_only"] = True + return program + + monkeypatch.setattr(export_utils, "replace_execute_engine", mutate_staged_program) + monkeypatch.setattr( + executorch.exir, + "to_edge_transform_and_lower", + MagicMock(side_effect=RuntimeError("lowering failed")), + ) + + with pytest.raises(RuntimeError, match="lowering failed"): + executorch_api.export(source) + + assert source.graph_module.code == original_code + assert set(source.state_dict) == original_state_keys + assert "staged_only" not in source.graph_module.meta + + +@pytest.mark.unit +def test_export_accepts_false_retrace_for_exported_program(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + program = FakeExportedProgram() + + export_module.export(program, retrace=False) + + assert lower.call_args.args == (("rewritten", program),) + + +@pytest.mark.unit +def test_export_rejects_true_retrace_for_exported_program(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + program = FakeExportedProgram() + + with pytest.raises(ValueError, match="already-exported program"): + export_module.export(program, retrace=True) + + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_normalizes_none_per_method_values(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + prefill = FakeExportedProgram() + decode = FakeExportedProgram() + + export_module.export( + {"prefill": prefill, "decode": decode}, + partitioners={"prefill": None}, + compile_specs={"decode": None}, + ) + + pipelines = lower.call_args.kwargs["partitioner"] + assert len(pipelines["prefill"]) == 1 + assert len(pipelines["decode"]) == 1 + assert pipelines["prefill"][0].compile_specs == [] + assert pipelines["decode"][0].compile_specs == [] + + +@pytest.mark.unit +@pytest.mark.parametrize("container_type", [list, tuple]) +def test_prepare_graph_module_infers_nested_dynamic_shapes(monkeypatch, container_type): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + + class NestedModule(torch.nn.Module): + def forward(self, nested): + return nested[0] + + graph_module = torch.fx.symbolic_trace(NestedModule()) + exported = FakeExportedProgram() + export_graph_module = MagicMock(return_value=exported) + dynamic_input = torch_tensorrt.Input( + min_shape=(1, 2), + opt_shape=(2, 2), + max_shape=(4, 2), + name="nested", + ) + monkeypatch.setattr("torch_tensorrt.dynamo._exporter.export", export_graph_module) + + assert ( + export_module._prepare_graph_module( + graph_module, + arg_inputs=(container_type([dynamic_input]),), + kwarg_inputs={}, + dynamic_shapes=None, + retrace=False, + ) + is exported + ) + nested_shapes = export_graph_module.call_args.kwargs["dynamic_shapes"]["nested"] + assert isinstance(nested_shapes, container_type) + assert len(nested_shapes) == 1 + assert set(nested_shapes[0]) == {0} + + +@pytest.mark.unit +def test_prepare_graph_module_preserves_shared_dynamic_dimensions(monkeypatch): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + + class SharedBatchModule(torch.nn.Module): + def forward(self, left, right): + return left + right + + graph_module = torch.fx.symbolic_trace(SharedBatchModule()) + exported = FakeExportedProgram() + export_graph_module = MagicMock(return_value=exported) + monkeypatch.setattr("torch_tensorrt.dynamo._exporter.export", export_graph_module) + inputs = tuple( + torch_tensorrt.Input( + min_shape=(1, 2), + opt_shape=(2, 2), + max_shape=(4, 2), + name=name, + shared_dims={0: "batch"}, + ) + for name in ("left", "right") + ) + + assert ( + export_module._prepare_graph_module( + graph_module, + arg_inputs=inputs, + kwarg_inputs={}, + dynamic_shapes=None, + retrace=False, + ) + is exported + ) + dynamic_shapes = export_graph_module.call_args.kwargs["dynamic_shapes"] + assert dynamic_shapes["left"][0] is dynamic_shapes["right"][0] + + +@pytest.mark.unit +def test_prepare_graph_module_requires_shapes_for_mixed_dynamic_inputs(monkeypatch): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + + class MixedModule(torch.nn.Module): + def forward(self, dynamic, static): + return dynamic + static + + graph_module = torch.fx.symbolic_trace(MixedModule()) + dynamic_input = torch_tensorrt.Input( + min_shape=(1,), opt_shape=(2,), max_shape=(4,), name="dynamic" + ) + static_input = torch.ones(2) + + with pytest.raises(ValueError, match="require explicit dynamic_shapes"): + export_module._prepare_graph_module( + graph_module, + arg_inputs=(dynamic_input, static_input), + kwarg_inputs={}, + dynamic_shapes=None, + retrace=False, + ) + + exported = FakeExportedProgram() + export_graph_module = MagicMock(return_value=exported) + monkeypatch.setattr("torch_tensorrt.dynamo._exporter.export", export_graph_module) + explicit_shapes = ({0: torch.export.Dim("batch", min=1, max=4)}, None) + assert ( + export_module._prepare_graph_module( + graph_module, + arg_inputs=(dynamic_input, static_input), + kwarg_inputs={}, + dynamic_shapes=explicit_shapes, + retrace=False, + ) + is exported + ) + assert export_graph_module.call_args.kwargs["dynamic_shapes"] is explicit_shapes + + +@pytest.mark.unit +def test_prepare_graph_module_does_not_infer_shapes_without_inputs(monkeypatch): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + graph_module = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + exported = FakeExportedProgram() + export_graph_module = MagicMock(return_value=exported) + infer_args = MagicMock() + infer_kwargs = MagicMock() + monkeypatch.setattr("torch_tensorrt.dynamo._exporter.export", export_graph_module) + monkeypatch.setattr( + "torch_tensorrt.dynamo._tracer.get_dynamic_shapes_args", infer_args + ) + monkeypatch.setattr( + "torch_tensorrt.dynamo._tracer.get_dynamic_shapes_kwargs", infer_kwargs + ) + + assert ( + export_module._prepare_graph_module( + graph_module, + arg_inputs=(), + kwarg_inputs={}, + dynamic_shapes=None, + retrace=False, + ) + is exported + ) + infer_args.assert_not_called() + infer_kwargs.assert_not_called() + assert export_graph_module.call_args.kwargs["dynamic_shapes"] is None + + +@pytest.mark.unit +def test_export_rejects_ambiguous_sources_and_options(monkeypatch): + export_module, _ = _patch_lowering(monkeypatch) + program = FakeExportedProgram() + + with pytest.raises(ValueError, match="already-exported program"): + export_module.export(program, arg_inputs=[torch.ones(1)]) + with pytest.raises(TypeError, match="Compile nn.Module inputs"): + export_module.export(torch.nn.Linear(1, 1)) + with pytest.raises(ValueError, match="unknown methods"): + export_module.export({"forward": program}, partitioners={"missing": [object()]}) + with pytest.raises(ValueError, match="collide"): + export_module.export({"forward": program}, constant_methods={"forward": 1}) + + +@pytest.mark.unit +@pytest.mark.parametrize("option_name", ["partitioners", "compile_specs"]) +def test_export_rejects_string_option_sequences(monkeypatch, option_name): + export_module, _ = _patch_lowering(monkeypatch) + program = FakeExportedProgram() + + with pytest.raises(TypeError, match=option_name): + export_module.export(program, **{option_name: "invalid"}) + + with pytest.raises(TypeError, match=option_name): + export_module.export( + {"forward": program}, **{option_name: {"forward": "invalid"}} + ) + + +@pytest.mark.unit +def test_prepare_graph_module_rejects_argument_mismatch(monkeypatch): + export_module = importlib.import_module("torch_tensorrt.executorch.export") + + class OneInput(torch.nn.Module): + def forward(self, x): + return x + + graph_module = torch.fx.symbolic_trace(OneInput()) + with pytest.raises(TypeError): + export_module._prepare_graph_module( + graph_module, + arg_inputs=( + torch_tensorrt.Input((1,), name="x"), + torch_tensorrt.Input((1,), name="extra"), + ), + kwarg_inputs={}, + dynamic_shapes=None, + retrace=False, + ) + + +@pytest.mark.unit +def test_export_warns_for_multiple_engines(monkeypatch, caplog): + program = FakeExportedProgram() + export_module, _ = _patch_lowering(monkeypatch, {program: 2}) + + export_module.export(program) + assert "contains 2 TRT engines" in caplog.text From d0e5aabb77c7fa1a507555013f15d102353c657b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 6 Aug 2026 13:18:29 -0700 Subject: [PATCH 2/4] fix(executorch): require per-method partitioners for multi-method export A partitioner can carry method-specific state. ExecuTorch backends bake the method name into the DelegationSpec built in the partitioner constructor, so reusing one instance across methods tags every method with the first method's name, and each delegate then looks up the wrong compiled method at runtime. export() broadcast a flat partitioners= sequence to every method, which made that mistake easy and silent. Reject it when there is more than one method and point at the per-method mapping form instead. A flat sequence still works for a single method, and compile_specs (plain data, not stateful objects) still broadcasts unchanged. Tests: reject a shared instance across two methods, allow a flat sequence for one method, and confirm a per-method mapping preserves each instance. --- py/torch_tensorrt/executorch/export.py | 17 ++++++++ tests/py/dynamo/executorch/test_export.py | 51 +++++++++++++++++++++++ 2 files changed, 68 insertions(+) diff --git a/py/torch_tensorrt/executorch/export.py b/py/torch_tensorrt/executorch/export.py index 66e0ae5785..91d8548203 100644 --- a/py/torch_tensorrt/executorch/export.py +++ b/py/torch_tensorrt/executorch/export.py @@ -218,6 +218,19 @@ def _per_method_values( "sequence." ) shared = list(value or ()) + # A partitioner carries per-method state: ExecuTorch backends bake the method + # name into the DelegationSpec built in their constructor, so one instance + # reused across methods would tag every method with the first method's name. + if shared and option_name == "partitioners" and len(method_names) > 1: + raise ValueError( + "partitioners must be a mapping of method name to sequence when " + f"exporting multiple methods ({', '.join(method_names)}). A single " + "partitioner instance cannot be shared across methods because a " + "partitioner may carry method-specific state, such as a method name " + "baked into its delegation spec. Construct one per method, for " + 'example partitioners={"prefill": [MyPartitioner(...)], "decode": ' + "[MyPartitioner(...)]}." + ) return {name: list(shared) for name in method_names} @@ -252,6 +265,10 @@ def export( TensorRT engines. Transform passes must treat shared payload values as immutable. Method mappings preserve independent entry points but do not imply shared mutable state between them. + + When exporting more than one method, pass ``partitioners`` as a mapping of + method name to its own partitioner instances. A partitioner may carry + method-specific state, so sharing one instance across methods is rejected. """ from torch_tensorrt._features import ENABLED_FEATURES diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index d1eae88b27..1876d717b4 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -676,6 +676,57 @@ def test_export_normalizes_none_per_method_values(monkeypatch): assert pipelines["decode"][0].compile_specs == [] +@pytest.mark.unit +def test_export_rejects_shared_partitioners_across_methods(monkeypatch): + """One partitioner instance cannot serve several methods. + + ExecuTorch backends bake the method name into the delegation spec built in + the partitioner constructor, so a shared instance would tag every method + with the first method's name. + """ + export_module, lower = _patch_lowering(monkeypatch) + + with pytest.raises(ValueError, match="must be a mapping of method name"): + export_module.export( + {"prefill": FakeExportedProgram(), "decode": FakeExportedProgram()}, + partitioners=[object()], + ) + + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_allows_shared_partitioners_for_single_method(monkeypatch): + """A flat sequence stays valid when there is only one method to tag.""" + export_module, lower = _patch_lowering(monkeypatch) + partitioner = object() + + export_module.export(FakeExportedProgram(), partitioners=[partitioner]) + + pipeline = lower.call_args.kwargs["partitioner"] + assert pipeline[-1] is partitioner + + +@pytest.mark.unit +def test_export_accepts_per_method_partitioner_instances(monkeypatch): + """Each method keeps the partitioner instance it was given.""" + export_module, lower = _patch_lowering(monkeypatch) + prefill_partitioner = object() + decode_partitioner = object() + + export_module.export( + {"prefill": FakeExportedProgram(), "decode": FakeExportedProgram()}, + partitioners={ + "prefill": [prefill_partitioner], + "decode": [decode_partitioner], + }, + ) + + pipelines = lower.call_args.kwargs["partitioner"] + assert pipelines["prefill"][-1] is prefill_partitioner + assert pipelines["decode"][-1] is decode_partitioner + + @pytest.mark.unit @pytest.mark.parametrize("container_type", [list, tuple]) def test_prepare_graph_module_infers_nested_dynamic_shapes(monkeypatch, container_type): From 6c9cb0a522670b765861e35bbe434e4534b7e81e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 6 Aug 2026 14:28:22 -0700 Subject: [PATCH 3/4] fix(executorch): stop sharing metadata containers and harden option validation Three fixes found while reviewing the composable export path. 1. Staging shared a node.meta value wholesale whenever it contained a symbolic size or fake tensor anywhere inside it. A multi-output op stores a list of fake tensors in meta["val"], so that list stayed shared with the caller's program and an Edge transform mutating it would corrupt the caller's input. Seed only the shape-bound leaves into the deepcopy memo, so the leaves are still shared (they must be, they belong to the live ShapeEnv) while the container around them is copied. 2. A partitioner can carry method-specific state: ExecuTorch backends bake the method name into the DelegationSpec built in the constructor, so reusing one instance across methods tags every method with the first method's name and each delegate then looks up the wrong compiled method at runtime. Reject a reused instance whether it arrives as a flat sequence broadcast to every method or as the same object under two method keys. 3. ExecuTorch dispatches per-method transform passes on isinstance(passes, dict), so a Mapping that is not a dict silently ran no passes at all. Normalize it. Tests: a container holding symbolic leaves is copied while the leaves stay shared; a reused partitioner is rejected in both forms while distinct instances and the single-method flat form still work; a non-dict Mapping of transform passes is normalized to dict. --- py/torch_tensorrt/executorch/_export_utils.py | 44 +++++--- py/torch_tensorrt/executorch/export.py | 49 ++++++--- tests/py/dynamo/executorch/test_export.py | 101 +++++++++++++++++- 3 files changed, 159 insertions(+), 35 deletions(-) diff --git a/py/torch_tensorrt/executorch/_export_utils.py b/py/torch_tensorrt/executorch/_export_utils.py index ebacd4bbce..d6752cefb4 100644 --- a/py/torch_tensorrt/executorch/_export_utils.py +++ b/py/torch_tensorrt/executorch/_export_utils.py @@ -8,19 +8,26 @@ from torch.export.graph_signature import InputKind -def _is_graph_bound_metadata(value: Any) -> bool: - """True if a metadata value is tied to the exported program's shape environment. +def _seed_graph_bound_leaves(value: Any, memo: dict[int, Any]) -> bool: + """Seed shape-environment-bound leaves into ``memo`` so deepcopy shares them. Fake tensors and symbolic sizes reach back to a live ``ShapeEnv`` that owns tensors deepcopy refuses to touch, and cloning them would also detach the copy - from the symbols the graph is guarded on. + from the symbols the graph is guarded on. Seeding only the leaves lets the + surrounding container still be copied, so a container such as the ``list`` a + multi-output op stores in ``meta["val"]`` is not shared with the caller. + + Returns True if the value has to be shared wholesale because deepcopy cannot + reach its leaves through pytree. """ + found = False for leaf in torch.utils._pytree.tree_leaves(value): - if isinstance(leaf, (torch.SymInt, torch.SymFloat, torch.SymBool)): - return True - if isinstance(leaf, torch.Tensor) and is_fake(leaf): - return True - return False + if isinstance(leaf, (torch.SymInt, torch.SymFloat, torch.SymBool)) or ( + isinstance(leaf, torch.Tensor) and is_fake(leaf) + ): + memo[id(leaf)] = leaf + found = True + return found def get_engine_info_from_state(engine_obj: Any) -> list[Any]: @@ -150,14 +157,19 @@ def _stage_graph_module( f"Staged GraphModule {name or ''!r} changed node identities." ) for node_name, source_node in source_nodes.items(): - staged_nodes[node_name].meta = { - key: ( - value - if _is_graph_bound_metadata(value) - else copy.deepcopy(value, payload_memo) - ) - for key, value in source_node.meta.items() - } + staged_meta = {} + for key, value in source_node.meta.items(): + # Seed the shape-bound leaves first so the container around them + # is still copied while the leaves stay shared. + _seed_graph_bound_leaves(value, payload_memo) + try: + staged_meta[key] = copy.deepcopy(value, payload_memo) + except Exception: + # A value deepcopy cannot reach through pytree (for example a + # live ShapeEnv) has to be shared rather than lose the graph + # its guards refer to. + staged_meta[key] = value + staged_nodes[node_name].meta = staged_meta return staged diff --git a/py/torch_tensorrt/executorch/export.py b/py/torch_tensorrt/executorch/export.py index 91d8548203..da09c2b4d6 100644 --- a/py/torch_tensorrt/executorch/export.py +++ b/py/torch_tensorrt/executorch/export.py @@ -189,6 +189,32 @@ def _prepare_programs( ) +def _reject_shared_partitioners(per_method: dict[str, list[Any]]) -> None: + """Reject one partitioner instance serving several methods. + + A partitioner can carry method-specific state: ExecuTorch backends bake the + method name into the DelegationSpec built in the constructor, so a reused + instance tags every method with the first method's name and each delegate + then looks up the wrong compiled method at runtime. + """ + if len(per_method) < 2: + return + owner_by_id: dict[int, str] = {} + for name, partitioners in per_method.items(): + for partitioner in partitioners: + previous = owner_by_id.setdefault(id(partitioner), name) + if previous != name: + raise ValueError( + f"partitioners reuses the same " + f"{type(partitioner).__name__} instance for {previous!r} and " + f"{name!r}. A partitioner may carry method-specific state, " + "such as a method name baked into its delegation spec, so " + "each method needs its own instance. For example: " + 'partitioners={"prefill": [MyPartitioner(...)], "decode": ' + "[MyPartitioner(...)]}." + ) + + def _per_method_values( value: Sequence[Any] | Mapping[str, Sequence[Any] | None] | None, method_names: tuple[str, ...], @@ -218,19 +244,6 @@ def _per_method_values( "sequence." ) shared = list(value or ()) - # A partitioner carries per-method state: ExecuTorch backends bake the method - # name into the DelegationSpec built in their constructor, so one instance - # reused across methods would tag every method with the first method's name. - if shared and option_name == "partitioners" and len(method_names) > 1: - raise ValueError( - "partitioners must be a mapping of method name to sequence when " - f"exporting multiple methods ({', '.join(method_names)}). A single " - "partitioner instance cannot be shared across methods because a " - "partitioner may carry method-specific state, such as a method name " - "baked into its delegation spec. Construct one per method, for " - 'example partitioners={"prefill": [MyPartitioner(...)], "decode": ' - "[MyPartitioner(...)]}." - ) return {name: list(shared) for name in method_names} @@ -266,9 +279,9 @@ def export( immutable. Method mappings preserve independent entry points but do not imply shared mutable state between them. - When exporting more than one method, pass ``partitioners`` as a mapping of - method name to its own partitioner instances. A partitioner may carry - method-specific state, so sharing one instance across methods is rejected. + When exporting more than one method, give each method its own partitioner + instances via ``partitioners={"method": [...]}``. A partitioner may carry + method-specific state, so reusing one instance across methods is rejected. """ from torch_tensorrt._features import ENABLED_FEATURES @@ -304,6 +317,7 @@ def export( ) program_map = {"forward": programs} if not isinstance(programs, dict) else programs extra_partitioners = _per_method_values(partitioners, method_names, "partitioners") + _reject_shared_partitioners(extra_partitioners) method_compile_specs = _per_method_values( compile_specs, method_names, "compile_specs" ) @@ -321,6 +335,9 @@ def export( raise ValueError( f"transform_passes contains unknown methods: {sorted(unknown)}" ) + # ExecuTorch dispatches per-method passes on isinstance(passes, dict), so a + # Mapping that is not a dict silently runs no passes at all. + transform_passes = dict(transform_passes) engine_counts = { name: validate_engine_program(program) for name, program in program_map.items() diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index 1876d717b4..5315c776d3 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -9,6 +9,7 @@ import torch # noqa: E402 import torch_tensorrt # noqa: E402 +from torch._subclasses.fake_tensor import is_fake # noqa: E402 from torch.export.graph_signature import ( # noqa: E402 CustomObjArgument, ExportGraphSignature, @@ -353,14 +354,25 @@ def forward(self, x): # Symbolic metadata stays bound to the original ShapeEnv, so it must be shared # rather than copied for the staged program to remain consistent. + # Symbolic leaves stay bound to the original ShapeEnv, so they must be shared + # rather than copied for the staged program to remain consistent. The container + # holding them is still copied so the caller's metadata cannot be mutated. source_nodes = {node.name: node for node in program.graph.nodes} symbolic_shared = 0 for node in staged.graph.nodes: source_node = source_nodes[node.name] for key, source_value in source_node.meta.items(): - if not export_utils._is_graph_bound_metadata(source_value): + source_leaves = [ + leaf + for leaf in torch.utils._pytree.tree_leaves(source_value) + if isinstance(leaf, (torch.SymInt, torch.SymFloat, torch.SymBool)) + or (isinstance(leaf, torch.Tensor) and is_fake(leaf)) + ] + if not source_leaves: continue - assert node.meta[key] is source_value + staged_leaves = torch.utils._pytree.tree_leaves(node.meta[key]) + for leaf in source_leaves: + assert any(leaf is staged_leaf for staged_leaf in staged_leaves) symbolic_shared += 1 assert symbolic_shared @@ -368,6 +380,39 @@ def forward(self, x): assert staged.state_dict["lin.weight"] is program.state_dict["lin.weight"] +@pytest.mark.unit +def test_stage_exported_program_copies_containers_holding_symbolic_leaves(): + """A container in node.meta must not be shared just because it holds a leaf. + + A multi-output op stores a list of fake tensors in meta["val"]. Sharing that + list wholesale would let an Edge transform mutate the caller's program. + """ + export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") + + class Model(torch.nn.Module): + def forward(self, x): + first, second = torch.chunk(x, 2, dim=1) + return first + second + + program = torch.export.export(Model(), (torch.randn(8, 4),)) + container_nodes = [ + node + for node in program.graph.nodes + if isinstance(node.meta.get("val"), (list, tuple)) + ] + assert container_nodes, "expected a multi-output op with a list-valued meta" + + staged = export_utils.stage_exported_program(program) + staged_nodes = {node.name: node for node in staged.graph.nodes} + + for source_node in container_nodes: + source_value = source_node.meta["val"] + staged_value = staged_nodes[source_node.name].meta["val"] + assert staged_value is not source_value + for source_leaf, staged_leaf in zip(source_value, staged_value): + assert staged_leaf is source_leaf + + @pytest.mark.unit def test_stage_exported_program_clones_nested_graph_modules(): export_utils = importlib.import_module("torch_tensorrt.executorch._export_utils") @@ -686,7 +731,7 @@ def test_export_rejects_shared_partitioners_across_methods(monkeypatch): """ export_module, lower = _patch_lowering(monkeypatch) - with pytest.raises(ValueError, match="must be a mapping of method name"): + with pytest.raises(ValueError, match="reuses the same"): export_module.export( {"prefill": FakeExportedProgram(), "decode": FakeExportedProgram()}, partitioners=[object()], @@ -695,6 +740,21 @@ def test_export_rejects_shared_partitioners_across_methods(monkeypatch): lower.assert_not_called() +@pytest.mark.unit +def test_export_rejects_shared_partitioner_inside_mapping(monkeypatch): + """The same instance listed under two method keys is rejected too.""" + export_module, lower = _patch_lowering(monkeypatch) + shared = object() + + with pytest.raises(ValueError, match="reuses the same"): + export_module.export( + {"prefill": FakeExportedProgram(), "decode": FakeExportedProgram()}, + partitioners={"prefill": [shared], "decode": [shared]}, + ) + + lower.assert_not_called() + + @pytest.mark.unit def test_export_allows_shared_partitioners_for_single_method(monkeypatch): """A flat sequence stays valid when there is only one method to tag.""" @@ -727,6 +787,41 @@ def test_export_accepts_per_method_partitioner_instances(monkeypatch): assert pipelines["decode"][-1] is decode_partitioner +@pytest.mark.unit +def test_export_normalizes_mapping_transform_passes_to_dict(monkeypatch): + """ExecuTorch dispatches per-method passes on isinstance(passes, dict). + + A Mapping that is not a dict would silently run no passes at all, so it must + be normalized before being forwarded. + """ + from collections.abc import Mapping as AbcMapping + + class CustomMapping(AbcMapping): + def __init__(self, data): + self._data = data + + def __getitem__(self, key): + return self._data[key] + + def __iter__(self): + return iter(self._data) + + def __len__(self): + return len(self._data) + + export_module, lower = _patch_lowering(monkeypatch) + passes = ["a-pass"] + + export_module.export( + {"prefill": FakeExportedProgram(), "decode": FakeExportedProgram()}, + transform_passes=CustomMapping({"prefill": passes}), + ) + + forwarded = lower.call_args.kwargs["transform_passes"] + assert type(forwarded) is dict + assert forwarded == {"prefill": passes} + + @pytest.mark.unit @pytest.mark.parametrize("container_type", [list, tuple]) def test_prepare_graph_module_infers_nested_dynamic_shapes(monkeypatch, container_type): From 1b09190076548912906129bcfe84e918418ce230 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 6 Aug 2026 20:48:58 -0700 Subject: [PATCH 4/4] docs(executorch): document the shared-payload contract; validate constant methods Two review findings from adversarial testing of the composable export path. The staged Edge program shares weight and TensorRT engine storage with the programs the caller passed in. Only structure is copied. That is deliberate (it keeps a multi-gigabyte engine from being duplicated) but it was documented only in the export() docstring, while the saving guide showed transform_passes= with no caveat. An in-place edit in a transform pass, or a change to a source program after export, silently changes the caller's program and any sibling exported from it, with no warning. Document the contract where callers read it, and say plainly that neither case raises. Also document the several-methods-in-one-pte form, which was described in prose but had no example, including the requirement that each method gets its own partitioner instances. Constant-method names are baked into the .pte as method names, so reject keys that are not valid Python identifiers instead of forwarding them. Tests: invalid constant-method keys (spaces, leading digit, empty, non-string) are rejected; a valid key is forwarded unchanged. --- .../runtime_performance/saving_models.rst | 47 +++++++++++++++++-- py/torch_tensorrt/executorch/export.py | 9 ++++ tests/py/dynamo/executorch/test_export.py | 26 ++++++++++ 3 files changed, 79 insertions(+), 3 deletions(-) diff --git a/docsrc/user_guide/runtime_performance/saving_models.rst b/docsrc/user_guide/runtime_performance/saving_models.rst index d32351d293..e4dbc36c92 100644 --- a/docsrc/user_guide/runtime_performance/saving_models.rst +++ b/docsrc/user_guide/runtime_performance/saving_models.rst @@ -270,9 +270,50 @@ inspect the delegated graph, add metadata, or control final memory planning: engine-bearing ``ExportedProgram``, or a mapping of independently exported methods. It always applies the TensorRT partitioner first, followed by caller ``partitioners`` in order, and returns ExecuTorch's native -``EdgeProgramManager``. Method mappings preserve independent entry points but -do not imply shared mutable state. Use ``save`` for the one-shot path where -Torch-TensorRT manages Edge lowering, finalization, and persistence. +``EdgeProgramManager``. Use ``save`` for the one-shot path where Torch-TensorRT +manages Edge lowering, finalization, and persistence. + +**Several methods in one .pte** + +Pass a mapping of method name to ``ExportedProgram`` to keep independent entry +points, such as a separate prefill and decode, in one program. Give each method +its own partitioner instances: a partitioner can carry method-specific state (an +ExecuTorch backend bakes the method name into the delegation spec it builds in +its constructor), so reusing one instance across methods is rejected. + +.. code-block:: python + + edge = torch_tensorrt.executorch.export( + {"prefill": prefill_program, "decode": decode_program}, + partitioners={ + "prefill": [CudaPartitioner([])], + "decode": [CudaPartitioner([])], + }, + ) + +Each method is validated before any program is rewritten, so an error in one +method leaves the others untouched. A method mapping preserves independent entry +points but does not by itself give them shared mutable state. + +.. warning:: + + **The returned Edge program shares weight and engine storage with the + programs you passed in.** Only structure is copied: the graph, the graph + signature, the ``state_dict`` keys, and node metadata. Tensor and TensorRT + engine contents are shared by reference, which is what keeps a multi-gigabyte + engine from being duplicated. + + Two consequences to plan for: + + * A transform pass must not modify a shared payload in place. An in-place + edit such as ``weight.data.mul_(scale)`` changes the program you passed in, + and every other Edge program exported from it. Build a new tensor and + rebind it instead of mutating the existing one. + * Modifying a source program after calling ``export`` is also visible in the + Edge program. Finish preparing a program before exporting it. + + Neither case raises an error or a warning, so treat every shared payload as + read-only. **Coalesced TensorRT + CUDA .pte** diff --git a/py/torch_tensorrt/executorch/export.py b/py/torch_tensorrt/executorch/export.py index da09c2b4d6..5d26d3024f 100644 --- a/py/torch_tensorrt/executorch/export.py +++ b/py/torch_tensorrt/executorch/export.py @@ -323,6 +323,15 @@ def export( ) if constant_methods is not None: + invalid = [ + name + for name in constant_methods + if not isinstance(name, str) or not name.isidentifier() + ] + if invalid: + raise ValueError( + f"constant_methods keys must be valid Python identifiers: {invalid}" + ) collisions = set(constant_methods) & set(method_names) if collisions: raise ValueError( diff --git a/tests/py/dynamo/executorch/test_export.py b/tests/py/dynamo/executorch/test_export.py index 5315c776d3..2aea13426f 100644 --- a/tests/py/dynamo/executorch/test_export.py +++ b/tests/py/dynamo/executorch/test_export.py @@ -787,6 +787,32 @@ def test_export_accepts_per_method_partitioner_instances(monkeypatch): assert pipelines["decode"][-1] is decode_partitioner +@pytest.mark.unit +def test_export_rejects_invalid_constant_method_names(monkeypatch): + """Constant-method names are baked into the .pte, so reject unusable ones.""" + export_module, lower = _patch_lowering(monkeypatch) + + for bad_name in ("not an identifier", "1_leading_digit", ""): + with pytest.raises(ValueError, match="valid Python identifiers"): + export_module.export(FakeExportedProgram(), constant_methods={bad_name: 1}) + + with pytest.raises(ValueError, match="valid Python identifiers"): + export_module.export(FakeExportedProgram(), constant_methods={42: 1}) + + lower.assert_not_called() + + +@pytest.mark.unit +def test_export_accepts_valid_constant_method_names(monkeypatch): + export_module, lower = _patch_lowering(monkeypatch) + + export_module.export( + FakeExportedProgram(), constant_methods={"get_vocab_size": 256} + ) + + assert lower.call_args.kwargs["constant_methods"] == {"get_vocab_size": 256} + + @pytest.mark.unit def test_export_normalizes_mapping_transform_passes_to_dict(monkeypatch): """ExecuTorch dispatches per-method passes on isinstance(passes, dict).