Enhance displaying and debuggability in dump comparator (#19466)

2026-02-27 09:06:19 +08:00
parent 23cbbd6d41
commit 4add6ec0f6
17 changed files with 879 additions and 31 deletions
@@ -124,7 +124,6 @@ class TestExecuteAlignerPlan:
                y=[self._make_step_plan(step=0, indices=[0])],
            ),
            token_aligner_plan=None,
-            token_dims=Pair(x=0, y=0),
        )

        tensors_pair: Pair[list[torch.Tensor]] = Pair(
@@ -146,7 +145,6 @@ class TestExecuteAlignerPlan:
                y=[self._make_step_plan(step=0, indices=[0, 1])],
            ),
            token_aligner_plan=None,
-            token_dims=Pair(x=0, y=0),
        )

        tensors_pair: Pair[list[torch.Tensor]] = Pair(
@@ -168,7 +166,6 @@ class TestExecuteAlignerPlan:
                y=[self._make_step_plan(step=0, indices=[0])],
            ),
            token_aligner_plan=None,
-            token_dims=Pair(x=0, y=0),
        )

        t_x: torch.Tensor = torch.tensor([1.0, 2.0])
@@ -191,7 +188,6 @@ class TestExecuteAlignerPlan:
                y=[self._make_step_plan(step=0, indices=[0])],
            ),
            token_aligner_plan=None,
-            token_dims=Pair(x=0, y=0),
        )

        tensors_pair: Pair[list[torch.Tensor]] = Pair(
@@ -0,0 +1,386 @@
+import sys
+from pathlib import Path
+from typing import Any, Optional
+
+import polars as pl
+import pytest
+import torch
+
+from sglang.srt.debug_utils.comparator.display import (
+    _collect_input_ids_and_positions,
+    _collect_rank_info,
+    _extract_parallel_info,
+    _render_polars_as_text,
+)
+from sglang.srt.debug_utils.comparator.output_types import (
+    InputIdsRecord,
+    RankInfoRecord,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="default", nightly=True)
+
+
+def _save_dump_file(
+    directory: Path,
+    *,
+    name: str,
+    step: int,
+    rank: int,
+    dump_index: int,
+    value: torch.Tensor,
+    meta: dict,
+) -> str:
+    filename = f"name={name}___step={step}___rank={rank}___dump_index={dump_index}.pt"
+    torch.save({"value": value, "meta": meta}, directory / filename)
+    return filename
+
+
+def _make_df(rows: list[dict]) -> pl.DataFrame:
+    df = pl.DataFrame(rows)
+    df = df.with_columns(
+        pl.col("step").cast(int),
+        pl.col("rank").cast(int),
+        pl.col("dump_index").cast(int),
+    )
+    return df
+
+
+class TestRenderPolarsAsText:
+    def test_renders_table(self) -> None:
+        df = pl.DataFrame({"col_a": [1, 2], "col_b": ["x", "y"]})
+        text: str = _render_polars_as_text(df, title="test table")
+
+        assert "test table" in text
+        assert "col_a" in text
+        assert "col_b" in text
+
+    def test_renders_empty_dataframe(self) -> None:
+        df = pl.DataFrame({"a": [], "b": []})
+        text: str = _render_polars_as_text(df, title="empty")
+        assert "empty" in text
+
+
+class TestCollectRankInfo:
+    def test_collects_rank_info(self, tmp_path: Path) -> None:
+        sglang_info = {
+            "tp_rank": 0,
+            "tp_size": 2,
+            "pp_rank": 0,
+            "pp_size": 1,
+        }
+        filename: str = _save_dump_file(
+            tmp_path,
+            name="input_ids",
+            step=0,
+            rank=0,
+            dump_index=0,
+            value=torch.tensor([1, 2, 3]),
+            meta={"sglang_parallel_info": sglang_info},
+        )
+        df = _make_df(
+            [
+                {
+                    "filename": filename,
+                    "name": "input_ids",
+                    "step": 0,
+                    "rank": 0,
+                    "dump_index": 0,
+                }
+            ]
+        )
+
+        rows: Optional[list[dict[str, Any]]] = _collect_rank_info(df, dump_dir=tmp_path)
+
+        assert rows is not None
+        assert len(rows) == 1
+        assert rows[0]["rank"] == 0
+        assert rows[0]["tp"] == "0/2"
+        assert rows[0]["pp"] == "0/1"
+
+    def test_returns_none_when_no_input_ids(self, tmp_path: Path) -> None:
+        df = _make_df(
+            [
+                {
+                    "filename": "f.pt",
+                    "name": "some_other",
+                    "step": 0,
+                    "rank": 0,
+                    "dump_index": 0,
+                }
+            ]
+        )
+        result = _collect_rank_info(df, dump_dir=tmp_path)
+        assert result is None
+
+    def test_deduplicates_ranks(self, tmp_path: Path) -> None:
+        meta = {"sglang_parallel_info": {"tp_rank": 0, "tp_size": 1}}
+        f1: str = _save_dump_file(
+            tmp_path,
+            name="input_ids",
+            step=0,
+            rank=0,
+            dump_index=0,
+            value=torch.tensor([1]),
+            meta=meta,
+        )
+        f2: str = _save_dump_file(
+            tmp_path,
+            name="input_ids",
+            step=1,
+            rank=0,
+            dump_index=1,
+            value=torch.tensor([2]),
+            meta=meta,
+        )
+        df = _make_df(
+            [
+                {
+                    "filename": f1,
+                    "name": "input_ids",
+                    "step": 0,
+                    "rank": 0,
+                    "dump_index": 0,
+                },
+                {
+                    "filename": f2,
+                    "name": "input_ids",
+                    "step": 1,
+                    "rank": 0,
+                    "dump_index": 1,
+                },
+            ]
+        )
+
+        rows = _collect_rank_info(df, dump_dir=tmp_path)
+
+        assert rows is not None
+        assert len(rows) == 1
+
+
+class TestCollectInputIdsAndPositions:
+    def test_collects_ids_and_positions(self, tmp_path: Path) -> None:
+        f_ids: str = _save_dump_file(
+            tmp_path,
+            name="input_ids",
+            step=0,
+            rank=0,
+            dump_index=0,
+            value=torch.tensor([10, 20, 30]),
+            meta={},
+        )
+        f_pos: str = _save_dump_file(
+            tmp_path,
+            name="positions",
+            step=0,
+            rank=0,
+            dump_index=1,
+            value=torch.tensor([0, 1, 2]),
+            meta={},
+        )
+        df = _make_df(
+            [
+                {
+                    "filename": f_ids,
+                    "name": "input_ids",
+                    "step": 0,
+                    "rank": 0,
+                    "dump_index": 0,
+                },
+                {
+                    "filename": f_pos,
+                    "name": "positions",
+                    "step": 0,
+                    "rank": 0,
+                    "dump_index": 1,
+                },
+            ]
+        )
+
+        rows = _collect_input_ids_and_positions(df, dump_dir=tmp_path)
+
+        assert rows is not None
+        assert len(rows) == 1
+        assert rows[0]["step"] == 0
+        assert rows[0]["rank"] == 0
+        assert rows[0]["num_tokens"] == 3
+        assert "10" in rows[0]["input_ids"]
+        assert "0" in rows[0]["positions"]
+
+    def test_returns_none_when_empty(self, tmp_path: Path) -> None:
+        df = _make_df(
+            [
+                {
+                    "filename": "f.pt",
+                    "name": "weight",
+                    "step": 0,
+                    "rank": 0,
+                    "dump_index": 0,
+                }
+            ]
+        )
+        result = _collect_input_ids_and_positions(df, dump_dir=tmp_path)
+        assert result is None
+
+    def test_with_mock_tokenizer(self, tmp_path: Path) -> None:
+        f_ids: str = _save_dump_file(
+            tmp_path,
+            name="input_ids",
+            step=0,
+            rank=0,
+            dump_index=0,
+            value=torch.tensor([1, 2]),
+            meta={},
+        )
+        df = _make_df(
+            [
+                {
+                    "filename": f_ids,
+                    "name": "input_ids",
+                    "step": 0,
+                    "rank": 0,
+                    "dump_index": 0,
+                }
+            ]
+        )
+
+        class _MockTokenizer:
+            def decode(self, ids: list[int], skip_special_tokens: bool = False) -> str:
+                return f"decoded:{ids}"
+
+        rows = _collect_input_ids_and_positions(
+            df, dump_dir=tmp_path, tokenizer=_MockTokenizer()
+        )
+
+        assert rows is not None
+        assert "decoded_text" in rows[0]
+        assert "decoded:" in rows[0]["decoded_text"]
+
+
+class TestRankInfoRecordSnapshot:
+    def test_to_text_snapshot(self) -> None:
+        record = RankInfoRecord(
+            label="baseline",
+            rows=[
+                {"rank": 0, "tp": "0/2", "pp": "0/1"},
+                {"rank": 1, "tp": "1/2", "pp": "0/1"},
+            ],
+        )
+        text: str = record.to_text()
+
+        assert "baseline ranks" in text
+        assert "rank" in text
+        assert "tp" in text
+        assert "pp" in text
+        assert "0/2" in text
+        assert "1/2" in text
+        assert "0/1" in text
+
+    def test_json_roundtrip(self) -> None:
+        record = RankInfoRecord(
+            label="target",
+            rows=[{"rank": 0, "tp": "0/4"}],
+        )
+        json_str: str = record.model_dump_json()
+
+        assert '"type":"rank_info"' in json_str
+        assert '"label":"target"' in json_str
+        assert '"tp":"0/4"' in json_str
+
+
+class TestInputIdsRecordSnapshot:
+    def test_to_text_snapshot(self) -> None:
+        record = InputIdsRecord(
+            label="target",
+            rows=[
+                {
+                    "step": 0,
+                    "rank": 0,
+                    "num_tokens": 3,
+                    "input_ids": "[10, 20, 30]",
+                    "positions": "[0, 1, 2]",
+                },
+            ],
+        )
+        text: str = record.to_text()
+
+        assert "target input_ids & positions" in text
+        assert "step" in text
+        assert "num_tokens" in text
+        assert "10, 20, 30" in text
+        assert "0, 1, 2" in text
+
+    def test_json_roundtrip(self) -> None:
+        record = InputIdsRecord(
+            label="baseline",
+            rows=[
+                {
+                    "step": 0,
+                    "rank": 0,
+                    "num_tokens": 2,
+                    "input_ids": "[1, 2]",
+                    "positions": "[0, 1]",
+                    "decoded_text": "'hello'",
+                },
+            ],
+        )
+        json_str: str = record.model_dump_json()
+
+        assert '"type":"input_ids"' in json_str
+        assert '"label":"baseline"' in json_str
+        assert '"decoded_text"' in json_str
+
+    def test_to_text_with_decoded(self) -> None:
+        record = InputIdsRecord(
+            label="test",
+            rows=[
+                {
+                    "step": 0,
+                    "rank": 0,
+                    "num_tokens": 2,
+                    "input_ids": "[1, 2]",
+                    "positions": "[0, 1]",
+                    "decoded_text": "'hello world'",
+                },
+            ],
+        )
+        text: str = record.to_text()
+
+        assert "decoded_text" in text
+        assert "hello world" in text
+
+
+class TestExtractParallelInfo:
+    def test_extracts_rank_size_pairs(self) -> None:
+        info: dict = {
+            "tp_rank": 1,
+            "tp_size": 4,
+            "pp_rank": 0,
+            "pp_size": 2,
+        }
+        row_data: dict = {}
+        _extract_parallel_info(row_data=row_data, info=info)
+
+        assert row_data["tp"] == "1/4"
+        assert row_data["pp"] == "0/2"
+
+    def test_skips_error_info(self) -> None:
+        row_data: dict = {}
+        _extract_parallel_info(
+            row_data=row_data, info={"error": True, "tp_rank": 0, "tp_size": 1}
+        )
+        assert row_data == {}
+
+    def test_skips_empty_info(self) -> None:
+        row_data: dict = {}
+        _extract_parallel_info(row_data=row_data, info={})
+        assert row_data == {}
+
+    def test_ignores_rank_without_size(self) -> None:
+        row_data: dict = {}
+        _extract_parallel_info(row_data=row_data, info={"tp_rank": 0})
+        assert "tp" not in row_data
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))
@@ -0,0 +1,62 @@
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+from sglang.srt.debug_utils.dump_loader import read_tokenizer_path
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="default", nightly=True)
+
+
+def _save_pt(
+    directory: Path, filename: str, *, value: torch.Tensor, meta: dict
+) -> None:
+    torch.save({"value": value, "meta": meta}, directory / filename)
+
+
+class TestReadTokenizerPath:
+    def test_finds_tokenizer_path(self, tmp_path: Path) -> None:
+        _save_pt(
+            tmp_path,
+            "name=x___step=0___rank=0___dump_index=0.pt",
+            value=torch.tensor([1.0]),
+            meta={"tokenizer_path": "/models/llama-3"},
+        )
+        result = read_tokenizer_path(tmp_path)
+        assert result == "/models/llama-3"
+
+    def test_returns_none_when_no_tokenizer_path(self, tmp_path: Path) -> None:
+        _save_pt(
+            tmp_path,
+            "name=x___step=0___rank=0___dump_index=0.pt",
+            value=torch.tensor([1.0]),
+            meta={},
+        )
+        result = read_tokenizer_path(tmp_path)
+        assert result is None
+
+    def test_returns_none_for_empty_directory(self, tmp_path: Path) -> None:
+        result = read_tokenizer_path(tmp_path)
+        assert result is None
+
+    def test_skips_files_without_tokenizer_path(self, tmp_path: Path) -> None:
+        _save_pt(
+            tmp_path,
+            "name=a___step=0___rank=0___dump_index=0.pt",
+            value=torch.tensor([1.0]),
+            meta={},
+        )
+        _save_pt(
+            tmp_path,
+            "name=b___step=0___rank=0___dump_index=1.pt",
+            value=torch.tensor([2.0]),
+            meta={"tokenizer_path": "/models/deepseek"},
+        )
+        result = read_tokenizer_path(tmp_path)
+        assert result == "/models/deepseek"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))
@@ -1,8 +1,13 @@
+import json
 import sys

 import pytest
 from pydantic import ValidationError

+from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
+    AlignerPerStepPlan,
+    AlignerPlan,
+)
 from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
    PositionalSeqId,
    TokenAlignerPlan,
@@ -10,13 +15,18 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.types import (
    TokenAlignerStepAux,
    TokenLocator,
 )
-from sglang.srt.debug_utils.comparator.aligner.unsharder.types import AxisInfo
-from sglang.srt.debug_utils.comparator.dims import TokenLayout
+from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
+    AxisInfo,
+    ConcatParams,
+    UnsharderPlan,
+)
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis, TokenLayout
 from sglang.srt.debug_utils.comparator.output_types import (
    ComparisonRecord,
    GeneralWarning,
    SkipRecord,
    SummaryRecord,
+    parse_record_json,
 )
 from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
    DiffInfo,
@@ -241,5 +251,79 @@ class TestOutputRecordCategories:
        assert record.category == "passed"


+def _make_aligner_plan() -> AlignerPlan:
+    unsharder = UnsharderPlan(
+        axis=ParallelAxis.TP,
+        params=ConcatParams(dim_name="h"),
+        groups=[[0, 1]],
+    )
+    return AlignerPlan(
+        per_step_plans=Pair(
+            x=[
+                AlignerPerStepPlan(
+                    step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
+                )
+            ],
+            y=[
+                AlignerPerStepPlan(
+                    step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
+                )
+            ],
+        ),
+    )
+
+
+class TestAlignerPlanInComparisonRecord:
+    def test_comparison_record_with_aligner_plan(self) -> None:
+        plan: AlignerPlan = _make_aligner_plan()
+        record: ComparisonRecord = _make_comparison_record(
+            diff=_make_diff_info(passed=True),
+        )
+        record_with_plan = record.model_copy(update={"aligner_plan": plan})
+        assert record_with_plan.aligner_plan is not None
+        assert record_with_plan.aligner_plan.per_step_plans.x[0].step == 0
+
+    def test_aligner_plan_json_roundtrip(self) -> None:
+        plan: AlignerPlan = _make_aligner_plan()
+        record: ComparisonRecord = _make_comparison_record(
+            diff=_make_diff_info(passed=True),
+        )
+        record_with_plan = record.model_copy(update={"aligner_plan": plan})
+
+        json_str: str = record_with_plan.model_dump_json()
+        parsed = json.loads(json_str)
+        assert "aligner_plan" in parsed
+        assert (
+            parsed["aligner_plan"]["per_step_plans"]["x"][0]["sub_plans"][0]["type"]
+            == "unsharder"
+        )
+
+        roundtripped: ComparisonRecord = parse_record_json(json_str)
+        assert roundtripped.aligner_plan is not None
+        assert (
+            roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type
+            == "unsharder"
+        )
+
+    def test_comparison_record_without_aligner_plan(self) -> None:
+        record: ComparisonRecord = _make_comparison_record(
+            diff=_make_diff_info(passed=True),
+        )
+        json_str: str = record.model_dump_json()
+        roundtripped: ComparisonRecord = parse_record_json(json_str)
+        assert roundtripped.aligner_plan is None
+
+    def test_aligner_plan_text_format(self) -> None:
+        plan: AlignerPlan = _make_aligner_plan()
+        record: ComparisonRecord = _make_comparison_record(
+            diff=_make_diff_info(passed=True),
+        )
+        record_with_plan = record.model_copy(update={"aligner_plan": plan})
+
+        text: str = record_with_plan.to_text()
+        assert "Aligner Plan:" in text
+        assert "unsharder" in text
+
+
 if __name__ == "__main__":
    sys.exit(pytest.main([__file__]))