diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 00add6ddb..e7f69f69f 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -1,19 +1,17 @@ import argparse from pathlib import Path -from typing import Optional import polars as pl -import torch from sglang.srt.debug_utils.comparator.output_types import ( - ComparisonRecord, ConfigRecord, - SkipRecord, SummaryRecord, print_record, ) -from sglang.srt.debug_utils.comparator.tensor_comparison import compare_tensors -from sglang.srt.debug_utils.dump_loader import ValueWithMeta, find_row, read_meta +from sglang.srt.debug_utils.comparator.pipeline import process_tensor_group +from sglang.srt.debug_utils.dump_loader import filter_rows, read_meta + +_NON_KEY_COLS = {"dump_index", "filename", "duplicate_index"} def main() -> None: @@ -22,6 +20,8 @@ def main() -> None: def run(args: argparse.Namespace) -> None: + df_baseline = read_meta(args.baseline_path) + df_target = read_meta(args.target_path) df_target = df_target.filter( (pl.col("step") >= args.start_step) & (pl.col("step") <= args.end_step) @@ -30,8 +30,6 @@ def run(args: argparse.Namespace) -> None: df_target = df_target.filter(pl.col("filename").str.contains(args.filter)) assert all(c in df_target.columns for c in ["rank", "step", "dump_index", "name"]) - df_baseline = read_meta(args.baseline_path) - print_record( ConfigRecord( baseline_path=args.baseline_path, @@ -44,60 +42,27 @@ def run(args: argparse.Namespace) -> None: ) counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0} + grouping: str = args.grouping - for row in df_target.iter_rows(named=True): - path_target = Path(args.target_path) / row["filename"] - baseline_step = row["step"] + non_key_cols = _NON_KEY_COLS | ({"rank"} if grouping == "logical" else set()) + key_cols = [c for c in df_target.columns if c not in non_key_cols] + tensor_group_keys = df_target.unique(subset=key_cols) - row_baseline = find_row( - df_baseline, - conditions=dict( - step=baseline_step, - **{ - k: v - for k, v in row.items() - if k not in ["step", "dump_index", "filename"] - }, - ), - ) + for tensor_group_key in tensor_group_keys.iter_rows(named=True): + conditions = {k: tensor_group_key[k] for k in key_cols} + baseline_rows = filter_rows(df_baseline, conditions=conditions) + target_rows = filter_rows(df_target, conditions=conditions) - if row_baseline is None: - counts["skipped"] += 1 - print_record( - SkipRecord(name=row["name"], reason="no_baseline"), - output_format=args.output_format, - ) - continue - - path_baseline = Path(args.baseline_path) / row_baseline["filename"] - - x_baseline = _load_tensor(path_baseline) - x_target = _load_tensor(path_target) - - if x_baseline is None or x_target is None: - counts["skipped"] += 1 - print_record( - SkipRecord(name=row["name"], reason="load_failed"), - output_format=args.output_format, - ) - continue - - info = compare_tensors( - x_baseline=x_baseline, - x_target=x_target, - name=row["name"], + record = process_tensor_group( + name=tensor_group_key["name"], + baseline_filenames=[r["filename"] for r in baseline_rows], + target_filenames=[r["filename"] for r in target_rows], + baseline_path=Path(args.baseline_path), + target_path=Path(args.target_path), diff_threshold=args.diff_threshold, ) - - if info.diff is not None and info.diff.passed: - counts["passed"] += 1 - else: - counts["failed"] += 1 - - print_record( - ComparisonRecord(**info.model_dump()), - output_format=args.output_format, - ) + counts[record.category] += 1 + print_record(record, output_format=args.output_format) print_record( SummaryRecord(total=sum(counts.values()), **counts), @@ -105,15 +70,7 @@ def run(args: argparse.Namespace) -> None: ) -def _load_tensor(path: Path) -> Optional[torch.Tensor]: - loaded = ValueWithMeta.load(path) - if not isinstance(loaded.value, torch.Tensor): - return None - return loaded.value - - def _parse_args() -> argparse.Namespace: - # python -m sglang.srt.debug_utils.comparator --baseline-path ... --target-path ... parser = argparse.ArgumentParser() parser.add_argument("--baseline-path", type=str) parser.add_argument("--target-path", type=str) @@ -130,4 +87,11 @@ def _parse_args() -> argparse.Namespace: default="text", help="Output format: text (default) or json (JSONL, one JSON object per line)", ) + parser.add_argument( + "--grouping", + type=str, + choices=["logical", "raw"], + default="logical", + help="Grouping mode: logical (cross-rank unshard) or raw (rank-by-rank)", + ) return parser.parse_args() diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index 7ee29a629..28f325ab2 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -38,6 +38,10 @@ class SkipRecord(_OutputRecord): name: str reason: str + @property + def category(self): + return "skipped" + def to_text(self) -> str: return f"Skip: {self.name} ({self.reason})" @@ -45,6 +49,10 @@ class SkipRecord(_OutputRecord): class ComparisonRecord(TensorComparisonInfo, _OutputRecord): type: Literal["comparison"] = "comparison" + @property + def category(self): + return "passed" if self.diff is not None and self.diff.passed else "failed" + def to_text(self) -> str: return format_comparison(self) diff --git a/python/sglang/srt/debug_utils/comparator/pipeline.py b/python/sglang/srt/debug_utils/comparator/pipeline.py new file mode 100644 index 000000000..20715a09a --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/pipeline.py @@ -0,0 +1,119 @@ +from pathlib import Path +from typing import Any, Optional + +import torch + +from sglang.srt.debug_utils.comparator.dims import parse_dims +from sglang.srt.debug_utils.comparator.output_types import ( + ComparisonRecord, + SkipRecord, +) +from sglang.srt.debug_utils.comparator.tensor_comparison.compare import compare_tensors +from sglang.srt.debug_utils.comparator.unshard.executor import execute_unshard_plan +from sglang.srt.debug_utils.comparator.unshard.parallel_info import ( + normalize_parallel_info, +) +from sglang.srt.debug_utils.comparator.unshard.planner import compute_unshard_plan +from sglang.srt.debug_utils.comparator.unshard.types import Plan, UnshardPlan +from sglang.srt.debug_utils.dump_loader import ValueWithMeta + + +def process_tensor_group( + *, + name: str, + baseline_filenames: list[str], + target_filenames: list[str], + baseline_path: Path, + target_path: Path, + diff_threshold: float, +) -> ComparisonRecord | SkipRecord: + b_tensors = _load_tensors(baseline_filenames, baseline_path) + t_tensors = _load_tensors(target_filenames, target_path) + + b_plans, t_plans = _compute_plans( + baseline_metas=[item.meta for item in b_tensors], + target_metas=[item.meta for item in t_tensors], + ) + + b_extracted = _extract_tensors(b_tensors) + t_extracted = _extract_tensors(t_tensors) + del b_tensors, t_tensors + + b_tensor = _execute_plans(b_extracted, b_plans) + t_tensor = _execute_plans(t_extracted, t_plans) + + if b_tensor is None or t_tensor is None: + reason = "baseline_load_failed" if b_tensor is None else "target_load_failed" + return SkipRecord(name=name, reason=reason) + + info = compare_tensors( + x_baseline=b_tensor, + x_target=t_tensor, + name=name, + diff_threshold=diff_threshold, + ) + + return ComparisonRecord(**info.model_dump()) + + +def _load_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]: + return [ValueWithMeta.load(base_path / f) for f in filenames] + + +def _compute_plans( + *, + baseline_metas: list[dict[str, Any]], + target_metas: list[dict[str, Any]], +) -> tuple[list[Plan], list[Plan]]: + """This function deliberately takes metadata, since plan computation must never inspect actual tensor data.""" + return ( + _compute_plans_for_group(baseline_metas), + _compute_plans_for_group(target_metas), + ) + + +def _compute_plans_for_group(metas: list[dict[str, Any]]) -> list[Plan]: + if not metas or len(metas) == 1: + return [] + + dims_str = metas[0].get("dims") + if dims_str is None: + return [] + + dim_specs = parse_dims(dims_str) + parallel_infos = [normalize_parallel_info(meta) for meta in metas] + plan = compute_unshard_plan(dim_specs=dim_specs, parallel_infos=parallel_infos) + + return [plan] if plan is not None else [] + + +def _extract_tensors( + loaded: list[ValueWithMeta], +) -> Optional[list[torch.Tensor]]: + return [value for item in loaded if isinstance(value := item.value, torch.Tensor)] + + +def _execute_plans( + tensors: list[torch.Tensor], + plans: list[Plan], +) -> Optional[torch.Tensor]: + if not tensors: + return None + + if not plans: + if len(tensors) != 1: + return None + return tensors[0] + + assert len(plans) <= 1, "multi-plan not supported yet" + + for plan in plans: + if isinstance(plan, UnshardPlan): + # TODO: incorrect `tensors_by_world_rank` if multi UnshardPlan + tensors = execute_unshard_plan( + plan, tensors_by_world_rank=dict(enumerate(tensors)) + ) + else: + raise NotImplementedError(f"Unknown {plan=}") + + return tensors diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index ee218a8a5..d92be4ca7 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -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__":