Support TP unification and enhance tests in dump comparator (#19278)

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()

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)
 

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

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__":