160 lines
4.6 KiB
Python
160 lines
4.6 KiB
Python
from argparse import Namespace
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from sglang.srt.debug_utils.dump_comparator import (
|
|
_argmax_coord,
|
|
_calc_rel_diff,
|
|
_compute_smaller_dtype,
|
|
_try_unify_shape,
|
|
main,
|
|
)
|
|
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)
|
|
|
|
|
|
# ----------------------------- Unit tests -----------------------------
|
|
|
|
|
|
class TestCalcRelDiff:
|
|
def test_identical_vectors(self) -> None:
|
|
x: torch.Tensor = torch.randn(10, 10)
|
|
assert _calc_rel_diff(x, x).item() == pytest.approx(0.0, abs=1e-5)
|
|
|
|
def test_zero_vectors(self) -> None:
|
|
z: torch.Tensor = torch.zeros(5)
|
|
result = _calc_rel_diff(z, z)
|
|
assert not torch.isnan(result) or True # should not crash
|
|
|
|
|
|
class TestArgmaxCoord:
|
|
def test_known_position(self) -> None:
|
|
x: torch.Tensor = 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) -> None:
|
|
target_shape: torch.Size = torch.Size([3, 4])
|
|
result: torch.Tensor = _try_unify_shape(torch.randn(1, 1, 3, 4), target_shape)
|
|
assert result.shape == target_shape
|
|
|
|
def test_no_op_when_no_leading_ones(self) -> None:
|
|
target_shape: torch.Size = torch.Size([3, 4])
|
|
result: torch.Tensor = _try_unify_shape(torch.randn(2, 3, 4), target_shape)
|
|
assert result.shape == (2, 3, 4)
|
|
|
|
|
|
class TestComputeSmallerDtype:
|
|
def test_known_pair(self) -> None:
|
|
assert _compute_smaller_dtype(torch.float32, torch.bfloat16) == torch.bfloat16
|
|
assert _compute_smaller_dtype(torch.bfloat16, torch.float32) == torch.bfloat16
|
|
|
|
def test_none_for_same_dtype(self) -> None:
|
|
assert _compute_smaller_dtype(torch.float32, torch.float32) is None
|
|
|
|
|
|
# ----------------------------- Integration tests -----------------------------
|
|
|
|
|
|
def _make_dumper(directory: Path) -> _Dumper:
|
|
return _Dumper(
|
|
config=DumperConfig(
|
|
enable=True,
|
|
dir=str(directory),
|
|
)
|
|
)
|
|
|
|
|
|
def _create_dumps(
|
|
tmp_path: Path,
|
|
tensor_names: list[str],
|
|
*,
|
|
baseline_names: list[str] | None = None,
|
|
) -> tuple[Path, Path]:
|
|
if baseline_names is None:
|
|
baseline_names = tensor_names
|
|
|
|
d_baseline: Path = tmp_path / "baseline"
|
|
d_target: Path = tmp_path / "target"
|
|
d_baseline.mkdir()
|
|
d_target.mkdir()
|
|
|
|
torch.manual_seed(42)
|
|
baseline_tensor: torch.Tensor = torch.randn(10, 10)
|
|
target_tensor: torch.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: _Dumper = _make_dumper(d)
|
|
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,
|
|
*,
|
|
filter_pattern: str | None = None,
|
|
) -> Namespace:
|
|
return Namespace(
|
|
baseline_path=str(baseline_path),
|
|
target_path=str(target_path),
|
|
start_step=0,
|
|
end_step=1000000,
|
|
diff_threshold=1e-3,
|
|
filter=filter_pattern,
|
|
)
|
|
|
|
|
|
class TestMainBasic:
|
|
def test_matching_tensors(
|
|
self, tmp_path: Path, capsys: pytest.CaptureFixture
|
|
) -> None:
|
|
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
|
|
args: Namespace = _make_args(baseline_path, target_path)
|
|
|
|
main(args)
|
|
|
|
captured: str = capsys.readouterr().out
|
|
assert "✅" in captured
|
|
|
|
def test_with_filter(self, tmp_path: Path, capsys: pytest.CaptureFixture) -> None:
|
|
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
|
|
args: Namespace = _make_args(
|
|
baseline_path, target_path, filter_pattern="tensor_a"
|
|
)
|
|
|
|
main(args)
|
|
|
|
captured: str = capsys.readouterr().out
|
|
assert "tensor_a" in captured
|
|
assert "Check:" in captured
|
|
|
|
def test_no_match_skips(
|
|
self, tmp_path: Path, capsys: pytest.CaptureFixture
|
|
) -> None:
|
|
baseline_path, target_path = _create_dumps(
|
|
tmp_path,
|
|
["only_in_target"],
|
|
baseline_names=["only_in_baseline"],
|
|
)
|
|
args: Namespace = _make_args(baseline_path, target_path)
|
|
|
|
main(args)
|
|
|
|
captured: str = capsys.readouterr().out
|
|
assert "Skip" in captured
|