Merged main

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

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,

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

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