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

[Conarnar]

Description

torch_tensorrt.save(..., retrace=False) runs the legacy exporter, which calls inline_torch_modules to collapse submodules into the parent graph. That inlining wired submodule inputs by placeholder name. When a submodule input placeholder's name collides with an unrelated node in the parent graph, graph_copy maps the submodule body onto that node through a pre-seeded val_map — so the original producer is consumed and disappears from the graph entirely. This is silent: it surfaces only later as a delegate (or TensorRT engine) reporting the wrong argument count at runtime, or as a wrong numerical result rather than a crash.

A related failure hit multi-output submodules: a mis-wired input orphaned one submodule output, dead-code elimination then pruned it, and a downstream consumer (in the hybrid case, a TensorRT engine) was left short an output.

This path is on by default for composable hybrid TensorRT + CUDA .pte export (retrace=False), so it matters for that flow.

Fix

• Wire submodule inputs positionally from the call_module node's args (the authoritative ordered list of what the node actually consumes) instead of matching by placeholder name, and remove the name-based get_duplicate_nodes matching. Names are incidental and can legally collide after graph surgery; there is no correct name-matching here.
• create_trt_exp_program kwargs ordering:pytree flattens kwargs in dict-insertion order while the graph consumes placeholders positionally, so kwargs passed out of signature order silently bound each value to the wrong input (the num_leaves assert only checks the count). Reorder kwargs into placeholder order before flattening.
• persistent=True on lifted BUFFER specs: torch >= 2.3 asserts an explicit persistent flag on BUFFER-kind InputSpecs. A buffer only reaches this branch when it is in state_dict, which excludes non-persistent buffers by construction, so True is always correct here.

Note: the kwargs reorder handles flat, top-level kwargs passed out of signature order. An out-of-order nested kwarg (one whose value is itself a dict/list that flattens to multiple placeholders) falls through the set(kwarg_targets) == set(example_kwargs) guard unchanged — a pre-existing, narrower edge this PR doesn't address.

Tests

tests/py/dynamo/models/test_exporter_inlining.py (all CPU-only, no GPU/TensorRT build required):

• test_inline_torch_modules_preserves_all_submodule_outputs — a submodule placeholder is deliberately named to collide with a parent input, so the test fails on the old name-matching path and passes with positional wiring.
• test_inline_torch_modules_wires_inputs_by_position — pins the positional-wiring behavior / no leaked placeholder.
• test_create_trt_exp_program_reorders_kwargs_to_placeholder_order — exports a - b, passes kwargs reversed, asserts in_spec records placeholder order and the result is 7.0 (not -7.0).
• test_lift_sets_persistent_true_on_buffer_spec — pins the persistent=True flag on lifted BUFFER specs.
• test_create_trt_exp_program_rebuilds_in_spec_without_inputs — the no-input in_spec path.

[shoumikhin]

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:

• This is effectively a prerequisite for hybrid TensorRT plus CUDA .pte export via
  retrace=False, which is the default in feat(executorch): expose composable Edge export API #4440. I checked and there is no merge
  conflict between the two, they touch different files.
• The multi-output preservation fix matters directly for fix(executorch): support KV-cache aliased I/O in the TensorRT delegate #4445: the failure mode you
  describe (an orphaned output pruned by dead-code elimination, leaving a delegate short
  an argument) is the same class of arg-count error that PR reports, so it would be
  worth confirming whether some of that 0x12 mismatch is actually this bug rather than
  missing aliasing metadata. If the two overlap, landing this first may shrink what
  fix(executorch): support KV-cache aliased I/O in the TensorRT delegate #4445 has to do.

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

[Conarnar]

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.

[shoumikhin]

The new branch can silently return wrong numbers when keyword arguments are passed in a different order than the placeholders.

_exporter.py:402 builds the input spec from the caller's kwargs dict, and pytree (PyTorch's helper for flattening nested inputs) keeps dict keys in insertion order, while the graph consumes the flattened values by position. Same graph computing a - b with a=10, b=3:

kwargs {'a': 10, 'b': 3} (placeholder order) -> 7.0 correct
kwargs {'b': 3, 'a': 10} (reversed) -> -7.0 wrong, no error

The num_leaves assert at 404 does not catch it, because it compares counts (2 == 2) and not which value goes where. This is not a regression, since on main this branch raised AttributeError before it got this far, but it is a new sharp edge. Sorting kwargs into placeholder order, or asserting the spec's key order matches the placeholder names, would close it.

[shoumikhin]

test_inline_torch_modules_preserves_all_submodule_outputs (test_exporter_inlining.py:63) passes both with and without your fix, so right now it does not test the fix.

The submodule placeholders are named a and b while the parent uses x and y, so there is no name collision and the old code takes its already-correct path. Renaming the submodule's first placeholder from a to y makes the test discriminate:

as written now, on this PR: PASS on main: PASS (proves nothing)
after a -> y, on this PR: PASS on main: TypeError: forward() missing 1 required positional argument: 'b'

Worth doing, since this is the test covering the multi-output half of the change. The name-collision test in the same file does genuinely fail on main, so the core claim is well covered, it is only this one that is inert.

One thing that may be worth adding to the description: the bug this fixes is a wrong answer, not a crash. With a parent node named getitem and a submodule placeholder also named getitem, main returns -17.0 where this PR returns the correct 17.0.

[shoumikhin]

In support of the persistent= line at _exporter.py:242, in case anyone asks you to split it out as unrelated: it is required, not optional.

InputSpec(kind=InputKind.BUFFER, arg=TensorArgument(name="b"), target="b")
-> AssertionError: Failed to specify persistent flag on BUFFER.

Main builds InputSpec with a variable kind and no persistent at line 233, and grep persistent on main finds nothing, so any buffer reaching that branch asserts. That makes the fix belong with a change that lets buffers reach it. The only gap is coverage: nothing in the new test file uses a BUFFER-kind spec, so a small case with a registered buffer would pin this line.