Update some scripts in subgraph decomposer (#432)

* Fix

* Optimize typical_sequence_decomposer_test

* change the entry of naive_graph_decomposer from graph_net.torch.run_model to graph_net.model_path_handler

* update test_compiler and validator backend to support config and model_list

* Tidy model lists in repo

* Fix relative script

* Update the process of typical_sequence_decomposer

* renamed: naive_graph_decomposer -> graph_decomposer

* remove antique test code

* Fix: process torch.device and torch.dtype in inputs

* change the key of split_results.json from model_name to model_path

* Remove range_decomposer_backend for graph_net.torch.test_compiler, which is replaced by graph_net.torch.graph_decomposer

* change the debug info of fx.GraphModule replay error

* revert some change in decompose_util

* revert some change in decompose_util

Author: JewelRoam

graph_net/test/decomposer_validator_test.sh

@@ -248,6 +248,10 @@ def get_args_node(arg):
             yield arg.start
             yield arg.stop
             yield arg.step
+        elif isinstance(arg, torch.device):
+            pass
+        elif isinstance(arg, torch.dtype):
+            pass
         else:
             assert isinstance(
                 arg, (int, bool, float, str, type(...), type(None))

graph_net/torch/backend/range_decomposer_backend.py

@@ -15,6 +15,7 @@ def get_torch_module_and_inputs(model_path):
 def _get_torch_module(model_path):
     py_module = load_module(f"{model_path}/model.py")
     torch_module_cls = py_module.GraphModule
+    torch_module_cls.__graph_net_file_path__ = model_path
     return torch_module_cls()
 
 

graph_net/torch/decompose_util.py

@@ -221,10 +221,11 @@ def handle_underscore_suffix_difference():
 
     zip_filter_names = get_zip_filter_names()
 
-    def zip_filter_names_str():
+    def zip_filter_names_error_str():
         for triple in zip_filter_names:
             print(triple)
-        return "<printed before>"
+        error_model_path = module.__graph_net_file_path__
+        return f"{error_model_path=}"
 
     from pathlib import Path
 

graph_net/torch/fx_graph_module_util.py

@@ -193,9 +193,7 @@ def _make_config(
     def __call__(self, rel_model_path):
         model_path = os.path.join(self.config["model_path_prefix"], rel_model_path)
         split_results = load_json(self.config["split_results_path"])
-        split_positions = split_results[os.path.basename(rel_model_path)][
-            "split_points"
-        ]
+        split_positions = split_results[rel_model_path]["split_positions"]
         config = {
             "split_positions": split_positions,
             "group_head_and_tail": self.config.get("group_head_and_tail", False),

graph_net/torch/fx_graph_parse_util.py

@@ -21,7 +21,6 @@
 from graph_net.torch.backend.blade_disc_backend import BladeDISCBackend
 from graph_net.torch.backend.nope_backend import NopeBackend
 from graph_net.torch.backend.unstable_to_stable_backend import UnstableToStableBackend
-from graph_net.torch.backend.range_decomposer_backend import RangeDecomposerBackend
 from graph_net.torch.backend.range_decomposer_validator_backend import (
     RangeDecomposerValidatorBackend,
 )
@@ -37,7 +36,6 @@
     "bladedisc": BladeDISCBackend(),
     "nope": NopeBackend(),
     "unstable_to_stable": UnstableToStableBackend(),
-    "range_decomposer": RangeDecomposerBackend(),
     "range_decomposer_validator": RangeDecomposerValidatorBackend(),
 }
 

graph_net/torch/graph_decomposer.py

@@ -203,7 +203,7 @@ def analyze(self, model_paths_file: str, device: str) -> Dict[str, Dict]:
                 )
 
             current_idx = 0
-            split_points_set = set()
+            split_positions = set()
             total_len = sum(token2len.get(t, 1) for t in seq_tokens)
 
             for token_id in seq_tokens:
@@ -212,22 +212,22 @@ def analyze(self, model_paths_file: str, device: str) -> Dict[str, Dict]:
 
                 if is_pattern:
                     if current_idx > 0:
-                        split_points_set.add(current_idx)
+                        split_positions.add(current_idx)
                     end_idx = current_idx + length
                     if end_idx < total_len:
-                        split_points_set.add(end_idx)
+                        split_positions.add(end_idx)
 
                 current_idx += length
 
-            sorted_splits = sorted(list(split_points_set))
+            sorted_splits = sorted(list(split_positions))
 
             self._print_analysis(
                 model_name, str(original_path), sorted_splits, total_len, full_model_ops
             )
 
-            results[model_name] = {
-                "path": str(original_path),
-                "split_points": sorted_splits,
+            results[str(original_path)] = {
+                "model_name": model_name,
+                "split_positions": sorted_splits,
                 "total_length": total_len,
             }