From 464d756af3ae69dc35fc79c6e0e0c9c64da27c51 Mon Sep 17 00:00:00 2001 From: Conan Jeffrey Truong Date: Wed, 29 Jul 2026 16:09:56 -0700 Subject: [PATCH] fix(dynamo): correct legacy exporter (retrace=False) submodule inlining 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). --- py/torch_tensorrt/dynamo/_exporter.py | 108 +++++++-------- .../dynamo/models/test_exporter_inlining.py | 127 ++++++++++++++++++ 2 files changed, 175 insertions(+), 60 deletions(-) create mode 100644 tests/py/dynamo/models/test_exporter_inlining.py diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index dacbea140a..059f32075d 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -4,6 +4,7 @@ from typing import Any, Dict, Optional, Sequence, Tuple import torch +import torch.utils._pytree as pytree from torch._export.non_strict_utils import make_constraints from torch._guards import detect_fake_mode from torch._library.fake_class_registry import FakeScriptObject @@ -19,6 +20,7 @@ OutputSpec, TensorArgument, ) +from torch.fx.graph import _PyTreeCodeGen from torch_tensorrt._features import ENABLED_FEATURES from torch_tensorrt.dynamo.runtime._TorchTensorRTModule import ENGINE_IDX, NAME_IDX @@ -230,6 +232,12 @@ def lift( kind=input_kind, arg=input_spec_arg, target=node.target, + # torch>=2.3 requires an explicit persistent flag on BUFFER + # specs. state_dict() excludes non-persistent buffers by + # construction, so any buffer reaching this in-state_dict + # branch is persistent (non-persistent buffers take the + # not-in-state_dict path above and are lifted as constants). + persistent=(True if input_kind == InputKind.BUFFER else None), ), ) non_user_input_idx += 1 @@ -245,29 +253,6 @@ def lift( return gm, graph_signature, state_dict, constants -def get_duplicate_nodes( - gm: torch.fx.GraphModule, submodule: torch.fx.GraphModule -) -> Tuple[Sequence[Any], Sequence[Any]]: - """ - We check if there are duplicate nodes when we copy submodule graph into gm. - Handle the case where the subgraph input placeholders are same as - gm placeholders. This happens when the first submodule in the graph is - a pytorch submodule - """ - 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 - ] - return submodule_duplicate_inputs, gm_duplicate_inputs - - def inline_torch_modules(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: """ Inline a submodule within the parent graph (gm). All `call_module` nodes @@ -285,43 +270,31 @@ def inline_torch_modules(gm: torch.fx.GraphModule) -> torch.fx.GraphModule: # or a placeholder of the main graph submodule_inputs = gm_node.args - 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] - - # Copy all nodes in the submodule into gm and - # store the output node of this submodule which is now present in gm + # Copy the submodule's nodes into gm, then wire its inputs POSITIONALLY. + # + # We deliberately do NOT pre-seed val_map by matching submodule input + # placeholders to gm nodes by NAME. Name matching silently binds an + # input to the WRONG node when names collide (e.g. a _run_on_gpu input + # placeholder whose name matches 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. It also leaks a mixed submodule's + # computed-intermediate inputs as spurious graph inputs. gm_node.args is + # the authoritative, ordered list of the real inputs, so we let + # graph_copy create a fresh (auto-renamed on collision) placeholder for + # every submodule input and rewire each to submodule_inputs[i] by + # position, then erase it. + val_map: Dict[Any, Any] = {} submodule_output = gm.graph.graph_copy(submodule.graph, val_map) - # Get their references (since we copied) in the parent graph (gm) - if len(submodule_duplicate_inputs) == 0: - submodule_placeholder_input_names = [ - node.name - for node in submodule.graph.nodes - if node.op == "placeholder" - ] - gm_added_placeholder_inputs = [ - node - for node in gm.graph.nodes - if node.name in submodule_placeholder_input_names - ] - - assert len(submodule_inputs) == len(gm_added_placeholder_inputs) - - # Replace the added placeholder inputs with original inputs to this submodule node - for idx in range(len(gm_added_placeholder_inputs)): - gm_added_placeholder_inputs[idx].replace_all_uses_with( - submodule_inputs[idx] - ) - - # Erase the placeholder input nodes in the gm - for idx in range(len(gm_added_placeholder_inputs)): - gm.graph.erase_node(gm_added_placeholder_inputs[idx]) + submodule_placeholders = [ + node for node in submodule.graph.nodes if node.op == "placeholder" + ] + assert len(submodule_placeholders) == len(submodule_inputs) + for idx, submodule_placeholder in enumerate(submodule_placeholders): + copied_placeholder = val_map[submodule_placeholder] + copied_placeholder.replace_all_uses_with(submodule_inputs[idx]) + gm.graph.erase_node(copied_placeholder) # Replace the pytorch submodule node (call_module) with the inlined subgraph output # Special handling when submodule returns multiple outputs (tuple) @@ -403,14 +376,29 @@ def create_trt_exp_program( input_specs=input_specs, output_specs=output_specs ) + # A hybrid TRT+CUDA GraphModule from dynamo.compile carries a plain fx.CodeGen + # (no pytree_info): the module already returns a flat tuple, so rebuild + # in_spec/out_spec from the example inputs and the flat graph outputs. This + # matches retrace=True, which traces the same flat module and likewise cannot + # re-nest -- so the fallback never diverges from it. + codegen = gm.graph._codegen + if isinstance(codegen, _PyTreeCodeGen): + in_spec = codegen.pytree_info.in_spec + out_spec = codegen.pytree_info.out_spec + else: + example_args = tuple(arg_inputs) if arg_inputs is not None else () + example_kwargs = kwarg_inputs or {} + in_spec = pytree.tree_flatten((example_args, example_kwargs))[1] + out_spec = pytree.tree_flatten(tuple(output_nodes))[1] + module_call_graph = [ ModuleCallEntry( "", ModuleCallSignature( inputs=[], outputs=[], - in_spec=gm.graph._codegen.pytree_info.in_spec, - out_spec=gm.graph._codegen.pytree_info.out_spec, + in_spec=in_spec, + out_spec=out_spec, ), ) ] diff --git a/tests/py/dynamo/models/test_exporter_inlining.py b/tests/py/dynamo/models/test_exporter_inlining.py new file mode 100644 index 0000000000..0793c24213 --- /dev/null +++ b/tests/py/dynamo/models/test_exporter_inlining.py @@ -0,0 +1,127 @@ +"""Unit tests for the legacy dynamo exporter's submodule inlining +(torch_tensorrt.dynamo._exporter). These run on plain fx graphs and need neither a +GPU nor a TensorRT build.""" + +import operator + +import pytest +import torch +from torch_tensorrt.dynamo._exporter import inline_torch_modules + + +@pytest.mark.unit +def test_inline_torch_modules_wires_inputs_by_position(): + """inline_torch_modules must wire a _run_on_gpu submodule's inputs from the + call_module args by POSITION, not by matching placeholder names to graph nodes. + + Regression: the old name-matching path bound a submodule input to a same-named + but unrelated graph node, rewiring a consumer to the wrong producer (and, for a + submodule mixing graph-input and computed-intermediate inputs, leaking the + latter as spurious graph placeholders). Here the submodule's first input + placeholder is named "y", colliding with the parent's second input "y" even + though the first *argument* is the parent's "x"; positional wiring must ignore + the collision. Subtraction makes the input order observable. + """ + # Submodule: out = first - second. First placeholder deliberately named "y". + sub_graph = torch.fx.Graph() + first = sub_graph.placeholder("y") + second = sub_graph.placeholder("z") + sub_graph.output(sub_graph.call_function(torch.sub, (first, second))) + submodule = torch.fx.GraphModule(torch.nn.Module(), sub_graph) + + # Parent: inputs (x, y); call _run_on_gpu_0(x, y) -> expected x - y. + parent_graph = torch.fx.Graph() + x = parent_graph.placeholder("x") + y = parent_graph.placeholder("y") + root = torch.nn.Module() + root.add_module("_run_on_gpu_0", submodule) + call = parent_graph.call_module("_run_on_gpu_0", (x, y)) + parent_graph.output(call) + parent = torch.fx.GraphModule(root, parent_graph) + + n_placeholders_before = sum(1 for n in parent.graph.nodes if n.op == "placeholder") + + inline_torch_modules(parent) + parent.recompile() + + # No spurious placeholders leaked by the inlining. + assert ( + sum(1 for n in parent.graph.nodes if n.op == "placeholder") + == n_placeholders_before + ) + # No call_module node survives (the submodule was inlined). + assert not any(n.op == "call_module" for n in parent.graph.nodes) + # Positional wiring: first input <- x, second input <- y, so out == x - y. + out = parent(torch.tensor(5.0), torch.tensor(3.0)) + assert torch.allclose(out, torch.tensor(2.0)) + + +@pytest.mark.unit +def test_inline_torch_modules_preserves_all_submodule_outputs(): + """A multi-output _run_on_gpu submodule must keep every output wired to its + consumer after inlining. Regression: a mis-wired input orphaned one submodule + output, which dead-code elimination then pruned, leaving a downstream consumer + (or, in the hybrid case, a TensorRT engine) short an output at runtime. + """ + # Submodule returns (a + b, a - b); both outputs are consumed downstream. + sub_graph = torch.fx.Graph() + a = sub_graph.placeholder("a") + b = sub_graph.placeholder("b") + add = sub_graph.call_function(torch.add, (a, b)) + sub = sub_graph.call_function(torch.sub, (a, b)) + sub_graph.output((add, sub)) + submodule = torch.fx.GraphModule(torch.nn.Module(), sub_graph) + + parent_graph = torch.fx.Graph() + x = parent_graph.placeholder("x") + y = parent_graph.placeholder("y") + root = torch.nn.Module() + root.add_module("_run_on_gpu_0", submodule) + call = parent_graph.call_module("_run_on_gpu_0", (x, y)) + o0 = parent_graph.call_function(operator.getitem, (call, 0)) + o1 = parent_graph.call_function(operator.getitem, (call, 1)) + # Consume both outputs: (a+b) * (a-b). + parent_graph.output(parent_graph.call_function(torch.mul, (o0, o1))) + parent = torch.fx.GraphModule(root, parent_graph) + + inline_torch_modules(parent) + parent.recompile() + + # (x+y)*(x-y) == x^2 - y^2 ; with x=5, y=3 -> 25 - 9 = 16. + out = parent(torch.tensor(5.0), torch.tensor(3.0)) + assert torch.allclose(out, torch.tensor(16.0)) + + +@pytest.mark.unit +def test_inline_torch_modules_computed_intermediate_inputs(): + """A _run_on_gpu submodule whose inputs are computed intermediates (not top-level + graph placeholders, and not name-matching any graph node) must inline correctly. + This is the case the old zero-duplicate path handled; positional wiring preserves + it (and it is the shape that leaked spurious placeholders in the mixed case). + """ + # Submodule: out = m + n. Names don't collide with anything in the parent. + sub_graph = torch.fx.Graph() + m = sub_graph.placeholder("m") + n = sub_graph.placeholder("n") + sub_graph.output(sub_graph.call_function(torch.add, (m, n))) + submodule = torch.fx.GraphModule(torch.nn.Module(), sub_graph) + + # Parent: x -> c0 = x*2, c1 = x+1; call _run_on_gpu_0(c0, c1) -> c0 + c1. + parent_graph = torch.fx.Graph() + x = parent_graph.placeholder("x") + c0 = parent_graph.call_function(torch.mul, (x, 2)) + c1 = parent_graph.call_function(torch.add, (x, 1)) + root = torch.nn.Module() + root.add_module("_run_on_gpu_0", submodule) + call = parent_graph.call_module("_run_on_gpu_0", (c0, c1)) + parent_graph.output(call) + parent = torch.fx.GraphModule(root, parent_graph) + + inline_torch_modules(parent) + parent.recompile() + + # No spurious placeholders leaked; the computed intermediates stay in-graph. + assert sum(1 for node in parent.graph.nodes if node.op == "placeholder") == 1 + # out = (x*2) + (x+1); x=5 -> 10 + 6 = 16. + out = parent(torch.tensor(5.0)) + assert torch.allclose(out, torch.tensor(16.0))