[Feature Enhancement] Support multiple incorrect subgraphs

Author: TelGome

graph_net/subgraph_decompose_and_evaluation_step.py

@@ -201,11 +201,18 @@ def run_decomposer_for_multi_models(
     print(
         f"[Decomposition] max_subgraph_size: {max_subgraph_size}, log_path: {log_path}"
     )
+
     for model_name, task_info in tasks_map.items():
         original_path = task_info["original_path"]
-        split_positions = calculate_split_positions_for_subgraph(
-            task_info["subgraph_size"], max_subgraph_size
-        )
+        split_positions = []
+        ranges = task_info["subgraph_sizes"]
+
+        for rng in ranges:
+            splits = calculate_split_positions_for_subgraph(rng, max_subgraph_size)
+            split_positions.extend(splits)
+
+        # Deduplicate and sort
+        split_positions = sorted(list(set(split_positions)))
         task_info["split_positions"] = split_positions
 
         rectified_model_path = get_rectfied_model_path(original_path)
@@ -222,6 +229,7 @@ def run_decomposer_for_multi_models(
         )
         if not success:
             failed_decomposition.append(rectified_model_path)
+
     return tasks_map, failed_decomposition
 
 
@@ -290,7 +298,7 @@ def generate_initial_tasks(args):
         tasks_map[model_name] = {
             "subgraph_path": model_path,
             "original_path": model_path,
-            "subgraph_size": [0, kMaxGraphSize],
+            "subgraph_sizes": [[0, kMaxGraphSize]],
             "split_positions": set(),
         }
 
@@ -307,7 +315,6 @@ def generate_refined_tasks(base_output_dir, current_pass_id):
     prev_incorrect_subgraphs = prev_config.get("incorrect_models", [])
     prev_tasks_map = prev_config.get("tasks_map", {})
 
-    # Load previous max size as fallback
     prev_max_subgraph_size = prev_config.get("max_subgraph_size")
     max_subgraph_size = prev_max_subgraph_size // 2
 
@@ -324,20 +331,24 @@ def generate_refined_tasks(base_output_dir, current_pass_id):
         assert model_name in prev_tasks_map
         pre_task_for_model = prev_tasks_map[model_name]
 
-        # Reconstruct previous subgraph size to locate the failing segment
         prev_split_positions = pre_task_for_model.get("split_positions", [])
-        subgraph_size = reconstruct_subgraph_size(prev_split_positions)
+        subgraph_sizes = reconstruct_subgraph_size(prev_split_positions)
+
         assert subgraph_idx < len(
-            subgraph_size
+            subgraph_sizes
         ), f"subgraph_idx {subgraph_idx} is out of bounds for {model_name} (previous split_positions: {prev_split_positions})"
 
+        current_fail_range = subgraph_sizes[subgraph_idx]
+
         if model_name not in tasks_map:
             tasks_map[model_name] = {
                 "subgraph_path": subgraph_path,
                 "original_path": pre_task_for_model["original_path"],
-                "subgraph_size": subgraph_size[subgraph_idx],
+                "subgraph_sizes": [current_fail_range],
                 "split_positions": set(),
             }
+        else:
+            tasks_map[model_name]["subgraph_sizes"].append(current_fail_range)
 
     return tasks_map, max_subgraph_size
 
@@ -403,9 +414,11 @@ def execute_decomposition_phase(max_subgraph_size, tasks_map, framework, workspa
             shutil.rmtree(decomposed_samples_dir)
             os.makedirs(decomposed_samples_dir, exist_ok=True)
             for model_name, task_info in tasks_map.items():
-                task_info["subgraph_size"][1] = (
-                    task_info["subgraph_size"][0] + max_subgraph_size
-                )
+                for i in range(len(task_info["subgraph_sizes"])):
+                    # Attempt to expand the end position for retry
+                    task_info["subgraph_sizes"][i][1] = (
+                        task_info["subgraph_sizes"][i][0] + max_subgraph_size
+                    )
             max_subgraph_size = max(1, max_subgraph_size // 2)
         else:
             need_decompose = False
@@ -476,6 +489,12 @@ def main(args):
         print("\n--- Phase 3: Analysis ---")
         next_round_models = get_incorrect_models(args.tolerance, pass_log_path)
         print(f"[Analysis] Found {len(next_round_models)} incorrect subgraphs.\n")
+        if len(next_round_models) > 0:
+            print("[DEBUG] List of detected incorrect models:")
+            for idx, model_path in enumerate(sorted(list(next_round_models))):
+                print(f"  [{idx}] {model_path}")
+        else:
+            print("[DEBUG] No incorrect models detected.")
         print_summary_and_suggestion(next_round_models, max_subgraph_size)
 
     # --- Step 5: Save States ---

graph_net/test/error_model/graph_net.json

@@ -0,0 +1,7 @@
+{
+    "framework": "torch",
+    "num_devices_required": 1,
+    "num_nodes_required": 1,
+    "dynamic": false,
+    "model_name": "error_model"
+}

graph_net/test/error_model/input_meta.py

@@ -0,0 +1,153 @@
+import torch
+
+from torch import device
+
+
+class GraphModule(torch.nn.Module):
+    def forward(
+        self,
+        add_22,
+        extended_attention_mask_2,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_3_modules_output_modules_layer_norm_parameters_bias_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_3_modules_output_modules_layer_norm_parameters_weight_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_dense_parameters_bias_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_dense_parameters_weight_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_layer_norm_parameters_bias_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_layer_norm_parameters_weight_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_key_parameters_bias_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_key_parameters_weight_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_query_parameters_bias_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_query_parameters_weight_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_value_parameters_bias_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_value_parameters_weight_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_intermediate_modules_dense_parameters_bias_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_intermediate_modules_dense_parameters_weight_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_dense_parameters_bias_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_dense_parameters_weight_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_layer_norm_parameters_bias_,
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_layer_norm_parameters_weight_,
+    ):
+        hidden_states_66 = torch.nn.functional.layer_norm(
+            add_22,
+            (32,),
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_3_modules_output_modules_layer_norm_parameters_weight_,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_3_modules_output_modules_layer_norm_parameters_bias_,
+            1e-12,
+        )
+        add_22 = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_3_modules_output_modules_layer_norm_parameters_weight_ = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_3_modules_output_modules_layer_norm_parameters_bias_ = (None)
+        linear_44 = torch.nn.functional.linear(
+            hidden_states_66,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_query_parameters_weight_,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_query_parameters_bias_,
+        )
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_query_parameters_weight_ = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_query_parameters_bias_ = (None)
+        view_16 = linear_44.view(2, -1, 4, 8)
+        linear_44 = None
+        query_layer_4 = view_16.transpose(1, 2)
+        view_16 = None
+        linear_45 = torch.nn.functional.linear(
+            hidden_states_66,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_key_parameters_weight_,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_key_parameters_bias_,
+        )
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_key_parameters_weight_ = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_key_parameters_bias_ = (None)
+        view_17 = linear_45.view(2, -1, 4, 8)
+        linear_45 = None
+        key_layer_4 = view_17.transpose(1, 2)
+        view_17 = None
+        linear_46 = torch.nn.functional.linear(
+            hidden_states_66,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_value_parameters_weight_,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_value_parameters_bias_,
+        )
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_value_parameters_weight_ = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_self_modules_value_parameters_bias_ = (None)
+        view_18 = linear_46.view(2, -1, 4, 8)
+        linear_46 = None
+        value_layer_4 = view_18.transpose(1, 2)
+        view_18 = None
+        transpose_25 = key_layer_4.transpose(-1, -2)
+        key_layer_4 = None
+        attention_scores_22 = torch.matmul(query_layer_4, transpose_25)
+        query_layer_4 = transpose_25 = None
+        attention_scores_23 = attention_scores_22 / 2.8284271247461903
+        attention_scores_22 = None
+        eps = torch.tensor(1e-8, device=attention_scores_23.device)
+        nan_val = eps / (eps - eps)
+        attention_scores_23 = attention_scores_23 + nan_val
+        nan_val = None
+        to_8 = extended_attention_mask_2.to(device(type="cuda", index=0))
+        extended_attention_mask_2 = None
+        attention_scores_24 = attention_scores_23 + to_8
+        attention_scores_23 = to_8 = None
+        _log_api_usage_once_4 = torch._C._log_api_usage_once("python.nn_module")
+        _log_api_usage_once_4 = None
+        attention_probs_14 = torch.nn.functional.softmax(
+            attention_scores_24, -1, _stacklevel=5
+        )
+        attention_scores_24 = None
+        attention_probs_dropped_4 = torch.nn.functional.dropout(
+            attention_probs_14, 0.0, False, False
+        )
+        attention_probs_14 = None
+        context_layer_22 = torch.matmul(attention_probs_dropped_4, value_layer_4)
+        attention_probs_dropped_4 = value_layer_4 = None
+        permute_14 = context_layer_22.permute(0, 2, 1, 3)
+        context_layer_22 = None
+        context_layer_23 = permute_14.contiguous()
+        permute_14 = None
+        context_layer_24 = context_layer_23.view(2, 14, 32)
+        context_layer_23 = None
+        hidden_states_67 = torch.nn.functional.linear(
+            context_layer_24,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_dense_parameters_weight_,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_dense_parameters_bias_,
+        )
+        context_layer_24 = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_dense_parameters_weight_ = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_dense_parameters_bias_ = (None)
+        hidden_states_68 = torch.nn.functional.dropout(
+            hidden_states_67, 0.0, False, False
+        )
+        hidden_states_67 = None
+        add_24 = hidden_states_68 + hidden_states_66
+        hidden_states_68 = hidden_states_66 = None
+        hidden_states_69 = torch.nn.functional.layer_norm(
+            add_24,
+            (32,),
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_layer_norm_parameters_weight_,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_layer_norm_parameters_bias_,
+            1e-12,
+        )
+        add_24 = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_layer_norm_parameters_weight_ = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_attention_modules_output_modules_layer_norm_parameters_bias_ = (None)
+        hidden_states_70 = torch.nn.functional.linear(
+            hidden_states_69,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_intermediate_modules_dense_parameters_weight_,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_intermediate_modules_dense_parameters_bias_,
+        )
+        l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_intermediate_modules_dense_parameters_weight_ = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_intermediate_modules_dense_parameters_bias_ = (None)
+        hidden_states_71 = torch.nn.functional.gelu(hidden_states_70)
+        hidden_states_70 = None
+        hidden_states_72 = torch.nn.functional.linear(
+            hidden_states_71,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_dense_parameters_weight_,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_dense_parameters_bias_,
+        )
+        hidden_states_71 = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_dense_parameters_weight_ = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_dense_parameters_bias_ = (None)
+        hidden_states_73 = torch.nn.functional.dropout(
+            hidden_states_72, 0.0, False, False
+        )
+        hidden_states_72 = None
+        nan_val = torch.tensor(0.0, device=hidden_states_73.device) / torch.tensor(
+            0.0, device=hidden_states_73.device
+        )
+        hidden_states_73 = hidden_states_73 + nan_val
+        nan_val = None
+        add_25 = hidden_states_73 + hidden_states_69
+        hidden_states_73 = hidden_states_69 = None
+        hidden_states_74 = torch.nn.functional.layer_norm(
+            add_25,
+            (32,),
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_layer_norm_parameters_weight_,
+            l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_layer_norm_parameters_bias_,
+            1e-12,
+        )
+        add_25 = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_layer_norm_parameters_weight_ = l_l_self_modules_text_model_modules_encoder_modules_layer_modules_4_modules_output_modules_layer_norm_parameters_bias_ = (None)
+        return (hidden_states_74,)