diff --git a/py/torch_tensorrt/dynamo/_compiler.py b/py/torch_tensorrt/dynamo/_compiler.py index 3ebac2a21f..2b5695a9fd 100644 --- a/py/torch_tensorrt/dynamo/_compiler.py +++ b/py/torch_tensorrt/dynamo/_compiler.py @@ -5,7 +5,6 @@ import os import platform import warnings - from typing import Any, Collection, Dict, List, Optional, Sequence, Tuple, Union import sympy @@ -48,6 +47,9 @@ inline_lifted_buffers_into_gm, lift_mutated_buffers, ) +from torch_tensorrt.dynamo.lowering.passes.reset_folded_constructors import ( + reset_folded_constructors, +) from torch_tensorrt.dynamo.partitioning._resource_partitioner import ( resource_partition, ) @@ -1353,6 +1355,12 @@ def preserve_module_specs( f"node_name: {name} does not exist in the submodule node dictionary" ) + # Partitioning can expose an internal folded constructor as a new TRT + # subgraph output. Give that compiler-owned value fresh storage on each + # invocation before downstream eager code can mutate it. + submodule = reset_folded_constructors(submodule, settings) + setattr(partitioned_module, name, submodule) + # set the submodule metadata back to the parent trt_module_node metadata_list = get_output_metadata(submodule) assert len(metadata_list) > 0 diff --git a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py index b92165f83f..d008d7d264 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/_aten_lowering_pass.py @@ -26,6 +26,7 @@ from .repair_input_as_output import repair_input_as_output from .replace_fused_rms_norm import replace_fused_rms_norm from .replace_max_pool_with_indices import replace_max_pool_with_indices +from .reset_folded_constructors import reset_folded_constructors from .rule_based_autocast import rule_based_autocast pre_lowering_pass_list = [ @@ -39,6 +40,7 @@ remove_input_alias_fixing_clones, constant_fold, repair_input_as_output, + reset_folded_constructors, fuse_prims_broadcast, replace_max_pool_with_indices, remove_assert_nodes, diff --git a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py index 60dec56f3b..d7e4ea6bd4 100644 --- a/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py +++ b/py/torch_tensorrt/dynamo/lowering/passes/constant_folding.py @@ -7,6 +7,9 @@ from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( clean_up_graph_after_modifications, ) +from torch_tensorrt.dynamo.lowering.passes.reset_folded_constructors import ( + FOLDED_CONSTRUCTOR_META, +) from packaging import version @@ -58,6 +61,7 @@ def constant_fold( gm.graph.erase_node(node) gm = clean_up_graph_after_modifications(gm) + # Delete the constant folder instance which holds GPU memory del cf @@ -92,6 +96,9 @@ def replace_node_with_constant( new_input_node = g.create_node("get_attr", qualname, (), {}) node.replace_all_uses_with(new_input_node) new_input_node.meta.update(node.meta) + # Distinguishes this from module state that existed before folding: the + # value came from an op in the graph body, so eager rebuilds it per call. + new_input_node.meta[FOLDED_CONSTRUCTOR_META] = True g.erase_node(node) # Needed to suppress `does not reference an nn.Module, nn.Parameter, or buffer` warning diff --git a/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py b/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py new file mode 100644 index 0000000000..9cd50c58fa --- /dev/null +++ b/py/torch_tensorrt/dynamo/lowering/passes/reset_folded_constructors.py @@ -0,0 +1,76 @@ +import logging + +import torch +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.lowering.passes.pass_utils import ( + clean_up_graph_after_modifications, +) + +logger = logging.getLogger(__name__) + +# Set by constant folding on every attribute it creates. Such a value was built +# by an op in the graph body, so eager semantics allocate it again on every +# call. Attributes that already existed on the module are never folded, so they +# never carry this tag and keep their persistent, caller-visible identity. +FOLDED_CONSTRUCTOR_META = "folded_constructor" + +# Marks the copy inserted below, so running this pass again (once per TensorRT +# submodule after partitioning) does not stack redundant copies. +_FRESH_COPY_META = "folded_constructor_reset" + + +def _mutates_its_input(user: torch.fx.Node) -> bool: + target = user.target + if not isinstance(target, torch._ops.OpOverload): + return False + return bool(target._schema.is_mutable) + + +def reset_folded_constructors( + gm: torch.fx.GraphModule, settings: CompilationSettings +) -> torch.fx.GraphModule: + """Rebuild folded constructors that must not persist between calls. + + Constant folding hoists ops out of the graph body into module attributes, + which turns per-call values into state shared by every invocation. That is + only observable when the value escapes as an output or is mutated in place; + a folded value that is merely read stays a genuine constant and is left as + is, so weights keep converting to TensorRT constants. + + The copy is inserted at the point of construction rather than at the return, + so an in-place mutation inside the graph also sees fresh storage. Every user + reads the same copy, which preserves aliasing when a value is returned more + than once. + + Values the caller owns, such as placeholders and pre-existing module + attributes, are untagged and deliberately untouched: Python passes those by + reference and mutations are expected to persist. + """ + modified = False + + for node in list(gm.graph.nodes): + if node.op != "get_attr" or not node.meta.get(FOLDED_CONSTRUCTOR_META): + continue + + users = list(node.users) + if not users or any(user.meta.get(_FRESH_COPY_META) for user in users): + continue + if not any(user.op == "output" or _mutates_its_input(user) for user in users): + continue + + with gm.graph.inserting_after(node): + fresh = gm.graph.call_function(torch.ops.aten.clone.default, args=(node,)) + fresh.meta.update(node.meta) + fresh.meta.pop(FOLDED_CONSTRUCTOR_META, None) + fresh.meta[_FRESH_COPY_META] = True + + for user in users: + user.replace_input_with(node, fresh) + + modified = True + + if modified: + gm = clean_up_graph_after_modifications(gm) + logger.debug("Graph after resetting folded constructors:\n%s", gm.graph) + + return gm diff --git a/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py new file mode 100644 index 0000000000..1ac07394ab --- /dev/null +++ b/tests/py/dynamo/lowering/test_constant_folding_mutable_outputs.py @@ -0,0 +1,251 @@ +# Regression tests for https://github.com/pytorch/TensorRT/issues/4466 +# +# Constant folding may replace input-independent factories (e.g. torch.zeros) +# with persistent module attributes. If those are returned across a graph break +# and mutated in eager code, later calls must not observe that mutation. + +import torch +import torch_tensorrt # noqa: F401 # Registers the "tensorrt" backend +from torch.fx import Graph, GraphModule +from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt.dynamo._settings import CompilationSettings +from torch_tensorrt.dynamo.lowering.passes.constant_folding import constant_fold +from torch_tensorrt.dynamo.lowering.passes.reset_folded_constructors import ( + FOLDED_CONSTRUCTOR_META, + reset_folded_constructors, +) + + +@torch._dynamo.disable +def mutate_accumulators(values, weight, sample): + # Runs eagerly after the graph break. + values += sample + weight += 1 + return values / weight + + +def accumulate_with_fresh_state(sample): + # Input-independent factories; constant folding replaces these with + # persistent module attributes. + values = torch.zeros((8,), dtype=torch.float32, device="cuda") + weight = torch.zeros((8,), dtype=torch.float32, device="cuda") + return mutate_accumulators(values, weight, sample) + + +class TestFoldedConstantGraphBreak(TestCase): + def tearDown(self): + torch._dynamo.reset() + + def test_state_is_fresh_across_calls(self): + """End-to-end repro from issue #4466.""" + compiled = torch.compile( + accumulate_with_fresh_state, + backend="tensorrt", + dynamic=False, + options={"min_block_size": 1}, + ) + + first = compiled(torch.ones(8, device="cuda")) + second = compiled(torch.full((8,), 3.0, device="cuda")) + + # Eager semantics: each call starts from fresh zeros. + # Before the fix, second would be 2.0 from reused mutated state. + torch.testing.assert_close(first, torch.ones_like(first)) + torch.testing.assert_close(second, torch.full_like(second, 3.0)) + + def test_folded_constructor_outputs_are_reset(self): + """Folded constants returned as graph outputs must be cloned.""" + g = Graph() + with g.inserting_after(): + values = g.call_function( + torch.ops.aten.zeros.default, + args=((8,),), + kwargs={"dtype": torch.float32}, + ) + weight = g.call_function( + torch.ops.aten.zeros.default, + args=((8,),), + kwargs={"dtype": torch.float32}, + ) + g.output((values, weight)) + + gm = GraphModule(torch.nn.Module(), g) + gm = constant_fold(gm, CompilationSettings()) + gm = reset_folded_constructors(gm, CompilationSettings()) + + # Outputs should be clones of the folded constants, not the attrs themselves. + output_node = next(n for n in gm.graph.nodes if n.op == "output") + outs = output_node.args[0] + self.assertEqual(len(outs), 2) + for out in outs: + self.assertEqual(out.op, "call_function") + self.assertIn(out.target, (torch.clone, torch.ops.aten.clone.default)) + + out0, out1 = gm() + frozen = [ + getattr(gm, name) + for name in dir(gm) + if name.startswith("_frozen_param") + and isinstance(getattr(gm, name), torch.Tensor) + ] + self.assertGreaterEqual(len(frozen), 2) + + # Mutating returned tensors must not change stored folded constants. + out0.add_(1) + out1.add_(1) + for t in frozen: + self.assertTrue(torch.equal(t.detach().cpu(), torch.zeros(8))) + + def test_repeated_folded_outputs_preserve_alias(self): + """Same folded value returned twice must still share storage (eager semantics).""" + g = Graph() + with g.inserting_after(): + values = g.call_function( + torch.ops.aten.zeros.default, + args=((8,),), + kwargs={"dtype": torch.float32}, + ) + g.output((values, values)) # same value twice + + gm = GraphModule(torch.nn.Module(), g) + gm = constant_fold(gm, CompilationSettings()) + gm = reset_folded_constructors(gm, CompilationSettings()) + + output_node = next(n for n in gm.graph.nodes if n.op == "output") + out0, out1 = output_node.args[0] + # One clone node reused for both outputs. + self.assertIs(out0, out1) + self.assertEqual(out0.op, "call_function") + self.assertIn(out0.target, (torch.clone, torch.ops.aten.clone.default)) + + a, b = gm() + self.assertEqual(a.data_ptr(), b.data_ptr()) + a.add_(1) + self.assertTrue(torch.equal(a, b)) # mutation shared across aliases + + frozen = [ + getattr(gm, name) + for name in dir(gm) + if name.startswith("_frozen_param") + and isinstance(getattr(gm, name), torch.Tensor) + ] + self.assertGreaterEqual(len(frozen), 1) + for t in frozen: + self.assertTrue(torch.equal(t.detach().cpu(), torch.zeros(8))) + + def test_post_partition_frozen_output_is_reset(self): + """A folded value exposed as a subgraph output is reset by the late pass.""" + root = torch.nn.Module() + root.register_parameter( + "_frozen_param0", + torch.nn.Parameter(torch.zeros(8), requires_grad=False), + ) + g = Graph() + frozen = g.get_attr("_frozen_param0") + frozen.meta[FOLDED_CONSTRUCTOR_META] = True + g.output({"first": frozen, "nested": (frozen,)}) + + gm = GraphModule(root, g) + gm = reset_folded_constructors(gm, CompilationSettings()) + + first = gm() + self.assertEqual(first["first"].data_ptr(), first["nested"][0].data_ptr()) + first["first"].add_(1) + self.assertTrue(torch.equal(first["first"], first["nested"][0])) + self.assertTrue(torch.equal(gm._frozen_param0, torch.zeros(8))) + + second = gm() + self.assertTrue(torch.equal(second["first"], torch.zeros(8))) + + def test_user_owned_output_is_not_reset(self): + """Graph inputs remain caller-owned and retain copy-by-reference semantics.""" + g = Graph() + value = g.placeholder("value") + g.output(value) + + gm = GraphModule(torch.nn.Module(), g) + gm = reset_folded_constructors(gm, CompilationSettings()) + + supplied = torch.zeros(8) + returned = gm(supplied) + self.assertEqual(returned.data_ptr(), supplied.data_ptr()) + returned.add_(1) + self.assertTrue(torch.equal(supplied, torch.ones(8))) + + def _mutate_and_return_attr(self, root, attr, folded): + """Graph for `self. += 1; return self.`.""" + g = Graph() + state = g.get_attr(attr) + if folded: + state.meta[FOLDED_CONSTRUCTOR_META] = True + g.output(g.call_function(torch.ops.aten.add_.Tensor, args=(state, 1))) + gm = GraphModule(root, g) + return reset_folded_constructors(gm, CompilationSettings()) + + def test_module_state_keeps_mutation_between_calls(self): + """State assigned in __init__ is caller-visible, so mutation must persist.""" + root = torch.nn.Module() + root.register_buffer("weight", torch.zeros(8)) + + gm = self._mutate_and_return_attr(root, "weight", folded=False) + + self.assertTrue(torch.equal(gm(), torch.ones(8))) + self.assertTrue(torch.equal(gm(), torch.full((8,), 2.0))) + self.assertTrue(torch.equal(gm.weight, torch.full((8,), 2.0))) + + def test_function_state_is_rebuilt_between_calls(self): + """A constructor folded out of forward must not accumulate across calls.""" + root = torch.nn.Module() + root.register_parameter( + "_frozen_param0", + torch.nn.Parameter(torch.zeros(8), requires_grad=False), + ) + + gm = self._mutate_and_return_attr(root, "_frozen_param0", folded=True) + + # The copy precedes the in-place op, so the stored value never changes. + self.assertTrue(torch.equal(gm(), torch.ones(8))) + self.assertTrue(torch.equal(gm(), torch.ones(8))) + self.assertTrue(torch.equal(gm._frozen_param0.detach(), torch.zeros(8))) + + def test_read_only_folded_constant_is_left_alone(self): + """A folded constant that only feeds an op stays a constant for TensorRT.""" + root = torch.nn.Module() + root.register_parameter( + "_frozen_param0", + torch.nn.Parameter(torch.ones(8), requires_grad=False), + ) + g = Graph() + weight = g.get_attr("_frozen_param0") + weight.meta[FOLDED_CONSTRUCTOR_META] = True + value = g.placeholder("value") + g.output(g.call_function(torch.ops.aten.add.Tensor, args=(value, weight))) + + gm = GraphModule(root, g) + gm = reset_folded_constructors(gm, CompilationSettings()) + + clones = [n for n in gm.graph.nodes if n.target is torch.ops.aten.clone.default] + self.assertEqual(clones, []) + + def test_repeated_runs_do_not_stack_copies(self): + """Re-running per TensorRT submodule must not add a copy each time.""" + root = torch.nn.Module() + root.register_parameter( + "_frozen_param0", + torch.nn.Parameter(torch.zeros(8), requires_grad=False), + ) + g = Graph() + frozen = g.get_attr("_frozen_param0") + frozen.meta[FOLDED_CONSTRUCTOR_META] = True + g.output(frozen) + + gm = GraphModule(root, g) + for _ in range(3): + gm = reset_folded_constructors(gm, CompilationSettings()) + + clones = [n for n in gm.graph.nodes if n.target is torch.ops.aten.clone.default] + self.assertEqual(len(clones), 1) + + +if __name__ == "__main__": + run_tests()