From d7578ce279a2695939b11b5bbeb30536281e834e Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Wed, 25 Feb 2026 09:37:21 +0800 Subject: [PATCH] Implement simplest dump comparator v2 (#19274) --- .../srt/debug_utils/comparator/__init__.py | 0 .../srt/debug_utils/comparator/__main__.py | 4 + .../srt/debug_utils/comparator/entrypoint.py | 93 +++++ .../comparator/tensor_comparison/__init__.py | 4 + .../comparator/tensor_comparison/compare.py | 115 ++++++ .../comparator/tensor_comparison/printer.py | 88 +++++ .../comparator/tensor_comparison/types.py | 46 +++ .../srt/debug_utils/comparator/utils.py | 55 +++ .../sglang/srt/debug_utils/dump_comparator.py | 328 ------------------ .../tensor_comparison/test_compare.py | 150 ++++++++ .../tensor_comparison/test_printer.py | 262 ++++++++++++++ .../debug_utils/comparator/test_entrypoint.py | 125 +++++++ .../debug_utils/comparator/test_utils.py | 119 +++++++ .../debug_utils/test_dump_comparator.py | 139 -------- 14 files changed, 1061 insertions(+), 467 deletions(-) create mode 100644 python/sglang/srt/debug_utils/comparator/__init__.py create mode 100644 python/sglang/srt/debug_utils/comparator/__main__.py create mode 100644 python/sglang/srt/debug_utils/comparator/entrypoint.py create mode 100644 python/sglang/srt/debug_utils/comparator/tensor_comparison/__init__.py create mode 100644 python/sglang/srt/debug_utils/comparator/tensor_comparison/compare.py create mode 100644 python/sglang/srt/debug_utils/comparator/tensor_comparison/printer.py create mode 100644 python/sglang/srt/debug_utils/comparator/tensor_comparison/types.py create mode 100644 python/sglang/srt/debug_utils/comparator/utils.py delete mode 100644 python/sglang/srt/debug_utils/dump_comparator.py create mode 100644 test/registered/debug_utils/comparator/tensor_comparison/test_compare.py create mode 100644 test/registered/debug_utils/comparator/tensor_comparison/test_printer.py create mode 100644 test/registered/debug_utils/comparator/test_entrypoint.py create mode 100644 test/registered/debug_utils/comparator/test_utils.py delete mode 100644 test/registered/debug_utils/test_dump_comparator.py diff --git a/python/sglang/srt/debug_utils/comparator/__init__.py b/python/sglang/srt/debug_utils/comparator/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/debug_utils/comparator/__main__.py b/python/sglang/srt/debug_utils/comparator/__main__.py new file mode 100644 index 000000000..511d5f6d1 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/__main__.py @@ -0,0 +1,4 @@ +from sglang.srt.debug_utils.comparator.entrypoint import main + +if __name__ == "__main__": + main() diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py new file mode 100644 index 000000000..4f4018804 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -0,0 +1,93 @@ +import argparse +from pathlib import Path + +import polars as pl + +from sglang.srt.debug_utils.comparator.tensor_comparison import ( + compare_tensors, + print_comparison, +) +from sglang.srt.debug_utils.comparator.utils import load_object +from sglang.srt.debug_utils.dump_loader import find_row, read_meta +from sglang.srt.debug_utils.dumper import get_truncated_value + + +def main() -> None: + args = _parse_args() + run(args) + + +def run(args: argparse.Namespace) -> None: + 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) + + for row in df_target.iter_rows(named=True): + path_target = Path(args.target_path) / row["filename"] + baseline_step = row["step"] + + row_baseline = find_row( + df_baseline, + conditions=dict( + step=baseline_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']})" + ) + + 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: " + f"x_baseline={x_baseline}, x_target={x_target}" + ) + continue + + info = compare_tensors( + x_baseline=x_baseline, + x_target=x_target, + name=row["name"], + ) + print_comparison(info=info, diff_threshold=args.diff_threshold) + print() + + +def _parse_args() -> argparse.Namespace: + # python -m sglang.srt.debug_utils.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" + ) + return parser.parse_args() diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparison/__init__.py b/python/sglang/srt/debug_utils/comparator/tensor_comparison/__init__.py new file mode 100644 index 000000000..82c7788fa --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparison/__init__.py @@ -0,0 +1,4 @@ +from sglang.srt.debug_utils.comparator.tensor_comparison.compare import compare_tensors +from sglang.srt.debug_utils.comparator.tensor_comparison.printer import ( + print_comparison, +) diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparison/compare.py b/python/sglang/srt/debug_utils/comparator/tensor_comparison/compare.py new file mode 100644 index 000000000..de2280813 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparison/compare.py @@ -0,0 +1,115 @@ +from typing import Optional + +import torch + +from sglang.srt.debug_utils.comparator.tensor_comparison.types import ( + DiffInfo, + TensorComparisonInfo, + TensorInfo, + TensorStats, +) +from sglang.srt.debug_utils.comparator.utils import ( + argmax_coord, + calc_rel_diff, + compute_smaller_dtype, + try_unify_shape, +) +from sglang.srt.debug_utils.dumper import get_truncated_value + +QUANTILE_NUMEL_THRESHOLD = 10_000_000 +SAMPLE_DIFF_THRESHOLD = 1e-3 + + +def compare_tensors( + x_baseline: torch.Tensor, + x_target: torch.Tensor, + name: str = "", +) -> TensorComparisonInfo: + baseline_info = TensorInfo( + shape=x_baseline.shape, + dtype=x_baseline.dtype, + stats=_compute_tensor_stats(x_baseline.float()), + sample=None, + ) + target_info = TensorInfo( + shape=x_target.shape, + dtype=x_target.dtype, + stats=_compute_tensor_stats(x_target.float()), + sample=None, + ) + + x_baseline = try_unify_shape(x_baseline, target_shape=x_target.shape) + unified_shape = x_baseline.shape + + baseline_original_dtype = x_baseline.dtype + target_original_dtype = x_target.dtype + + x_baseline_f = x_baseline.float() + x_target_f = x_target.float() + + shape_mismatch = x_baseline_f.shape != x_target_f.shape + + diff: Optional[DiffInfo] = None + diff_downcast: Optional[DiffInfo] = None + downcast_dtype: Optional[torch.dtype] = None + + if not shape_mismatch: + diff = _compute_diff(x_baseline=x_baseline_f, x_target=x_target_f) + + needs_sample = diff.max_abs_diff > SAMPLE_DIFF_THRESHOLD + if needs_sample: + baseline_info.sample = get_truncated_value(x_baseline_f) + target_info.sample = get_truncated_value(x_target_f) + + if baseline_original_dtype != target_original_dtype: + downcast_dtype = compute_smaller_dtype( + baseline_original_dtype, target_original_dtype + ) + if downcast_dtype is not None: + diff_downcast = _compute_diff( + x_baseline=x_baseline_f.to(downcast_dtype), + x_target=x_target_f.to(downcast_dtype), + ) + + return TensorComparisonInfo( + name=name, + baseline=baseline_info, + target=target_info, + unified_shape=unified_shape, + shape_mismatch=shape_mismatch, + diff=diff, + diff_downcast=diff_downcast, + downcast_dtype=downcast_dtype, + ) + + +def _compute_tensor_stats(x: torch.Tensor) -> TensorStats: + include_quantiles = x.numel() < QUANTILE_NUMEL_THRESHOLD + return TensorStats( + mean=torch.mean(x).item(), + std=torch.std(x).item(), + min=torch.min(x).item(), + max=torch.max(x).item(), + p1=_quantile_or_none(x, q=0.01, include=include_quantiles), + p5=_quantile_or_none(x, q=0.05, include=include_quantiles), + p95=_quantile_or_none(x, q=0.95, include=include_quantiles), + p99=_quantile_or_none(x, q=0.99, include=include_quantiles), + ) + + +def _quantile_or_none(x: torch.Tensor, *, q: float, include: bool) -> Optional[float]: + return torch.quantile(x, q).item() if include else None + + +def _compute_diff(x_baseline: torch.Tensor, x_target: torch.Tensor) -> DiffInfo: + raw_abs_diff = (x_target - x_baseline).abs() + max_diff_coord = argmax_coord(raw_abs_diff) + + return DiffInfo( + rel_diff=calc_rel_diff(x_target, x_baseline).item(), + max_abs_diff=raw_abs_diff.max().item(), + mean_abs_diff=raw_abs_diff.mean().item(), + max_diff_coord=max_diff_coord, + baseline_at_max=x_baseline[max_diff_coord].item(), + target_at_max=x_target[max_diff_coord].item(), + ) diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparison/printer.py b/python/sglang/srt/debug_utils/comparator/tensor_comparison/printer.py new file mode 100644 index 000000000..51dccc723 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparison/printer.py @@ -0,0 +1,88 @@ +from dataclasses import fields + +from sglang.srt.debug_utils.comparator.tensor_comparison.types import ( + DiffInfo, + TensorComparisonInfo, + TensorStats, +) + + +def print_comparison(info: TensorComparisonInfo, diff_threshold: float) -> None: + baseline = info.baseline + target = info.target + + dtype_marker = "" if baseline.dtype == target.dtype else "🟠" + print( + f"Raw " + f"[shape] {baseline.shape} vs {target.shape}\t" + f"[{dtype_marker}dtype] {baseline.dtype} vs {target.dtype}" + ) + + if info.unified_shape != baseline.shape: + print( + f"Unify shape: {baseline.shape} -> {info.unified_shape} " + f"(to match {target.shape})" + ) + + print( + f"After unify " + f"[shape] {info.unified_shape} vs {target.shape}\t" + f"[dtype] {baseline.dtype} vs {target.dtype}" + ) + + _print_stats_comparison(baseline=baseline.stats, target=target.stats) + + if info.shape_mismatch: + print("⚠️ Shape mismatch") + return + + if info.diff is not None: + _print_diff( + diff=info.diff, + diff_threshold=diff_threshold, + ) + + if info.diff_downcast is not None and info.downcast_dtype is not None: + _print_diff( + diff=info.diff_downcast, + diff_threshold=diff_threshold, + prefix_text=f"When downcast to {info.downcast_dtype}: ", + ) + + if baseline.sample is not None: + print(f"x_baseline(sample)={baseline.sample}") + if target.sample is not None: + print(f"x_target(sample)={target.sample}") + + +def _print_stats_comparison(baseline: TensorStats, target: TensorStats) -> None: + stat_names = [f.name for f in fields(TensorStats)] + for stat_name in stat_names: + value_baseline = getattr(baseline, stat_name) + value_target = getattr(target, stat_name) + if value_baseline is None or value_target is None: + continue + print( + f"[{stat_name}] {value_baseline:.4f} vs {value_target:.4f} " + f"(diff: {value_target - value_baseline:.4f})" + ) + + +def _print_diff(diff: DiffInfo, diff_threshold: float, prefix_text: str = "") -> None: + print( + prefix_text + + "\t".join( + f"{'❌' if value > diff_threshold else '✅'} {name}={value}" + for name, value in [ + ("rel_diff", diff.rel_diff), + ("max_abs_diff", diff.max_abs_diff), + ("mean_abs_diff", diff.mean_abs_diff), + ] + ) + ) + + print( + f"max_abs_diff happens at coord={diff.max_diff_coord} with " + f"baseline={diff.baseline_at_max} " + f"target={diff.target_at_max}" + ) diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparison/types.py b/python/sglang/srt/debug_utils/comparator/tensor_comparison/types.py new file mode 100644 index 000000000..644f81f94 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparison/types.py @@ -0,0 +1,46 @@ +from dataclasses import dataclass +from typing import Optional, Tuple + +import torch + + +@dataclass +class TensorStats: + mean: float + std: float + min: float + max: float + p1: Optional[float] = None + p5: Optional[float] = None + p95: Optional[float] = None + p99: Optional[float] = None + + +@dataclass +class TensorInfo: + shape: torch.Size + dtype: torch.dtype + stats: TensorStats + sample: Optional[str] = None + + +@dataclass +class DiffInfo: + rel_diff: float + max_abs_diff: float + mean_abs_diff: float + max_diff_coord: Tuple[int, ...] + baseline_at_max: float + target_at_max: float + + +@dataclass +class TensorComparisonInfo: + name: str + baseline: TensorInfo + target: TensorInfo + unified_shape: Optional[torch.Size] + shape_mismatch: bool + diff: Optional[DiffInfo] = None + diff_downcast: Optional[DiffInfo] = None + downcast_dtype: Optional[torch.dtype] = None diff --git a/python/sglang/srt/debug_utils/comparator/utils.py b/python/sglang/srt/debug_utils/comparator/utils.py new file mode 100644 index 000000000..667081fe1 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/utils.py @@ -0,0 +1,55 @@ +import functools +from pathlib import Path +from typing import Optional, Tuple + +import torch + + +def argmax_coord(x: torch.Tensor) -> Tuple[int, ...]: + flat_idx = x.argmax() + return tuple(idx.item() for idx in torch.unravel_index(flat_idx, x.shape)) + + +def compute_smaller_dtype( + dtype_a: torch.dtype, dtype_b: torch.dtype +) -> Optional[torch.dtype]: + 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: torch.Size) -> torch.Tensor: + 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] + ): + return functools.reduce(lambda a, _: a.squeeze(0), range(num_dim_to_remove), x) + + return x + + +# Copied from DeepGEMM +def calc_rel_diff(x: torch.Tensor, y: torch.Tensor) -> 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: Path) -> Optional[torch.Tensor]: + 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() diff --git a/python/sglang/srt/debug_utils/dump_comparator.py b/python/sglang/srt/debug_utils/dump_comparator.py deleted file mode 100644 index 05af80d77..000000000 --- a/python/sglang/srt/debug_utils/dump_comparator.py +++ /dev/null @@ -1,328 +0,0 @@ -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 - -from sglang.srt.debug_utils.dump_loader import find_row, read_meta -from sglang.srt.debug_utils.dumper import get_truncated_value - - -def main(args): - df_target = read_meta(args.target_path) - df_target = df_target.filter( - (pl.col("step") >= args.start_id) & (pl.col("step") <= 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", "step", "dump_index", "name"]) - - df_baseline = read_meta(args.baseline_path) - 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["step"]) - if location_info is None: - continue - baseline_step = location_info.baseline_step - baseline_token_slice = location_info.baseline_token_slice - else: - baseline_step = row["step"] - 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( - step=baseline_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"], - baseline_token_slice=baseline_token_slice, - tensor_dim_desc=tensor_dim_desc, - ) - print() - - -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) - - 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: - 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) - - 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) - - print( - prefix_text - + "\t".join( - f"{'❌' if value > diff_threshold else '✅'} {name}={value}" - for name, value in [ - ("rel_diff", rel_diff), - ("max_abs_diff", max_abs_diff), - ("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() - - -# 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_step: int - baseline_token_slice: slice - - -def _get_location_info_of_target_pass_id() -> Optional[Dict[int, LocationInfo]]: - """Customization endpoint.""" - return None - - -@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) - parser.add_argument( - "--filter", type=str, default=None, help="Regex to filter filenames" - ) - args = parser.parse_args() - main(args) diff --git a/test/registered/debug_utils/comparator/tensor_comparison/test_compare.py b/test/registered/debug_utils/comparator/tensor_comparison/test_compare.py new file mode 100644 index 000000000..15a5b14ab --- /dev/null +++ b/test/registered/debug_utils/comparator/tensor_comparison/test_compare.py @@ -0,0 +1,150 @@ +import sys + +import pytest +import torch + +from sglang.srt.debug_utils.comparator.tensor_comparison.compare import ( + QUANTILE_NUMEL_THRESHOLD, + SAMPLE_DIFF_THRESHOLD, + _compute_diff, + _compute_tensor_stats, + compare_tensors, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=20, suite="default", nightly=True) + + +class TestComputeTensorStats: + def test_basic_stats(self): + x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0]) + stats = _compute_tensor_stats(x) + + assert stats.mean == pytest.approx(3.0, abs=1e-4) + assert stats.std == pytest.approx(1.5811, abs=1e-3) + assert stats.min == pytest.approx(1.0, abs=1e-4) + assert stats.max == pytest.approx(5.0, abs=1e-4) + + def test_quantile_values(self): + x = torch.linspace(0.0, 100.0, steps=1000) + stats = _compute_tensor_stats(x) + + assert stats.p1 == pytest.approx(1.0, abs=0.5) + assert stats.p5 == pytest.approx(5.0, abs=0.5) + assert stats.p95 == pytest.approx(95.0, abs=0.5) + assert stats.p99 == pytest.approx(99.0, abs=0.5) + + def test_large_tensor_skips_quantiles(self): + x = torch.randn(QUANTILE_NUMEL_THRESHOLD + 1) + stats = _compute_tensor_stats(x) + + assert stats.mean is not None + assert stats.p1 is None + assert stats.p5 is None + assert stats.p95 is None + assert stats.p99 is None + + +class TestComputeDiff: + def test_identical_tensors(self): + x = torch.ones(10, 10) + diff = _compute_diff(x_baseline=x, x_target=x) + + assert diff.rel_diff == pytest.approx(0.0, abs=1e-5) + assert diff.max_abs_diff == pytest.approx(0.0, abs=1e-5) + assert diff.mean_abs_diff == pytest.approx(0.0, abs=1e-5) + + def test_known_offset(self): + x = torch.ones(10, 10) + y = x.clone() + y[3, 7] = 1.5 + + diff = _compute_diff(x_baseline=x, x_target=y) + + assert diff.max_abs_diff == pytest.approx(0.5, abs=1e-4) + assert diff.max_diff_coord == (3, 7) + assert diff.baseline_at_max == pytest.approx(1.0, abs=1e-4) + assert diff.target_at_max == pytest.approx(1.5, abs=1e-4) + assert diff.mean_abs_diff == pytest.approx(0.5 / 100, abs=1e-4) + + def test_rel_diff_value(self): + x = torch.tensor([1.0, 0.0]) + y = torch.tensor([0.0, 1.0]) + diff = _compute_diff(x_baseline=x, x_target=y) + + assert diff.rel_diff == pytest.approx(1.0, abs=1e-5) + + +class TestCompareTensors: + def test_normal(self): + x = torch.randn(5, 5) + y = x + torch.randn(5, 5) * 0.001 + + info = compare_tensors(x_baseline=x, x_target=y, name="test") + + assert info.name == "test" + assert info.baseline.shape == torch.Size([5, 5]) + assert info.target.shape == torch.Size([5, 5]) + assert info.shape_mismatch is False + assert info.diff is not None + assert info.diff_downcast is None + + def test_shape_mismatch(self): + x = torch.randn(3, 4) + y = torch.randn(5, 6) + + info = compare_tensors(x_baseline=x, x_target=y, name="mismatch") + + assert info.shape_mismatch is True + assert info.diff is None + + def test_dtype_mismatch(self): + x = torch.randn(5, 5, dtype=torch.float32) + y = torch.randn(5, 5, dtype=torch.bfloat16) + + info = compare_tensors(x_baseline=x, x_target=y, name="dtype_test") + + assert info.shape_mismatch is False + assert info.diff is not None + assert info.diff_downcast is not None + assert info.downcast_dtype == torch.bfloat16 + + def test_shape_unification(self): + torch.manual_seed(0) + core = torch.randn(4, 8) + x = core.unsqueeze(0).unsqueeze(0) # [1, 1, 4, 8] + y = core.clone() # [4, 8] + + info = compare_tensors(x_baseline=x, x_target=y, name="unify") + + assert info.baseline.shape == torch.Size([1, 1, 4, 8]) + assert info.unified_shape == torch.Size([4, 8]) + assert info.shape_mismatch is False + assert info.diff is not None + assert info.diff.max_abs_diff == pytest.approx(0.0, abs=1e-5) + + def test_sample_generated_when_large_diff(self): + x = torch.zeros(5, 5) + y = torch.ones(5, 5) + + info = compare_tensors(x_baseline=x, x_target=y, name="big_diff") + + assert info.diff is not None + assert info.diff.max_abs_diff > SAMPLE_DIFF_THRESHOLD + assert info.baseline.sample is not None + assert info.target.sample is not None + + def test_no_sample_when_small_diff(self): + x = torch.ones(5, 5) + y = x + 1e-5 + + info = compare_tensors(x_baseline=x, x_target=y, name="tiny_diff") + + assert info.diff is not None + assert info.diff.max_abs_diff < SAMPLE_DIFF_THRESHOLD + assert info.baseline.sample is None + assert info.target.sample is None + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/tensor_comparison/test_printer.py b/test/registered/debug_utils/comparator/tensor_comparison/test_printer.py new file mode 100644 index 000000000..447dc3747 --- /dev/null +++ b/test/registered/debug_utils/comparator/tensor_comparison/test_printer.py @@ -0,0 +1,262 @@ +import sys + +import pytest +import torch + +from sglang.srt.debug_utils.comparator.tensor_comparison.printer import ( + print_comparison, +) +from sglang.srt.debug_utils.comparator.tensor_comparison.types import ( + DiffInfo, + TensorComparisonInfo, + TensorInfo, + TensorStats, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="default", nightly=True) + + +def _make_stats( + mean: float = 0.0, + std: float = 1.0, + min: float = -2.0, + max: float = 2.0, + p1: float | None = -1.8, + p5: float | None = -1.5, + p95: float | None = 1.5, + p99: float | None = 1.8, +) -> TensorStats: + return TensorStats( + mean=mean, std=std, min=min, max=max, p1=p1, p5=p5, p95=p95, p99=p99 + ) + + +def _make_diff( + rel_diff: float = 0.0001, + max_abs_diff: float = 0.0005, + mean_abs_diff: float = 0.0002, +) -> DiffInfo: + return DiffInfo( + rel_diff=rel_diff, + max_abs_diff=max_abs_diff, + mean_abs_diff=mean_abs_diff, + max_diff_coord=(2, 3), + baseline_at_max=1.0, + target_at_max=1.0005, + ) + + +def _make_tensor_info( + shape: torch.Size = torch.Size([4, 8]), + dtype: torch.dtype = torch.float32, + stats: TensorStats | None = None, + sample: str | None = None, +) -> TensorInfo: + return TensorInfo( + shape=shape, + dtype=dtype, + stats=stats if stats is not None else _make_stats(), + sample=sample, + ) + + +# Snapshot strings below are intentionally spelled out in full per test. +# The shared skeleton (stats block, diff block) looks duplicated, but keeping +# each test self-contained makes failures immediately readable without chasing +# helper functions. Do not extract common fragments. +class TestPrintComparison: + def test_normal(self, capsys): + info = TensorComparisonInfo( + name="test", + baseline=_make_tensor_info( + stats=_make_stats(mean=0.1, std=1.0, min=-2.0, max=2.0), + ), + target=_make_tensor_info( + stats=_make_stats(mean=0.1001, std=1.0001, min=-2.0001, max=2.0001), + ), + unified_shape=torch.Size([4, 8]), + shape_mismatch=False, + diff=_make_diff(), + ) + + print_comparison(info=info, diff_threshold=1e-3) + + assert capsys.readouterr().out == ( + "Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "[dtype] torch.float32 vs torch.float32\n" + "[mean] 0.1000 vs 0.1001 (diff: 0.0001)\n" + "[std] 1.0000 vs 1.0001 (diff: 0.0001)\n" + "[min] -2.0000 vs -2.0001 (diff: -0.0001)\n" + "[max] 2.0000 vs 2.0001 (diff: 0.0001)\n" + "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" + "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" + "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" + "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" + "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=(2, 3) with " + "baseline=1.0 target=1.0005\n" + ) + + def test_shape_mismatch(self, capsys): + info = TensorComparisonInfo( + name="mismatch", + baseline=_make_tensor_info(shape=torch.Size([3, 4])), + target=_make_tensor_info(shape=torch.Size([5, 6])), + unified_shape=torch.Size([3, 4]), + shape_mismatch=True, + ) + + print_comparison(info=info, diff_threshold=1e-3) + + assert capsys.readouterr().out == ( + "Raw [shape] torch.Size([3, 4]) vs torch.Size([5, 6])\t" + "[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] torch.Size([3, 4]) vs torch.Size([5, 6])\t" + "[dtype] torch.float32 vs torch.float32\n" + "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" + "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" + "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" + "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" + "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" + "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" + "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" + "⚠️ Shape mismatch\n" + ) + + def test_with_downcast(self, capsys): + info = TensorComparisonInfo( + name="downcast", + baseline=_make_tensor_info(), + target=_make_tensor_info(dtype=torch.bfloat16), + unified_shape=torch.Size([4, 8]), + shape_mismatch=False, + diff=_make_diff(rel_diff=0.002, max_abs_diff=0.005, mean_abs_diff=0.001), + diff_downcast=_make_diff( + rel_diff=0.0001, max_abs_diff=0.0005, mean_abs_diff=0.0002 + ), + downcast_dtype=torch.bfloat16, + ) + + print_comparison(info=info, diff_threshold=1e-3) + + assert capsys.readouterr().out == ( + "Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "[🟠dtype] torch.float32 vs torch.bfloat16\n" + "After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "[dtype] torch.float32 vs torch.bfloat16\n" + "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" + "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" + "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" + "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" + "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" + "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" + "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" + "❌ rel_diff=0.002\t❌ max_abs_diff=0.005\t✅ mean_abs_diff=0.001\n" + "max_abs_diff happens at coord=(2, 3) with " + "baseline=1.0 target=1.0005\n" + "When downcast to torch.bfloat16: " + "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=(2, 3) with " + "baseline=1.0 target=1.0005\n" + ) + + def test_with_shape_unification(self, capsys): + info = TensorComparisonInfo( + name="unify", + baseline=_make_tensor_info(shape=torch.Size([1, 1, 4, 8])), + target=_make_tensor_info(), + unified_shape=torch.Size([4, 8]), + shape_mismatch=False, + diff=_make_diff(), + ) + + print_comparison(info=info, diff_threshold=1e-3) + + assert capsys.readouterr().out == ( + "Raw [shape] torch.Size([1, 1, 4, 8]) vs torch.Size([4, 8])\t" + "[dtype] torch.float32 vs torch.float32\n" + "Unify shape: torch.Size([1, 1, 4, 8]) -> torch.Size([4, 8]) " + "(to match torch.Size([4, 8]))\n" + "After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "[dtype] torch.float32 vs torch.float32\n" + "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" + "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" + "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" + "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" + "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" + "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" + "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" + "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=(2, 3) with " + "baseline=1.0 target=1.0005\n" + ) + + def test_with_samples(self, capsys): + info = TensorComparisonInfo( + name="samples", + baseline=_make_tensor_info(sample="tensor([0.1, 0.2, ...])"), + target=_make_tensor_info(sample="tensor([0.1, 0.3, ...])"), + unified_shape=torch.Size([4, 8]), + shape_mismatch=False, + diff=_make_diff(), + ) + + print_comparison(info=info, diff_threshold=1e-3) + + assert capsys.readouterr().out == ( + "Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "[dtype] torch.float32 vs torch.float32\n" + "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" + "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" + "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" + "[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n" + "[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n" + "[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n" + "[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n" + "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=(2, 3) with " + "baseline=1.0 target=1.0005\n" + "x_baseline(sample)=tensor([0.1, 0.2, ...])\n" + "x_target(sample)=tensor([0.1, 0.3, ...])\n" + ) + + def test_none_quantiles(self, capsys): + stats_no_quantiles = _make_stats(p1=None, p5=None, p95=None, p99=None) + + info = TensorComparisonInfo( + name="no_quantiles", + baseline=_make_tensor_info(stats=stats_no_quantiles), + target=_make_tensor_info(stats=stats_no_quantiles), + unified_shape=torch.Size([4, 8]), + shape_mismatch=False, + diff=_make_diff(), + ) + + print_comparison(info=info, diff_threshold=1e-3) + + assert capsys.readouterr().out == ( + "Raw [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] torch.Size([4, 8]) vs torch.Size([4, 8])\t" + "[dtype] torch.float32 vs torch.float32\n" + "[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n" + "[std] 1.0000 vs 1.0000 (diff: 0.0000)\n" + "[min] -2.0000 vs -2.0000 (diff: 0.0000)\n" + "[max] 2.0000 vs 2.0000 (diff: 0.0000)\n" + "✅ rel_diff=0.0001\t✅ max_abs_diff=0.0005\t✅ mean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=(2, 3) with " + "baseline=1.0 target=1.0005\n" + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py new file mode 100644 index 000000000..c2a7b4835 --- /dev/null +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -0,0 +1,125 @@ +import sys +from argparse import Namespace +from pathlib import Path + +import pytest +import torch + +from sglang.srt.debug_utils.comparator.entrypoint import run +from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=30, suite="default", nightly=True) + + +def _make_dumper(directory: Path) -> _Dumper: + return _Dumper( + config=DumperConfig(enable=True, dir=str(directory), enable_http_server=False) + ) + + +def _create_dumps( + tmp_path: Path, + tensor_names: list[str], + *, + baseline_names: list[str] | None = None, + num_steps: int = 1, +) -> tuple[Path, Path]: + """Create baseline and target dump directories with given tensor names. + + If baseline_names is None, uses the same names as tensor_names. + Each step dumps all names with the same tensor (different per baseline/target). + """ + if baseline_names is None: + baseline_names = tensor_names + + d_baseline = tmp_path / "baseline" + d_target = tmp_path / "target" + d_baseline.mkdir() + d_target.mkdir() + + torch.manual_seed(42) + baseline_tensor = torch.randn(10, 10) + target_tensor = baseline_tensor + torch.randn(10, 10) * 0.01 + + exp_paths: list[Path] = [] + for d, names, tensor in [ + (d_baseline, baseline_names, baseline_tensor), + (d_target, tensor_names, target_tensor), + ]: + dumper = _make_dumper(d) + for _ in range(num_steps): + for name in names: + dumper.dump(name, tensor) + dumper.step() + exp_paths.append(d / dumper._config.exp_name) + + return exp_paths[0], exp_paths[1] + + +def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace: + defaults = dict( + baseline_path=str(baseline_path), + target_path=str(target_path), + start_step=0, + end_step=1000000, + diff_threshold=1e-3, + filter=None, + ) + defaults.update(overrides) + return Namespace(**defaults) + + +class TestEntrypoint: + def test_run_basic(self, tmp_path, capsys): + baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"]) + args = _make_args(baseline_path, target_path) + + run(args) + + output = capsys.readouterr().out + assert "df_target" in output + assert "df_baseline" in output + assert output.count("Check:") == 2 + assert "tensor_a" in output + assert "tensor_b" in output + assert "rel_diff" in output + assert "Skip" not in output + + def test_filter(self, tmp_path, capsys): + baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"]) + args = _make_args(baseline_path, target_path, filter="tensor_a") + + run(args) + + output = capsys.readouterr().out + assert output.count("Check:") == 1 + assert "tensor_a" in output + + def test_no_baseline_skip(self, tmp_path, capsys): + baseline_path, target_path = _create_dumps( + tmp_path, + tensor_names=["tensor_a", "tensor_extra"], + baseline_names=["tensor_a"], + ) + args = _make_args(baseline_path, target_path) + + run(args) + + output = capsys.readouterr().out + assert output.count("Check:") == 1 + assert "Skip:" in output + assert "since no baseline" in output + + def test_step_range(self, tmp_path, capsys): + baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=3) + args = _make_args(baseline_path, target_path, start_step=1, end_step=1) + + run(args) + + output = capsys.readouterr().out + assert output.count("Check:") == 1 + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/test_utils.py b/test/registered/debug_utils/comparator/test_utils.py new file mode 100644 index 000000000..824d645d0 --- /dev/null +++ b/test/registered/debug_utils/comparator/test_utils.py @@ -0,0 +1,119 @@ +import sys +from pathlib import Path + +import pytest +import torch + +from sglang.srt.debug_utils.comparator.utils import ( + argmax_coord, + calc_rel_diff, + compute_smaller_dtype, + load_object, + try_unify_shape, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=10, suite="default", nightly=True) + + +class TestCalcRelDiff: + def test_identical_tensors(self): + x = torch.randn(10, 10) + assert calc_rel_diff(x, x).item() == pytest.approx(0.0, abs=1e-5) + + def test_orthogonal_tensors(self): + result = calc_rel_diff( + torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0]) + ).item() + assert result == pytest.approx(1.0, abs=1e-5) + + def test_similar_tensors(self): + x = torch.tensor([1.0, 2.0, 3.0]) + y = torch.tensor([1.01, 2.01, 3.01]) + result = calc_rel_diff(x, y).item() + assert 0.0 < result < 0.01 + + def test_negated_tensors(self): + x = torch.tensor([1.0, 2.0]) + result = calc_rel_diff(x, -x).item() + assert result == pytest.approx(2.0, abs=1e-5) + + +class TestArgmaxCoord: + def test_1d_tensor(self): + x = torch.tensor([0.0, 0.0, 5.0, 0.0]) + assert argmax_coord(x) == (2,) + + def test_2d_tensor(self): + x = torch.zeros(3, 4) + x[1, 2] = 10.0 + assert argmax_coord(x) == (1, 2) + + def test_3d_tensor(self): + x = torch.zeros(2, 3, 4) + x[1, 2, 3] = 10.0 + assert argmax_coord(x) == (1, 2, 3) + + +class TestTryUnifyShape: + def test_squeeze_leading_ones(self): + target = torch.Size([3, 4]) + assert try_unify_shape(torch.randn(1, 1, 3, 4), target).shape == target + + def test_no_squeeze_when_leading_dim_not_one(self): + target = torch.Size([3, 4]) + assert try_unify_shape(torch.randn(2, 3, 4), target).shape == (2, 3, 4) + + def test_same_shape_noop(self): + target = torch.Size([3, 4]) + x = torch.randn(3, 4) + result = try_unify_shape(x, target) + assert result.shape == target + assert result.data_ptr() == x.data_ptr() + + def test_trailing_dims_mismatch(self): + target = torch.Size([5, 6]) + x = torch.randn(1, 3, 4) + result = try_unify_shape(x, target) + assert result.shape == (1, 3, 4) + + +class TestComputeSmallerDtype: + def test_float32_bfloat16(self): + assert compute_smaller_dtype(torch.float32, torch.bfloat16) == torch.bfloat16 + + def test_reverse_order(self): + assert compute_smaller_dtype(torch.bfloat16, torch.float32) == torch.bfloat16 + + def test_same_dtype_returns_none(self): + assert compute_smaller_dtype(torch.float32, torch.float32) is None + + def test_unknown_pair_returns_none(self): + assert compute_smaller_dtype(torch.int32, torch.int64) is None + + +class TestLoadObject: + def test_load_tensor(self, tmp_path): + path = tmp_path / "tensor.pt" + torch.save(torch.randn(5, 5), path) + assert load_object(path).shape == (5, 5) + + def test_load_dict_with_value_key(self, tmp_path): + path = tmp_path / "wrapped.pt" + tensor = torch.randn(3, 3) + torch.save({"value": tensor}, path) + result = load_object(path) + assert result is not None + assert result.shape == (3, 3) + + def test_non_tensor_returns_none(self, tmp_path): + path = tmp_path / "tensor.pt" + torch.save({"dict": 1}, path) + assert load_object(path) is None + + def test_nonexistent_returns_none(self): + assert load_object(Path("/nonexistent.pt")) is None + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/test_dump_comparator.py b/test/registered/debug_utils/test_dump_comparator.py deleted file mode 100644 index 3cab0eae7..000000000 --- a/test/registered/debug_utils/test_dump_comparator.py +++ /dev/null @@ -1,139 +0,0 @@ -import os -import tempfile -import unittest -from contextlib import contextmanager -from pathlib import Path - -import torch - -from sglang.test.ci.ci_register import register_cpu_ci -from sglang.test.test_utils import CustomTestCase - -register_cpu_ci(est_time=60, suite="default", nightly=True) - - -class TestDumpComparator(CustomTestCase): - def test_calc_rel_diff(self): - from sglang.srt.debug_utils.dump_comparator import _calc_rel_diff - - x = torch.randn(10, 10) - self.assertAlmostEqual(_calc_rel_diff(x, x).item(), 0.0, places=5) - self.assertAlmostEqual( - _calc_rel_diff(torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])).item(), - 1.0, - places=5, - ) - - def test_argmax_coord(self): - from sglang.srt.debug_utils.dump_comparator import _argmax_coord - - x = torch.zeros(2, 3, 4) - x[1, 2, 3] = 10.0 - self.assertEqual(_argmax_coord(x), (1, 2, 3)) - - def test_try_unify_shape(self): - from sglang.srt.debug_utils.dump_comparator import _try_unify_shape - - target = torch.Size([3, 4]) - self.assertEqual( - _try_unify_shape(torch.randn(1, 1, 3, 4), target).shape, target - ) - self.assertEqual( - _try_unify_shape(torch.randn(2, 3, 4), target).shape, (2, 3, 4) - ) - - def test_compute_smaller_dtype(self): - from sglang.srt.debug_utils.dump_comparator import _compute_smaller_dtype - - self.assertEqual( - _compute_smaller_dtype(torch.float32, torch.bfloat16), torch.bfloat16 - ) - self.assertIsNone(_compute_smaller_dtype(torch.float32, torch.float32)) - - def test_einops_pattern(self): - from sglang.srt.debug_utils.dump_comparator import ( - _get_einops_dim_index, - _split_einops_pattern, - ) - - self.assertEqual(_split_einops_pattern("a (b c) d"), ["a", "(b c)", "d"]) - self.assertEqual(_get_einops_dim_index("a b c", "b"), 1) - - def test_load_object(self): - from sglang.srt.debug_utils.dump_comparator import _load_object - - with tempfile.TemporaryDirectory() as tmpdir: - path = Path(tmpdir) / "tensor.pt" - torch.save(torch.randn(5, 5), path) - self.assertEqual(_load_object(path).shape, (5, 5)) - - torch.save({"dict": 1}, path) - self.assertIsNone(_load_object(path)) - - self.assertIsNone(_load_object("/nonexistent.pt")) - - def test_compute_and_print_diff(self): - from sglang.srt.debug_utils.dump_comparator import _compute_and_print_diff - - x = torch.ones(10, 10) - self.assertAlmostEqual( - _compute_and_print_diff(x, x, 1e-3)["max_abs_diff"], 0.0, places=5 - ) - self.assertAlmostEqual( - _compute_and_print_diff(x, x + 0.5, 1e-3)["max_abs_diff"], 0.5, places=4 - ) - - -class TestEndToEnd(CustomTestCase): - def test_main(self): - from argparse import Namespace - - from sglang.srt.debug_utils.dump_comparator import main - from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper - - with tempfile.TemporaryDirectory() as d1, tempfile.TemporaryDirectory() as d2: - baseline_tensor = torch.randn(10, 10) - target_tensor = baseline_tensor + torch.randn(10, 10) * 0.01 - - dump_dirs = [] - for d, tensor in [(d1, baseline_tensor), (d2, target_tensor)]: - dumper = _Dumper( - config=DumperConfig( - enable=True, - dir=d, - enable_http_server=False, - ) - ) - dumper.dump("tensor_a", tensor) - dumper.step() - dumper.dump("tensor_b", tensor * 2) - dumper.step() - dump_dirs.append(Path(d) / dumper._config.exp_name) - - args = Namespace( - baseline_path=str(dump_dirs[0]), - target_path=str(dump_dirs[1]), - start_id=0, - end_id=1, - baseline_start_id=0, - diff_threshold=1e-3, - filter=None, - ) - main(args) - - -@contextmanager -def _with_env(name: str, value: str): - old = os.environ.get(name) - os.environ[name] = value - try: - yield - finally: - if old is None: - os.environ.pop(name, None) - else: - os.environ[name] = old - - -if __name__ == "__main__": - unittest.main()