Uh oh!
There was an error while loading. Please reload this page.
feat(executorch): expose composable Edge export API - #4440
Merged
lanluo-nvidia merged 44 commits intoAug 19, 2026
Conversation
shoumikhinforce-pushed
the
executorch-composable-export
branch
from
July 29, 2026 18:48
cd55b95 to
4dc2d25Compareshoumikhinforce-pushed
the
executorch-composable-export
branch
from
July 31, 2026 06:51
4dc2d25 to
cb3633fCompareUh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
Uh oh!
There was an error while loading. Please reload this page.
shoumikhinforce-pushed
the
executorch-composable-export
branch
from
July 31, 2026 23:34
cb3633f to
93efd54Compare| retrace=False, arg_inputs=inputs, | ||
| ) | ||
| ``save`` writes both the ``.pte`` and any external ``.ptd`` tensor-data files. |
Collaborator
There was a problem hiding this comment.
I would suggest here first add the torch_tensorrt.save() as the default method for user wants to export a executorch pte file.
then add the advanced method, when user needs a customization before serialize into a pte file.
So that we explain user both scenario clearly.
Collaborator
There was a problem hiding this comment.
Also please wait for PR: 4398 to be merged to main(I will push by the end of today.)
please rebase on top of that.
7 tasks
shoumikhinforce-pushed
the
executorch-composable-export
branch
2 times, most recently
from
August 18, 2026 02:49
6606f9f to
2fd1609Compare
`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)
print(edge.exported_program().graph)
program = edge.to_executorch()
with open("model.pte", "wb") as output:
program.write_to_file(output)
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.
Use `save()` when you just want a `.pte`. Use
`torch_tensorrt.executorch.export()` when you need control before the `.pte` is
created.
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()
```
```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.
- 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.
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.
- 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.
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.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.
…alidation 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.
…tant 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.
…each engine once
The multi-method example passed CudaPartitioner([]) for every method. The AOTI backend
reads the method name from the partitioner's compile specs while building the named data
store, and CudaPartitioner does not add one, so running the example fails with "Could not
find method name in compile specs". This is the example that teaches the rule that each
method needs its own partitioner state, so it now carries that state:
partitioners={
"prefill": [
CudaPartitioner([CudaBackend.generate_method_name_compile_spec("prefill")])
],
}
The example also used CudaBackend and CudaPartitioner while the only imports for them
appeared in a later section, so it could not be copied and run. Both imports are now part
of the example.
Separately, every engine was serialized twice. Reading an engine's state calls serialize()
and base64 encodes the result at 1.33x, and both the validation pre-pass and the later
rewrite resolved the same engine. Validation now records what it resolved and the rewrite
reuses it, so the work happens once.
The record is keyed by node name rather than by program or engine identity, because
staging returns a new ExportedProgram and the rewrite runs on that copy, so identity
cannot bridge the two steps.
Testing: the whole executorch test directory passes, 110 tests. A new test drives the two
steps with an engine that counts serializations and asserts exactly one; it fails against
the previous code.Passing transform_passes={} raised KeyError: 'forward' from ExecuTorch. ExecuTorch
decides between global and per-method passes by checking whether the value is a dict,
and then looks up every exported method, so an empty dict is read as "per-method passes
were supplied" and the first lookup fails. Measured against the installed ExecuTorch:
transform_passes=None -> ok
transform_passes={} -> KeyError: 'forward'
transform_passes={"forward": []} -> ok
An empty mapping now normalizes to None, which is what constant_methods already does a
few lines further down.
The docstring also gained three limits that were not written down: generate_etrecord
copies the whole program so it sits outside the payload sharing the rest of the
docstring promises, constant_methods keys are restricted to Python identifiers here
which is narrower than ExecuTorch accepts, and a per-method partitioner needs the
compile spec for the method it serves.
Testing: the whole executorch test directory passes, 111 tests. The new test asserts
None reaches ExecuTorch for an empty mapping and fails against the previous code.…write The test added for single-engine serialization never called the rewrite. It counted resolutions across validation, looked up one dictionary key, then asserted the count was still one. Nothing between the two assertions could change it, so deleting the reuse path in replace_execute_engine left the test passing while the engine was serialized twice again. Confirmed by removing the reuse path: the test still passed. The test now drives both steps. It validates one program into the resolved mapping, then calls the rewrite on a separate program object holding the same node names, which is the condition staging produces, and asserts the engine was resolved only once. It also asserts the node names match, since that is what lets the rewrite find the earlier work on a different object. Verified the test fails when the reuse path is removed, so it can detect the regression it was written for. Testing: the whole executorch test directory passes, 111 tests.
…only when it names a method Passing a single partitioner for several methods was an unconditional error, so the documented sequence form was unusable with a method mapping of two or more methods. That shape is not a misuse: ExecuTorch's own multi-method examples pass one partitioner for several methods, in the CoreML test-resource generator and in its static LLM example. The justification was also too broad. It claimed ExecuTorch backends bake the method name into the delegation spec built in the constructor. AotiPartitioner takes a backend name and a compile spec list and derives no method name; the name only appears when the caller puts it there with generate_method_name_compile_spec. So the hazard is real only when the specs actually carry a method name. The check now warns instead of raising, and only fires when the shared instance's compile specs name a method, which is the case where two methods would be tagged with the same name and the delegates would look up the wrong one. Testing: the whole executorch test directory passes, 111 tests. One test asserts the warning fires for a shared instance whose specs name a method, and another asserts a shared instance without one is accepted and still lowers.
Four small corrections found while reviewing the export path. The multi-method memory cost was only written in the pull request description. It now sits in the export() docstring next to the generate_etrecord note, since that is what a caller reads. Staging caught every exception while deep copying node metadata, including programming errors, and then silently shared the value with the caller's program. Sharing a mutable value there is the one case that breaks the isolation the rest of that function provides, so the handler now catches only the failure modes a value that cannot be copied raises, and logs at debug when it happens. The payload sharing memo seeded every value from the state dict and the constants without checking its type, while the two loops just below it only shared tensors and script objects. So a custom-object constant was shared even though the documentation promises sharing for tensors and engines. The first loop now applies the same check. torch_tensorrt.executorch had no page in the Python API reference, so the user guide was its only documentation. Added one covering export(), get_edge_compile_config(), TensorRTPartitioner and TensorRTBackend, which are the four names the module exports. Testing: the whole executorch test directory passes, 111 tests. Verified the four documented names all resolve on the installed module.
Moving this code out of the compile module dropped three explanations that are not recoverable from the code itself, so anyone changing it would have to rediscover them. Why the engine node is replaced at all: ExecuTorch's lowering runs passes that dispatch through the C++ schema validator, and the validator rejects the engine argument because it arrives as a custom-object placeholder rather than a real script object. The replacement node carries the same information as plain strings, so the passes never see a script object. Why the payload becomes a buffer plus a get_attr reference instead of an inline value: CPython's tokenizer cannot parse a string literal larger than about 2 GB, so an inline base64 string breaks recompilation for any engine past that size. Why the existing fake mode is reused rather than a new one: a fresh mode fails downstream with a mismatch error as soon as a pass mixes tensors from the two modes. Testing: the whole executorch test directory passes, 111 tests. Comment-only change.
Staging makes a private copy of a program's structure so ExecuTorch lowering cannot mutate the caller's program, while the big payloads (weights and TensorRT engines) stay shared. Copying an engine is expensive: it serializes and then deserializes the whole engine, which can be gigabytes. That sharing was not happening for a lifted engine. A lifted engine does not sit on the graph module as an attribute; it arrives as a FakeScriptObject that holds the real engine on .real_obj. FakeScriptObject is not a torch.ScriptObject subclass, so the two isinstance checks that decide what to share skipped it, and deepcopy walked straight into the engine. It happened twice per program: once for the wrapper in node.meta["val"], and once through the graph signature, which was deep-copied with no memo at all and holds the same wrapper on CustomObjArgument.fake_val. Measured on a stand-in engine that counts serialize and deserialize calls: two of each before this change, zero after. The fix adds FakeScriptObject to the shared payload types, and passes the same memo to the graph signature copy so the graph and the signature keep pointing at one engine. Testing: the new unit test stages a lifted-engine program and asserts the wrapper is the same object in the staged node metadata, in the staged signature, and in the staged constants, and that the engine was never serialized. It fails before this change and passes after. The executorch test directory was run before and after with no other change in results.
`validate_engine_program` resolved every execute_engine node separately and keyed the result by node name. When one engine feeds several nodes, which is what the shared-engine cleanup path in `replace_execute_engine` already handles, the engine was serialized once per node and every copy of its base64 bytes stayed resident until the export finished. The rewrite then discarded all but one of them. Resolve once per engine node and hand the same engine-info list to every node that uses it. The `resolved` mapping keeps its node-name keys, so the rewrite is unchanged. Test plan: - New unit test `test_validate_engine_program_serializes_a_shared_engine_once` builds a program with three execute_engine nodes on one engine. On the parent commit the engine is serialized 3 times; with this change it is serialized once and all three entries point at the same list. - Full `tests/py/dynamo/executorch/test_export.py` run locally: 49 passed.
Two engine copies were held longer than needed. Every method's resolved engine info stayed in `resolved_engines` until the last method was rewritten, and inside the rewrite the read-only decoded bytes stayed alive next to the writable copy `torch.frombuffer` needs, for the rest of the loop iteration. Hand each method its own resolved mapping with `pop`, so the payloads go away with the rewrite that consumed them, and rebind the decoded bytes to the bytearray so only one of the two survives the conversion. The docstring claimed peak memory of roughly 1.33x the engine bytes. That counted only the base64 text. Replace it with what the code actually does. Test plan: - New unit test `test_export_releases_a_methods_engines_once_it_is_rewritten` weak-references one payload per method and records which are still live at each rewrite. On the parent commit the first method's payload is still live during the second rewrite; with this change it is gone. - Full `tests/py/dynamo/executorch/test_export.py` run locally: 50 passed.
…e buffer The engine buffer name came from `get_unique_attr_name_in_module`, which only looks at module attributes. torch.export lifts buffers out of the module and into `state_dict`, so a model carrying a buffer named `_trt_engine_0` keeps that name in `state_dict` while the module reports nothing. The rewrite then picked the same name and its `state_dict[buffer_name] = engine_tensor` line replaced the model's own buffer with the engine bytes. Search for a name that is free in both places. Test plan: - New unit test `test_replace_execute_engine_keeps_a_lifted_buffer_of_the_same_name` puts a `_trt_engine_0` entry in state_dict before the rewrite. On the parent commit that entry comes back as the engine tensor; with this change it is untouched and the engine lands under a different name. - Full `tests/py/dynamo/executorch/test_export.py` run locally: 51 passed.
`from .export import export` rebinds `torch_tensorrt.executorch.export` from the submodule to the function, so the same name means two things depending on what has been imported. It also lets `import torch_tensorrt.executorch.export` install the module as a package attribute, which skips the `__getattr__` guard that is supposed to explain that ExecuTorch is not installed, and leaves `from torch_tensorrt.executorch import export` handing back a module that cannot be called. Name the module `_export`, next to `_export_utils`, so the public name belongs to the function alone. The public API `torch_tensorrt.executorch.export(...)` is unchanged, and so are the docs that reference it. Test plan: - Full `tests/py/dynamo/executorch/test_export.py` run locally: 51 passed. - Checked the tree for other references: the remaining `executorch.export` mentions in docsrc, `_compile.py`, and the composition test all name the function, not the module.
`import importlib` does not import the `util` submodule. `python -S -c "import importlib; importlib.util"` raises AttributeError, which the guard's `except ModuleNotFoundError` would not catch. It only works today because something else imported it first. Two comments no longer match the code they describe: - A non-dict Mapping does not make ExecuTorch silently run no passes. It matches none of the dispatch branches and reaches a KeyError on the first method, the same way an empty dict does. - `_seed_graph_bound_leaves` documented its return as "the value has to be shared wholesale", which is the opposite of what it returns, and the only caller ignores it. Drop the return. Test plan: - `python -S -c "import importlib; importlib.util"` reproduces the AttributeError. - Full `tests/py/dynamo/executorch/test_export.py` run locally: 51 passed.
_prepare_graph_module materializes Input specs on the default device, which is a GPU, so five tests that check dynamic-shape inference errored on a CPU host instead of skipping. Four of them do not care about the device, so point them at the CPU. The fifth needs a tensor already on the GPU, so gate it at test runtime, matching test_api.py and test_cuda_partitioner_composition.py. Test plan: pytest tests/py/dynamo/executorch/test_export.py with the GPU hidden goes from 5 failed, 46 passed to 50 passed, 1 skipped; unchanged with a GPU.
_TRTEngine.__getstate__ base64-encodes the engine into a str, so the str branch of replace_execute_engine is the only one a real engine takes, and every test fed it raw bytes instead. Parametrize the round trip over both shapes. Test plan: pytest tests/py/dynamo/executorch/test_export.py, 53 passed. With the str branch deleted, both new base64 cases fail with TypeError: string argument without an encoding.
Validation resolves every method before any rewrite, so the per-method handoff
could be misrouted and nothing would notice: the multi-method end-to-end test
only checks that the two delegate payloads differ, which a swap preserves.
Test plan: pytest tests/py/dynamo/executorch/test_export.py, 54 passed. With
the handoff swapped to pop another method's entry, the new test fails with
{'prefill': 'decode', 'decode': 'prefill'}.The PR replaced _count_executorch_engine_nodes with validate_engine_program but dropped the test that pinned its two accepted node shapes, so nothing covered the no_op placeholder branch or the plain-args resolve path any more. Test plan: pytest tests/py/dynamo/executorch/test_export.py, 55 passed. With the no_op branch removed from validate_engine_program, the restored test fails with assert 1 == 2.
The sharing assertion in this test reads a state dict that stage_exported_program copies shallowly, so it held no matter what the memo did. A real tensor in node metadata is copied without the memo, so assert on that instead. Test plan: pytest tests/py/dynamo/executorch/test_export.py, 55 passed. With _payload_sharing_memo stubbed to return an empty dict, this test now fails alongside the two that already caught it.
ExecuTorch's EdgeProgramManager.transform deep-copies every method missing from
a per-method pass dict, and after the rewrite that copy duplicates the method's
whole engine buffer. Measured on executorch 1.1: a two-method manager given
{'a': []} clones b's state dict tensors, and given {'a': [], 'b': []} shares
both, because an empty PassManager reports modified=False and _transform then
returns the program itself.
Test plan: pytest tests/py/dynamo/executorch/test_export.py, 56 passed.The test docstring still said it silently runs no passes. ExecuTorch leaves the entry unassigned and raises KeyError on the next line, which is what the source comment now says. Test plan: pytest tests/py/dynamo/executorch/test_export.py, 56 passed.
executorch.rst publishes export() with autofunction, but the docstring described only behavior, so the rendered page listed no parameters at all. Every sibling public API has a napoleon Arguments block. Add Arguments, Keyword Arguments and Returns blocks covering the source and all twelve keyword parameters, including which of them a pre-exported source rejects.
torch_tensorrt/__init__.py never imports torch_tensorrt.executorch, and a subpackage only becomes an attribute of its parent once something imports it, so the example as published raises AttributeError on its first line.
torch_tensorrt/__init__.py never imports torch_tensorrt.executorch, so the attribute exists only after something else imports it. These two tests reach it after `import torch_tensorrt` alone, which works in a full directory run because sibling test modules import the subpackage during collection, and raises AttributeError when the file runs on its own.
Every method is serialized before the first one is rewritten, so the base64 text of all of them is resident at that point, not just the text of the method being rewritten. The old wording scoped the cost to one method.
…its own
A partitioner filed under one method while its compile specs name another passed every
check and produced a program whose delegate carried the wrong label. Measured what the
delegates actually received:
correct wiring, each method names itself
prefill {'method_name': b'prefill', 'target_device': b'cuda:0'}
decode {'method_name': b'decode', 'target_device': b'cuda:0'}
prefill given the decode name
prefill {'method_name': b'decode', 'target_device': b'cuda:0'}
decode {'method_name': b'decode', 'target_device': b'cuda:0'}
At runtime the prefill delegate then looks up the compiled decode method. This is the same
failure the shared-instance check exists to prevent, reached another way: that check
compares instances against each other and never compares an instance's declared method
name against the key it is filed under.
The helper now returns the declared name instead of a boolean, and a mismatch against the
key is rejected. The early return for fewer than two methods is gone, because a single
method can be mis-wired too.
Testing: the whole executorch test directory passes, 115 tests. Two new tests cover the
mismatch, one for a two-method mapping and one for a single method, and both fail against
the previous code. Also verified against a real CUDA partitioner that the correct
per-method shape is still accepted and that a shared instance carrying no method name is
still allowed.The guide showed both ways to produce a .pte without saying which one a reader should reach for, so the choice had to be inferred from the order the examples appeared in. It now states plainly that torch_tensorrt.save() is the default and is the right choice whenever a .pte file is all that is needed, and that torch_tensorrt.executorch.export() is the advanced path, for when the program has to be changed before it is written to disk. Each example also gets a heading, so a reader scanning the page can tell them apart.
shoumikhinforce-pushed
the
executorch-composable-export
branch
from
August 18, 2026 18:59
bbca11b to
8103e57Compareshoumikhin
commented
Aug 18, 2026
ContributorAuthor
Both suggestions are done. Documenting save() as the default. The guide showed both ways to write a
Waiting for the runtime wheel PR. That has merged, and this branch is rebased on a main |
Exporting a model whose forward argument is named after a Python builtin, for example `sum`, failed with "Staged GraphModule '<root>' changed node identities." Copying an FX graph renames any node whose name shadows a builtin, so a placeholder named `sum` comes back as `sum_1`. Staging restored the source names for placeholders only, but the rename cascades: `sum_1` was already taken by the `torch.sum` node, which the copy then pushed to `sum_2`. The identity check that runs right after saw the mismatch and refused to export. Graph copying preserves node order, so restore the source names position by position over all nodes instead of just the placeholders. Restoring every name is also what makes the identity check below pass, so that check can no longer notice a copy that reordered nodes, and the pairing this loop does would then be wrong for every node. Pair on `op` and `target`, which copying keeps verbatim, and raise if a pair does not line up. Example that used to fail and now exports: class Model(torch.nn.Module): def forward(self, sum, x): return torch.sum(x) + sum Tested by extending the staging test to a model that hits the cascade. It fails on the previous code and passes now. The rest of the export test file is unchanged and still passes.
The docs promised that the Edge program returned by `export()` shares TensorRT engine storage with the programs passed in, and that this is what keeps a multi-gigabyte engine from being duplicated. Only the first half is true for tensors, and the engine half is not true at all. What actually happens: staging never deep copies an engine object, because copying one means serializing and deserializing it. But the rewrite that follows has to put the engine somewhere ExecuTorch can read, so it decodes the engine into a byte buffer that the returned program owns. Measured on a single-engine program, the returned program retains one full private copy of the engine bytes while the source engine stays alive and independent. Weights and other tensors really are shared by reference, so the warning about not mutating a shared payload in place still stands. Doc-only change, no behavior change. Verified by measuring retained memory across the validate, stage, and rewrite sequence.
The docstring said peak memory adds about 1.33x the engine bytes, the size of the base64 text. Measured, a single-engine export peaks at about 3.7x. The missing part is the decode. Turning one engine's base64 text back into bytes holds about 2.3x that engine's bytes at once: the text has to be converted to ASCII first, then decoded, then copied into the writable buffer the tensor wraps. That cost is paid per engine while the text of every method not yet rewritten is still resident, so it hurts most with a single method, which is exactly what `save()` produces. Measured with 128 MB engines, peak added over the program itself: 1 method 3.67x 2 methods 2.50x 4 methods 1.92x Doc-only change, no behavior change.
reStructuredText needs a blank line before a section title. Without it the title is not a title at all, so the "Advanced: export()" section vanished from the page and its table of contents entry, and its underline row of 22 quote characters was rendered as ordinary text in the middle of the previous paragraph. Nothing warned about this. Parsing the page shows the section title missing before the fix and present after it, with the 22 stray quote characters gone.
…vers Three problems in the same check, in the docs, the code, and the name. The docs said sharing one partitioner instance across methods is fine as long as its compile specs name no method. That is true for a backend built for it, such as CoreML, whose own multi-method example shares one instance. It is not true for a backend that reads its method name from those specs. Passing one unnamed `CudaPartitioner` for two methods passes this check and then fails inside ExecuTorch with "Could not find method name in compile specs". Say so, and point at the one-instance-per-method form that works. The second error message in the check could never be reached. It fires when one instance whose specs name a method is reused for two methods, but the mismatch check above it already raises for whichever of the two the name does not match, and a name can match at most one. Remove the unreachable branch and the bookkeeping it needed. With that branch gone, nothing about the check is specific to sharing: it rejects any partitioner whose specs name a method other than the one it was given, including a single-method mapping. Rename it to `_reject_misnamed_partitioners` and rewrite the docstring to match. The eight partitioner tests still pass, including the one that asserts an unnamed shared instance is accepted.
The three new `save(output_format="executorch")` tests failed on a build without the Torch-TensorRT runtime extension. They mock out the lowering, but `save` checks for the runtime before it gets there, so the call raises RuntimeError instead of doing what the test asserts. Guard them with the same skipif the engine tests in this directory already use. Without the runtime the file now reports 4 passed and 3 skipped instead of 4 passed and 3 failed, and with the runtime all 7 pass.
…ine buffer The engine buffer name has to be free in three places, and only two were checked. `torch.export` lifts buffers into `state_dict` and script objects and non-persistent tensors into `constants`, and `hasattr` on the module sees neither. A program carrying a constant literally named `_trt_engine_0` got an engine buffer with that same name, so the two entries then referred to different objects under one name. Also check `constants`, the same way `state_dict` is already checked. Covered by a new test that fails on the previous code.
The multi-engine warning counted execute_engine calls but called them engines, so one engine feeding three calls reported three engines. It also said "Multi-engine .pte exports", but `export()` returns an Edge program and never writes a `.pte`. Say calls, and say what each extra call actually costs. The comment above the string arguments said `str()` on the engine slot would cost four times the engine bytes. Nothing puts raw bytes in that slot: the engine info always arrives with base64 text there, and `str()` on a str returns the same object. The slot is skipped because `engine_attr_node` takes its place in `no_op_args`, not to avoid a copy, so say that instead.
Uh oh!
There was an error while loading. Please reload this page.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for freeto join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why this PR is needed
torch_tensorrt.save(..., output_format="executorch")is the easiest way to create an ExecuTorch.ptefile. It completes the whole export pipeline immediately:That remains the recommended API for common one-step exports.
Some workflows need to stop before final ExecuTorch memory planning and serialization. For example, users may want to:
prefillanddecode;This PR adds
torch_tensorrt.executorch.export(). It returns ExecuTorch's standardEdgeProgramManager, which is the supported boundary for inspecting and customizing an Edge program before callingto_executorch().Which API should I use?
Use
torch_tensorrt.save()when you want a.ptein one step.Use
torch_tensorrt.executorch.export()when you need control before the final.pteis created:save()now delegates to the same implementation, so the simple and advanced paths do not maintain separate lowering logic.Example: TensorRT with a fallback backend
TensorRT always receives the first chance to claim its prebuilt engine nodes. Additional partitioners run afterward and can claim operations that TensorRT does not support.
Assume
cuda_partitioneris an ExecuTorch partitioner configured for your target:This can produce one ExecuTorch program containing both TensorRT and CUDA delegates.
Example: multiple methods
Assume
prefill_programanddecode_programare independent, engine-bearingExportedProgramobjects, and the two spec lists are configured for their methods:The mapping preserves method names and allows per-method partitioner and compile-spec pipelines. It does not imply that mutable state is shared between methods.
Example: transforms, metadata, and ETRecord
Assume the transform passes and Edge configuration below are provided by the caller:
Accepted inputs
The new API accepts:
torch.fx.GraphModule;torch.export.ExportedProgram;ExportedProgramobjects.A plain
torch.nn.Modulemust be compiled with Torch-TensorRT first.Safety and memory behavior
Engine rewriting changes the exported graph. The new API stages an independent copy of graph structure, signatures, state containers, constants, and metadata before rewriting anything. Rewrite and lowering failures therefore do not consume or structurally modify the caller's original program.
Tensor and TensorRT engine payload storage is treated as immutable and shared with the staged structure. This avoids deep-copying potentially multi-gigabyte engines. Custom transform passes must not mutate shared payload objects.
Metadata describing symbolic shapes is also shared rather than copied. It points back into the live shape environment the exported program is guarded on, and that environment holds fake tensors which cannot be copied. Sharing it is what keeps dynamic-shape models working; copying it would also detach the copy from the shape symbols the graph depends on.
When a lifted TensorRT engine is replaced, the old placeholder, graph-signature entry, and constant are removed. An engine feeding several call nodes is serialized once, and those nodes reuse one materialized
uint8payload buffer.Compatibility
save(..., output_format="executorch")behavior is preserved, with onedifference: a falsy non-list
partitionersorcompile_specs, such as{}, now raisesTypeErrorinstead of being read as an empty list.save()lowers a single method, so itkeeps rejecting the per-method mappings that
export()accepts..pteserialization and external.ptdpersistence are preserved.Input(shared_dims=...)remain shared.Validation
Built from source with the Torch-TensorRT C++ runtime active, against TensorRT and
ExecuTorch, and exercised on both x86_64 and aarch64 CUDA GPUs.
Checked end to end with a real compiled TensorRT engine:
export()writes a loadable.pte;export()writes a loadable.pteof the same size the previousone-step save path produced;
GraphModulestill holds its original engine nodes afterexport();save(..., output_format="executorch")still writes a.pte;save()still raisesTypeErrorfor a non-listpartitionersorcompile_specs.The whole ExecuTorch test directory passes on a CUDA machine, including the five tests
in
test_cuda_partitioner_composition.py.These tests do run on a pull request. The standard test matrix lists this directory in the
nightly lane, but the ExecuTorch runtime channel of the Linux x86_64 workflow runs
pytest executorch/on every lane exceptskipandrtx, and that channel feeds therequired gate.
One of those tests drives validation and then the rewrite, and asserts the engine was
resolved only once. Reading an engine's state serializes it and base64 encodes the
result, so resolving the same engine in both steps would double that cost on programs
whose purpose is to avoid copying large engines. The test was verified to fail when the
reuse path is removed from the rewrite.
Note on memory: each engine is serialized once, and a method's engine payloads are
released as soon as that method is rewritten, so a multi-method export does not hold
every method's payloads at once. While a method is being rewritten its engines are
resident both as the base64 text the serialization produced and as the decoded buffer,
so peak memory is roughly 2.3x that method's engine bytes, and briefly 3.3x for the
engine being decoded.
Static checks on the changed files: import ordering, formatting, lint, and type
checking.
The end-to-end run with a real compiled engine above covers the first revision of this
branch. The later commits change placeholder staging, the deepcopy memo, engine
serialization, the engine buffer name, and the per-method pass dictionary. Each of them
is covered by a unit test in
test_export.pythat was checked against its parent commit,so the test fails without the fix and passes with it. The end-to-end run will be repeated
on this head before merge.
Related work
PR #4433 forwards additional Edge-lowering options through one-step
save(). After this PR lands, #4433 can keepsave()as a thin wrapper overtorch_tensorrt.executorch.export()instead of extending the previous private lowering path.