[Feature Enhancement] Support multiple incorrect subgraphs. (#407)

* [Feature Enhancement] Support multiple incorrect subgraphs

* support tolerance range

* fix

Author: TelGome

graph_net/subgraph_decompose_and_evaluation_step.py

@@ -18,6 +18,29 @@ def convert_b64_string_to_json(b64str):
     return json.loads(base64.b64decode(b64str).decode("utf-8"))
 
 
+def get_ranged_incorrect_models(tolerance_args: List[int], log_path: str) -> set:
+    if not os.path.exists(log_path):
+        return set()
+
+    t_start = tolerance_args[0]
+    models_start = set(get_incorrect_models(t_start, log_path))
+
+    if len(tolerance_args) == 1:
+        return models_start
+
+    t_end = tolerance_args[1]
+    models_end = set(get_incorrect_models(t_end, log_path))
+
+    print(f"[Filter] Tolerance Range: {t_start} -> {t_end}")
+    print(
+        f"[Filter] Fail({t_start}): {len(models_start)}, Fail({t_end}): {len(models_end)}"
+    )
+
+    diff_set = models_start - models_end
+
+    return diff_set
+
+
 class TaskController:
     def __init__(self, args):
         self.root_output_dir = os.path.abspath(args.output_dir)
@@ -203,10 +226,10 @@ def run_decomposer_for_multi_models(
     )
     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
-        )
-        task_info["split_positions"] = split_positions
+
+        split_positions = task_info["split_positions"]
+        if isinstance(split_positions, set):
+            split_positions = sorted(list(split_positions))
 
         rectified_model_path = get_rectfied_model_path(original_path)
         assert os.path.exists(
@@ -282,19 +305,28 @@ def calculate_split_positions_for_subgraph(subgraph_size, max_subgraph_size):
 def generate_initial_tasks(args):
     """Generates tasks for Pass 0 based on the initial log file."""
     print(f"[Init] Pass 0: Reading from log file: {args.log_file}")
-    initial_failures = get_incorrect_models(args.tolerance, args.log_file)
+    initial_failures = get_ranged_incorrect_models(args.tolerance, args.log_file)
 
     tasks_map = {}
+    max_subgraph_size = args.max_subgraph_size
+
     for model_path in initial_failures:
         model_name = get_model_name_with_subgraph_tag(model_path)
+
+        initial_range = [0, kMaxGraphSize]
+        initial_splits = calculate_split_positions_for_subgraph(
+            initial_range, max_subgraph_size
+        )
+
         tasks_map[model_name] = {
             "subgraph_path": model_path,
             "original_path": model_path,
-            "subgraph_size": [0, kMaxGraphSize],
-            "split_positions": set(),
+            "split_positions": set(initial_splits),
         }
 
-    max_subgraph_size = args.max_subgraph_size
+    for task in tasks_map.values():
+        task["split_positions"] = sorted(list(task["split_positions"]))
+
     return tasks_map, max_subgraph_size
 
 
@@ -307,7 +339,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 +355,30 @@ 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_ranges = reconstruct_subgraph_size(prev_split_positions)
+
         assert subgraph_idx < len(
-            subgraph_size
+            subgraph_ranges
         ), f"subgraph_idx {subgraph_idx} is out of bounds for {model_name} (previous split_positions: {prev_split_positions})"
 
+        current_fail_range = subgraph_ranges[subgraph_idx]
+
+        new_splits = calculate_split_positions_for_subgraph(
+            current_fail_range, max_subgraph_size
+        )
+
         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],
-                "split_positions": set(),
+                "split_positions": set(new_splits),
             }
+        else:
+            tasks_map[model_name]["split_positions"].update(new_splits)
+
+    for task in tasks_map.values():
+        task["split_positions"] = sorted(list(task["split_positions"]))
 
     return tasks_map, max_subgraph_size
 
@@ -402,11 +443,23 @@ def execute_decomposition_phase(max_subgraph_size, tasks_map, framework, workspa
             need_decompose = True
             shutil.rmtree(decomposed_samples_dir)
             os.makedirs(decomposed_samples_dir, exist_ok=True)
+            max_subgraph_size = max(1, max_subgraph_size // 2)
             for model_name, task_info in tasks_map.items():
-                task_info["subgraph_size"][1] = (
-                    task_info["subgraph_size"][0] + max_subgraph_size
+                splits = task_info["split_positions"]
+                if not splits or len(splits) < 2:
+                    continue
+                if isinstance(splits, set):
+                    splits = sorted(list(splits))
+                start_pos = splits[0]
+                first_segment_end = splits[1]
+                new_splits = list(
+                    range(start_pos, first_segment_end + 1, max_subgraph_size)
                 )
-            max_subgraph_size = max(1, max_subgraph_size // 2)
+
+                if new_splits[-1] != first_segment_end:
+                    new_splits.append(first_segment_end)
+
+                task_info["split_positions"] = sorted(list(set(new_splits)))
         else:
             need_decompose = False
         print()
@@ -474,8 +527,18 @@ def main(args):
     next_round_models = set()
     if task_controller.task_scheduler["post_analysis"]:
         print("\n--- Phase 3: Analysis ---")
-        next_round_models = get_incorrect_models(args.tolerance, pass_log_path)
+        analysis_tolerance = (
+            args.tolerance[0] if isinstance(args.tolerance, list) else args.tolerance
+        )
+        next_round_models = get_incorrect_models(analysis_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 ---
@@ -497,7 +560,11 @@ def main(args):
         "--test-config", type=str, required=True, help="Base64 encoded test config"
     )
     parser.add_argument(
-        "--tolerance", type=int, required=True, help="Tolerance level range [-10, 5)"
+        "--tolerance",
+        type=int,
+        nargs="+",
+        required=True,
+        help="Tolerance level range [-10, 5)",
     )
     parser.add_argument("--max-subgraph-size", type=int, default=4096)
     args = parser.parse_args()

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,)