Change the entry of naive_graph_decomposer from graph_net.torch.run_model to graph_net.model_path_handler (#413)

JewelRoam · web-flow · commit 8ba1b8f4e89b · 2025-12-04T22:48:14.000+08:00
* 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
diff --git a/graph_net/test/naive_graph_decomposer_test.sh b/graph_net/test/naive_graph_decomposer_test.sh
@@ -6,23 +6,20 @@ os.path.dirname(graph_net.__file__))")
 # input model path
 MODEL_NAME=resnet18
 MODEL_PATH_IN_SAMPLES=/timm/$MODEL_NAME
-decorator_config_json_str=$(cat <<EOF
+config_json_str=$(cat <<EOF
 {
-    "decorator_path": "$GRAPH_NET_ROOT/torch/extractor.py",
-    "decorator_config": {
-        "name": "$MODEL_NAME",
-        "custom_extractor_path": "$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py",
-        "custom_extractor_config": {
-            "output_dir": "/tmp/naive_decompose_workspace",
-            "split_positions": [8, 16, 32],
-            "group_head_and_tail": true,
-            "filter_path":"$GRAPH_NET_ROOT/torch/naive_subgraph_filter.py",
-            "filter_config": {}
-        }
+    "handler_path": "$GRAPH_NET_ROOT/torch/naive_graph_decomposer.py",
+    "handler_class_name": "NaiveDecomposerExtractor",
+    "handler_config": {
+        "output_dir": "/tmp/naive_decompose_workspace",
+        "split_positions": [8, 16, 32],
+        "group_head_and_tail": true,
+        "filter_path":"$GRAPH_NET_ROOT/torch/naive_subgraph_filter.py",
+        "filter_config": {}
     }
 }
 EOF
 )
-DECORATOR_CONFIG=$(echo $decorator_config_json_str | base64 -w 0)
+CONFIG=$(echo $config_json_str | base64 -w 0)
 
-python3 -m graph_net.torch.run_model --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --decorator-config=$DECORATOR_CONFIG
+python3 -m graph_net.model_path_handler --model-path $GRAPH_NET_ROOT/../samples/$MODEL_PATH_IN_SAMPLES --handler-config=$CONFIG
diff --git a/graph_net/test/typical_sequence_decomposer_test.sh b/graph_net/test/typical_sequence_decomposer_test.sh
@@ -1,46 +1,60 @@
 #!/bin/bash
 
 GRAPH_NET_ROOT=$(python3 -c "import graph_net; import os; print(os.path.dirname(os.path.dirname(graph_net.__file__)))")
+DECOMPOSE_PATH=$GRAPH_NET_ROOT/decompose_workspace
+
+mkdir -p "$DECOMPOSE_PATH"
 
-MODEL1="$GRAPH_NET_ROOT/samples/torchvision/resnet18"
-MODEL2="$GRAPH_NET_ROOT/samples/torchvision/resnet34"
-MODEL_LIST_FILE=$(mktemp)
-echo "$MODEL1" > "$MODEL_LIST_FILE"
-echo "$MODEL2" >> "$MODEL_LIST_FILE"
+temp_model_list=$(mktemp)
+cat "$GRAPH_NET_ROOT/graph_net/config/torch_samples_list.txt" > "$temp_model_list"
 
 python3 -m graph_net.torch.typical_sequence_split_points \
-    --model-list "$MODEL_LIST_FILE" \
+    --model-list "$temp_model_list" \
     --device "cuda" \
     --window-size 10 \
-    --output-json "$GRAPH_NET_ROOT/split_results.json"
+    --output-json "$DECOMPOSE_PATH/split_results.json"
 
-rm -f "$MODEL_LIST_FILE"
+while IFS= read -r MODEL_PATH_IN_SAMPLES; do
+    if [[ -n "$MODEL_PATH_IN_SAMPLES" ]]; then
+        MODEL_FULL_PATH="$GRAPH_NET_ROOT/$MODEL_PATH_IN_SAMPLES"
+        MODEL_NAME=$(basename "$MODEL_PATH_IN_SAMPLES")
 
+        echo "== Decomposing $MODEL_PATH_IN_SAMPLES. =="
 
-MODEL_PATH_IN_SAMPLES=/torchvision/resnet18
-MODEL_NAME=$(basename "$MODEL_PATH_IN_SAMPLES")
-
-decomposer_config_json_str=$(cat <<EOF
+        decomposer_config_json_str=$(cat <<EOF
 {
-    "split_results_path": "$GRAPH_NET_ROOT/split_results.json",
-    "workspace_path": "$GRAPH_NET_ROOT/decompose_workspace",
-    "chain_style": "True"
+    "split_results_path": "$DECOMPOSE_PATH/split_results.json",
+    "workspace_path": "$DECOMPOSE_PATH",
+    "chain_style": true,
+    "target_model_name": "$MODEL_NAME"
 }
 EOF
-)
-DECOMPOSER_CONFIG=$(echo $decomposer_config_json_str | base64 -w 0)
+        )
+        DECOMPOSER_CONFIG=$(echo $decomposer_config_json_str | base64 -w 0)
 
-python3 -m graph_net.torch.test_compiler --model-path $GRAPH_NET_ROOT/samples/$MODEL_PATH_IN_SAMPLES --compiler range_decomposer --device cuda --config=$DECOMPOSER_CONFIG
+        python3 -m graph_net.torch.test_compiler \
+            --model-path "$MODEL_FULL_PATH" \
+            --compiler range_decomposer \
+            --device cuda \
+            --config="$DECOMPOSER_CONFIG"
 
+        cp -r "$MODEL_FULL_PATH" "$DECOMPOSE_PATH/"
 
-DECOMPOSE_PATH=$GRAPH_NET_ROOT/decompose_workspace
-cp -r "$GRAPH_NET_ROOT/samples/$MODEL_PATH_IN_SAMPLES" "$DECOMPOSE_PATH/"
+        echo "== Validating $MODEL_PATH_IN_SAMPLES. =="
+
+        python3 -m graph_net.torch.test_compiler \
+            --model-path "$DECOMPOSE_PATH/$MODEL_NAME" \
+            --compiler range_decomposer_validator \
+            --device cuda > "$DECOMPOSE_PATH/${MODEL_NAME}_validation.log" 2>&1
+
+        echo "== Finished processing $MODEL_PATH_IN_SAMPLES. =="
+    fi
+done < $temp_model_list
+
+rm -f "$temp_model_list"
 
-python3 -m graph_net.torch.test_compiler \
-    --model-path $DECOMPOSE_PATH/$MODEL_NAME \
-    --compiler range_decomposer_validator \
-    --device cuda > "$DECOMPOSE_PATH/log.log" 2>&1
+cat $DECOMPOSE_PATH/*_validation.log >> $DECOMPOSE_PATH/combined.log
 
 python3 -m graph_net.plot_ESt \
-  --benchmark-path $DECOMPOSE_PATH/log.log \
-  --output-dir $DECOMPOSE_PATH \
+    --benchmark-path "$DECOMPOSE_PATH/combined.log" \
+    --output-dir "$DECOMPOSE_PATH"
diff --git a/graph_net/torch/fx_graph_cache_util.py b/graph_net/torch/fx_graph_cache_util.py
@@ -1,51 +1,15 @@
-import logging
-import torch
 import copy
 import os
-import inspect
-from graph_net.tensor_meta import TensorMeta
 from graph_net.torch.fx_graph_parse_util import parse_sole_graph_module
-from graph_net.imp_util import load_module
-from dataclasses import asdict
+from graph_net.torch.fx_graph_module_util import get_torch_module_and_inputs
 
 
 def parse_immutable_model_path_into_sole_graph_module(model_path):
     model_path = os.path.realpath(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")
+        module, inputs = get_torch_module_and_inputs(model_path)
         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])
 
 
-def _get_torch_module(model_path):
-    py_module = load_module(f"{model_path}/model.py")
-    torch_module_cls = py_module.GraphModule
-    return torch_module_cls()
-
-
-def _get_tensor_metas(model_path):
-    make = TensorMeta.unserialize_from_py_file
-    return [
-        *make(os.path.join(model_path, "input_meta.py")),
-        *make(os.path.join(model_path, "weight_meta.py")),
-    ]
-
-
-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 get_dummy_named_tensors
-
-    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
-    )
-
-
 g_model_path2graph_module = {}
diff --git a/graph_net/torch/fx_graph_module_util.py b/graph_net/torch/fx_graph_module_util.py
@@ -0,0 +1,37 @@
+import os
+import inspect
+from graph_net.tensor_meta import TensorMeta
+from graph_net.imp_util import load_module
+from dataclasses import asdict
+
+
+def get_torch_module_and_inputs(model_path):
+    module = _get_torch_module(model_path)
+    tensor_metas = _get_tensor_metas(model_path)
+    inputs = _create_inputs_by_metas(module, tensor_metas)
+    return module, inputs
+
+
+def _get_torch_module(model_path):
+    py_module = load_module(f"{model_path}/model.py")
+    torch_module_cls = py_module.GraphModule
+    return torch_module_cls()
+
+
+def _get_tensor_metas(model_path):
+    make = TensorMeta.unserialize_from_py_file
+    return [
+        *make(os.path.join(model_path, "input_meta.py")),
+        *make(os.path.join(model_path, "weight_meta.py")),
+    ]
+
+
+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 get_dummy_named_tensors
+
+    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
@@ -76,12 +76,45 @@ def get_zip_filter_names():
             if name_from_signature != name_from_placeholder
         )
 
+    if len(get_zip_filter_names()) > 0 and set(get_input_names_from_signature()) == set(
+        get_input_names_from_placeholder()
+    ):
+        traced_module = _reorder_placeholders(
+            traced_module, get_input_names_from_signature()
+        )
+
     zip_filter_names = get_zip_filter_names()
 
     def zip_filter_names_str():
         for triple in zip_filter_names:
             print(triple)
         return "<printed before>"
 
+    from pathlib import Path
+
+    Path("/tmp/a.py").write_text(traced_module.code)
     assert len(zip_filter_names) == 0, f"{zip_filter_names_str()=}"
     return traced_module
+
+
+def _reorder_placeholders(gm, sorted_names):
+    sorted_names = list(sorted_names)
+    name2placeholder = {
+        node.name: node for node in gm.graph.nodes if node.op == "placeholder"
+    }
+    for i, current_placeholder_name in enumerate(sorted_names):
+        if i == 0:
+            continue
+        prev_node = name2placeholder[sorted_names[i - 1]]
+        current_node = name2placeholder[current_placeholder_name]
+        with gm.graph.inserting_after(prev_node):
+            new_node = gm.graph.placeholder(current_node.name)
+            # force rename
+            new_node.name = current_node.name
+            new_node.target = current_node.target
+            current_node.replace_all_uses_with(new_node)
+            name2placeholder[current_placeholder_name] = new_node
+            gm.graph.erase_node(current_node)
+
+    gm.recompile()
+    return gm
diff --git a/graph_net/torch/naive_graph_decomposer.py b/graph_net/torch/naive_graph_decomposer.py
