🐛 [Bug] inline_torch_modules matches submodule placeholders to parent nodes by name, mis-wiring the graph

[SrivastavaKshitij]

Bug Description

When inline_torch_modules inlines a torch-executed (_run_on_gpu_*) submodule into the
parent graph, it decides which parent node each submodule placeholder corresponds to by
comparing node names. Node names are only unique within one graph, so a name that the
submodule uses for a placeholder can also belong to a completely unrelated node of the parent
graph. The pairing takes that coincidence for identity.

py/torch_tensorrt/dynamo/_exporter.py :: get_duplicate_nodes (line numbers from the build
tested, see Environment):

    submodule_placeholder_inputs = [
        node for node in submodule.graph.nodes if node.op == "placeholder"
    ]
    submodule_input_node_names = [node.name for node in submodule_placeholder_inputs]
    gm_node_names = [node.name for node in gm.graph.nodes]
    submodule_duplicate_inputs = [
        node for node in submodule_placeholder_inputs if node.name in gm_node_names
    ]
    gm_duplicate_inputs = [
        node for node in gm.graph.nodes if node.name in submodule_input_node_names
    ]

inline_torch_modules zips those two lists into the val_map handed to graph_copy, so the
copied submodule body reads gm_duplicate_inputs[i] wherever it should read the i-th argument
of the call_module node:

                submodule_duplicate_inputs, gm_duplicate_inputs = get_duplicate_nodes(
                    gm, submodule
                )
                assert len(submodule_duplicate_inputs) == len(gm_duplicate_inputs)
                # Avoid creating new copies of duplicate inputs by creating a mapping
                val_map = {}
                for i in range(len(submodule_duplicate_inputs)):
                    val_map[submodule_duplicate_inputs[i]] = gm_duplicate_inputs[i]

Two things make this worse than a mis-ordered mapping:

1. The correct mapping is already in hand. A submodule's placeholders are, in order, the
   args of its call_module node. inline_torch_modules even computes them
   (submodule_inputs = gm_node.args) a few lines earlier -- but only uses them inside the
   len(submodule_duplicate_inputs) == 0 branch.
2. One collision disables placeholder replacement for the whole submodule, because that
   replacement is guarded all-or-nothing:

                # Get their references (since we copied) in the parent graph (gm)
                if len(submodule_duplicate_inputs) == 0:

So the observable outcome depends on where the colliding parent node happens to sit:

• Colliding node defined AFTER the call_module node (what the script below builds): the
  inlined body is inserted before the node it now reads, and the graph is no longer
  topologically ordered. transform aborts with
  RuntimeError: Argument 'getitem_1' of Node 'mul_1' was used before it has been defined!
• Colliding node defined BEFORE the call_module node: no error at all. The graph passes
  lint(), the inlined body silently reads the wrong node, and the node that was actually
  passed as the argument is left with num_users=0. This is a silent wrong-answer case.
• Partial collision (a submodule with several placeholders, only some of whose names
  collide): the duplicate list is non-empty, so replacement is skipped entirely, and the
  non-colliding placeholders are copied into the middle of the parent graph and left dangling
  while their real arguments are orphaned at num_users=0.

To Reproduce

docker run --rm --gpus all --ipc=host -v "$PWD":/w -w /w \
        nvcr.io/nvidia/pytorch:26.07-py3 python repro.py

repro.py

repro.py

import sys
import traceback

import torch
import torch_tensorrt
from torch import nn
from torch_tensorrt.dynamo._exporter import get_duplicate_nodes, transform

COLLIDING_NAME = "getitem_1"


def build_submodule() -> torch.fx.GraphModule:
    """A torch-executed submodule whose single placeholder is named `getitem_1`."""
    graph = torch.fx.Graph()
    placeholder = graph.placeholder(COLLIDING_NAME)
    doubled = graph.call_function(torch.mul, (placeholder, 2.0))
    graph.output(doubled)
    return torch.fx.GraphModule(nn.Module(), graph)


def build_parent(submodule: torch.fx.GraphModule) -> torch.fx.GraphModule:
    """A parent graph that also holds a node named `getitem_1`, defined later."""
    root = nn.Module()
    root.add_module("_run_on_gpu_0", submodule)
    graph = torch.fx.Graph()
    x = graph.placeholder("x")
    real_input = graph.call_function(torch.mul, (x, 3.0))
    submodule_call = graph.call_module("_run_on_gpu_0", (real_input,))
    collision = graph.create_node("call_function", torch.neg, (x,), name=COLLIDING_NAME)
    graph.output(graph.call_function(torch.add, (submodule_call, collision)))
    return torch.fx.GraphModule(root, graph)


def main() -> int:
    print("torch", torch.__version__, "torch_tensorrt", torch_tensorrt.__version__)

    submodule = build_submodule()
    gm = build_parent(submodule)
    print(f"parent graph:\n{gm.graph}")
    print(f"submodule graph:\n{submodule.graph}")

    submodule_call = next(node for node in gm.graph.nodes if node.op == "call_module")
    real_arguments = list(submodule_call.args)
    submodule_duplicates, gm_duplicates = get_duplicate_nodes(gm, submodule)
    print(f"call_module args (the correct mapping): {real_arguments}")
    print(f"get_duplicate_nodes paired {submodule_duplicates} with {gm_duplicates}")

    paired_wrong_node = (
        len(gm_duplicates) == 1
        and gm_duplicates[0].name == COLLIDING_NAME
        and gm_duplicates[0] not in real_arguments
    )
    print(f"placeholder paired with a node that is not its argument: {paired_wrong_node}")

    formatted = ""
    try:
        patched = transform(gm)
        print(f"--- transform succeeded ---\n{patched.graph}")
    except BaseException:
        formatted = traceback.format_exc()
        print(f"--- transform raised ---\n{formatted}")

    graph_invalid = "used before it has been defined" in formatted
    print(f"transform produced a graph that is not topologically ordered: {graph_invalid}")

    reproduced = paired_wrong_node and graph_invalid
    print(f"reproduced: {reproduced} (constructed collision)")
    return 0 if reproduced else 1


if __name__ == "__main__":
    sys.exit(main())

output

torch 2.13.0a0+9186a08b2c.nv26.07 torch_tensorrt 2.14.0a0
parent graph:
graph():
    %x : [num_users=2] = placeholder[target=x]
    %mul : [num_users=1] = call_function[target=torch.mul](args = (%x, 3.0), kwargs = {})
    %_run_on_gpu_0 : [num_users=1] = call_module[target=_run_on_gpu_0](args = (%mul,), kwargs = {})
    %getitem_1 : [num_users=1] = call_function[target=torch.neg](args = (%x,), kwargs = {})
    %add : [num_users=1] = call_function[target=torch.add](args = (%_run_on_gpu_0, %getitem_1), kwargs = {})
    return add
submodule graph:
graph():
    %getitem_1 : [num_users=1] = placeholder[target=getitem_1]
    %mul : [num_users=1] = call_function[target=torch.mul](args = (%getitem_1, 2.0), kwargs = {})
    return mul
call_module args (the correct mapping): [mul]
get_duplicate_nodes paired [getitem_1] with [getitem_1]
placeholder paired with a node that is not its argument: True
--- transform raised ---
Traceback (most recent call last):
  File "/w/repro.py", line 97, in main
    patched = transform(gm)
              ^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch_tensorrt/dynamo/_exporter.py", line 92, in transform
    gm.graph.eliminate_dead_code()
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph.py", line 2249, in eliminate_dead_code
    self.lint()
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph.py", line 2156, in lint
    check_arg(arg, node)
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph.py", line 2141, in check_arg
    raise RuntimeError(
RuntimeError: Argument 'getitem_1' of Node 'mul_1' was used before it has been defined! Please check that Nodes in the graph are topologically ordered
graph():
    %x : [num_users=2] = placeholder[target=x]
    %mul : [num_users=0] = call_function[target=torch.mul](args = (%x, 3.0), kwargs = {})
    %mul_1 : [num_users=1] = call_function[target=torch.mul](args = (%getitem_1, 2.0), kwargs = {})
    %getitem_1 : [num_users=2] = call_function[target=torch.neg](args = (%x,), kwargs = {})
    %add : [num_users=1] = call_function[target=torch.add](args = (%mul_1, %getitem_1), kwargs = {})
    return add

transform produced a graph that is not topologically ordered: True
reproduced: True (constructed collision)

Expected behavior

node's args**, never by name. submodule.graph placeholders are in the same order as
gm_node.args, so the mapping is:

val_map = dict(zip(submodule_placeholder_inputs, gm_node.args))

With that, graph_copy wires the body to the real arguments directly, no placeholders are
copied into the parent graph, and the "duplicate node" concept -- and the all-or-nothing
len(submodule_duplicate_inputs) == 0 guard around placeholder replacement -- can go away
entirely.

Separately, and worth fixing even if the pairing rule stays: the placeholder-replacement step
should be per-placeholder rather than skipped wholesale when the duplicate list is non-empty.
As written, a partial name collision leaves stray placeholders in the parent graph and orphans
the real arguments without any error.

Environment

Build information about Torch-TensorRT can be found by turning on debug messages

• Pytorch NGC container : 26.07-py3

[kvnloo]

I think the strongest fix is to remove node-name matching from placeholder identity entirely.

For a call_module node, the correct mapping is already available structurally:

submodule placeholders in order -> call_module.args in order.

Names are graph-local presentation details and should not participate in cross-graph identity.

The regression should cover all three failure modes from the issue:

• collision with a parent node defined after the call;
• collision with one defined before the call;
• partial collision with several placeholders.

For every case:

• graph remains topologically valid;
• each inlined placeholder reads the actual call argument;
• no real argument becomes orphaned;
• outputs match the pre-inline graph.

If unowned, this seems like a clean exporter fix.