Implement simplest dump comparator v2 (#19274)

2026-02-25 09:37:21 +08:00
parent 9cec98b445
commit d7578ce279
14 changed files with 1061 additions and 467 deletions
@@ -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__]))