Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Conarnar wants to merge 1 commit into
base: main
Choose a base branch
from
+175 −60
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
108 changes: 48 additions & 60 deletions py/torch_tensorrt/dynamo/_exporter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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,
),
)
]
Expand Down
127 changes: 127 additions & 0 deletions tests/py/dynamo/models/test_exporter_inlining.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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))
Loading