Support TP unification and enhance tests in dump comparator (#19278)
This commit is contained in:
fzyzcjy
@@ -5,6 +5,7 @@ from pathlib import Path
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import sglang.srt.debug_utils.dumper as _dumper_module
|
||||
from sglang.srt.debug_utils.comparator.entrypoint import run
|
||||
from sglang.srt.debug_utils.comparator.output_types import (
|
||||
AnyRecord,
|
||||
@@ -20,6 +21,285 @@ from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=30, suite="default", nightly=True)
|
||||
|
||||
_FIXED_EXP_NAME = "my_exp_name"
|
||||
|
||||
# Each test has a one-line docstring describing the scenario it covers.
|
||||
|
||||
|
||||
class TestEntrypointGroupingRaw:
|
||||
"""Test `--grouping raw` scenarios"""
|
||||
|
||||
def test_run_basic(self, tmp_path, capsys):
|
||||
"""Two matching tensors produce ConfigRecord, 2 ComparisonRecords, and SummaryRecord."""
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
|
||||
args = _make_args(baseline_path, target_path, grouping="raw")
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
assert isinstance(records[0], ConfigRecord)
|
||||
|
||||
assert len(_get_comparisons(records)) == 2
|
||||
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.total == 2
|
||||
assert summary.skipped == 0
|
||||
|
||||
def test_filter(self, tmp_path, capsys):
|
||||
"""--filter selects only the matching tensor, producing 1 ComparisonRecord."""
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
|
||||
args = _make_args(baseline_path, target_path, filter="tensor_a", grouping="raw")
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
assert len(_get_comparisons(records)) == 1
|
||||
|
||||
def test_no_baseline_skip(self, tmp_path, capsys):
|
||||
"""Target tensor missing from baseline emits a SkipRecord with reason baseline_load_failed."""
|
||||
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, grouping="raw")
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
skips = [r for r in records if isinstance(r, SkipRecord)]
|
||||
assert len(skips) == 1
|
||||
assert skips[0].reason == "baseline_load_failed"
|
||||
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.skipped == 1
|
||||
|
||||
def test_step_range(self, tmp_path, capsys):
|
||||
"""--start_step/--end_step restricts comparison to a single step out of three."""
|
||||
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, grouping="raw"
|
||||
)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.total == 1
|
||||
|
||||
def test_all_valid_records(self, tmp_path, capsys):
|
||||
"""Every emitted JSON record is a valid _OutputRecord subclass."""
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=2)
|
||||
args = _make_args(baseline_path, target_path, grouping="raw")
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
assert all(isinstance(r, _OutputRecord) for r in records)
|
||||
|
||||
def test_text_output_smoke(self, tmp_path, capsys):
|
||||
"""Text output format renders without errors and contains Config/Summary sections."""
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
|
||||
args = _make_args(
|
||||
baseline_path, target_path, output_format="text", grouping="raw"
|
||||
)
|
||||
capsys.readouterr()
|
||||
|
||||
run(args)
|
||||
|
||||
output = capsys.readouterr().out
|
||||
assert "Config:" in output
|
||||
assert "Summary:" in output
|
||||
|
||||
|
||||
class TestEntrypointGroupingLogical:
|
||||
"""Test `--grouping logical` scenarios"""
|
||||
|
||||
def test_no_dims_single_rank(self, tmp_path, capsys):
|
||||
"""Single-rank dumps without dims fall back to raw loading."""
|
||||
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
|
||||
args = _make_args(baseline_path, target_path)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
assert len(_get_comparisons(records)) == 2
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.total == 2
|
||||
assert summary.skipped == 0
|
||||
|
||||
def test_tp_unshard_same_size(self, tmp_path, capsys):
|
||||
"""Both sides TP=2: shards are concatenated before comparison."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8)
|
||||
full_target = full_baseline + torch.randn(4, 8) * 0.001
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
baseline_path = _create_tp_sharded_dumps(
|
||||
baseline_dir,
|
||||
full_tensor=full_baseline,
|
||||
name="hidden",
|
||||
tp_size=2,
|
||||
shard_dim=1,
|
||||
dims_str="b h(tp)",
|
||||
)
|
||||
target_path = _create_tp_sharded_dumps(
|
||||
target_dir,
|
||||
full_tensor=full_target,
|
||||
name="hidden",
|
||||
tp_size=2,
|
||||
shard_dim=1,
|
||||
dims_str="b h(tp)",
|
||||
)
|
||||
|
||||
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
comp = _assert_single_comparison_passed(records)
|
||||
assert comp.name == "hidden"
|
||||
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.total == 1
|
||||
assert summary.passed == 1
|
||||
|
||||
def test_tp_unshard_different_sizes(self, tmp_path, capsys):
|
||||
"""Baseline TP=4 vs target TP=2: different shard counts are handled correctly."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8)
|
||||
full_target = full_baseline + torch.randn(4, 8) * 0.001
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
baseline_path = _create_tp_sharded_dumps(
|
||||
baseline_dir,
|
||||
full_tensor=full_baseline,
|
||||
name="hidden",
|
||||
tp_size=4,
|
||||
shard_dim=1,
|
||||
dims_str="b h(tp)",
|
||||
)
|
||||
target_path = _create_tp_sharded_dumps(
|
||||
target_dir,
|
||||
full_tensor=full_target,
|
||||
name="hidden",
|
||||
tp_size=2,
|
||||
shard_dim=1,
|
||||
dims_str="b h(tp)",
|
||||
)
|
||||
|
||||
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
_assert_single_comparison_passed(records)
|
||||
|
||||
def test_one_side_dims_single_baseline(self, tmp_path, capsys):
|
||||
"""Baseline has no dims (single rank), target has TP shards: unshard target only."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8)
|
||||
target_full = full_tensor + torch.randn(4, 8) * 0.001
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
baseline_path = _create_rank_dump(
|
||||
baseline_dir, rank=0, name="hidden", tensor=full_tensor
|
||||
)
|
||||
|
||||
target_path = _create_tp_sharded_dumps(
|
||||
target_dir,
|
||||
full_tensor=target_full,
|
||||
name="hidden",
|
||||
tp_size=2,
|
||||
shard_dim=1,
|
||||
dims_str="b h(tp)",
|
||||
)
|
||||
|
||||
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
_assert_single_comparison_passed(records)
|
||||
|
||||
def test_ambiguous_baseline_no_dims(self, tmp_path, capsys):
|
||||
"""Multi-rank baseline without dims cannot be unsharded, so it is skipped."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8)
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
for rank, shard in [(0, full_tensor[:, :4]), (1, full_tensor[:, 4:])]:
|
||||
baseline_path = _create_rank_dump(
|
||||
baseline_dir, rank=rank, name="hidden", tensor=shard
|
||||
)
|
||||
|
||||
target_path = _create_tp_sharded_dumps(
|
||||
target_dir,
|
||||
full_tensor=full_tensor,
|
||||
name="hidden",
|
||||
tp_size=2,
|
||||
shard_dim=1,
|
||||
dims_str="b h(tp)",
|
||||
)
|
||||
|
||||
args = _make_args(baseline_path, target_path)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
skips = [r for r in records if isinstance(r, SkipRecord)]
|
||||
assert len(skips) == 1
|
||||
assert skips[0].reason == "baseline_load_failed"
|
||||
|
||||
def test_summary_counts_unshard(self, tmp_path, capsys):
|
||||
"""Two TP-sharded tensors: summary counts total=2, passed=2, skipped=0."""
|
||||
torch.manual_seed(42)
|
||||
full_a = torch.randn(4, 8)
|
||||
full_b = torch.randn(4, 8)
|
||||
|
||||
baseline_dir = tmp_path / "baseline"
|
||||
target_dir = tmp_path / "target"
|
||||
|
||||
for tensor_name, tensor in [("t_a", full_a), ("t_b", full_b)]:
|
||||
baseline_path = _create_tp_sharded_dumps(
|
||||
baseline_dir,
|
||||
full_tensor=tensor,
|
||||
name=tensor_name,
|
||||
tp_size=2,
|
||||
shard_dim=1,
|
||||
dims_str="b h(tp)",
|
||||
)
|
||||
target_tensor = tensor + torch.randn_like(tensor) * 0.0001
|
||||
target_path = _create_tp_sharded_dumps(
|
||||
target_dir,
|
||||
full_tensor=target_tensor,
|
||||
name=tensor_name,
|
||||
tp_size=2,
|
||||
shard_dim=1,
|
||||
dims_str="b h(tp)",
|
||||
)
|
||||
|
||||
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.total == 2
|
||||
assert summary.passed == 2
|
||||
assert summary.failed == 0
|
||||
assert summary.skipped == 0
|
||||
|
||||
|
||||
# --------------------------- Assertion helpers -------------------
|
||||
|
||||
|
||||
def _get_comparisons(records: list[AnyRecord]) -> list[ComparisonRecord]:
|
||||
return [r for r in records if isinstance(r, ComparisonRecord)]
|
||||
|
||||
|
||||
def _assert_single_comparison_passed(records: list[AnyRecord]) -> ComparisonRecord:
|
||||
comparisons = _get_comparisons(records)
|
||||
assert len(comparisons) == 1
|
||||
assert comparisons[0].diff is not None
|
||||
assert comparisons[0].diff.passed
|
||||
return comparisons[0]
|
||||
|
||||
|
||||
# --------------------------- Utils ------------------------------
|
||||
|
||||
|
||||
def _make_dumper(directory: Path) -> _Dumper:
|
||||
return _Dumper(
|
||||
@@ -74,114 +354,77 @@ def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace
|
||||
end_step=1000000,
|
||||
diff_threshold=1e-3,
|
||||
filter=None,
|
||||
output_format="text",
|
||||
output_format="json",
|
||||
grouping="logical",
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return Namespace(**defaults)
|
||||
|
||||
|
||||
def _run_and_parse(args: Namespace, capsys: pytest.CaptureFixture) -> list[AnyRecord]:
|
||||
capsys.readouterr()
|
||||
run(args)
|
||||
return _parse_jsonl(capsys.readouterr().out)
|
||||
|
||||
|
||||
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()
|
||||
def _create_rank_dump(
|
||||
directory: Path,
|
||||
*,
|
||||
rank: int,
|
||||
name: str,
|
||||
tensor: torch.Tensor,
|
||||
dims: str | None = None,
|
||||
parallel_info: dict | None = None,
|
||||
) -> Path:
|
||||
"""Create a dump file via the real dumper, as if running on the given rank."""
|
||||
with pytest.MonkeyPatch.context() as mp:
|
||||
mp.setattr(_dumper_module, "_get_rank", lambda: rank)
|
||||
|
||||
run(args)
|
||||
|
||||
output = capsys.readouterr().out
|
||||
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 "rel_diff" 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"],
|
||||
dumper = _Dumper(
|
||||
config=DumperConfig(
|
||||
enable=True,
|
||||
dir=str(directory),
|
||||
exp_name=_FIXED_EXP_NAME,
|
||||
enable_http_server=False,
|
||||
)
|
||||
)
|
||||
args = _make_args(baseline_path, target_path)
|
||||
capsys.readouterr()
|
||||
|
||||
run(args)
|
||||
static_meta: dict = {"world_rank": rank, "world_size": 1}
|
||||
if parallel_info is not None:
|
||||
static_meta["sglang_parallel_info"] = parallel_info
|
||||
dumper.__dict__["_static_meta"] = static_meta
|
||||
|
||||
output = capsys.readouterr().out
|
||||
assert "Skip:" in output
|
||||
assert "no_baseline" in output
|
||||
dumper.dump(name, tensor, dims=dims)
|
||||
dumper.step()
|
||||
|
||||
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 "Summary:" in output
|
||||
return directory / _FIXED_EXP_NAME
|
||||
|
||||
|
||||
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"],
|
||||
def _create_tp_sharded_dumps(
|
||||
directory: Path,
|
||||
*,
|
||||
full_tensor: torch.Tensor,
|
||||
name: str,
|
||||
tp_size: int,
|
||||
shard_dim: int,
|
||||
dims_str: str,
|
||||
) -> Path:
|
||||
"""Create TP-sharded dump files from a full tensor via the real dumper."""
|
||||
shards = list(full_tensor.chunk(tp_size, dim=shard_dim))
|
||||
for tp_rank in range(tp_size):
|
||||
_create_rank_dump(
|
||||
directory,
|
||||
rank=tp_rank,
|
||||
name=name,
|
||||
tensor=shards[tp_rank],
|
||||
dims=dims_str,
|
||||
parallel_info={"tp_rank": tp_rank, "tp_size": tp_size},
|
||||
)
|
||||
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)
|
||||
return directory / _FIXED_EXP_NAME
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user