Support agent-friendly output formats in dump comparator (#19275)
This commit is contained in:
fzyzcjy
@@ -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()
|
||||
|
||||
83
python/sglang/srt/debug_utils/comparator/output_types.py
Normal file
83
python/sglang/srt/debug_utils/comparator/output_types.py
Normal file
@@ -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())
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
),
|
||||
)
|
||||
|
||||
@@ -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}",
|
||||
]
|
||||
@@ -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
|
||||
|
||||
@@ -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, ...]:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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__]))
|
||||
@@ -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__]))
|
||||
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user