Tiny add filter, support duplications, add visualizations, fix error and robustness for dump comparator (#16262)

python/sglang/srt/debug_utils/dump_comparator.py

@@ -15,11 +15,12 @@ from sglang.srt.debug_utils.dumper import get_truncated_value
 
 def main(args):
     df_target = read_meta(args.target_path)
-    df_target = df_target.sort("rank", "dump_index")
     df_target = df_target.filter(
         (pl.col("forward_pass_id") >= args.start_id)
         & (pl.col("forward_pass_id") <= args.end_id)
     )
+    if args.filter:
+        df_target = df_target.filter(pl.col("filename").str.contains(args.filter))
     assert all(
         c in df_target.columns
         for c in ["rank", "forward_pass_id", "dump_index", "name"]
@@ -77,7 +78,11 @@ def main(args):
             continue
 
         path_baseline = Path(args.baseline_path) / row_baseline["filename"]
-        print(f"Check: target={str(path_target)} baseline={str(path_baseline)}")
+        print(
+            f"Check:\n"
+            f"target={str(path_target)} (duplicate_index={row['duplicate_index']})\n"
+            f"baseline={str(path_baseline)} (duplicate_index={row_baseline['duplicate_index']})"
+        )
         check_tensor_pair(
             path_baseline=path_baseline,
             path_target=path_target,
@@ -109,6 +114,12 @@ def check_tensor_pair(
     x_baseline = _load_object(path_baseline)
     x_target = _load_object(path_target)
 
+    if x_baseline is None or x_target is None:
+        print(
+            f"Skip comparison because of None: x_baseline={x_baseline}, x_target={x_target}"
+        )
+        return
+
     print(
         f"Raw "
         f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
@@ -218,9 +229,21 @@ def _compute_and_print_diff(
         )
     )
 
+    max_diff_coord = _argmax_coord(raw_abs_diff)
+    print(
+        f"max_abs_diff happens at coord={max_diff_coord} with "
+        f"baseline={x_baseline[max_diff_coord].item()} "
+        f"target={x_target[max_diff_coord].item()}"
+    )
+
     return dict(max_abs_diff=max_abs_diff)
 
 
+def _argmax_coord(x: torch.Tensor) -> tuple:
+    flat_idx = x.argmax()
+    return tuple(idx.item() for idx in torch.unravel_index(flat_idx, x.shape))
+
+
 def _compute_smaller_dtype(dtype_a, dtype_b):
     info_dict = {
         (torch.float32, torch.bfloat16): torch.bfloat16,
@@ -251,7 +274,12 @@ def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor):
 
 
 def _load_object(path):
-    x = torch.load(path, weights_only=False)
+    try:
+        x = torch.load(path, weights_only=False)
+    except Exception as e:
+        print(f"Skip load {path} since error {e}")
+        return None
+
     if not isinstance(x, torch.Tensor):
         print(f"Skip load {path} since {type(x)=} is not a Tensor ({x=})")
         return None
@@ -270,9 +298,9 @@ class LocationInfo:
     baseline_token_slice: slice
 
 
-def _get_location_info_of_target_pass_id() -> Dict[int, LocationInfo]:
+def _get_location_info_of_target_pass_id() -> Optional[Dict[int, LocationInfo]]:
     """Customization endpoint."""
-    return {}
+    return None
 
 
 @dataclass
@@ -296,5 +324,8 @@ if __name__ == "__main__":
     parser.add_argument("--end-id", type=int, default=1000000)
     parser.add_argument("--baseline-start-id", type=int, default=0)
     parser.add_argument("--diff-threshold", type=float, default=1e-3)
+    parser.add_argument(
+        "--filter", type=str, default=None, help="Regex to filter filenames"
+    )
     args = parser.parse_args()
     main(args)

python/sglang/srt/debug_utils/dump_loader.py

@@ -47,16 +47,19 @@ def read_meta(directory):
 
     rows = []
     for p in directory.glob("*.pt"):
-        full_kwargs = {}
-        for kv in p.stem.split("___"):
-            k, v = kv.split("=")
-            full_kwargs[k] = v
-        rows.append(
-            {
-                "filename": str(p.name),
-                **full_kwargs,
-            }
-        )
+        try:
+            full_kwargs = {}
+            for kv in p.stem.split("___"):
+                k, v = kv.split("=")
+                full_kwargs[k] = v
+            rows.append(
+                {
+                    "filename": str(p.name),
+                    **full_kwargs,
+                }
+            )
+        except Exception as e:
+            print(f"[DumpLoader] skip loading {p} due to error {e}")
 
     df = pl.DataFrame(rows)
     df = df.with_columns(
@@ -64,6 +67,17 @@ def read_meta(directory):
         pl.col("rank").cast(int),
         pl.col("dump_index").cast(int),
     )
+    df = _add_duplicate_index(df)
+    df = df.sort("rank", "dump_index")
+    return df
+
+
+def _add_duplicate_index(df: pl.DataFrame) -> pl.DataFrame:
+    group_cols = [c for c in df.columns if c not in ["filename", "dump_index"]]
+    df = df.sort(group_cols + ["dump_index"])
+    df = df.with_columns(
+        pl.cum_count("dump_index").over(group_cols).sub(1).alias("duplicate_index")
+    )
     return df