more dimension generalization passes for token dimension

Author: lixinqi

graph_net/config/small_torch_samples_list.txt

@@ -1,3 +1,3 @@
 # samples/timm/resnetaa50d.d_in12k
-samples/transformers-auto-model/opus-mt-en-gmw
-# samples/transformers-auto-model/Michielo_mt5-small_nl-en_translation
+# samples/transformers-auto-model/opus-mt-en-gmw
+samples/transformers-auto-model/Michielo_mt5-small_nl-en_translation

graph_net/constraint_util.py

@@ -12,6 +12,7 @@
 import shutil
 from pathlib import Path
 import json
+from dataclasses import asdict
 
 
 class UpdateInputTensorConstraints:
@@ -75,7 +76,6 @@ def _make_config(
 
     def __call__(self, model_path):
         model_path = os.path.join(self.config["model_path_prefix"], model_path)
-        print(f"{model_path=}")
         cstr_path = os.path.join(model_path, "input_tensor_constraints.py")
         if (
             self.config["resume"]
@@ -143,15 +143,13 @@ def _try_dimension_generalization(self, dim_axes_pairs, model_path, tensor_metas
         )
         dim_generalizer = decorator_cls(self.config["dimension_generalizer_config"])
         dim_gen_pass = dim_generalizer(model, dim_axes_pairs)
-        if not dim_gen_pass.need_rewrite():
+        tensor_meta_attrs_list = [asdict(tensor_meta) for tensor_meta in tensor_metas]
+        inputs = dim_gen_pass.create_inputs_by_metas(tensor_meta_attrs_list)
+        if not dim_gen_pass.need_rewrite(inputs):
             yield model_path, ()
             return
-        from dataclasses import asdict
 
-        tensor_meta_attrs_list = [asdict(tensor_meta) for tensor_meta in tensor_metas]
-        graph_module = dim_gen_pass.rewrite_with_tensor_meta_attrs_list(
-            tensor_meta_attrs_list=tensor_meta_attrs_list,
-        )
+        graph_module = dim_gen_pass.rewrite(inputs)
         with tempfile.TemporaryDirectory() as tmp_dir:
             shutil.copytree(Path(model_path), Path(tmp_dir), dirs_exist_ok=True)
             dim_gen_pass.save_graph_module(graph_module, tmp_dir)

graph_net/imp_util.py

@@ -4,6 +4,7 @@
 def load_module(path, name="unamed"):
     spec = imp.spec_from_file_location(name, path)
     module = imp.module_from_spec(spec)
+    module.__file__ = path
     spec.loader.exec_module(module)
     module.__graph_net_file_path__ = path
     return module

graph_net/test/shape_prop_test.sh

@@ -0,0 +1,15 @@
+#!/bin/bash
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
+os.path.dirname(graph_net.__file__))")
+
+config_json_str=$(cat <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/torch/constraint_util.py",
+    "handler_class_name": "ShapePropagatablePredicator"
+}
+EOF
+)
+CONFIG=$(echo $config_json_str | base64 -w 0)
+
+python3 -m graph_net.model_path_handler --model-path $1 --handler-config=$CONFIG

graph_net/torch/dim_gen_passes/naive_call_method_expand_pass.py

@@ -1,12 +1,16 @@
 import torch
 import torch.fx as fx
 from graph_net.torch.dim_gen_passes import DimensionGeneralizationPass
+import os
 
 
 class ConcretePass(DimensionGeneralizationPass):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+    def get_pass_name(cls) -> bool:
+        return os.path.basename(__file__)[:-3]
+
     def need_rewrite(self, traced_module: fx.GraphModule) -> bool:
         if 0 not in self.axes:
             return False

graph_net/torch/dim_gen_passes/naive_call_method_reshape_pass.py

@@ -1,12 +1,16 @@
 import torch
 import torch.fx as fx
 from graph_net.torch.dim_gen_passes import DimensionGeneralizationPass
+import os
 
 
 class ConcretePass(DimensionGeneralizationPass):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+    def get_pass_name(cls) -> bool:
+        return os.path.basename(__file__)[:-3]
+
     def need_rewrite(self, traced_module: fx.GraphModule) -> bool:
         if 0 not in self.axes:
             return False

graph_net/torch/dim_gen_passes/naive_call_method_view_pass.py

@@ -1,12 +1,16 @@
 import torch
 import torch.fx as fx
 from graph_net.torch.dim_gen_passes import DimensionGeneralizationPass
+import os
 
 
 class ConcretePass(DimensionGeneralizationPass):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+    def get_pass_name(cls) -> bool:
+        return os.path.basename(__file__)[:-3]
+
     def need_rewrite(self, traced_module: fx.GraphModule) -> bool:
         if 0 not in self.axes:
             return False

graph_net/torch/dim_gen_passes/non_batch_call_function_arange_pass.py

@@ -2,32 +2,38 @@
 import torch.fx as fx
 from graph_net.torch.dim_gen_passes import DimensionGeneralizationPass
 from collections import namedtuple
+import os
 
 
 class ConcretePass(DimensionGeneralizationPass):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+    def get_pass_name(cls) -> bool:
+        return os.path.basename(__file__)[:-3]
+
     def need_rewrite(self, traced_module: fx.GraphModule) -> bool:
         # non batch
         if 0 in self.axes:
             return False
         return any(self._node_need_rewrite(node) for node in traced_module.graph.nodes)
 
+    def node_target(self):
+        return torch.arange
+
     def _node_need_rewrite(self, node) -> bool:
-        return (
-            node.op == "call_function"
-            and node.target == torch.arange
-            and len(node.args) >= 2
-            and node.args[0] == 0
-            and node.args[1] == self.dim
-        )
+        if not (node.op == "call_function"):
+            return False
+        if not (node.target == self.node_target()):
+            return False
+        if len(node.args) == 1:
+            return node.args[0] == self.dim
+        elif len(node.args) == 2:
+            return node.args[0] == 0 and node.args[1] == self.dim
+        else:
+            return False
 
     def rewrite(self, traced_module: fx.GraphModule) -> fx.GraphModule:
-        """
-        Fx Pass: Replaces hardcoded constants in 'torch.arange' ops that match an input tensor dimension
-        with a dynamic 'size()' call. The primary goal is to dynamicize the batch size.
-        """
         # Create a new graph to hold the rewritten nodes
         new_graph = fx.Graph()
 
@@ -50,8 +56,8 @@ def try_reset_last_node_axis(node, new_node):
                 last_node_axis = NodeAxis(node=new_node, shape_axis=axis)
                 return
 
-        def get_new_node_arg(i, arg):
-            if i != 1:
+        def get_new_node_arg(i, arg, len_args):
+            if not ((len_args == 2 and i == 1) or (len_args == 1 and i == 0)):
                 return val_map[arg] if arg in val_map else arg
             # i == 1
             assert arg == self.dim
@@ -70,12 +76,14 @@ def create_new_node(node):
                 try_reset_last_node_axis(node=node, new_node=new_node)
                 return new_node
 
-            new_arange_args = tuple(
-                get_new_node_arg(i, arg) for i, arg in enumerate(node.args)
+            new_node_args = tuple(
+                get_new_node_arg(i, arg, len(node.args))
+                for i, arg in enumerate(node.args)
             )
 
-            # --- Rebuild the torch.arange node ---
-            new_node = new_graph.call_function(torch.arange, args=new_arange_args)
+            new_node = new_graph.call_function(
+                self.node_target(), args=new_node_args, kwargs=node.kwargs
+            )
 
             return new_node
 

graph_net/torch/dim_gen_passes/non_batch_call_function_arange_plus_one_pass.py

@@ -0,0 +1,118 @@
+import torch
+import torch.fx as fx
+from graph_net.torch.dim_gen_passes import DimensionGeneralizationPass
+from collections import namedtuple
+import os
+import operator
+
+
+class ConcretePass(DimensionGeneralizationPass):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def get_pass_name(cls) -> bool:
+        return os.path.basename(__file__)[:-3]
+
+    def need_rewrite(self, traced_module: fx.GraphModule) -> bool:
+        # non batch
+        if 0 in self.axes:
+            return False
+        if not any(self._is_pivote_node(node) for node in traced_module.graph.nodes):
+            return False
+        return any(self._node_need_rewrite(node) for node in traced_module.graph.nodes)
+
+    def _node_target(self):
+        return torch.arange
+
+    def _is_pivote_node(self, node):
+        if not (node.op == "call_function"):
+            return False
+        if not (self._is_pivote_target(node.target)):
+            return False
+        if not (len(node.args) == 1):
+            return False
+        input_node = node.args[0]
+        input_node_meta = input_node.meta.get("tensor_meta")
+        if not (input_node_meta is not None):
+            return False
+        if not (hasattr(input_node_meta, "shape")):
+            return False
+        shape = input_node_meta.shape
+        if not any(dim == self.dim + self._dyn_dim_delta() for dim in shape):
+            return False
+        return True
+
+    def _is_pivote_target(self, target):
+        return target == torch.triu
+
+    def _node_need_rewrite(self, node) -> bool:
+        if not (node.op == "call_function"):
+            return False
+        if not (node.target == self._node_target()):
+            return False
+        if not (len(node.args) == 1):
+            return False
+        if not (node.args[0] == self.dim + self._dyn_dim_delta()):
+            return False
+        return True
+
+    def _dyn_dim_delta(self):
+        return 1
+
+    def rewrite(self, traced_module: fx.GraphModule) -> fx.GraphModule:
+        # Create a new graph to hold the rewritten nodes
+        new_graph = fx.Graph()
+
+        # Create a map to link nodes from the old graph to nodes in the new graph
+        val_map = {}
+
+        NodeAxis = namedtuple("NodeAxis", ["node", "shape_axis"])
+        last_node_axis = None
+
+        def try_reset_last_node_axis(node, new_node):
+            nonlocal last_node_axis
+            node_meta = node.meta.get("tensor_meta")
+            if node_meta is None:
+                return
+            if not hasattr(node_meta, "shape"):
+                return
+            for axis, dim in enumerate(node_meta.shape):
+                if not (dim == self.dim and axis > 0):
+                    continue
+                last_node_axis = NodeAxis(node=new_node, shape_axis=axis)
+                return
+
+        def get_new_node_dim(dim):
+            if not (dim == self.dim + self._dyn_dim_delta()):
+                return val_map[dim] if dim in val_map else dim
+            assert dim == self.dim + self._dyn_dim_delta()
+
+            # Use the size() method to retrieve the dynamic dimension
+            size_node = new_graph.call_method(
+                "size", args=(last_node_axis.node, last_node_axis.shape_axis)
+            )
+            plus_one_node = new_graph.call_function(operator.add, args=(size_node, 1))
+            return plus_one_node
+
+        def create_new_node(node):
+            if not (self._node_need_rewrite(node) and last_node_axis is not None):
+                # Copy other nodes to the new graph
+                new_node = new_graph.node_copy(node, lambda x: val_map[x])
+                try_reset_last_node_axis(node=node, new_node=new_node)
+                return new_node
+
+            new_node_dim = get_new_node_dim(dim=node.args[0])
+
+            new_node = new_graph.call_function(
+                self._node_target(), args=(new_node_dim,), kwargs=node.kwargs
+            )
+
+            return new_node
+
+        for node in traced_module.graph.nodes:
+            val_map[node] = create_new_node(node)
+
+        # Replace the old graph with the new graph and return
+        traced_module.graph = new_graph
+        traced_module.recompile()
+        return traced_module