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

[Conarnar]

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).

[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]

Should fix: when no example inputs are given, this builds an in_spec that says the model takes zero inputs, while the graph still has its placeholders.

torch_tensorrt.save(..., retrace=False) passes arg_inputs=(), not None (_compile.py:831), and the API tells users inputs are not needed on this path (_compile.py:1012). So example_args is (), in_spec ends up with 0 leaves, and nothing compares it against the placeholders.

I reproduced this on torch 2.11. The ExportedProgram builds, torch.export.save and torch.export.load both succeed, and it only fails much later:

• ep.module()(x) raises Trying to flatten user inputs with exported input tree spec ...
• output_format="executorch" fails inside to_edge with leaves has length 2 but the spec refers to a pytree that holds 0 items

Before this PR the same call failed immediately with 'CodeGen' object has no attribute 'pytree_info', so this turns a loud failure into a silently broken artifact.

It is reachable in the normal flow: lowering/_buffer_lifting.py:166 does gm.graph.set_codegen(torch.fx.graph.CodeGen()) for any model with mutated buffers, and nothing puts the pytree codegen back.

One more thing: the comment above says this fallback matches retrace=True. In the no-input case it does not. retrace=True goes through torch.export.export(module, args=()) and raises TypeError: missing a required argument.

Suggestion: when no inputs are supplied, build the spec positionally from the graph placeholders instead of from an empty tuple, and assert that in_spec.num_leaves matches the placeholder count. The supplied-input path is already correct, including kwargs, so it does not need to change.