From 8b1ab4aaf90f06a9b6d5c7dc8ac82beac441aa0d Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Wed, 25 Feb 2026 09:40:33 +0800 Subject: [PATCH] Support agent-friendly output formats in dump comparator (#19275) --- .../srt/debug_utils/comparator/entrypoint.py | 72 +++-- .../debug_utils/comparator/output_types.py | 83 ++++++ .../comparator/tensor_comparison/__init__.py | 3 - .../comparator/tensor_comparison/compare.py | 49 ++-- .../comparator/tensor_comparison/formatter.py | 88 ++++++ .../comparator/tensor_comparison/types.py | 28 +- .../srt/debug_utils/comparator/utils.py | 5 + .../tensor_comparison/test_compare.py | 15 +- .../tensor_comparison/test_formatter.py | 253 ++++++++++++++++++ .../tensor_comparison/test_types.py | 122 +++++++++ .../debug_utils/comparator/test_entrypoint.py | 83 +++++- 11 files changed, 730 insertions(+), 71 deletions(-) create mode 100644 python/sglang/srt/debug_utils/comparator/output_types.py create mode 100644 python/sglang/srt/debug_utils/comparator/tensor_comparison/formatter.py create mode 100644 test/registered/debug_utils/comparator/tensor_comparison/test_formatter.py create mode 100644 test/registered/debug_utils/comparator/tensor_comparison/test_types.py diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 4f4018804..196439277 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -3,13 +3,16 @@ 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.output_types import ( + ComparisonRecord, + ConfigRecord, + SkipRecord, + SummaryRecord, + print_record, ) +from sglang.srt.debug_utils.comparator.tensor_comparison import compare_tensors 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: @@ -27,8 +30,19 @@ def run(args: argparse.Namespace) -> None: 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) + + print_record( + ConfigRecord( + baseline_path=args.baseline_path, + target_path=args.target_path, + diff_threshold=args.diff_threshold, + start_step=args.start_step, + end_step=args.end_step, + ), + output_format=args.output_format, + ) + + counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0} for row in df_target.iter_rows(named=True): path_target = Path(args.target_path) / row["filename"] @@ -47,26 +61,23 @@ def run(args: argparse.Namespace) -> None: ) 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)}") + counts["skipped"] += 1 + print_record( + SkipRecord(name=row["name"], reason="no_baseline"), + output_format=args.output_format, + ) 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}" + counts["skipped"] += 1 + print_record( + SkipRecord(name=row["name"], reason="load_failed"), + output_format=args.output_format, ) continue @@ -74,9 +85,23 @@ def run(args: argparse.Namespace) -> None: x_baseline=x_baseline, x_target=x_target, name=row["name"], + diff_threshold=args.diff_threshold, ) - print_comparison(info=info, diff_threshold=args.diff_threshold) - print() + + if info.diff is not None and info.diff.passed: + counts["passed"] += 1 + else: + counts["failed"] += 1 + + print_record( + ComparisonRecord(**info.model_dump()), + output_format=args.output_format, + ) + + print_record( + SummaryRecord(total=sum(counts.values()), **counts), + output_format=args.output_format, + ) def _parse_args() -> argparse.Namespace: @@ -90,4 +115,11 @@ def _parse_args() -> argparse.Namespace: parser.add_argument( "--filter", type=str, default=None, help="Regex to filter filenames" ) + parser.add_argument( + "--output-format", + type=str, + choices=["text", "json"], + default="text", + help="Output format: text (default) or json (JSONL, one JSON object per line)", + ) return parser.parse_args() diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py new file mode 100644 index 000000000..7ee29a629 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -0,0 +1,83 @@ +from abc import abstractmethod +from typing import Annotated, Literal, Union + +from pydantic import Discriminator, TypeAdapter + +from sglang.srt.debug_utils.comparator.tensor_comparison.formatter import ( + format_comparison, +) +from sglang.srt.debug_utils.comparator.tensor_comparison.types import ( + TensorComparisonInfo, +) +from sglang.srt.debug_utils.comparator.utils import _StrictBase + + +class _OutputRecord(_StrictBase): + @abstractmethod + def to_text(self) -> str: ... + + +class ConfigRecord(_OutputRecord): + type: Literal["config"] = "config" + baseline_path: str + target_path: str + diff_threshold: float + start_step: int + end_step: int + + def to_text(self) -> str: + return ( + f"Config: baseline={self.baseline_path} target={self.target_path}\n" + f"diff_threshold={self.diff_threshold} " + f"steps=[{self.start_step}, {self.end_step}]" + ) + + +class SkipRecord(_OutputRecord): + type: Literal["skip"] = "skip" + name: str + reason: str + + def to_text(self) -> str: + return f"Skip: {self.name} ({self.reason})" + + +class ComparisonRecord(TensorComparisonInfo, _OutputRecord): + type: Literal["comparison"] = "comparison" + + def to_text(self) -> str: + return format_comparison(self) + + +class SummaryRecord(_OutputRecord): + type: Literal["summary"] = "summary" + total: int + passed: int + failed: int + skipped: int + + def to_text(self) -> str: + return ( + f"Summary: {self.passed} passed, {self.failed} failed, " + f"{self.skipped} skipped (total {self.total})" + ) + + +AnyRecord = Annotated[ + Union[ConfigRecord, SkipRecord, ComparisonRecord, SummaryRecord], + Discriminator("type"), +] + + +_any_record_adapter = TypeAdapter(AnyRecord) + + +def parse_record_json(json_str: str | bytes) -> AnyRecord: + return _any_record_adapter.validate_json(json_str) + + +def print_record(record: _OutputRecord, output_format: str) -> None: + if output_format == "json": + print(record.model_dump_json()) + else: + print(record.to_text()) diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparison/__init__.py b/python/sglang/srt/debug_utils/comparator/tensor_comparison/__init__.py index 82c7788fa..14718ef55 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparison/__init__.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparison/__init__.py @@ -1,4 +1 @@ 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 index de2280813..43c056a5a 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparison/compare.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparison/compare.py @@ -24,22 +24,21 @@ def compare_tensors( x_baseline: torch.Tensor, x_target: torch.Tensor, name: str = "", + diff_threshold: float = 1e-3, ) -> TensorComparisonInfo: baseline_info = TensorInfo( - shape=x_baseline.shape, - dtype=x_baseline.dtype, + shape=list(x_baseline.shape), + dtype=str(x_baseline.dtype), stats=_compute_tensor_stats(x_baseline.float()), - sample=None, ) target_info = TensorInfo( - shape=x_target.shape, - dtype=x_target.dtype, + shape=list(x_target.shape), + dtype=str(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 + unified_shape = list(x_baseline.shape) baseline_original_dtype = x_baseline.dtype target_original_dtype = x_target.dtype @@ -54,12 +53,16 @@ def compare_tensors( downcast_dtype: Optional[torch.dtype] = None if not shape_mismatch: - diff = _compute_diff(x_baseline=x_baseline_f, x_target=x_target_f) + diff = _compute_diff( + x_baseline=x_baseline_f, + x_target=x_target_f, + diff_threshold=diff_threshold, + ) 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) + baseline_info.sample = str(get_truncated_value(x_baseline_f)) + target_info.sample = str(get_truncated_value(x_target_f)) if baseline_original_dtype != target_original_dtype: downcast_dtype = compute_smaller_dtype( @@ -69,6 +72,7 @@ def compare_tensors( diff_downcast = _compute_diff( x_baseline=x_baseline_f.to(downcast_dtype), x_target=x_target_f.to(downcast_dtype), + diff_threshold=diff_threshold, ) return TensorComparisonInfo( @@ -79,7 +83,7 @@ def compare_tensors( shape_mismatch=shape_mismatch, diff=diff, diff_downcast=diff_downcast, - downcast_dtype=downcast_dtype, + downcast_dtype=str(downcast_dtype) if downcast_dtype is not None else None, ) @@ -101,15 +105,28 @@ def _quantile_or_none(x: torch.Tensor, *, q: float, include: bool) -> Optional[f return torch.quantile(x, q).item() if include else None -def _compute_diff(x_baseline: torch.Tensor, x_target: torch.Tensor) -> DiffInfo: +def _compute_diff( + x_baseline: torch.Tensor, + x_target: torch.Tensor, + diff_threshold: float = 1e-3, +) -> DiffInfo: raw_abs_diff = (x_target - x_baseline).abs() max_diff_coord = argmax_coord(raw_abs_diff) + 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() + 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, + rel_diff=rel_diff, + max_abs_diff=max_abs_diff, + mean_abs_diff=mean_abs_diff, + max_diff_coord=list(max_diff_coord), baseline_at_max=x_baseline[max_diff_coord].item(), target_at_max=x_target[max_diff_coord].item(), + passed=( + rel_diff <= diff_threshold + and max_abs_diff <= diff_threshold + and mean_abs_diff <= diff_threshold + ), ) diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparison/formatter.py b/python/sglang/srt/debug_utils/comparator/tensor_comparison/formatter.py new file mode 100644 index 000000000..57fbbd8e2 --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparison/formatter.py @@ -0,0 +1,88 @@ +from sglang.srt.debug_utils.comparator.tensor_comparison.types import ( + DiffInfo, + TensorComparisonInfo, + TensorStats, +) + + +def format_comparison(info: TensorComparisonInfo) -> str: + lines: list[str] = [] + baseline = info.baseline + target = info.target + + dtype_marker = "" if baseline.dtype == target.dtype else "🟠" + lines.append( + 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: + lines.append( + f"Unify shape: {baseline.shape} -> {info.unified_shape} " + f"(to match {target.shape})" + ) + + lines.append( + f"After unify " + f"[shape] {info.unified_shape} vs {target.shape}\t" + f"[dtype] {baseline.dtype} vs {target.dtype}" + ) + + lines.extend(_format_stats_comparison(baseline=baseline.stats, target=target.stats)) + + if info.shape_mismatch: + lines.append("⚠️ Shape mismatch") + return "\n".join(lines) + + if info.diff is not None: + lines.extend(_format_diff(diff=info.diff)) + + if info.diff_downcast is not None and info.downcast_dtype is not None: + lines.extend( + _format_diff( + diff=info.diff_downcast, + prefix_text=f"When downcast to {info.downcast_dtype}: ", + ) + ) + + if baseline.sample is not None: + lines.append(f"x_baseline(sample)={baseline.sample}") + if target.sample is not None: + lines.append(f"x_target(sample)={target.sample}") + + return "\n".join(lines) + + +def _format_stats_comparison(baseline: TensorStats, target: TensorStats) -> list[str]: + lines: list[str] = [] + for stat_name in TensorStats.model_fields: + value_baseline = getattr(baseline, stat_name) + value_target = getattr(target, stat_name) + if value_baseline is None or value_target is None: + continue + lines.append( + f"[{stat_name}] {value_baseline:.4f} vs {value_target:.4f} " + f"(diff: {value_target - value_baseline:.4f})" + ) + return lines + + +def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]: + marker = "✅" if diff.passed else "❌" + return [ + prefix_text + + marker + + " " + + "\t".join( + f"{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), + ] + ), + 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 index 644f81f94..a71402030 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparison/types.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparison/types.py @@ -1,11 +1,9 @@ -from dataclasses import dataclass -from typing import Optional, Tuple +from typing import Optional -import torch +from sglang.srt.debug_utils.comparator.utils import _StrictBase -@dataclass -class TensorStats: +class TensorStats(_StrictBase): mean: float std: float min: float @@ -16,31 +14,29 @@ class TensorStats: p99: Optional[float] = None -@dataclass -class TensorInfo: - shape: torch.Size - dtype: torch.dtype +class TensorInfo(_StrictBase): + shape: list[int] + dtype: str stats: TensorStats sample: Optional[str] = None -@dataclass -class DiffInfo: +class DiffInfo(_StrictBase): rel_diff: float max_abs_diff: float mean_abs_diff: float - max_diff_coord: Tuple[int, ...] + max_diff_coord: list[int] baseline_at_max: float target_at_max: float + passed: bool -@dataclass -class TensorComparisonInfo: +class TensorComparisonInfo(_StrictBase): name: str baseline: TensorInfo target: TensorInfo - unified_shape: Optional[torch.Size] + unified_shape: Optional[list[int]] shape_mismatch: bool diff: Optional[DiffInfo] = None diff_downcast: Optional[DiffInfo] = None - downcast_dtype: Optional[torch.dtype] = None + downcast_dtype: Optional[str] = None diff --git a/python/sglang/srt/debug_utils/comparator/utils.py b/python/sglang/srt/debug_utils/comparator/utils.py index 667081fe1..63a310fa7 100644 --- a/python/sglang/srt/debug_utils/comparator/utils.py +++ b/python/sglang/srt/debug_utils/comparator/utils.py @@ -3,6 +3,11 @@ from pathlib import Path from typing import Optional, Tuple import torch +from pydantic import BaseModel, ConfigDict + + +class _StrictBase(BaseModel): + model_config = ConfigDict(extra="forbid") def argmax_coord(x: torch.Tensor) -> Tuple[int, ...]: diff --git a/test/registered/debug_utils/comparator/tensor_comparison/test_compare.py b/test/registered/debug_utils/comparator/tensor_comparison/test_compare.py index 15a5b14ab..e42bd20f2 100644 --- a/test/registered/debug_utils/comparator/tensor_comparison/test_compare.py +++ b/test/registered/debug_utils/comparator/tensor_comparison/test_compare.py @@ -53,6 +53,7 @@ class TestComputeDiff: 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) + assert diff.passed is True def test_known_offset(self): x = torch.ones(10, 10) @@ -62,10 +63,11 @@ class TestComputeDiff: 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.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) + assert diff.passed is False def test_rel_diff_value(self): x = torch.tensor([1.0, 0.0]) @@ -73,6 +75,7 @@ class TestComputeDiff: diff = _compute_diff(x_baseline=x, x_target=y) assert diff.rel_diff == pytest.approx(1.0, abs=1e-5) + assert diff.passed is False class TestCompareTensors: @@ -83,8 +86,8 @@ class TestCompareTensors: 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.baseline.shape == [5, 5] + assert info.target.shape == [5, 5] assert info.shape_mismatch is False assert info.diff is not None assert info.diff_downcast is None @@ -107,7 +110,7 @@ class TestCompareTensors: 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 + assert info.downcast_dtype == "torch.bfloat16" def test_shape_unification(self): torch.manual_seed(0) @@ -117,8 +120,8 @@ class TestCompareTensors: 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.baseline.shape == [1, 1, 4, 8] + assert info.unified_shape == [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) diff --git a/test/registered/debug_utils/comparator/tensor_comparison/test_formatter.py b/test/registered/debug_utils/comparator/tensor_comparison/test_formatter.py new file mode 100644 index 000000000..33c44920c --- /dev/null +++ b/test/registered/debug_utils/comparator/tensor_comparison/test_formatter.py @@ -0,0 +1,253 @@ +import sys + +import pytest + +from sglang.srt.debug_utils.comparator.tensor_comparison.formatter import ( + format_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, + passed: bool = True, +) -> 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, + passed=passed, + ) + + +def _make_tensor_info( + shape: list[int] | None = None, + dtype: str = "torch.float32", + stats: TensorStats | None = None, + sample: str | None = None, +) -> TensorInfo: + return TensorInfo( + shape=shape if shape is not None else [4, 8], + 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 TestFormatComparison: + def test_normal(self): + 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=[4, 8], + shape_mismatch=False, + diff=_make_diff(), + ) + + assert format_comparison(info) == ( + "Raw [shape] [4, 8] vs [4, 8]\t" + "[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] [4, 8] vs [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\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=[2, 3] with " + "baseline=1.0 target=1.0005" + ) + + def test_shape_mismatch(self): + info = TensorComparisonInfo( + name="mismatch", + baseline=_make_tensor_info(shape=[3, 4]), + target=_make_tensor_info(shape=[5, 6]), + unified_shape=[3, 4], + shape_mismatch=True, + ) + + assert format_comparison(info) == ( + "Raw [shape] [3, 4] vs [5, 6]\t" + "[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] [3, 4] vs [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" + ) + + def test_with_downcast(self): + info = TensorComparisonInfo( + name="downcast", + baseline=_make_tensor_info(), + target=_make_tensor_info(dtype="torch.bfloat16"), + unified_shape=[4, 8], + shape_mismatch=False, + diff=_make_diff( + rel_diff=0.002, max_abs_diff=0.005, mean_abs_diff=0.001, passed=False + ), + diff_downcast=_make_diff( + rel_diff=0.0001, max_abs_diff=0.0005, mean_abs_diff=0.0002, passed=True + ), + downcast_dtype="torch.bfloat16", + ) + + assert format_comparison(info) == ( + "Raw [shape] [4, 8] vs [4, 8]\t" + "[🟠dtype] torch.float32 vs torch.bfloat16\n" + "After unify [shape] [4, 8] vs [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\tmax_abs_diff=0.005\tmean_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\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=[2, 3] with " + "baseline=1.0 target=1.0005" + ) + + def test_with_shape_unification(self): + info = TensorComparisonInfo( + name="unify", + baseline=_make_tensor_info(shape=[1, 1, 4, 8]), + target=_make_tensor_info(), + unified_shape=[4, 8], + shape_mismatch=False, + diff=_make_diff(), + ) + + assert format_comparison(info) == ( + "Raw [shape] [1, 1, 4, 8] vs [4, 8]\t" + "[dtype] torch.float32 vs torch.float32\n" + "Unify shape: [1, 1, 4, 8] -> [4, 8] " + "(to match [4, 8])\n" + "After unify [shape] [4, 8] vs [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\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=[2, 3] with " + "baseline=1.0 target=1.0005" + ) + + def test_with_samples(self): + 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=[4, 8], + shape_mismatch=False, + diff=_make_diff(), + ) + + assert format_comparison(info) == ( + "Raw [shape] [4, 8] vs [4, 8]\t" + "[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] [4, 8] vs [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\tmax_abs_diff=0.0005\tmean_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, ...])" + ) + + def test_none_quantiles(self): + 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=[4, 8], + shape_mismatch=False, + diff=_make_diff(), + ) + + assert format_comparison(info) == ( + "Raw [shape] [4, 8] vs [4, 8]\t" + "[dtype] torch.float32 vs torch.float32\n" + "After unify [shape] [4, 8] vs [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\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n" + "max_abs_diff happens at coord=[2, 3] with " + "baseline=1.0 target=1.0005" + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/tensor_comparison/test_types.py b/test/registered/debug_utils/comparator/tensor_comparison/test_types.py new file mode 100644 index 000000000..cffd08d28 --- /dev/null +++ b/test/registered/debug_utils/comparator/tensor_comparison/test_types.py @@ -0,0 +1,122 @@ +import json +import sys + +import pytest + +from sglang.srt.debug_utils.comparator.output_types import ( + ComparisonRecord, + ConfigRecord, + SkipRecord, + SummaryRecord, + parse_record_json, +) +from sglang.srt.debug_utils.comparator.tensor_comparison.types import ( + DiffInfo, + 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(**overrides: float) -> TensorStats: + defaults = dict( + mean=0.5, + std=1.0, + min=-2.0, + max=3.0, + p1=-1.8, + p5=-1.5, + p95=2.5, + p99=2.8, + ) + defaults.update(overrides) + return TensorStats(**defaults) + + +def _make_diff(**overrides) -> DiffInfo: + defaults = dict( + rel_diff=1e-4, + max_abs_diff=5e-4, + mean_abs_diff=2e-4, + max_diff_coord=[2, 3], + baseline_at_max=1.0, + target_at_max=1.0005, + passed=True, + ) + defaults.update(overrides) + return DiffInfo(**defaults) + + +def _make_tensor_info(**overrides) -> TensorInfo: + defaults = dict( + shape=[4, 8], + dtype="torch.float32", + stats=_make_stats(), + ) + defaults.update(overrides) + return TensorInfo(**defaults) + + +class TestStrictBase: + def test_rejects_extra_fields(self): + with pytest.raises(Exception): + TensorStats(mean=0.0, std=1.0, min=-1.0, max=1.0, bogus=42) + + def test_rejects_extra_fields_on_diff(self): + with pytest.raises(Exception): + DiffInfo( + rel_diff=0.0, + max_abs_diff=0.0, + mean_abs_diff=0.0, + max_diff_coord=[0], + baseline_at_max=0.0, + target_at_max=0.0, + passed=True, + extra_field=123, + ) + + +class TestRecordTypes: + def test_comparison_record_inherits_tensor_fields(self): + record = ComparisonRecord( + name="hidden_states", + baseline=_make_tensor_info(), + target=_make_tensor_info(), + unified_shape=[4, 8], + shape_mismatch=False, + diff=_make_diff(), + ) + parsed = json.loads(record.model_dump_json()) + assert parsed["type"] == "comparison" + assert parsed["name"] == "hidden_states" + assert "baseline" in parsed + assert "diff" in parsed + + def test_discriminated_union_parsing(self): + for record in [ + ConfigRecord( + baseline_path="/a", + target_path="/b", + diff_threshold=1e-3, + start_step=0, + end_step=100, + ), + SkipRecord(name="attn", reason="no_baseline"), + ComparisonRecord( + name="mlp", + baseline=_make_tensor_info(), + target=_make_tensor_info(), + unified_shape=[4, 8], + shape_mismatch=False, + ), + SummaryRecord(total=10, passed=8, failed=1, skipped=1), + ]: + restored = parse_record_json(record.model_dump_json()) + assert type(restored) is type(record) + assert restored == record + + +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 index c2a7b4835..ee218a8a5 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -6,6 +6,15 @@ import pytest import torch from sglang.srt.debug_utils.comparator.entrypoint import run +from sglang.srt.debug_utils.comparator.output_types import ( + AnyRecord, + ComparisonRecord, + ConfigRecord, + SkipRecord, + SummaryRecord, + _OutputRecord, + parse_record_json, +) from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper from sglang.test.ci.ci_register import register_cpu_ci @@ -65,36 +74,39 @@ def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace end_step=1000000, diff_threshold=1e-3, filter=None, + output_format="text", ) defaults.update(overrides) return Namespace(**defaults) +def _parse_jsonl(output: str) -> list[AnyRecord]: + return [parse_record_json(line) for line in output.strip().splitlines()] + + 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) + capsys.readouterr() 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 "Config:" in output assert "rel_diff" in output + assert "Summary:" 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") + capsys.readouterr() run(args) output = capsys.readouterr().out - assert output.count("Check:") == 1 - assert "tensor_a" in output + assert "rel_diff" in output def test_no_baseline_skip(self, tmp_path, capsys): baseline_path, target_path = _create_dumps( @@ -103,22 +115,73 @@ class TestEntrypoint: baseline_names=["tensor_a"], ) args = _make_args(baseline_path, target_path) + capsys.readouterr() run(args) output = capsys.readouterr().out - assert output.count("Check:") == 1 assert "Skip:" in output - assert "since no baseline" in output + assert "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) + capsys.readouterr() run(args) output = capsys.readouterr().out - assert output.count("Check:") == 1 + assert "Summary:" in output + + +class TestEntrypointJsonl: + def test_jsonl_basic(self, tmp_path, capsys): + baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"]) + args = _make_args(baseline_path, target_path, output_format="json") + capsys.readouterr() + + run(args) + + records = _parse_jsonl(capsys.readouterr().out) + assert isinstance(records[0], ConfigRecord) + + comparisons = [r for r in records if isinstance(r, ComparisonRecord)] + assert len(comparisons) == 2 + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.total == 2 + assert summary.skipped == 0 + + def test_jsonl_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, output_format="json") + capsys.readouterr() + + run(args) + + records = _parse_jsonl(capsys.readouterr().out) + skips = [r for r in records if isinstance(r, SkipRecord)] + assert len(skips) == 1 + assert skips[0].reason == "no_baseline" + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.skipped == 1 + + def test_jsonl_all_valid_records(self, tmp_path, capsys): + baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=2) + args = _make_args(baseline_path, target_path, output_format="json") + capsys.readouterr() + + run(args) + + records = _parse_jsonl(capsys.readouterr().out) + assert all(isinstance(r, _OutputRecord) for r in records) if __name__ == "__main__":