Skip to content

fix(dynamo): correct legacy exporter (retrace=False) submodule inlining for hybrid graphs - #4446

Open
Conarnar wants to merge 1 commit into
pytorch:mainfrom
Conarnar:fix/executorch-retrace-false-hybrid
Open

fix(dynamo): correct legacy exporter (retrace=False) submodule inlining for hybrid graphs#4446
Conarnar wants to merge 1 commit into
pytorch:mainfrom
Conarnar:fix/executorch-retrace-false-hybrid

Conversation

@Conarnar

Copy link
Copy Markdown
Contributor

torch_tensorrt.save(retrace=False) uses the legacy dynamo exporter, which inlines the partitioned _run_on_gpu (non-TensorRT) submodules back into the graph before building an ExportedProgram. For a hybrid graph interleaving TensorRT engines with a CUDA/pytorch delegated op, inline_torch_modules wired each submodule's inputs by MATCHING placeholder names to graph nodes (get_duplicate_nodes). Name matching binds an input to a same-named but unrelated node on a collision (e.g. a submodule input placeholder name-matching a different engine's getitem), which:

  • rewires a consumer to the wrong producer and orphans the real one; the orphan is then pruned by dead-code elimination, leaving a delegate short an output at runtime (an aliased engine reports "expected N args, got N-1"); and
  • for a submodule mixing graph-input and computed-intermediate inputs, leaks the computed intermediates as spurious graph placeholders (misclassified USER_INPUTs).

Wire submodule inputs POSITIONALLY from the call_module args (gm_node.args, which is authoritative) instead of by name: let graph_copy create a fresh placeholder for each submodule input, then rewire each to submodule_inputs[i] by position and erase it. Drop get_duplicate_nodes (now unused).

Also fix two torch-version-compat gaps this path hits on recent torch:

  • lift(): pass an explicit persistent= flag on BUFFER InputSpecs (required since 2.3).
  • create_trt_exp_program(): an inlined GraphModule may carry a plain fx.CodeGen (no pytree_info); fall back to specs rebuilt from the example inputs + graph outputs.

With these, retrace=False export of a hybrid TensorRT+CUDA program is bit-identical to retrace=True (validated on a 2-layer int4 MoE decode: per-step argmax + logits match).

Tests: tests/py/dynamo/models/test_exporter_inlining.py -- positional input wiring under a name collision, and multi-output preservation (GPU-free fx unit tests).

@meta-cla meta-cla Bot added the cla signed label Jul 30, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Jul 30, 2026
@github-actions
github-actions Bot requested a review from cehongwang July 30, 2026 00:09
@shoumikhin

Copy link
Copy Markdown
Contributor

Reviewed this one carefully because #4440 lands right on top of it: the composable
ExecuTorch export path defaults to retrace=False, so it goes through
create_trt_exp_program and inline_torch_modules on every call. Direction and
approach both look right to me, and I reproduced the bug locally to be sure.

The bug is real, and it is worse than mis-wiring

I built a small parent graph whose submodule input placeholder name collides with an
unrelated node, then ran inline_torch_modules on current main. The colliding node was
not merely re-wired, it disappeared from the graph entirely: graph_copy mapped the
submodule's body onto it through the pre-seeded val_map, so the original producer was
consumed. Silent node loss in an exporter is about as bad as it gets, and it is
invisible until a delegate reports the wrong arg count at runtime.

Your two unit tests are genuine regression tests, not decoration. I ran them both ways:

  • with this patch: 3 passed
  • with _exporter.py reverted to main: test_inline_torch_modules_wires_inputs_by_position
    fails on assert 3 == 2, the spurious leaked placeholder

That is exactly the property that matters, and it is nice that they need neither a GPU
nor a TensorRT build.

On the approach

Wiring positionally from gm_node.args is the right call. args is the authoritative
ordered list of what the call_module node actually consumes; placeholder names are
incidental and, as this shows, can collide with anything. Deleting get_duplicate_nodes
rather than trying to make name matching smarter is also the right instinct: there is no
version of name matching that is correct here, because two unrelated nodes are allowed
to share a name after graph surgery.

Letting graph_copy create a fresh placeholder and then erasing it is slightly more
churn than pre-seeding, but it is obviously correct, which I would take over clever.

Questions on the two compat fixes

These are separate from the inlining fix and I would like to understand them a bit
better, mostly because they are easy to get subtly wrong.

1. persistent=True on BUFFER specs. The reasoning in the comment (a buffer only
reaches this branch when it is in state_dict, so it is persistent) reads as correct to
me. Worth double checking one case: a buffer registered with persistent=False that
some earlier pass has nonetheless placed in state_dict. If that cannot happen, could
the comment say so directly, so a future reader does not have to re-derive it?

2. The _PyTreeCodeGen fallback. in_spec rebuilt from (example_args, example_kwargs) looks equivalent to what a real pytree_info.in_spec would carry. My
question is about out_spec = pytree.tree_flatten(tuple(output_nodes))[1], which always
produces a flat tuple spec even when the module's real out_spec describes a nested
structure.

I tried to reach that path and could not, which is reassuring but leaves the question
open. Both a dict-returning model and a nested-tuple model still carry a
_PyTreeCodeGen after inlining, so they take the normal branch and their nesting is
preserved exactly:

dict_output    -> codegen _PyTreeCodeGen, out_spec TreeSpec(dict, ['logits','aux'], [*, *])
nested_tuple   -> codegen _PyTreeCodeGen, out_spec TreeSpec(tuple, None, [*, TreeSpec(tuple, ...)])

So the fallback did not fire for either, and in both cases the exported program returned
the original structure.

That leaves two possibilities, and it would help to know which you saw. If the fallback
only triggers for graphs whose outputs really are a flat tuple, a one-line comment saying
so would save the next reader the same investigation. If it can trigger for a nested-output
graph, then that case would silently flatten and disagree with retrace=True, and
deriving out_spec from the pre-inlining module would be safer. What was the graph shape
that led you to add this?

Minor: the comment says "torch>=2.13" while the persistent comment says "since 2.3".
Worth confirming both version numbers, since the first looks like it might be a typo.

Composition with the ExecuTorch work

For anyone tracking how these fit together:

Thanks for tracking this down. The name-collision path has probably been quietly wrong
for a while.

@Conarnar
Conarnar force-pushed the fix/executorch-retrace-false-hybrid branch from d504244 to 60e39a4 Compare August 1, 2026 02:20
@Conarnar

Conarnar commented Aug 1, 2026

Copy link
Copy Markdown
Contributor Author

Thank you for your suggestions. I have added clarifications in the comments.

Minor: the comment says "torch>=2.13" while the persistent comment says "since 2.3". Worth confirming both version numbers, since the first looks like it might be a typo.

Verified that persistent=True was required since 2.3 but surfaced as a result of buffers.
Could not verify version of _PyTreeCodeGen but surfaced as a result of hybrid TRT + CUDA backends. Comment was removed.

…ng for hybrid graphs

torch_tensorrt.save(retrace=False) uses the legacy dynamo exporter, which inlines the
partitioned _run_on_gpu (non-TensorRT) submodules back into the graph before building an
ExportedProgram. For a hybrid graph interleaving TensorRT engines with a CUDA/pytorch
delegated op, inline_torch_modules wired each submodule's inputs by MATCHING placeholder
names to graph nodes (get_duplicate_nodes). Name matching binds an input to a same-named
but unrelated node on a collision (e.g. a submodule input placeholder name-matching a
different engine's getitem), which:
  - rewires a consumer to the wrong producer and orphans the real one; the orphan is then
    pruned by dead-code elimination, leaving a delegate short an output at runtime (an
    aliased engine reports "expected N args, got N-1"); and
  - for a submodule mixing graph-input and computed-intermediate inputs, leaks the
    computed intermediates as spurious graph placeholders (misclassified USER_INPUTs).

Wire submodule inputs POSITIONALLY from the call_module args (gm_node.args, which is
authoritative) instead of by name: let graph_copy create a fresh placeholder for each
submodule input, then rewire each to submodule_inputs[i] by position and erase it. Drop
get_duplicate_nodes (now unused).

Also fix two torch-version-compat gaps this path hits on recent torch:
  - lift(): pass an explicit persistent= flag on BUFFER InputSpecs (required since 2.3).
  - create_trt_exp_program(): an inlined GraphModule may carry a plain fx.CodeGen (no
    pytree_info); fall back to specs rebuilt from the example inputs + graph outputs.

With these, retrace=False export of a hybrid TensorRT+CUDA program is bit-identical to
retrace=True (validated on a 2-layer int4 MoE decode: per-step argmax + logits match).

Tests: tests/py/dynamo/models/test_exporter_inlining.py -- positional input wiring under a
name collision, and multi-output preservation (GPU-free fx unit tests).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants