Repository: tenstorrent/pytorch2.0_ttnn
Pull request #1015: Load Weights Once
Author: login: jmalone-tt
Created: 2025-05-14T00:33:44Z | Merged: 2025-05-21T18:11:41Z

This adds changes to only load weights once by adding a wrapper function with logic to load weights on device and use cached tensors for future iterations.

Cached tensors should be deallocated when the calling function ends. Since Native Device should have the same behavior, this is ok for now. If we start running into issues with running out of DRAM, we may need to be more strategic about which tensors are cached on device.

Note that this PR does not enable cached model parameters for training runs, only inference. This is accomplished through an additional GraphModuleAnalysisPass. This pass looks for characteristics that are typical of backward computation graphs (calls to backward aten ops), forward training graphs (outputting one of the inputs unchanged), and marks everything else as forward inference. This seems fragile, but any of the normal torch methods of determining whether a model is in training or eval mode appear to not work by the time the GraphModule gets to our backend.

This PR also fixes an issue where inputs for data parallel were not hoisted to the top of the forward function.

This PR also opens up the opportunity for preprocessing model weights. This is needed for convolution ops to be able to run metal trace, and will help speed up several other functions once implemented. Preprocessing will be addressed in a separate PR.

Local tests indicate the following performance improvements:
MNIST Batch 1, 5th iteration changes **5.29 ms -> 3.32 ms**
BERT Batch 8, 5th iteration now hits **23.3 sen/sec** (batch 1 on main today hits ~1.7 sen/sec, not sure about batch 8)


Diff:
diff --git a/torch_ttnn/backend.py b/torch_ttnn/backend.py
index 990b99b62d..78240e4764 100644
--- a/torch_ttnn/backend.py
+++ b/torch_ttnn/backend.py
@@ -67,6 +67,9 @@ def __init__(
         self._n_buffers = None
         self._n_arguments = None
 
+        # Used for pre-loading model params
+        self._is_end_to_end = False
+
     def reset_containers(self):
         self._out_fx_graphs = list()
         self.original_schema_list = list()
@@ -140,6 +143,7 @@ def aten_backend(
 
     # Run analysis passes to help with ttnn ops
     from torch.fx.passes.infra.pass_manager import PassManager
+    from torch_ttnn.passes.analysis.graph_module_analysis_pass import GraphModuleAnalysisPass
     from torch_ttnn.passes.analysis.input_analysis_pass import InputAnalysisPass
     from torch_ttnn.passes.analysis.multi_device_shard_analysis_pass import MultiDeviceShardAnalysisPass
 
@@ -157,13 +161,14 @@ def aten_backend(
     from torch_ttnn.passes.deallocation_pass import DeallocationPass
 
     passes = [
+        GraphModuleAnalysisPass(),
         InputAnalysisPass(option._n_parameters, option._n_buffers, option._n_arguments),
         MultiDeviceShardAnalysisPass(option.device),
         ConstantFoldingPass(),
         MultiDevicePass(option.device, example_inputs),
         ToTtPass(option.device, option.use_less_ttnn_op_types),
         FusionPass(),
-        AddDataMovePass(option.device),
+        AddDataMovePass(option.device, option._is_end_to_end),
         EliminateCoreopsPass(),
         CSEPass(),
         PermuteReshapeTuple(),
@@ -277,6 +282,9 @@ def ttnn_backend(
     options._n_buffers = len(list(gm.buffers()))
     options._n_arguments = len(example_inputs)
 
+    # Currently, we only support preprocessing weights for end-to-end converted models
+    options._is_end_to_end = gm.compile_subgraph_reason.graph_break == False
+
     tracer_option = options.tracer_option
     if tracer_option is not None:
         from ..tracer import Tracer
diff --git a/torch_ttnn/passes/analysis/graph_module_analysis_pass.py b/torch_ttnn/passes/analysis/graph_module_analysis_pass.py
new file mode 100644
index 0000000000..80273f7a08
--- /dev/null
+++ b/torch_ttnn/passes/analysis/graph_module_analysis_pass.py
@@ -0,0 +1,80 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from enum import Enum
+
+
+class ModelType(Enum):
+    """Enumeration of model types differentiating between inference and training, forward and backward.
+
+    :param INFERENCE: Model with this tag is for the forward pass of inference.
+    :param TRAIN_FORWARD: Model wih this tag is for the forward pass of a training run.
+    :param TRAIN_BACKWARD: Model with this tag is for the backward pass of a training run.
+    """
+
+    INFERENCE = 1
+    TRAIN_FORWARD = 2
+    TRAIN_BACKWARD = 3
+
+
+# this list seems small, but is based off all the backward calls from the Core Aten IR:
+# https://docs.pytorch.org/docs/stable/torch.compiler_ir.html
+aten_backward_ops = {
+    torch.ops.aten._adaptive_avg_pool2d_backward.default,
+    torch.ops.aten.avg_pool2d_backward.default,
+    torch.ops.aten.convolution_backward.default,
+    torch.ops.aten.embedding_dense_backward.default,
+    torch.ops.aten.max_pool2d_with_indices_backward.default,
+    torch.ops.aten.native_group_norm_backward.default,
+    torch.ops.aten.native_layer_norm_backward.default,
+}
+
+
+def is_train_backward(gm):
+    node_list = list(gm.graph.nodes)
+    for node in node_list:
+        if node.op == "call_function" and node.target in aten_backward_ops:
+            return True
+
+    return False
+
+
+def is_train_forward(gm):
+    # Assume training forward calls are differentiated by directly returning one of the inputs to be used for calculating gradients
+    # If this assumption fails in the future, we will need to update this function
+    outputs = [node for node in gm.graph.nodes if node.op == "output"]
+    for node in outputs:
+        placeholder_args = [arg for arg in node.args[0] if arg.op == "placeholder"]
+        if len(placeholder_args) > 0:
+            return True
+
+    return False
+
+
+class GraphModuleAnalysisPass(PassBase):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, gm: torch.fx.GraphModule):
+        """Marks the GraphModule as either training forward, training backward, or inference (forward).
+
+        This relies on commonalities between training forward, backward, and inference graphs. Namely, backward passes call backward versions of the forward functions to calculate gradients. Training forward passes return inputs unchanged. Inference forward functions do neither of these. It would be cleaner if we could just use something like `torch.is_grad_enabled()` or `gm.training` instead, but these appear to be inaccurate by the time the GraphModule is passed to our backend.
+
+        :param gm: Graph module for the function being compiled.
+        :return: Pass result with the updated graph module with metadata indicating the type of graph being compiled.
+        :rtype: PassResult[torch.fx.GraphModule, bool]
+        """
+
+        modified = False
+
+        # check nodes for backward function call
+        if is_train_backward(gm):
+            gm.meta["graph_type"] = ModelType.TRAIN_BACKWARD
+        elif is_train_forward(gm):
+            gm.meta["graph_type"] = ModelType.TRAIN_FORWARD
+        else:
+            gm.meta["graph_type"] = ModelType.INFERENCE
+
+        return PassResult(gm, modified)
diff --git a/torch_ttnn/passes/deallocation_pass.py b/torch_ttnn/passes/deallocation_pass.py
index 7351e44928..aa64438709 100644
--- a/torch_ttnn/passes/deallocation_pass.py
+++ b/torch_ttnn/passes/deallocation_pass.py
@@ -97,6 +97,9 @@ def call(self, gm: torch.fx.GraphModule):
                     # We don't want to delete these too early
                     if node.target in [target_wrappers.pack_to_tuple, operator.getitem]:
                         continue
+                    # Skip nodes that are cached between runs
+                    elif n.meta.get("is_cached", False):
+                        continue
                     with graph.inserting_after(node):
                         new_node = graph.call_function(deallocate, args=(n,))
                         modified = True
diff --git a/torch_ttnn/passes/lowering/add_data_move_pass.py b/torch_ttnn/passes/lowering/add_data_move_pass.py
index ecea1aa230..92ecdee65b 100644
--- a/torch_ttnn/passes/lowering/add_data_move_pass.py
+++ b/torch_ttnn/passes/lowering/add_data_move_pass.py
@@ -4,6 +4,7 @@
 import logging
 import torch
 import ttnn
+from torch_ttnn.passes.analysis.graph_module_analysis_pass import ModelType
 from torch_ttnn.passes.analysis.input_analysis_pass import PrimalTag
 from torch_ttnn.utils import (
     GraphCleanup,
@@ -269,7 +270,12 @@ class NodeInputAligner:
     def __init__(self, graph, device):
         self.graph = graph
         self.device = device
+        # aligned_node_dict maps DataMoveSpec to aligned version of the node to prevent calling the same data movement twice
         self.aligned_node_dict = {}
+        # marshaled_node_dict maps DataMoveSpec to index in the load_weights function that runs once. This is consumed once when adding data movement, and the DataMoveSpec will
+        # then be populated in the aligned_node_dict for further usage
+        self.marshaled_node_dict = {}
+        self.input_idx = 0
 
         # fields for data parallel
         self.shard_to_mesh = dict()
@@ -417,7 +423,7 @@ def _change_layout(self, spec):
 
         return input_node
 
-    def _create_aligned_node(self, spec):
+    def _extract_args_kwargs_from_spec(self, spec):
         if isinstance(spec, self.AlignSpecFromTorch):
             kwargs = {}
             args = (spec.input_node,)
@@ -450,7 +456,7 @@ def _create_aligned_node(self, spec):
                     args = spec.input_node.args
                 kwargs["mesh_mapper"] = mesh_mapper
 
-            return self.graph.call_function(ttnn.from_torch, args, kwargs)
+            return args, kwargs
 
         elif isinstance(spec, self.AlignSpecToTorch):
             kwargs = {"dtype": spec.dtype}
@@ -470,6 +476,21 @@ def _create_aligned_node(self, spec):
                     kwargs["mesh_composer"] = composer
                     args = (actual_inp_node,)
 
+            return args, kwargs
+
+        elif isinstance(spec, self.AlignSpecInTtnn):
+            return (), {}
+
+        else:
+            raise RuntimeError(f"Cannot create aligned node for unknown spec ({spec})")
+
+    def _create_aligned_node(self, spec):
+        args, kwargs = self._extract_args_kwargs_from_spec(spec)
+
+        if isinstance(spec, self.AlignSpecFromTorch):
+            return self.graph.call_function(ttnn.from_torch, args, kwargs)
+
+        elif isinstance(spec, self.AlignSpecToTorch):
             return self.graph.call_function(ttnn.to_torch, args, kwargs)
 
         elif isinstance(spec, self.AlignSpecInTtnn):
@@ -500,7 +521,38 @@ def _connect_aligned_node(self, node, aligned_node, input_site, input_site_type:
             new_arg[tuple_idx] = aligned_node
             node.update_kwarg(key, tuple(new_arg))
 
-    def align(self, node, input_node, input_site, input_site_type: InputSiteType, first_node):
+    def marshal_params(self, node, input_node, input_site, input_site_type, first_node):
+        if not isinstance(input_node, torch.fx.node.Node):
+            return 0
+
+        # examine first arg for multi device case
+        check_constant = input_node
+        if input_node.op == "call_function" and input_node.target in [
+            target_wrappers.shard_tensor,
+            target_wrappers.replicate_tensor,
+        ]:
+            check_constant = input_node.args[0]
+        is_constant_for_inference = (check_constant.op == "placeholder") and (
+            check_constant.meta.get("primal_tag") in [PrimalTag.PARAMETER, PrimalTag.BUFFER]
+        )
+        if not is_constant_for_inference:
+            return 0
+
+        data_move_spec = self._get_align_spec(node, input_node, input_site, input_site_type)
+        if not isinstance(data_move_spec, self.AlignSpecFromTorch):
+            # No need to align input_node
+            return 0
+
+        if data_move_spec in self.marshaled_node_dict:
+            # already handled
+            return 0
+
+        self.marshaled_node_dict[data_move_spec] = self.input_idx
+        self.input_idx += 1
+
+        return 1
+
+    def align(self, node, input_node, input_site, input_site_type: InputSiteType, first_node, ttnn_inputs):
         # assert input_site_type in ["args", "kwargs", "args_tuple", "kwargs_tuple"]
         data_move_spec = self._get_align_spec(node, input_node, input_site, input_site_type)
         if data_move_spec is None:
@@ -509,11 +561,26 @@ def align(self, node, input_node, input_site, input_site_type: InputSiteType, fi
 
         if data_move_spec in self.aligned_node_dict:
             aligned_node = self.aligned_node_dict[data_move_spec]
+        elif ttnn_inputs is not None and (input_idx := self.marshaled_node_dict.get(data_move_spec)) is not None:
+            with self.graph.inserting_before(node):
+                aligned_node = self.graph.call_function(getitem, (ttnn_inputs, input_idx))
+                # mark node cached so it isn't deallocated by DeallocationPass
+                aligned_node.meta["is_cached"] = True
+            # update aligned_node_dict
+            self.aligned_node_dict[data_move_spec] = aligned_node
         else:
+            # push from_torch calls to top of forward function if they are due to a placeholder
+            maybe_forward_input = input_node
+            # We have to test the first arg of shard_tensor and replicate_tensor calls instead
+            if input_node.op == "call_function" and input_node.target in [
+                target_wrappers.shard_tensor,
+                target_wrappers.replicate_tensor,
+            ]:
+                maybe_forward_input = input_node.args[0]
             if (
                 isinstance(data_move_spec, self.AlignSpecFromTorch)
-                and input_node.op == "placeholder"
-                and input_node.meta.get("primal_tag") != PrimalTag.ARGUMENT
+                and maybe_forward_input.op == "placeholder"
+                and maybe_forward_input.meta.get("primal_tag") != PrimalTag.ARGUMENT
             ):
                 # This will push all from_torch calls to the top of the forward function. This shouldn't impact performance, but it may impact memory usage since variables will be
                 # live longer than they would if from_torch calls occurred right before usage. If we start running out of DRAM or need to be more careful about memory usage, this
@@ -530,6 +597,46 @@ def align(self, node, input_node, input_site, input_site_type: InputSiteType, fi
         return 1
 
 
+def insert_load_params_once(gm, first_node, nodes, node_input_aligner):
+    SiteType = NodeInputAligner.InputSiteType
+    modifications_count = 0
+
+    for node in nodes:
+        args = node.args
+        for idx, arg in enumerate(args):
+            if isinstance(arg, (tuple, list, torch.fx.immutable_collections.immutable_list)):
+                for tuple_idx, tuple_arg in enumerate(arg):
+                    modifications_count += node_input_aligner.marshal_params(
+                        node, tuple_arg, [idx, tuple_idx], SiteType.ARGS_TUPLE, first_node
+                    )
+            else:
+                modifications_count += node_input_aligner.marshal_params(node, arg, idx, SiteType.ARGS, first_node)
+
+        kwargs = node.kwargs
+        for key, arg in kwargs.items():
+            if isinstance(arg, (tuple, list, torch.fx.immutable_collections.immutable_list)):
+                for tuple_idx, tuple_arg in enumerate(arg):
+                    modifications_count += node_input_aligner.marshal_params(
+                        node, tuple_arg, [key, tuple_idx], SiteType.KWARGS_TUPLE, first_node
+                    )
+            else:
+                modifications_count += node_input_aligner.marshal_params(node, arg, key, SiteType.KWARGS, first_node)
+
+    # reset run_once_count so recompilation triggers loading weights
+    with gm.graph.inserting_before(first_node):
+        ttnn_inputs = gm.graph.call_function(
+            target_wrappers.run_once,
+            tuple(
+                [
+                    node_input_aligner._extract_args_kwargs_from_spec(spec)
+                    for spec in node_input_aligner.marshaled_node_dict.keys()
+                ]
+            ),
+        )
+
+    return modifications_count, ttnn_inputs
+
+
 class AddDataMovePass(PassBase):
     """Pass that adds instructions to move data between host and device and align tensor dtype and layout.
 
@@ -538,39 +645,52 @@ class AddDataMovePass(PassBase):
     :param device: The device on which a workload will run (either a MeshDevice or Device).
     """
 
-    def __init__(self, device):
+    def __init__(self, device, is_end_to_end):
         self.device = device
+        self.is_end_to_end = is_end_to_end
 
     def call(self, gm: torch.fx.GraphModule):
         SiteType = NodeInputAligner.InputSiteType
 
-        i = 0
+        modifications_count = 0
         node_input_aligner = NodeInputAligner(gm.graph, self.device)
         nodes = list(gm.graph.nodes)
 
         first_node = [node for node in nodes if node.op != "placeholder"][0]
 
+        # first load weights
+        ttnn_inputs = None
+        if gm.meta.get("graph_type") == ModelType.INFERENCE and self.is_end_to_end:
+            global run_once_count
+            target_wrappers.run_once_count = 0
+            modifications_count, ttnn_inputs = insert_load_params_once(gm, first_node, nodes, node_input_aligner)
+
+        # then handle rest of the args and kwargs
         for node in nodes:
             args = node.args
             for idx, arg in enumerate(args):
                 if isinstance(arg, (tuple, list, torch.fx.immutable_collections.immutable_list)):
                     for tuple_idx, tuple_arg in enumerate(arg):
-                        i += node_input_aligner.align(
-                            node, tuple_arg, [idx, tuple_idx], SiteType.ARGS_TUPLE, first_node
+                        modifications_count += node_input_aligner.align(
+                            node, tuple_arg, [idx, tuple_idx], SiteType.ARGS_TUPLE, first_node, ttnn_inputs
                         )
                 else:
-                    i += node_input_aligner.align(node, arg, idx, SiteType.ARGS, first_node)
+                    modifications_count += node_input_aligner.align(
+                        node, arg, idx, SiteType.ARGS, first_node, ttnn_inputs
+                    )
 
             kwargs = node.kwargs
             for key, arg in kwargs.items():
                 if isinstance(arg, (tuple, list, torch.fx.immutable_collections.immutable_list)):
                     for tuple_idx, tuple_arg in enumerate(arg):
-                        i += node_input_aligner.align(
-                            node, tuple_arg, [key, tuple_idx], SiteType.KWARGS_TUPLE, first_node
+                        modifications_count += node_input_aligner.align(
+                            node, tuple_arg, [key, tuple_idx], SiteType.KWARGS_TUPLE, first_node, ttnn_inputs
                         )
                 else:
-                    i += node_input_aligner.align(node, arg, key, SiteType.KWARGS, first_node)
+                    modifications_count += node_input_aligner.align(
+                        node, arg, key, SiteType.KWARGS, first_node, ttnn_inputs
+                    )
 
-        modified = i > 0
+        modified = modifications_count > 0
         GraphCleanup(gm)
         return PassResult(gm, modified)
diff --git a/torch_ttnn/passes/lowering/target_wrappers.py b/torch_ttnn/passes/lowering/target_wrappers.py
index ba23d7e23a..23bfaebf6c 100644
--- a/torch_ttnn/passes/lowering/target_wrappers.py
+++ b/torch_ttnn/passes/lowering/target_wrappers.py
@@ -6,6 +6,25 @@
 
 from torch_ttnn.utils import TtnnDevice
 
+run_once_count = 0
+run_once_ans = tuple()
+
+
+@torch.fx.wrap
+def run_once(*args):
+    global run_once_count
+    global run_once_ans
+
+    if run_once_count == 0:
+
+        def convert_input(spec):
+            return ttnn.from_torch(*spec[0], **spec[1])
+
+        run_once_ans = tuple([convert_input(arg) for arg in args])
+        run_once_count += 1
+
+    return run_once_ans
+
 
 @torch.fx.wrap
 def clone(t):
diff --git a/torch_ttnn/passes/lowering/to_tt_pass.py b/torch_ttnn/passes/lowering/to_tt_pass.py
index 1d00944303..4b18ac1822 100644
--- a/torch_ttnn/passes/lowering/to_tt_pass.py
+++ b/torch_ttnn/passes/lowering/to_tt_pass.py
@@ -166,6 +166,12 @@ def __init__(self, module, device, use_less_ttnn_op_types):
         self.device = device
         self.use_less_ttnn_op_types = use_less_ttnn_op_types
 
+    def transform(self):
+        old_meta = self.module.meta
+        result = super().transform()
+        result.meta = old_meta
+        return result
+
     def get_attr(self, target, args, kwargs):
         # Restore original metadata for get_attr nodes
         proxy = super().get_attr(target, args, kwargs)
diff --git a/torch_ttnn/passes/multi_device_pass.py b/torch_ttnn/passes/multi_device_pass.py
index 86f630c676..fb2b0e6d8b 100644
--- a/torch_ttnn/passes/multi_device_pass.py
+++ b/torch_ttnn/passes/multi_device_pass.py
@@ -28,14 +28,14 @@ def __init__(self, device, example_inputs):
         self.device = device
         self.example_inputs = example_inputs
 
-    def replicate_to_mesh(self, node):
-        with node.graph.inserting_after(node):
+    def replicate_to_mesh(self, node, first_node):
+        with node.graph.inserting_before(first_node):
             new_node = node.graph.call_function(target_wrappers.replicate_tensor, (node,))
             node.replace_all_uses_with(new_node, delete_user_cb=lambda node: node != new_node and node.op != "output")
 
-    def shard_to_mesh(self, node):
+    def shard_to_mesh(self, node, first_node):
         batch_dimension = 0
-        with node.graph.inserting_after(node):
+        with node.graph.inserting_before(first_node):
             new_node = node.graph.call_function(
                 target_wrappers.shard_tensor, (node, batch_dimension, self.device.get_num_devices())
             )
@@ -84,15 +84,16 @@ def call(self, graph_module: torch.fx.GraphModule):
             return PassResult(graph_module, modified)
 
         node_list = list(graph_module.graph.nodes)
+        first_node = [node for node in node_list if node.op != "placeholder"][0]
         for node in node_list:
             if node.op == "placeholder":
                 try:
                     if (node.meta["primal_tag"] == PrimalTag.PARAMETER) or (
                         node.meta["primal_tag"] == PrimalTag.BUFFER
                     ):
-                        self.replicate_to_mesh(node)
+                        self.replicate_to_mesh(node, first_node)
                     elif node.meta["primal_tag"] == PrimalTag.ARGUMENT:
-                        self.shard_to_mesh(node)
+                        self.shard_to_mesh(node, first_node)
                 except Exception as e:
                     logging.error(f"Could not distribute node {node.name}: {e}")
 


Reviews:
user: login: dgomezTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-15T13:47:23Z
user: login: dgomezTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-15T13:47:55Z
user: login: dgomezTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-15T13:49:06Z
user: login: dgomezTT
body: looks good to me, but I would wait for Kevin or Artem to check it up!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-15T13:52:11Z
user: login: jmalone-tt
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-15T14:56:34Z
user: login: jmalone-tt
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-15T14:58:06Z
user: login: kevinwuTT
body: LGTM! I'm glad the analysis pass for determining training/eval works well enough. 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-15T16:54:56Z
user: login: jmalone-tt
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-15T18:27:48Z
user: login: ayerofieiev-tt
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T12:26:19Z
user: login: ayerofieiev-tt
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-05-16T12:27:52Z
user: login: ayerofieiev-tt
body: In general looks alright. Maybe a tiny bit hacky, but that’s the only way I can see right now too. 
state: APPROVED
author_association: MEMBER
submitted_at: 2025-05-16T12:34:38Z
user: login: philei-tt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T10:02:31Z
user: login: philei-tt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T10:04:18Z
user: login: philei-tt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T10:07:17Z
user: login: philei-tt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T11:27:53Z
user: login: philei-tt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T11:38:26Z
user: login: philei-tt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T11:44:07Z
user: login: philei-tt
body: Looks good, just a few small comments.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-20T11:44:34Z
user: login: jmalone-tt
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-20T13:54:32Z
user: login: jmalone-tt
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-20T13:54:59Z
user: login: jmalone-tt
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-20T13:56:12Z
user: login: jmalone-tt
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-20T13:56:47Z
user: login: jmalone-tt
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-20T14:34:07Z
user: login: philei-tt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T10:27:42Z
user: login: philei-tt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T10:49:21Z
user: login: jmalone-tt
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-05-21T15:07:43Z

Review comments:
diff_hunk: @@ -67,6 +67,9 @@ def __init__(
         self._n_buffers = None
         self._n_arguments = None
 
+        # Used to speed up execution
+        self._load_weights_once = False
path: torch_ttnn/backend.py
user: login: dgomezTT
body: maybe this should go to the global config?
created_at: 2025-05-15T13:47:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -187,7 +192,8 @@ def aten_backend(
     pm = PassManager(passes=passes)
     gm, modified = pm(gm)
 
-    GraphCleanup(gm)
+    # GraphCleanup(gm)
path: torch_ttnn/backend.py
user: login: dgomezTT
body: any reason to leave the comment?
created_at: 2025-05-15T13:47:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from enum import Enum
+
+
+class ModelType(Enum):
+    """Enumeration of compiled model inputs.
+
+    It is expected that PARAMETER and BUFFER tensors do not change between inference runs, but ARGUMENT tensors may change.
+
+    :param PARAMETER: Tensor tracked by optimizer during training. Represents weights and biases.
+    :param BUFFER: Tensor not updated during training but still part of model state. Normally used for non-trainable data like fixed weights or running statistics in batch norm.
+    :param ARGUMENT: Tensors not tracked by model and not persistent between calls. Can be input data or configuration options to module's methods.
+    """
+
+    INFERENCE = 1
+    TRAIN_FORWARD = 2
+    TRAIN_BACKWARD = 3
+
+
+aten_backward_ops = {
+    torch.ops.aten._adaptive_avg_pool2d_backward.default,
+    torch.ops.aten.avg_pool2d_backward.default,
+    torch.ops.aten.convolution_backward.default,
+    torch.ops.aten.embedding_dense_backward.default,
+    torch.ops.aten.max_pool2d_with_indices_backward.default,
+    torch.ops.aten.native_group_norm_backward.default,
+    torch.ops.aten.native_layer_norm_backward.default,
+}
+
+
+def is_train_backward(gm):
+    node_list = list(gm.graph.nodes)
+    for node in node_list:
+        if node.op == "call_function" and node.target in aten_backward_ops:
+            return True
+
+    return False
+
+
+def is_train_forward(gm):
+    outputs = [node for node in gm.graph.nodes if node.op == "output"]
+    for node in outputs:
+        placeholder_args = [arg for arg in node.args[0] if arg.op == "placeholder"]
+        if len(placeholder_args) > 0:
+            return True
+
+    return False
+
+
+class GraphModuleAnalysisPass(PassBase):
+    def __init__(self):
+        super().__init__()
+
+    def call(self, gm: torch.fx.GraphModule):
+        """Marks inputs for the compiled function as parameters, buffers, or arguments.
+
+        This relies on how PyTorch behaves under the hood. Namely, PyTorch will lay out function arguments in the order of parameters, then buffers, then arguments. We can rely on this layout order and the number of parameters, buffers, and arguments passed in from AOTAutograd to our backend function.
path: torch_ttnn/passes/analysis/graph_module_analysis_pass.py
user: login: dgomezTT
body: this is great, it helps a lot to understand the logic behind it!
created_at: 2025-05-15T13:49:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -187,7 +192,8 @@ def aten_backend(
     pm = PassManager(passes=passes)
     gm, modified = pm(gm)
 
-    GraphCleanup(gm)
+    # GraphCleanup(gm)
path: torch_ttnn/backend.py
user: login: jmalone-tt
body: Thanks for catching, we actually want to call GraphCleanup here but one of my previous attempts was deleted by the DCE that runs. I will revert this change
created_at: 2025-05-15T14:56:34Z
author_association: COLLABORATOR
diff_hunk: @@ -67,6 +67,9 @@ def __init__(
         self._n_buffers = None
         self._n_arguments = None
 
+        # Used to speed up execution
+        self._load_weights_once = False
path: torch_ttnn/backend.py
user: login: jmalone-tt
body: I think this can actually be deleted. I tried using TtnnOption so I didn't have to maintain graph module metadata, but it was causing issues. Good catch
created_at: 2025-05-15T14:58:06Z
author_association: COLLABORATOR
diff_hunk: @@ -553,6 +620,46 @@ def align(self, node, input_node, input_site, input_site_type: InputSiteType, fi
         return 1
 
 
+def insert_load_params_once(gm, first_node, nodes, node_input_aligner):
+    SiteType = NodeInputAligner.InputSiteType
+    i = 0
+
+    for node in nodes:
+        args = node.args
path: torch_ttnn/passes/lowering/add_data_move_pass.py
user: login: kevinwuTT
body: I know args and kwargs are iterated differently, but the rest of the code seems like it could be combined. Let me know if it's more difficult than it looks.
created_at: 2025-05-15T16:23:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -553,6 +620,46 @@ def align(self, node, input_node, input_site, input_site_type: InputSiteType, fi
         return 1
 
 
+def insert_load_params_once(gm, first_node, nodes, node_input_aligner):
+    SiteType = NodeInputAligner.InputSiteType
+    i = 0
+
+    for node in nodes:
+        args = node.args
path: torch_ttnn/passes/lowering/add_data_move_pass.py
user: login: jmalone-tt
body: I'm not entirely sure what change you're asking for here. Can you clarify?
created_at: 2025-05-15T18:27:47Z
author_association: COLLABORATOR
diff_hunk: @@ -270,6 +271,8 @@ def __init__(self, graph, device):
         self.graph = graph
         self.device = device
         self.aligned_node_dict = {}
path: torch_ttnn/passes/lowering/add_data_move_pass.py
user: login: ayerofieiev-tt
body: Can you please leave a comment on what this member is about?
created_at: 2025-05-16T12:26:19Z
author_association: MEMBER
diff_hunk: @@ -529,11 +579,28 @@ def align(self, node, input_node, input_site, input_site_type: InputSiteType, fi
 
         if data_move_spec in self.aligned_node_dict:
             aligned_node = self.aligned_node_dict[data_move_spec]
+        elif ttnn_inputs is not None and data_move_spec in self.marshaled_node_dict:
+            # add getitem call
+            # update aligned_node_dict
+            input_idx = self.marshaled_node_dict[data_move_spec]
path: torch_ttnn/passes/lowering/add_data_move_pass.py
user: login: ayerofieiev-tt
body: Here you are looking in dict twice. In above if statement with IN and now with access. I assume it’s not a big deal, but want to highlight in case it matters
created_at: 2025-05-16T12:27:52Z
author_association: MEMBER
diff_hunk: @@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from enum import Enum
+
+
+class ModelType(Enum):
+    """Enumeration of compiled model inputs.
+
+    It is expected that PARAMETER and BUFFER tensors do not change between inference runs, but ARGUMENT tensors may change.
+
+    :param PARAMETER: Tensor tracked by optimizer during training. Represents weights and biases.
+    :param BUFFER: Tensor not updated during training but still part of model state. Normally used for non-trainable data like fixed weights or running statistics in batch norm.
+    :param ARGUMENT: Tensors not tracked by model and not persistent between calls. Can be input data or configuration options to module's methods.
+    """
+
+    INFERENCE = 1
+    TRAIN_FORWARD = 2
+    TRAIN_BACKWARD = 3
path: torch_ttnn/passes/analysis/graph_module_analysis_pass.py
user: login: philei-tt
body: I think docs are from other enum? From PrimalTag?
created_at: 2025-05-20T10:02:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from enum import Enum
+
+
+class ModelType(Enum):
+    """Enumeration of compiled model inputs.
+
+    It is expected that PARAMETER and BUFFER tensors do not change between inference runs, but ARGUMENT tensors may change.
+
+    :param PARAMETER: Tensor tracked by optimizer during training. Represents weights and biases.
+    :param BUFFER: Tensor not updated during training but still part of model state. Normally used for non-trainable data like fixed weights or running statistics in batch norm.
+    :param ARGUMENT: Tensors not tracked by model and not persistent between calls. Can be input data or configuration options to module's methods.
+    """
+
+    INFERENCE = 1
+    TRAIN_FORWARD = 2
+    TRAIN_BACKWARD = 3
+
+
+aten_backward_ops = {
+    torch.ops.aten._adaptive_avg_pool2d_backward.default,
+    torch.ops.aten.avg_pool2d_backward.default,
+    torch.ops.aten.convolution_backward.default,
+    torch.ops.aten.embedding_dense_backward.default,
+    torch.ops.aten.max_pool2d_with_indices_backward.default,
+    torch.ops.aten.native_group_norm_backward.default,
+    torch.ops.aten.native_layer_norm_backward.default,
+}
path: torch_ttnn/passes/analysis/graph_module_analysis_pass.py
user: login: philei-tt
body: Are you sure this list will always determine that graph is train backward? Feels too small)
created_at: 2025-05-20T10:04:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from enum import Enum
+
+
+class ModelType(Enum):
+    """Enumeration of compiled model inputs.
+
+    It is expected that PARAMETER and BUFFER tensors do not change between inference runs, but ARGUMENT tensors may change.
+
+    :param PARAMETER: Tensor tracked by optimizer during training. Represents weights and biases.
+    :param BUFFER: Tensor not updated during training but still part of model state. Normally used for non-trainable data like fixed weights or running statistics in batch norm.
+    :param ARGUMENT: Tensors not tracked by model and not persistent between calls. Can be input data or configuration options to module's methods.
+    """
+
+    INFERENCE = 1
+    TRAIN_FORWARD = 2
+    TRAIN_BACKWARD = 3
+
+
+aten_backward_ops = {
+    torch.ops.aten._adaptive_avg_pool2d_backward.default,
+    torch.ops.aten.avg_pool2d_backward.default,
+    torch.ops.aten.convolution_backward.default,
+    torch.ops.aten.embedding_dense_backward.default,
+    torch.ops.aten.max_pool2d_with_indices_backward.default,
+    torch.ops.aten.native_group_norm_backward.default,
+    torch.ops.aten.native_layer_norm_backward.default,
+}
+
+
+def is_train_backward(gm):
+    node_list = list(gm.graph.nodes)
+    for node in node_list:
+        if node.op == "call_function" and node.target in aten_backward_ops:
+            return True
+
+    return False
+
+
+def is_train_forward(gm):
path: torch_ttnn/passes/analysis/graph_module_analysis_pass.py
user: login: philei-tt
body: Is this documented in somewhere in torch docs that train forward will return inputs? If no, this might change in the future
created_at: 2025-05-20T10:07:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -520,7 +541,39 @@ def _connect_aligned_node_layer_norm(
         else:
             self._connect_aligned_node(node, aligned_node, input_site, input_site_type)
 
-    def align(self, node, input_node, input_site, input_site_type: InputSiteType, first_node):
+    def marshal_params(self, node, input_node, input_site, input_site_type, first_node):
+        if not isinstance(input_node, torch.fx.node.Node):
+            return 0
+
+        # examine first arg for multi device case
+        check_constant = input_node
+        if input_node.op == "call_function" and input_node.target in [
+            target_wrappers.shard_tensor,
+            target_wrappers.replicate_tensor,
+        ]:
+            check_constant = input_node.args[0]
+        is_constant_for_inference = (check_constant.op == "placeholder") and (
+            check_constant.meta.get("primal_tag") in [PrimalTag.PARAMETER, PrimalTag.BUFFER]
+        )
+        if not is_constant_for_inference:
+            return 0
+
+        data_move_spec = self._get_align_spec(node, input_node, input_site, input_site_type)
+        if not isinstance(data_move_spec, self.AlignSpecFromTorch):
+            # No need to align input_node
+            return 0
+
+        if data_move_spec in self.marshaled_node_dict:
+            # already handled
+            return 0
+
+        with self.graph.inserting_before(first_node):
+            self.marshaled_node_dict[data_move_spec] = self.input_idx
+            self.input_idx += 1
path: torch_ttnn/passes/lowering/add_data_move_pass.py
user: login: philei-tt
body: Why is this block wrapped with with self.graph.inserting_before? You aren't inserting node here, or am I missing something?
created_at: 2025-05-20T11:27:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -6,6 +6,25 @@
 
 from torch_ttnn.utils import TtnnDevice
 
+run_once_count = 0
+run_once_ans = tuple()
+
+
+@torch.fx.wrap
+def run_once(*args):
path: torch_ttnn/passes/lowering/target_wrappers.py
user: login: philei-tt
body: (nit) using closure or @lru_cache might be a bit better than global variable. 
created_at: 2025-05-20T11:38:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -573,26 +681,34 @@ def call(self, gm: torch.fx.GraphModule):
 
         first_node = [node for node in nodes if node.op != "placeholder"][0]
 
+        # first load weights
+        ttnn_inputs = None
+        if gm.meta.get("graph_type") == ModelType.INFERENCE:
+            global run_once_count
+            target_wrappers.run_once_count = 0
+            i, ttnn_inputs = insert_load_params_once(gm, first_node, nodes, node_input_aligner)
+
+        # then handle rest of the args and kwargs
         for node in nodes:
             args = node.args
             for idx, arg in enumerate(args):
                 if isinstance(arg, (tuple, list, torch.fx.immutable_collections.immutable_list)):
                     for tuple_idx, tuple_arg in enumerate(arg):
                         i += node_input_aligner.align(
-                            node, tuple_arg, [idx, tuple_idx], SiteType.ARGS_TUPLE, first_node
+                            node, tuple_arg, [idx, tuple_idx], SiteType.ARGS_TUPLE, first_node, ttnn_inputs
                         )
                 else:
-                    i += node_input_aligner.align(node, arg, idx, SiteType.ARGS, first_node)
+                    i += node_input_aligner.align(node, arg, idx, SiteType.ARGS, first_node, ttnn_inputs)
 
             kwargs = node.kwargs
             for key, arg in kwargs.items():
                 if isinstance(arg, (tuple, list, torch.fx.immutable_collections.immutable_list)):
                     for tuple_idx, tuple_arg in enumerate(arg):
                         i += node_input_aligner.align(
-                            node, tuple_arg, [key, tuple_idx], SiteType.KWARGS_TUPLE, first_node
+                            node, tuple_arg, [key, tuple_idx], SiteType.KWARGS_TUPLE, first_node, ttnn_inputs
                         )
                 else:
-                    i += node_input_aligner.align(node, arg, key, SiteType.KWARGS, first_node)
+                    i += node_input_aligner.align(node, arg, key, SiteType.KWARGS, first_node, ttnn_inputs)
 
         modified = i > 0
path: torch_ttnn/passes/lowering/add_data_move_pass.py
user: login: philei-tt
body: (nit) maybe use modified = modified or node_input_aligner.align(...) for readability purposes? Not sure it's better, but variable "i" doesn't tell much about it's purpose
created_at: 2025-05-20T11:44:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from enum import Enum
+
+
+class ModelType(Enum):
+    """Enumeration of compiled model inputs.
+
+    It is expected that PARAMETER and BUFFER tensors do not change between inference runs, but ARGUMENT tensors may change.
+
+    :param PARAMETER: Tensor tracked by optimizer during training. Represents weights and biases.
+    :param BUFFER: Tensor not updated during training but still part of model state. Normally used for non-trainable data like fixed weights or running statistics in batch norm.
+    :param ARGUMENT: Tensors not tracked by model and not persistent between calls. Can be input data or configuration options to module's methods.
+    """
+
+    INFERENCE = 1
+    TRAIN_FORWARD = 2
+    TRAIN_BACKWARD = 3
path: torch_ttnn/passes/analysis/graph_module_analysis_pass.py
user: login: jmalone-tt
body: Thanks for catching this!
created_at: 2025-05-20T13:54:32Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from enum import Enum
+
+
+class ModelType(Enum):
+    """Enumeration of compiled model inputs.
+
+    It is expected that PARAMETER and BUFFER tensors do not change between inference runs, but ARGUMENT tensors may change.
+
+    :param PARAMETER: Tensor tracked by optimizer during training. Represents weights and biases.
+    :param BUFFER: Tensor not updated during training but still part of model state. Normally used for non-trainable data like fixed weights or running statistics in batch norm.
+    :param ARGUMENT: Tensors not tracked by model and not persistent between calls. Can be input data or configuration options to module's methods.
+    """
+
+    INFERENCE = 1
+    TRAIN_FORWARD = 2
+    TRAIN_BACKWARD = 3
+
+
+aten_backward_ops = {
+    torch.ops.aten._adaptive_avg_pool2d_backward.default,
+    torch.ops.aten.avg_pool2d_backward.default,
+    torch.ops.aten.convolution_backward.default,
+    torch.ops.aten.embedding_dense_backward.default,
+    torch.ops.aten.max_pool2d_with_indices_backward.default,
+    torch.ops.aten.native_group_norm_backward.default,
+    torch.ops.aten.native_layer_norm_backward.default,
+}
path: torch_ttnn/passes/analysis/graph_module_analysis_pass.py
user: login: jmalone-tt
body: I agree that it seems too small, but it's based on the list of backward ops from the aten core IR. I will add a comment with a link
created_at: 2025-05-20T13:54:58Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from enum import Enum
+
+
+class ModelType(Enum):
+    """Enumeration of compiled model inputs.
+
+    It is expected that PARAMETER and BUFFER tensors do not change between inference runs, but ARGUMENT tensors may change.
+
+    :param PARAMETER: Tensor tracked by optimizer during training. Represents weights and biases.
+    :param BUFFER: Tensor not updated during training but still part of model state. Normally used for non-trainable data like fixed weights or running statistics in batch norm.
+    :param ARGUMENT: Tensors not tracked by model and not persistent between calls. Can be input data or configuration options to module's methods.
+    """
+
+    INFERENCE = 1
+    TRAIN_FORWARD = 2
+    TRAIN_BACKWARD = 3
+
+
+aten_backward_ops = {
+    torch.ops.aten._adaptive_avg_pool2d_backward.default,
+    torch.ops.aten.avg_pool2d_backward.default,
+    torch.ops.aten.convolution_backward.default,
+    torch.ops.aten.embedding_dense_backward.default,
+    torch.ops.aten.max_pool2d_with_indices_backward.default,
+    torch.ops.aten.native_group_norm_backward.default,
+    torch.ops.aten.native_layer_norm_backward.default,
+}
+
+
+def is_train_backward(gm):
+    node_list = list(gm.graph.nodes)
+    for node in node_list:
+        if node.op == "call_function" and node.target in aten_backward_ops:
+            return True
+
+    return False
+
+
+def is_train_forward(gm):
path: torch_ttnn/passes/analysis/graph_module_analysis_pass.py
user: login: jmalone-tt
body: I don't know if it's documented anywhere, but inputs are used in backward pass, so this should be reliable. Can add a comment that notes this assumption though
created_at: 2025-05-20T13:56:12Z
author_association: COLLABORATOR
diff_hunk: @@ -520,7 +541,39 @@ def _connect_aligned_node_layer_norm(
         else:
             self._connect_aligned_node(node, aligned_node, input_site, input_site_type)
 
-    def align(self, node, input_node, input_site, input_site_type: InputSiteType, first_node):
+    def marshal_params(self, node, input_node, input_site, input_site_type, first_node):
+        if not isinstance(input_node, torch.fx.node.Node):
+            return 0
+
+        # examine first arg for multi device case
+        check_constant = input_node
+        if input_node.op == "call_function" and input_node.target in [
+            target_wrappers.shard_tensor,
+            target_wrappers.replicate_tensor,
+        ]:
+            check_constant = input_node.args[0]
+        is_constant_for_inference = (check_constant.op == "placeholder") and (
+            check_constant.meta.get("primal_tag") in [PrimalTag.PARAMETER, PrimalTag.BUFFER]
+        )
+        if not is_constant_for_inference:
+            return 0
+
+        data_move_spec = self._get_align_spec(node, input_node, input_site, input_site_type)
+        if not isinstance(data_move_spec, self.AlignSpecFromTorch):
+            # No need to align input_node
+            return 0
+
+        if data_move_spec in self.marshaled_node_dict:
+            # already handled
+            return 0
+
+        with self.graph.inserting_before(first_node):
+            self.marshaled_node_dict[data_move_spec] = self.input_idx
+            self.input_idx += 1
path: torch_ttnn/passes/lowering/add_data_move_pass.py
user: login: jmalone-tt
body: Good point, earlier in the change I was adding the nodes here, but then removed that code. Will update
created_at: 2025-05-20T13:56:47Z
author_association: COLLABORATOR
diff_hunk: @@ -6,6 +6,25 @@
 
 from torch_ttnn.utils import TtnnDevice
 
+run_once_count = 0
+run_once_ans = tuple()
+
+
+@torch.fx.wrap
+def run_once(*args):
path: torch_ttnn/passes/lowering/target_wrappers.py
user: login: jmalone-tt
body: I really like the lru_cache here, but it hits an error with the kwargs inputs being unhashable since it's a dict. The closure path seems less readable to me in this case. I agree that the global variable isn't the best, but I'm leaning towards keeping it unless you feel strongly
created_at: 2025-05-20T14:34:07Z
author_association: COLLABORATOR
diff_hunk: @@ -6,6 +6,25 @@
 
 from torch_ttnn.utils import TtnnDevice
 
+run_once_count = 0
+run_once_ans = tuple()
+
+
+@torch.fx.wrap
+def run_once(*args):
path: torch_ttnn/passes/lowering/target_wrappers.py
user: login: philei-tt
body: of course, it's no a big issue
created_at: 2025-05-21T10:27:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from enum import Enum
+
+
+class ModelType(Enum):
+    """Enumeration of compiled model inputs.
+
+    It is expected that PARAMETER and BUFFER tensors do not change between inference runs, but ARGUMENT tensors may change.
+
+    :param PARAMETER: Tensor tracked by optimizer during training. Represents weights and biases.
+    :param BUFFER: Tensor not updated during training but still part of model state. Normally used for non-trainable data like fixed weights or running statistics in batch norm.
+    :param ARGUMENT: Tensors not tracked by model and not persistent between calls. Can be input data or configuration options to module's methods.
+    """
+
+    INFERENCE = 1
+    TRAIN_FORWARD = 2
+    TRAIN_BACKWARD = 3
+
+
+aten_backward_ops = {
+    torch.ops.aten._adaptive_avg_pool2d_backward.default,
+    torch.ops.aten.avg_pool2d_backward.default,
+    torch.ops.aten.convolution_backward.default,
+    torch.ops.aten.embedding_dense_backward.default,
+    torch.ops.aten.max_pool2d_with_indices_backward.default,
+    torch.ops.aten.native_group_norm_backward.default,
+    torch.ops.aten.native_layer_norm_backward.default,
+}
path: torch_ttnn/passes/analysis/graph_module_analysis_pass.py
user: login: philei-tt
body: To be honest, I'm not convinced that all models will have aten.*_backward.* calls in backward pass. Here is an example that produces forward+backward aten IR for small linear model, that doesn't have any functions from the list: 

```python
import torch
import torch.nn as nn
from torch.fx import symbolic_trace
from torch.fx.experimental.proxy_tensor import make_fx

# Define a simple model with 2 linear layers and ReLU
class SimpleModel(nn.Module):
    def __init__(self, in_features=10, hidden_features=5, out_features=2):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_features, out_features)

    def forward(self, x):
        print(f"Input shape: {x.shape}")
        x = self.fc1(x)
        print(f"After fc1: {x.shape}")
        x = self.relu(x)
        print(f"After relu: {x.shape}")
        x = self.fc2(x)
        print(f"After fc2: {x.shape}")
        return x

if __name__ == "__main__":
    # Instantiate model and dummy input
    model = SimpleModel()
    x = torch.randn(1, 10, requires_grad=True)

    # Forward pass with shape printing
    y = model(x)

    # Trace the forward graph with torch.fx
    traced = symbolic_trace(model)
    print("\n--- Forward FX Graph ---")
    print(traced.graph)
    print("\n--- Forward FX Graph (Python IR) ---")
    print(traced.code)

    # Define a function combining forward & backward
    def forward_and_backward(inp):
        out = model(inp)
        loss = out.sum()
        loss.backward()
        return out

    # Use make_fx to capture both forward and backward aten IR
    fx_g = make_fx(forward_and_backward)(x)
    print("\n--- Forward+Backward FX Graph (aten IR) ---")
    print(fx_g.graph)
    
    print("\n --- Forward+Backward FX Graph (Python IR) ---")
    print(fx_g.code)
```

Forward + backward generated python code:
```
 --- Forward+Backward FX Graph (Python IR) ---



def forward(self, inp_1):
    _param_constant0 = self._param_constant0
    t = torch.ops.aten.t.default(_param_constant0);  _param_constant0 = None
    _param_constant1 = self._param_constant1
    addmm = torch.ops.aten.addmm.default(_param_constant1, inp_1, t);  _param_constant1 = None
    relu = torch.ops.aten.relu.default(addmm);  addmm = None
    detach = torch.ops.aten.detach.default(relu)
    _param_constant2 = self._param_constant2
    t_1 = torch.ops.aten.t.default(_param_constant2);  _param_constant2 = None
    _param_constant3 = self._param_constant3
    addmm_1 = torch.ops.aten.addmm.default(_param_constant3, relu, t_1);  _param_constant3 = None
    sum_1 = torch.ops.aten.sum.default(addmm_1)
    ones_like = torch.ops.aten.ones_like.default(sum_1, pin_memory = False, memory_format = torch.preserve_format);  sum_1 = None
    expand = torch.ops.aten.expand.default(ones_like, [1, 2]);  ones_like = None
    t_2 = torch.ops.aten.t.default(t_1);  t_1 = None
    mm = torch.ops.aten.mm.default(expand, t_2);  t_2 = None
    t_3 = torch.ops.aten.t.default(expand)
    mm_1 = torch.ops.aten.mm.default(t_3, relu);  t_3 = relu = None
    t_4 = torch.ops.aten.t.default(mm_1);  mm_1 = None
    sum_2 = torch.ops.aten.sum.dim_IntList(expand, [0], True);  expand = None
    view = torch.ops.aten.view.default(sum_2, [2]);  sum_2 = None
    detach_1 = torch.ops.aten.detach.default(view);  view = None
    detach_2 = torch.ops.aten.detach.default(detach_1);  detach_1 = None
    t_5 = torch.ops.aten.t.default(t_4);  t_4 = None
    detach_3 = torch.ops.aten.detach.default(t_5);  t_5 = None
    detach_4 = torch.ops.aten.detach.default(detach_3);  detach_3 = None
    detach_5 = torch.ops.aten.detach.default(detach);  detach = None
    threshold_backward = torch.ops.aten.threshold_backward.default(mm, detach_5, 0);  mm = detach_5 = None
    t_6 = torch.ops.aten.t.default(t);  t = None
    mm_2 = torch.ops.aten.mm.default(threshold_backward, t_6);  t_6 = None
    t_7 = torch.ops.aten.t.default(threshold_backward)
    mm_3 = torch.ops.aten.mm.default(t_7, inp_1);  t_7 = inp_1 = None
    t_8 = torch.ops.aten.t.default(mm_3);  mm_3 = None
    sum_3 = torch.ops.aten.sum.dim_IntList(threshold_backward, [0], True);  threshold_backward = None
    view_1 = torch.ops.aten.view.default(sum_3, [5]);  sum_3 = None
    detach_6 = torch.ops.aten.detach.default(view_1);  view_1 = None
    detach_7 = torch.ops.aten.detach.default(detach_6);  detach_6 = None
    detach_8 = torch.ops.aten.detach.default(mm_2);  mm_2 = None
    detach_9 = torch.ops.aten.detach.default(detach_8);  detach_8 = None
    t_9 = torch.ops.aten.t.default(t_8);  t_8 = None
    detach_10 = torch.ops.aten.detach.default(t_9);  t_9 = None
    detach_11 = torch.ops.aten.detach.default(detach_10);  detach_10 = None
    return addmm_1
 ```
created_at: 2025-05-21T10:49:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+#
+# SPDX-License-Identifier: Apache-2.0
+import torch
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from enum import Enum
+
+
+class ModelType(Enum):
+    """Enumeration of compiled model inputs.
+
+    It is expected that PARAMETER and BUFFER tensors do not change between inference runs, but ARGUMENT tensors may change.
+
+    :param PARAMETER: Tensor tracked by optimizer during training. Represents weights and biases.
+    :param BUFFER: Tensor not updated during training but still part of model state. Normally used for non-trainable data like fixed weights or running statistics in batch norm.
+    :param ARGUMENT: Tensors not tracked by model and not persistent between calls. Can be input data or configuration options to module's methods.
+    """
+
+    INFERENCE = 1
+    TRAIN_FORWARD = 2
+    TRAIN_BACKWARD = 3
+
+
+aten_backward_ops = {
+    torch.ops.aten._adaptive_avg_pool2d_backward.default,
+    torch.ops.aten.avg_pool2d_backward.default,
+    torch.ops.aten.convolution_backward.default,
+    torch.ops.aten.embedding_dense_backward.default,
+    torch.ops.aten.max_pool2d_with_indices_backward.default,
+    torch.ops.aten.native_group_norm_backward.default,
+    torch.ops.aten.native_layer_norm_backward.default,
+}
path: torch_ttnn/passes/analysis/graph_module_analysis_pass.py
user: login: jmalone-tt
body: Ok, looks like that graph contains `torch.ops.aten.threshold_backward.default`, which isn't in the list. I will go through and check the models that we support and make sure each of the backward passes would be caught in practice
created_at: 2025-05-21T15:07:43Z
author_association: COLLABORATOR

PR comments:
user: login: jmalone-tt
created_at: 2025-05-14T00:34:47Z
body: Based on a quick test, this lowers 5th iteration MNIST batch 1 from 5.29 ms to 3.32 ms. I have not tested BERT, data parallel, or other models.
author_association: COLLABORATOR
user: login: jmalone-tt
created_at: 2025-05-14T17:26:24Z
body: With this change, BERT batch 8 hits 23.3 sen/sec locally
author_association: COLLABORATOR
user: login: jmalone-tt
created_at: 2025-05-14T20:41:52Z
body: I'm sure the code in here could use some additional cleaning, but at this point I'm a bit too deep in the weeds, so I'm sure I'm missing things. Please point out areas where the code could be more readable
author_association: COLLABORATOR
user: login: jmalone-tt
created_at: 2025-05-15T18:23:39Z
body: Passed tests:
https://github.com/tenstorrent/pytorch2.0_ttnn/actions/runs/15048228625
author_association: COLLABORATOR
user: login: jmalone-tt
created_at: 2025-05-17T12:44:17Z
body: This passes consistently in the `Run Tests` workflow, and passes locally, but fails consistently in `Before Merge` 😕. Need to investigate further.
author_association: COLLABORATOR