chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Enhance metrics in dump comparator (#19560)

Browse Source
This commit is contained in:
fzyzcjy
2026-02-28 18:05:19 +08:00
committed by GitHub
parent 706ab9296a
commit b73aa53d7e
7 changed files with 129 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
View File

@@ -3,6 +3,7 @@ from typing import Optional
import torch
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DEFAULT_PERCENTILES,
DiffInfo,
TensorComparisonInfo,
TensorInfo,
@@ -89,21 +90,22 @@ def compare_tensor_pair(
def _compute_tensor_stats(x: torch.Tensor) -> TensorStats:
include_quantiles = x.numel() < QUANTILE_NUMEL_THRESHOLD
include_quantiles: bool = x.numel() < QUANTILE_NUMEL_THRESHOLD
return TensorStats(
mean=torch.mean(x).item(),
abs_mean=torch.mean(x.abs()).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),
percentiles=_compute_percentiles(x, 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_percentiles(x: torch.Tensor, *, include: bool) -> dict[int, float]:
if not include:
return {}
x_float: torch.Tensor = x.float()
return {p: torch.quantile(x_float, p / 100.0).item() for p in DEFAULT_PERCENTILES}
def _compute_diff(
@@ -118,10 +120,15 @@ def _compute_diff(
max_abs_diff = raw_abs_diff.max().item()
mean_abs_diff = raw_abs_diff.mean().item()
include_quantiles: bool = raw_abs_diff.numel() < QUANTILE_NUMEL_THRESHOLD
return DiffInfo(
rel_diff=rel_diff,
max_abs_diff=max_abs_diff,
mean_abs_diff=mean_abs_diff,
abs_diff_percentiles=_compute_percentiles(
raw_abs_diff, include=include_quantiles
),
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(),

27
python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py
View File

@@ -56,20 +56,30 @@ def format_comparison(info: TensorComparisonInfo) -> str:
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:
if stat_name == "percentiles":
continue
value_baseline: float = getattr(baseline, stat_name)
value_target: float = getattr(target, stat_name)
lines.append(
f"[{stat_name}] {value_baseline:.4f} vs {value_target:.4f} "
f"(diff: {value_target - value_baseline:.4f})"
)
for p in sorted(set(baseline.percentiles) & set(target.percentiles)):
value_baseline = baseline.percentiles[p]
value_target = target.percentiles[p]
lines.append(
f"[p{p}] {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]:
return [
lines: list[str] = [
prefix_text
+ "\t".join(
f"{'❌' if value > diff.diff_threshold else '✅'} {name}={value}"
@@ -83,3 +93,12 @@ def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
f"baseline={diff.baseline_at_max} "
f"target={diff.target_at_max}",
]
if diff.abs_diff_percentiles:
quantile_parts: list[str] = [
f"p{p}={value:.4f}"
for p, value in sorted(diff.abs_diff_percentiles.items())
]
lines.append("[abs_diff] " + " ".join(quantile_parts))
return lines

9
python/sglang/srt/debug_utils/comparator/tensor_comparator/types.py
View File

@@ -2,16 +2,16 @@ from typing import Optional
from sglang.srt.debug_utils.comparator.utils import _StrictBase
DEFAULT_PERCENTILES: tuple[int, ...] = (1, 5, 50, 95, 99)
class TensorStats(_StrictBase):
mean: float
abs_mean: float
std: float
min: float
max: float
p1: Optional[float] = None
p5: Optional[float] = None
p95: Optional[float] = None
p99: Optional[float] = None
percentiles: dict[int, float] = {}
class TensorInfo(_StrictBase):
@@ -25,6 +25,7 @@ class DiffInfo(_StrictBase):
rel_diff: float
max_abs_diff: float
mean_abs_diff: float
abs_diff_percentiles: dict[int, float] = {}
max_diff_coord: list[int]
baseline_at_max: float
target_at_max: float

35
test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py
View File

@@ -21,28 +21,34 @@ class TestComputeTensorStats:
stats = _compute_tensor_stats(x)
assert stats.mean == pytest.approx(3.0, abs=1e-4)
assert stats.abs_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_abs_mean_with_negative_values(self):
x = torch.tensor([-3.0, -1.0, 1.0, 3.0])
stats = _compute_tensor_stats(x)
assert stats.mean == pytest.approx(0.0, abs=1e-4)
assert stats.abs_mean == pytest.approx(2.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)
assert stats.percentiles[1] == pytest.approx(1.0, abs=0.5)
assert stats.percentiles[5] == pytest.approx(5.0, abs=0.5)
assert stats.percentiles[50] == pytest.approx(50.0, abs=0.5)
assert stats.percentiles[95] == pytest.approx(95.0, abs=0.5)
assert stats.percentiles[99] == 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
assert stats.percentiles == {}
class TestComputeDiff:
@@ -53,6 +59,9 @@ 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.abs_diff_percentiles[50] == pytest.approx(0.0, abs=1e-5)
assert diff.abs_diff_percentiles[95] == pytest.approx(0.0, abs=1e-5)
assert diff.abs_diff_percentiles[99] == pytest.approx(0.0, abs=1e-5)
assert diff.passed is True
def test_known_offset(self):
@@ -67,8 +76,18 @@ class TestComputeDiff:
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.abs_diff_percentiles[1] == pytest.approx(0.0, abs=1e-4)
assert diff.abs_diff_percentiles[50] == pytest.approx(0.0, abs=1e-4)
assert diff.abs_diff_percentiles[99] > 0
assert diff.passed is False
def test_large_tensor_skips_diff_quantiles(self):
x = torch.randn(QUANTILE_NUMEL_THRESHOLD + 1)
y = x + 0.001
diff = _compute_diff(x_baseline=x, x_target=y)
assert diff.abs_diff_percentiles == {}
def test_rel_diff_value(self):
x = torch.tensor([1.0, 0.0])
y = torch.tensor([0.0, 1.0])

65
test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py
View File

@@ -16,25 +16,47 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
_DEFAULT_PERCENTILES: dict[int, float] = {
1: -1.8,
5: -1.5,
50: 0.0,
95: 1.5,
99: 1.8,
}
def _make_stats(
mean: float = 0.0,
abs_mean: float = 0.8,
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,
percentiles: dict[int, float] | None = None,
) -> TensorStats:
return TensorStats(
mean=mean, std=std, min=min, max=max, p1=p1, p5=p5, p95=p95, p99=p99
mean=mean,
abs_mean=abs_mean,
std=std,
min=min,
max=max,
percentiles=percentiles if percentiles is not None else _DEFAULT_PERCENTILES,
)
_DEFAULT_ABS_DIFF_PERCENTILES: dict[int, float] = {
1: 0.0001,
5: 0.0001,
50: 0.0002,
95: 0.0004,
99: 0.0005,
}
def _make_diff(
rel_diff: float = 0.0001,
max_abs_diff: float = 0.0005,
mean_abs_diff: float = 0.0002,
abs_diff_percentiles: dict[int, float] | None = None,
diff_threshold: float = 1e-3,
passed: bool = True,
) -> DiffInfo:
@@ -42,6 +64,11 @@ def _make_diff(
rel_diff=rel_diff,
max_abs_diff=max_abs_diff,
mean_abs_diff=mean_abs_diff,
abs_diff_percentiles=(
abs_diff_percentiles
if abs_diff_percentiles is not None
else _DEFAULT_ABS_DIFF_PERCENTILES
),
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
@@ -89,16 +116,19 @@ class TestFormatComparison:
"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"
"[abs_mean] 0.8000 vs 0.8000 (diff: 0.0000)\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"
"[p50] 0.0000 vs 0.0000 (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"
"baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
)
def test_shape_mismatch(self):
@@ -116,11 +146,13 @@ class TestFormatComparison:
"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"
"[abs_mean] 0.8000 vs 0.8000 (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"
"[p50] 0.0000 vs 0.0000 (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"
@@ -148,20 +180,24 @@ class TestFormatComparison:
"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"
"[abs_mean] 0.8000 vs 0.8000 (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"
"[p50] 0.0000 vs 0.0000 (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"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.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"
"baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
)
def test_with_shape_unification(self):
@@ -182,16 +218,19 @@ class TestFormatComparison:
"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"
"[abs_mean] 0.8000 vs 0.8000 (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"
"[p50] 0.0000 vs 0.0000 (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"
"baseline=1.0 target=1.0005\n"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.0005"
)
def test_with_samples(self):
@@ -210,22 +249,25 @@ class TestFormatComparison:
"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"
"[abs_mean] 0.8000 vs 0.8000 (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"
"[p50] 0.0000 vs 0.0000 (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"
"[abs_diff] p1=0.0001 p5=0.0001 p50=0.0002 p95=0.0004 p99=0.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)
def test_empty_percentiles(self):
stats_no_quantiles = _make_stats(percentiles={})
info = TensorComparisonInfo(
name="no_quantiles",
@@ -233,7 +275,7 @@ class TestFormatComparison:
target=_make_tensor_info(stats=stats_no_quantiles),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
diff=_make_diff(abs_diff_percentiles={}),
)
assert format_comparison(info) == (
@@ -242,6 +284,7 @@ class TestFormatComparison:
"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"
"[abs_mean] 0.8000 vs 0.8000 (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"

12
test/registered/debug_utils/comparator/tensor_comparator/test_types.py
View File

@@ -23,16 +23,14 @@ 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(
def _make_stats(**overrides) -> TensorStats:
defaults: dict = dict(
mean=0.5,
abs_mean=1.2,
std=1.0,
min=-2.0,
max=3.0,
p1=-1.8,
p5=-1.5,
p95=2.5,
p99=2.8,
percentiles={1: -1.8, 5: -1.5, 50: 0.0, 95: 2.5, 99: 2.8},
)
defaults.update(overrides)
return TensorStats(**defaults)
@@ -66,7 +64,7 @@ def _make_tensor_info(**overrides) -> TensorInfo:
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)
TensorStats(mean=0.0, abs_mean=0.5, std=1.0, min=-1.0, max=1.0, bogus=42)
def test_rejects_extra_fields_on_diff(self):
with pytest.raises(Exception):

2
test/registered/debug_utils/comparator/test_model_validation.py
View File

@@ -187,7 +187,7 @@ def _make_tensor_info() -> TensorInfo:
return TensorInfo(
shape=[4, 4],
dtype="float32",
stats=TensorStats(mean=0.0, std=1.0, min=-2.0, max=2.0),
stats=TensorStats(mean=0.0, abs_mean=0.8, std=1.0, min=-2.0, max=2.0),
)