Merge branch 'develop' into binary_decomposer_paddle

Xreki · Xreki · commit 53b0e4faec68 · 2025-11-28T17:57:43.000+08:00
diff --git a/graph_net/config/small_torch_samples_list.txt b/graph_net/config/small_torch_samples_list.txt
@@ -1,3 +1,4 @@
+samples/timm/davit_giant
 # 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/Michielo_mt5-small_nl-en_translation
diff --git a/graph_net/config/todo_torch_samples_list.txt b/graph_net/config/todo_torch_samples_list.txt
diff --git a/graph_net/constraint_util.py b/graph_net/constraint_util.py
@@ -1,3 +1,4 @@
+import logging
 from graph_net.dynamic_dim_constraints import DynamicDimConstraints
 from contextlib import AbstractContextManager
 from graph_net.imp_util import load_module
@@ -88,14 +89,21 @@ def __call__(self, model_path):
                 return
 
         tensor_metas = self._get_tensor_metas(model_path)
+        tensor_meta_attrs_list = [asdict(tensor_meta) for tensor_meta in tensor_metas]
+        logging.warning(f"before create_inputs_by_metas")
+        inputs = self.get_dimension_generalizer().create_inputs_by_metas(
+            module=self.get_model(model_path),
+            tensor_meta_attrs_list=tensor_meta_attrs_list,
+        )
+        logging.warning(f"after create_inputs_by_metas")
         dyn_dim_cstr = make_dyn_dim_cstr_from_tensor_metas(tensor_metas)
 
         def data_input_predicator(input_var_name):
             return self.data_input_predicator(model_path, input_var_name)
 
         def get_tmp_model_path_ctx_mgr(dim_axes_pairs):
             return self._try_dimension_generalization(
-                dim_axes_pairs, model_path, tensor_metas
+                dim_axes_pairs, model_path, inputs
             )
 
         def get_predicator_is_dyn_dim_cstr_feasible(tmp_model_path):
@@ -128,28 +136,44 @@ def is_dyn_dim_cstr_feasible(dyn_dim_cstr):
                     )
                     sys.exit(0)
 
-    @contextmanager
-    def _try_dimension_generalization(self, dim_axes_pairs, model_path, tensor_metas):
-        if self.config["dimension_generalizer_filepath"] is None:
-            yield model_path, ()
-            return
-        py_module = load_module(os.path.join(model_path, "model.py"))
-        GraphModule = getattr(py_module, "GraphModule")
-        GraphModule.__graph_net_file_path__ = py_module.__graph_net_file_path__
-        model = GraphModule()
+    def get_dimension_generalizer(self):
+        if hasattr(self, "_dim_generalizer"):
+            return self._dim_generalizer
+        assert self.config["dimension_generalizer_filepath"] is not None
         decorator_cls = getattr(
             load_module(self.config["dimension_generalizer_filepath"]),
             self.config["dimension_generalizer_class_name"],
         )
-        dim_generalizer = decorator_cls(self.config["dimension_generalizer_config"])
+        self._dim_generalizer = decorator_cls(
+            self.config["dimension_generalizer_config"]
+        )
+        return self._dim_generalizer
+
+    def get_model(self, model_path):
+        py_module = load_module(os.path.join(model_path, "model.py"))
+        GraphModule = getattr(py_module, "GraphModule")
+        GraphModule.__graph_net_file_path__ = py_module.__graph_net_file_path__
+        return GraphModule()
+
+    @contextmanager
+    def _try_dimension_generalization(self, dim_axes_pairs, model_path, inputs):
+        logging.warning(f"enter _try_dimension_generalization")
+        if self.config["dimension_generalizer_filepath"] is None:
+            yield model_path, ()
+            return
+        model = self.get_model(model_path)
+        dim_generalizer = self.get_dimension_generalizer()
         dim_gen_pass = dim_generalizer(model, dim_axes_pairs)
-        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):
+        logging.warning(f"before need_rewrite")
+        need_rewrite = dim_gen_pass.need_rewrite(inputs)
+        logging.warning(f"after need_rewrite")
+        if not need_rewrite:
             yield model_path, ()
             return
 
+        logging.warning(f"before rewrite")
         graph_module = dim_gen_pass.rewrite(inputs)
+        logging.warning(f"after rewrite")
         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)
@@ -300,6 +324,7 @@ def append_dim_gen_pass_names(dim_gen_pass_names):
         )
 
     for i, picked_dim in enumerate(unqiue_dims):
+        logging.warning(f"{i=} {picked_dim=}")
         cur_dyn_dim_cstr = copy.deepcopy(dyn_dim_cstr)
 
         def filter_fn(input_name, input_idx, axis, dim):
@@ -319,11 +344,17 @@ def filter_fn(input_name, input_idx, axis, dim):
             (dim, axes) for dim in unqiue_dims[: i + 1] for axes in [dim2axes[dim]]
         )
         ctx_mgr = dyn_dim_cstr_feasibility_ctx_mgr
+        logging.warning(f"before dyn_dim_cstr_feasibility_ctx_mgr")
         with ctx_mgr(dim_axes_pairs) as dyn_dim_cstr_feasibility:
+            logging.warning(f"enter dyn_dim_cstr_feasibility_ctx_mgr")
             tmp_dyn_dim_cstr = copy.deepcopy(cur_dyn_dim_cstr)
             tmp_dyn_dim_cstr.update_symbol2example_value(sym2example_value)
-            if not dyn_dim_cstr_feasibility(tmp_dyn_dim_cstr):
+            logging.warning(f"before dyn_dim_cstr_feasibility")
+            is_dyn_dim_cstr_feasible = dyn_dim_cstr_feasibility(tmp_dyn_dim_cstr)
+            logging.warning(f"after dyn_dim_cstr_feasibility")
+            if not is_dyn_dim_cstr_feasible:
                 continue
             dyn_dim_cstr = cur_dyn_dim_cstr
             append_dim_gen_pass_names(dyn_dim_cstr_feasibility.dim_gen_pass_names)
+            logging.warning(f"leave dyn_dim_cstr_feasibility_ctx_mgr")
     return dyn_dim_cstr, total_dim_gen_pass_names
diff --git a/graph_net/dynamic_dim_constraints.py b/graph_net/dynamic_dim_constraints.py
@@ -44,6 +44,8 @@ def symbolize(
         ]
         Returns created symbol.
         """
+        import logging
+
         InputDim = namedtuple("InputDim", ["input_idx", "axis", "dim"])
         input_dims = [
             InputDim(input_idx, axis, dim)
diff --git a/graph_net/model_path_handler.py b/graph_net/model_path_handler.py
@@ -10,6 +10,11 @@
 import json
 import base64
 from contextlib import contextmanager
+import logging
+
+logging.basicConfig(
+    level=logging.WARNING, format="%(asctime)s [%(levelname)s] %(message)s"
+)
 
 
 def _convert_to_dict(config_str):
diff --git a/graph_net/test/batch_init_input_tensor_constraints_test.sh b/graph_net/test/batch_init_input_tensor_constraints_test.sh
@@ -11,7 +11,7 @@ config_json_str=$(cat <<EOF
     "handler_path": "$GRAPH_NET_ROOT/constraint_util.py",
     "handler_class_name": "UpdateInputTensorConstraints",
     "handler_config": {
-        "resume": false,
+        "resume": true,
         "model_path_prefix": "$GRAPH_NET_ROOT/../",
         "data_input_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
         "data_input_predicator_class_name": "NaiveDataInputPredicator",
@@ -26,4 +26,4 @@ EOF
 )
 CONFIG=$(echo $config_json_str | base64 -w 0)
 
-python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/small_torch_samples_list.txt --handler-config=$CONFIG
+python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/torch_samples_list.txt --handler-config=$CONFIG
diff --git a/graph_net/test/subgraph_decompose_and_evaluation_step_test.sh b/graph_net/test/subgraph_decompose_and_evaluation_step_test.sh
@@ -79,6 +79,11 @@ python3 -m graph_net.subgraph_decompose_and_evaluation_step \
     --tolerance="$TOLERANCE" \
     --max-subgraph-size="$INITIAL_MAX_SIZE"
 
-echo ""
-echo ">>> Pass execution finished."
-echo ">>> Run this script again to execute the NEXT pass if needed."
+if [ $? -ne 0 ]; then
+    echo ""
+    echo "[ERROR] Task failed! Please check logs and fix bugs before proceeding."
+else
+    echo ""
+    echo ">>> Pass execution finished."
+    echo ">>> Run this script again to execute the NEXT pass if needed."
+fi
diff --git a/graph_net/tools/batch_init_input_tensor_constraints.sh b/graph_net/tools/batch_init_input_tensor_constraints.sh
@@ -0,0 +1,29 @@
+#!/bin/bash
+
+GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(
+os.path.dirname(graph_net.__file__))")
+
+# input model path
+# model_runnable_predicator=ShapePropagatablePredicator
+model_runnable_predicator=ModelRunnablePredicator
+config_json_str=$(cat <<EOF
+{
+    "handler_path": "$GRAPH_NET_ROOT/constraint_util.py",
+    "handler_class_name": "UpdateInputTensorConstraints",
+    "handler_config": {
+        "resume": true,
+        "model_path_prefix": "$GRAPH_NET_ROOT/../",
+        "data_input_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
+        "data_input_predicator_class_name": "NaiveDataInputPredicator",
+        "model_runnable_predicator_filepath": "$GRAPH_NET_ROOT/torch/constraint_util.py",
+        "model_runnable_predicator_class_name": "$model_runnable_predicator",
+        "dimension_generalizer_filepath": "$GRAPH_NET_ROOT/torch/static_to_dynamic.py",
+        "dimension_generalizer_class_name": "StaticToDynamic",
+        "last_model_log_file": "/tmp/a.py"
+    }
+}
+EOF
+)
+CONFIG=$(echo $config_json_str | base64 -w 0)
+
+python3 -m graph_net.model_path_handler --model-path-list $GRAPH_NET_ROOT/config/torch_samples_list.txt --handler-config=$CONFIG
diff --git a/graph_net/torch/constraint_util.py b/graph_net/torch/constraint_util.py
@@ -1,5 +1,6 @@
 import sys
 import os
+import graph_net
 
 
 class NaiveDataInputPredicator:
@@ -12,23 +13,27 @@ def __call__(self, model_path, input_var_name: str) -> bool:
 
 class ModelRunnablePredicator:
     def __init__(self, config):
-        self.config = config
+        if config is None:
+            config = {}
+
+        graph_net_root = os.path.dirname(graph_net.__file__)
+        decorator_config = {"use_dummy_inputs": True}
+        self.predicator = RunModelPredicator(decorator_config)
 
     def __call__(self, model_path):
-        cmd = f"{sys.executable} -m graph_net.torch.run_model --model-path {model_path}"
-        return os.system(cmd) == 0
+        return self.predicator(model_path)
 
 
 class ShapePropagatablePredicator:
     def __init__(self, config=None):
         if config is None:
             config = {}
-        import graph_net
 
         graph_net_root = os.path.dirname(graph_net.__file__)
         decorator_config = {
             "decorator_path": f"{graph_net_root}/torch/shape_prop.py",
             "decorator_class_name": "ShapePropagate",
+            "use_dummy_inputs": True,
         }
         self.predicator = RunModelPredicator(decorator_config)
 
diff --git a/graph_net/torch/fx_graph_cache_util.py b/graph_net/torch/fx_graph_cache_util.py
@@ -1,3 +1,4 @@
+import logging
 import torch
 import copy
 import os
@@ -13,8 +14,12 @@ def parse_immutable_model_path_into_sole_graph_module(model_path):
     if model_path not in g_model_path2graph_module:
         module = _get_torch_module(model_path)
         tensor_metas = _get_tensor_metas(model_path)
+        logging.warning("before _create_inputs_by_metas")
         inputs = _create_inputs_by_metas(module, tensor_metas)
+        logging.warning("after _create_inputs_by_metas")
+        logging.warning("before parse_sole_graph_module")
         g_model_path2graph_module[model_path] = parse_sole_graph_module(module, inputs)
+        logging.warning("after parse_sole_graph_module")
     return copy.deepcopy(g_model_path2graph_module[model_path])
 
 
@@ -34,11 +39,9 @@ def _get_tensor_metas(model_path):
 
 def _create_inputs_by_metas(module, tensor_metas):
     tensor_meta_attrs_list = [asdict(tensor_meta) for tensor_meta in tensor_metas]
-    from graph_net.torch.utils import convert_tensor_meta_attrs_list_to_named_tensors
+    from graph_net.torch.utils import get_dummy_named_tensors
 
-    named_tensors = convert_tensor_meta_attrs_list_to_named_tensors(
-        tensor_meta_attrs_list
-    )
+    named_tensors = get_dummy_named_tensors(tensor_meta_attrs_list)
     name2tensor = {k: v for k, v in named_tensors}
     return tuple(
         name2tensor[name] for name in inspect.signature(module.forward).parameters
diff --git a/graph_net/torch/fx_graph_parse_util.py b/graph_net/torch/fx_graph_parse_util.py
@@ -1,4 +1,5 @@
 import torch
+import inspect
 
 
 def parse_sole_graph_module(module, inputs):
@@ -14,12 +15,46 @@ def my_backend(gm, sample_inputs):
 
     torch.compile(module, backend=my_backend)(*inputs)
     assert traced_module is not None
-    assert all(id(a) == id(b) for a, b in zip(inputs, traced_sample_inputs))
     for node in traced_module.graph.nodes:
         if node.op != "placeholder":
             continue
         assert node.target[:2] == "L_" or node.target[:2] == "l_", f"{node.target=}"
         node.target = node.target[2:]
+        if node.target[0] == "l":
+            node.target = "L" + node.target[1:]
         assert node.name[:2] == "L_" or node.name[:2] == "l_", f"{node.name=}"
         node.name = node.name[2:]
+        if node.name[0] == "l":
+            node.name = "L" + node.name[1:]
+
+    def get_input_names_from_signature():
+        return inspect.signature(module.forward).parameters
+
+    def get_input_names_from_placeholder():
+        return [
+            node.name for node in traced_module.graph.nodes if node.op == "placeholder"
+        ]
+
+    def get_diff_input_names():
+        placeholder_names = set(get_input_names_from_placeholder())
+        return [
+            (i, name)
+            for i, name in enumerate(get_input_names_from_signature())
+            if name not in placeholder_names
+        ]
+
+    if len(inputs) == len(traced_sample_inputs) + 1:
+        diff_input_names = get_diff_input_names()
+        assert len(diff_input_names) == 1, f"{diff_input_names=}"
+        pos, name = diff_input_names[0]
+        for i, node in enumerate(traced_module.graph.nodes):
+            if i < pos:
+                assert node.op == "placeholder"
+            elif i == pos:
+                with traced_module.graph.inserting_before(node):
+                    traced_module.graph.placeholder(name)
+            else:
+                break
+        traced_module.recompile()
+    assert len(get_diff_input_names()) == 0, f"{get_diff_input_names()=}"
     return traced_module
diff --git a/graph_net/torch/run_model.py b/graph_net/torch/run_model.py
@@ -30,10 +30,7 @@ def _convert_to_dict(config_str):
     return config
 
 
-def _get_decorator(args):
-    if args.decorator_config is None:
-        return lambda model: model
-    decorator_config = _convert_to_dict(args.decorator_config)
+def _get_decorator(decorator_config):
     if "decorator_path" not in decorator_config:
         return lambda model: model
     class_name = decorator_config.get("decorator_class_name", "RunModelDecorator")
@@ -44,6 +41,17 @@ def _get_decorator(args):
     return decorator_class(decorator_config.get("decorator_config", {}))
 
 
+def get_flag_use_dummy_inputs(decorator_config):
+    return "use_dummy_inputs" in decorator_config
+
+
+def replay_tensor(info, use_dummy_inputs):
+    if use_dummy_inputs:
+        return utils.get_dummy_tensor(info)
+    else:
+        return utils.replay_tensor(info)
+
+
 def main(args):
     model_path = args.model_path
     model_class = load_class_from_file(
@@ -53,11 +61,15 @@ def main(args):
     model = model_class()
     print(f"{model_path=}")
 
-    model = _get_decorator(args)(model)
+    decorator_config = _convert_to_dict(args.decorator_config)
+    if "decorator_path" in decorator_config:
+        model = _get_decorator(decorator_config)(model)
 
     inputs_params = utils.load_converted_from_text(f"{model_path}")
     params = inputs_params["weight_info"]
-    state_dict = {k: utils.replay_tensor(v) for k, v in params.items()}
+    use_dummy_inputs = get_flag_use_dummy_inputs(decorator_config)
+    print(f"{use_dummy_inputs=}")
+    state_dict = {k: replay_tensor(v, use_dummy_inputs) for k, v in params.items()}
 
     model(**state_dict)
 
diff --git a/graph_net/torch/static_to_dynamic.py b/graph_net/torch/static_to_dynamic.py
diff --git a/graph_net/torch/subgraph_decompose_and_evaluation_step.py b/graph_net/torch/subgraph_decompose_and_evaluation_step.py
diff --git a/graph_net/torch/utils.py b/graph_net/torch/utils.py

Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,8 @@ def symbolize(`
`44`	`44`	`]`
`45`	`45`	`Returns created symbol.`
`46`	`46`	`"""`
	`47`	`+ import logging`
	`48`	`+`
`47`	`49`	`InputDim = namedtuple("InputDim", ["input_idx", "axis", "dim"])`
`48`	`50`	`input_dims = [`
`49`	`51`	`InputDim(input_idx, axis, dim)`