Enhance dumper comparator with tensor unifier and location finder (#12623)

python/sglang/srt/debug_utils/dump_comparator.py

@@ -1,7 +1,11 @@
 import argparse
 import functools
+import re
+from dataclasses import dataclass
 from pathlib import Path
+from typing import Callable, Dict, List, Optional
 
+import einops
 import polars as pl
 import torch
 
@@ -25,15 +29,38 @@ def main(args):
     print("df_target", df_target)
     print("df_baseline", df_baseline)
 
+    location_info_of_target_pass_id = _get_location_info_of_target_pass_id()
+    tensor_dim_descs = _get_tensor_dim_descs()
+
     for row in df_target.iter_rows(named=True):
         path_target = Path(args.target_path) / row["filename"]
 
+        if location_info_of_target_pass_id is not None:
+            location_info = location_info_of_target_pass_id.get(row["forward_pass_id"])
+            if location_info is None:
+                continue
+            baseline_forward_pass_id = location_info.baseline_forward_pass_id
+            baseline_token_slice = location_info.baseline_token_slice
+        else:
+            baseline_forward_pass_id = (
+                row["forward_pass_id"] - args.start_id + args.baseline_start_id
+            )
+            baseline_token_slice = None
+
+        tensor_dim_desc = None
+        if tensor_dim_descs is not None:
+            tensor_dim_descs_filtered = [
+                desc
+                for desc in tensor_dim_descs
+                if re.search(desc["pattern"], row["filename"]) is not None
+            ]
+            if tensor_dim_descs_filtered:
+                tensor_dim_desc = tensor_dim_descs_filtered[0]
+
         row_baseline = find_row(
             df_baseline,
             conditions=dict(
-                forward_pass_id=row["forward_pass_id"]
+                forward_pass_id=baseline_forward_pass_id,
-                - args.start_id
-                + args.baseline_start_id,
                 **{
                     k: v
                     for k, v in row.items()
@@ -52,21 +79,56 @@ def main(args):
         path_baseline = Path(args.baseline_path) / row_baseline["filename"]
         print(f"Check: target={str(path_target)} baseline={str(path_baseline)}")
         check_tensor_pair(
-            path_baseline=path_baseline, path_target=path_target, name=row["name"]
+            path_baseline=path_baseline,
+            path_target=path_target,
+            diff_threshold=args.diff_threshold,
+            name=row["name"],
+            baseline_token_slice=baseline_token_slice,
+            tensor_dim_desc=tensor_dim_desc,
         )
         print()
 
 
-def check_tensor_pair(path_baseline, path_target, name=""):
+def _split_einops_pattern(pattern):
+    return re.findall(r"\([^()]*\)|\S+", pattern)
+
+
+def _get_einops_dim_index(pattern: str, dim_name: str):
+    pattern_list = _split_einops_pattern(pattern)
+    return pattern_list.index(dim_name)
+
+
+def check_tensor_pair(
+    path_baseline,
+    path_target,
+    diff_threshold: float = 1e-3,
+    name="",
+    baseline_token_slice=None,
+    tensor_dim_desc: Optional["TensorDimDesc"] = None,
+):
     x_baseline = _load_object(path_baseline)
     x_target = _load_object(path_target)
 
     print(
         f"Raw "
         f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
-        f"[dtype] {x_baseline.dtype} vs {x_target.dtype}"
+        f"[{'' if x_baseline.dtype == x_target.dtype else '🟠'}dtype] {x_baseline.dtype} vs {x_target.dtype}"
     )
 
+    if tensor_dim_desc is not None:
+        if (s := baseline_token_slice) is not None:
+            dim = _get_einops_dim_index(tensor_dim_desc.baseline_desc, "num_tokens")
+            x_baseline = x_baseline.narrow(
+                dim=dim, start=s.start, length=s.stop - s.start
+            )
+        x_baseline = einops.rearrange(
+            x_baseline,
+            tensor_dim_desc.baseline_desc + " -> " + tensor_dim_desc.target_desc,
+        )
+        if (f := tensor_dim_desc.baseline_cropper) is not None:
+            print("Apply baseline_cropper")
+            x_baseline = f(x_baseline)
+
     x_baseline, x_target = _comparison_preprocessor(x_baseline, x_target, name=name)
     x_baseline = _try_unify_shape(x_baseline, target_shape=x_target.shape)
 
@@ -76,19 +138,28 @@ def check_tensor_pair(path_baseline, path_target, name=""):
         f"[dtype] {x_baseline.dtype} vs {x_target.dtype}"
     )
 
+    x_baseline_original_dtype = x_baseline.dtype
+    x_target_original_dtype = x_target.dtype
+
     x_target = x_target.float()
     x_baseline = x_baseline.float()
 
-    for name, fn in (
+    for name, fn in [
         ("mean", torch.mean),
         ("std", torch.std),
         ("min", torch.min),
         ("max", torch.max),
-        ("p1", functools.partial(torch.quantile, q=0.01)),
+        *(
-        ("p5", functools.partial(torch.quantile, q=0.05)),
+            [
-        ("p95", functools.partial(torch.quantile, q=0.95)),
+                ("p1", functools.partial(torch.quantile, q=0.01)),
-        ("p99", functools.partial(torch.quantile, q=0.99)),
+                ("p5", functools.partial(torch.quantile, q=0.05)),
-    ):
+                ("p95", functools.partial(torch.quantile, q=0.95)),
+                ("p99", functools.partial(torch.quantile, q=0.99)),
+            ]
+            if x_baseline.numel() < 10_000_000
+            else []
+        ),
+    ]:
         value_baseline = fn(x_baseline).item()
         value_target = fn(x_target).item()
         print(
@@ -99,17 +170,46 @@ def check_tensor_pair(path_baseline, path_target, name=""):
         print(f"⚠️ Shape mismatch")
         return
 
+    diff_info = _compute_and_print_diff(
+        x_baseline=x_baseline,
+        x_target=x_target,
+        diff_threshold=diff_threshold,
+    )
+    needs_print = diff_info["max_abs_diff"] > 1e-3
+
+    if (x_baseline_original_dtype != x_target_original_dtype) and (
+        (
+            downcast_dtype := _compute_smaller_dtype(
+                x_baseline_original_dtype, x_target_original_dtype
+            )
+        )
+        is not None
+    ):
+        _compute_and_print_diff(
+            x_baseline=x_baseline.to(downcast_dtype),
+            x_target=x_target.to(downcast_dtype),
+            diff_threshold=diff_threshold,
+            prefix_text=f"When downcast to {downcast_dtype}: ",
+        )
+
+    if needs_print:
+        print(f"x_baseline(sample)={get_truncated_value(x_baseline)}")
+        print(f"x_target(sample)={get_truncated_value(x_target)}")
+
+
+def _compute_and_print_diff(
+    x_baseline, x_target, diff_threshold: float, prefix_text=""
+):
     raw_abs_diff = (x_target - x_baseline).abs()
 
     max_abs_diff = raw_abs_diff.max().item()
     mean_abs_diff = raw_abs_diff.mean().item()
     rel_diff = _calc_rel_diff(x_target, x_baseline)
 
-    needs_print = max_abs_diff > 1e-3
-
     print(
-        "\t".join(
+        prefix_text
-            f"{'❌' if value > 1e-3 else '✅'} {name}={value}"
+        + "\t".join(
+            f"{'❌' if value > diff_threshold else '✅'} {name}={value}"
             for name, value in [
                 ("rel_diff", rel_diff),
                 ("max_abs_diff", max_abs_diff),
@@ -118,9 +218,15 @@ def check_tensor_pair(path_baseline, path_target, name=""):
         )
     )
 
-    if needs_print:
+    return dict(max_abs_diff=max_abs_diff)
-        print(f"x_baseline(sample)={get_truncated_value(x_baseline)}")
+
-        print(f"x_target(sample)={get_truncated_value(x_target)}")
+
+def _compute_smaller_dtype(dtype_a, dtype_b):
+    info_dict = {
+        (torch.float32, torch.bfloat16): torch.bfloat16,
+        # ... add more ...
+    }
+    return info_dict.get((dtype_a, dtype_b)) or info_dict.get((dtype_b, dtype_a))
 
 
 def _try_unify_shape(x: torch.Tensor, target_shape):
@@ -144,25 +250,51 @@ def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor):
     return 1 - sim
 
 
-def _comparison_preprocessor(x_baseline, x_target, name):
-    # can insert arbitrary adhoc postprocessing logic here
-    return x_baseline, x_target
-
-
 def _load_object(path):
     x = torch.load(path, weights_only=False)
     if not isinstance(x, torch.Tensor):
-        print(f"Skip load {path} since {type(x)=} is not a Tensor")
+        print(f"Skip load {path} since {type(x)=} is not a Tensor ({x=})")
         return None
     return x.cuda()
 
 
+# TODO may make customization endpoints configurable via args pointing to code file
+def _comparison_preprocessor(x_baseline, x_target, name):
+    """Customization endpoint. Can insert arbitrary adhoc postprocessing logic here."""
+    return x_baseline, x_target
+
+
+@dataclass
+class LocationInfo:
+    baseline_forward_pass_id: int
+    baseline_token_slice: slice
+
+
+def _get_location_info_of_target_pass_id() -> Dict[int, LocationInfo]:
+    """Customization endpoint."""
+    return {}
+
+
+@dataclass
+class TensorDimDesc:
+    baseline_desc: str
+    target_desc: str
+    baseline_cropper: Optional[Callable[[torch.Tensor], torch.Tensor]]
+
+
+def _get_tensor_dim_descs() -> List[TensorDimDesc]:
+    """Customization endpoint."""
+    return []
+
+
 if __name__ == "__main__":
+    # python -m sglang.srt.debug_utils.dump_comparator --baseline-path ... --target-path ...
     parser = argparse.ArgumentParser()
     parser.add_argument("--baseline-path", type=str)
     parser.add_argument("--target-path", type=str)
     parser.add_argument("--start-id", type=int, default=0)
     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)
     args = parser.parse_args()
     main(args)

python/sglang/srt/debug_utils/dump_loader.py

@@ -72,12 +72,20 @@ def find_row(df, conditions: Dict[str, Any]):
         functools.reduce(
             lambda a, b: a & b,
             [
-                pl.col(col) == _cast_to_polars_dtype(conditions[col], df.schema[col])
+                (
+                    pl.col(col)
+                    == _cast_to_polars_dtype(conditions[col], df.schema[col])
+                    if conditions[col] is not None
+                    else pl.col(col).is_null()
+                )
                 for col in conditions.keys()
+                if col in df.columns
             ],
         )
     )
-    assert len(df_sub) <= 1
+    if len(df_sub) > 1:
+        print(f"find_row find ambiguous results: {df_sub=}")
+        return None
     return df_sub.to_dicts()[0] if len(df_sub) > 0 else None