chore: import upstream snapshot with attribution

python/sglang/srt/debug_utils/dump_comparator.py

@@ -0,0 +1,296 @@
+"""Simplified dump comparator — a self-contained single-file script for comparing
+two dump directories tensor-by-tensor.
+
+For advanced features (unshard, token alignment, per-dimension annotations), see the
+full ``comparator/`` package: ``python -m sglang.srt.debug_utils.comparator``.
+"""
+
+import argparse
+import functools
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, List, Optional
+
+import torch
+
+from sglang.srt.debug_utils.dumper import get_truncated_value
+
+
+def main(args):
+    import polars as pl
+
+    from sglang.srt.debug_utils.dump_loader import find_row, read_meta
+
+    df_target = read_meta(args.target_path)
+    df_target = df_target.filter(
+        (pl.col("step") >= args.start_step) & (pl.col("step") <= args.end_step)
+    )
+    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", "step", "dump_index", "name"])
+
+    df_baseline = read_meta(args.baseline_path)
+    print("df_target", df_target)
+    print("df_baseline", df_baseline)
+
+    tensor_dim_descs: List[TensorDimDesc] = _get_tensor_dim_descs()
+
+    for row in df_target.iter_rows(named=True):
+        path_target = Path(args.target_path) / row["filename"]
+
+        tensor_dim_desc: Optional[TensorDimDesc] = None
+        if tensor_dim_descs:
+            matched: list[TensorDimDesc] = [
+                desc
+                for desc in tensor_dim_descs
+                if re.search(desc.pattern, row["filename"]) is not None
+            ]
+            if matched:
+                tensor_dim_desc = matched[0]
+
+        row_baseline = find_row(
+            df_baseline,
+            conditions=dict(
+                step=row["step"],
+                **{
+                    k: v
+                    for k, v in row.items()
+                    if k not in ["step", "dump_index", "filename"]
+                },
+            ),
+        )
+
+        if row_baseline is None:
+            print(f"Skip: target={str(path_target)} since no baseline")
+            x_target = _load_object(path_target)
+            if x_target is not None:
+                print(f"x_target(sample)={get_truncated_value(x_target)}")
+            continue
+
+        path_baseline = Path(args.baseline_path) / row_baseline["filename"]
+        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,
+            diff_threshold=args.diff_threshold,
+            name=row["name"],
+            tensor_dim_desc=tensor_dim_desc,
+        )
+        print()
+
+
+def check_tensor_pair(
+    path_baseline,
+    path_target,
+    diff_threshold: float = 1e-3,
+    name="",
+    tensor_dim_desc: Optional["TensorDimDesc"] = None,
+):
+    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"
+        f"[{'' if x_baseline.dtype == x_target.dtype else '🟠'}dtype] {x_baseline.dtype} vs {x_target.dtype}"
+    )
+
+    if tensor_dim_desc is not None:
+        import einops
+
+        x_baseline = einops.rearrange(
+            x_baseline,
+            tensor_dim_desc.baseline_desc + " -> " + tensor_dim_desc.target_desc,
+        )
+        if tensor_dim_desc.baseline_cropper is not None:
+            print("Apply baseline_cropper")
+            x_baseline = tensor_dim_desc.baseline_cropper(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)
+
+    print(
+        f"After preprocessor "
+        f"[shape] {x_baseline.shape} vs {x_target.shape}\t"
+        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 [
+        ("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)),
+                ("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(
+            f"[{name}] {value_baseline :.4f} vs {value_target:.4f} (diff: {value_target - value_baseline:.4f})"
+        )
+
+    if x_baseline.shape != x_target.shape:
+        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)
+
+    rel_diff_marker: str = "❌" if rel_diff > diff_threshold else "✅"
+    print(
+        prefix_text
+        + f"{rel_diff_marker} rel_diff={rel_diff}\t"
+        + f"max_abs_diff={max_abs_diff}\t"
+        + f"mean_abs_diff={mean_abs_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,
+        # ... 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):
+    x_shape = x.shape
+    num_dim_to_remove = len(x_shape) - len(target_shape)
+    if (x_shape[num_dim_to_remove:] == target_shape) and all(
+        val == 1 for val in x_shape[:num_dim_to_remove]
+    ):
+        out = functools.reduce(lambda a, _: a.squeeze(0), range(num_dim_to_remove), x)
+        print(f"Unify shape: {x_shape} -> {out.shape} (to match {target_shape})")
+        return out
+
+    return x
+
+
+# Copied from DeepGEMM
+def _calc_rel_diff(x: torch.Tensor, y: torch.Tensor):
+    x, y = x.double(), y.double()
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+def _load_object(path):
+    try:
+        x = torch.load(path, weights_only=False)
+    except Exception as e:
+        print(f"Skip load {path} since error {e}")
+        return None
+
+    if isinstance(x, dict) and "value" in x:
+        x = x["value"]
+
+    if not isinstance(x, torch.Tensor):
+        print(f"Skip load {path} since {type(x)=} is not a Tensor ({x=})")
+        return None
+    return x.cuda()
+
+
+def _comparison_preprocessor(x_baseline, x_target, name):
+    """Customization endpoint. Can insert arbitrary adhoc postprocessing logic here."""
+    return x_baseline, x_target
+
+
+@dataclass
+class TensorDimDesc:
+    pattern: str
+    baseline_desc: str
+    target_desc: str
+    baseline_cropper: Optional[Callable[[torch.Tensor], torch.Tensor]] = None
+
+
+def _get_tensor_dim_descs() -> List[TensorDimDesc]:
+    """Customization endpoint. Return a list of TensorDimDesc to rearrange baseline
+    dimensions to match target layout via einops before comparison."""
+    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-step", type=int, default=0)
+    parser.add_argument("--end-step", type=int, default=1000000)
+    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)