Implement simplest dump comparator v2 (#19274)
This commit is contained in:
fzyzcjy
@@ -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__]))
|
||||
@@ -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__]))
|
||||
@@ -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__]))
|
||||
@@ -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__]))
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user