Support context parallel zigzag reordering in dump comparator (#19281)

2026-02-25 09:46:17 +08:00
parent 0de1f4b07b
commit 2e2b18e870
13 changed files with 423 additions and 42 deletions
@@ -0,0 +1,158 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.srt.debug_utils.comparator.aligner.reorder import (
+    ReorderPlan,
+    _reorder_zigzag_to_natural,
+    compute_reorder_plans,
+    execute_reorder_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
+    execute_unshard_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
+    compute_unshard_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=10, suite="default", nightly=True)
+
+
+class TestZigzagToNatural:
+    def test_zigzag_to_natural_cp2(self) -> None:
+        """cp_size=2: zigzag order [0,3,1,2] -> natural [0,1,2,3]."""
+        natural = torch.arange(24).reshape(4, 6)
+        chunks = list(natural.chunk(4, dim=0))
+
+        zigzag_order: list[int] = [0, 3, 1, 2]
+        zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=0)
+
+        result = _reorder_zigzag_to_natural(zigzagged, dim=0, cp_size=2)
+        assert torch.equal(result, natural)
+
+    def test_zigzag_to_natural_cp3(self) -> None:
+        """cp_size=3: zigzag 162534 -> natural 123456 (1-indexed)."""
+        natural = torch.arange(60).reshape(6, 10)
+        chunks = list(natural.chunk(6, dim=0))
+
+        zigzag_order: list[int] = [0, 5, 1, 4, 2, 3]
+        zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=0)
+
+        result = _reorder_zigzag_to_natural(zigzagged, dim=0, cp_size=3)
+        assert torch.equal(result, natural)
+
+    def test_zigzag_to_natural_arbitrary_dim(self) -> None:
+        """Reorder along dim=1 instead of dim=0."""
+        natural = torch.arange(48).reshape(3, 4, 4)
+        chunks = list(natural.chunk(4, dim=1))
+
+        zigzag_order: list[int] = [0, 3, 1, 2]
+        zigzagged = torch.cat([chunks[i] for i in zigzag_order], dim=1)
+
+        result = _reorder_zigzag_to_natural(zigzagged, dim=1, cp_size=2)
+        assert torch.equal(result, natural)
+
+
+class TestComputeReorderPlans:
+    def test_compute_reorder_plans_zigzag(self) -> None:
+        """s(cp,zigzag) produces a ReorderPlan."""
+        dim_specs = parse_dims("b s(cp,zigzag) h(tp)")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+        ]
+        plans = compute_reorder_plans(
+            dim_specs=dim_specs, parallel_infos=parallel_infos
+        )
+
+        assert len(plans) == 1
+        assert plans[0].params.op == "zigzag_to_natural"
+        assert plans[0].params.dim == 1
+        assert plans[0].params.cp_size == 2
+
+    def test_compute_reorder_plans_non_seq_dim_raises(self) -> None:
+        """Zigzag on non-sequence dim (e.g. t(cp,zigzag)) raises ValueError."""
+        dim_specs = parse_dims("t(cp,zigzag) h(tp)")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+        ]
+        with pytest.raises(ValueError, match="only supported on sequence dims"):
+            compute_reorder_plans(dim_specs=dim_specs, parallel_infos=parallel_infos)
+
+    def test_compute_reorder_plans_natural(self) -> None:
+        """s(cp) and s(cp,natural) produce no reorder plans."""
+        for dims_str in ["b s(cp) h(tp)", "b s(cp,natural) h(tp)"]:
+            dim_specs = parse_dims(dims_str)
+            parallel_infos: list[dict[ParallelAxis, AxisInfo]] = [
+                {
+                    ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                    ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+                },
+            ]
+            plans = compute_reorder_plans(
+                dim_specs=dim_specs, parallel_infos=parallel_infos
+            )
+            assert plans == []
+
+
+class TestCpZigzagTpE2E:
+    def test_cp_zigzag_tp_e2e(self) -> None:
+        """CP=2 zigzag + TP=2: full pipeline round-trip."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16)
+
+        # Shard: first split seq dim (dim=1) into CP=2 with zigzag ordering,
+        # then split hidden dim (dim=2) into TP=2.
+        natural_cp_chunks = list(full_tensor.chunk(4, dim=1))
+        zigzag_order: list[int] = [0, 3, 1, 2]
+        zigzagged = torch.cat([natural_cp_chunks[i] for i in zigzag_order], dim=1)
+
+        cp_chunks = list(zigzagged.chunk(2, dim=1))
+        tensors: list[torch.Tensor] = []
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for cp_rank in range(2):
+            tp_chunks = list(cp_chunks[cp_rank].chunk(2, dim=2))
+            for tp_rank in range(2):
+                tensors.append(tp_chunks[tp_rank])
+                parallel_infos.append(
+                    {
+                        ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                        ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
+                    }
+                )
+
+        dim_specs = parse_dims("b s(cp,zigzag) h(tp)")
+
+        unshard_plans = compute_unshard_plan(
+            dim_specs=dim_specs, parallel_infos=parallel_infos
+        )
+        reorder_plans = compute_reorder_plans(
+            dim_specs=dim_specs, parallel_infos=parallel_infos
+        )
+        all_plans = [*unshard_plans, *reorder_plans]
+
+        assert len(unshard_plans) == 2
+        assert len(reorder_plans) == 1
+
+        current: list[torch.Tensor] = tensors
+        for plan in all_plans:
+            if isinstance(plan, ReorderPlan):
+                current = execute_reorder_plan(plan, current)
+            else:
+                current = execute_unshard_plan(plan, current)
+
+        assert len(current) == 1
+        assert torch.allclose(current[0], full_tensor)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))
@@ -3,13 +3,15 @@ import sys
 import pytest
 import torch

-from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
-from sglang.srt.debug_utils.comparator.unshard.executor import (
+from sglang.srt.debug_utils.comparator.aligner.unshard.executor import (
    _apply_unshard,
    execute_unshard_plan,
 )
-from sglang.srt.debug_utils.comparator.unshard.planner import compute_unshard_plan
-from sglang.srt.debug_utils.comparator.unshard.types import AxisInfo
+from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
+    compute_unshard_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
 from sglang.test.ci.ci_register import register_cpu_ci

 register_cpu_ci(est_time=10, suite="default", nightly=True)
@@ -2,11 +2,11 @@ import sys

 import pytest

-from sglang.srt.debug_utils.comparator.dims import ParallelAxis
-from sglang.srt.debug_utils.comparator.unshard.parallel_info import (
+from sglang.srt.debug_utils.comparator.aligner.unshard.parallel_info import (
    normalize_parallel_info,
 )
-from sglang.srt.debug_utils.comparator.unshard.types import AxisInfo
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis
 from sglang.test.ci.ci_register import register_cpu_ci

 register_cpu_ci(est_time=10, suite="default", nightly=True)
@@ -2,9 +2,11 @@ import sys

 import pytest

+from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
+    compute_unshard_plan,
+)
+from sglang.srt.debug_utils.comparator.aligner.unshard.types import AxisInfo
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
-from sglang.srt.debug_utils.comparator.unshard.planner import compute_unshard_plan
-from sglang.srt.debug_utils.comparator.unshard.types import AxisInfo
 from sglang.test.ci.ci_register import register_cpu_ci

 register_cpu_ci(est_time=10, suite="default", nightly=True)
@@ -172,13 +174,14 @@ class TestComputeUnshardPlan:
        with pytest.raises(NotImplementedError, match="reduction"):
            compute_unshard_plan(dim_specs, parallel_infos)

-    def test_ordering_not_natural_raises(self) -> None:
+    def test_ordering_zigzag_accepted(self) -> None:
        dim_specs = parse_dims("s(cp,zigzag)")
        parallel_infos = [
            {ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
        ]
-        with pytest.raises(NotImplementedError, match="ordering"):
-            compute_unshard_plan(dim_specs, parallel_infos)
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 1
+        assert plans[0].axis == ParallelAxis.CP

    def test_ordering_natural_accepted(self) -> None:
        dim_specs = parse_dims("s(cp,natural)")
@@ -809,6 +809,74 @@ class TestEntrypointGroupingLogical:
        comp = _assert_single_comparison_passed(records)
        assert comp.name == "hidden"

+    def test_cp_zigzag_unshard(self, tmp_path, capsys):
+        """CP=2 zigzag reorder is correctly undone through the full pipeline."""
+        torch.manual_seed(42)
+        full_baseline = torch.randn(4, 8, 6)
+        full_target = full_baseline + torch.randn(4, 8, 6) * 0.001
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        for side_dir, full_tensor in [
+            (baseline_dir, full_baseline),
+            (target_dir, full_target),
+        ]:
+            _create_cp_zigzag_tp_sharded_dumps(
+                side_dir,
+                full_tensor=full_tensor,
+                name="attn_out",
+                cp_size=2,
+                tp_size=1,
+                seq_dim=1,
+                head_dim=2,
+                dims_str="b s(cp,zigzag) h",
+            )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=0.01,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "attn_out"
+
+    def test_cp_zigzag_tp_unshard(self, tmp_path, capsys):
+        """CP=2 zigzag + TP=2: multi-axis unshard with reorder through full pipeline."""
+        torch.manual_seed(42)
+        full_baseline = torch.randn(4, 8, 16)
+        full_target = full_baseline + torch.randn(4, 8, 16) * 0.001
+
+        baseline_dir = tmp_path / "baseline"
+        target_dir = tmp_path / "target"
+
+        for side_dir, full_tensor in [
+            (baseline_dir, full_baseline),
+            (target_dir, full_target),
+        ]:
+            _create_cp_zigzag_tp_sharded_dumps(
+                side_dir,
+                full_tensor=full_tensor,
+                name="hidden",
+                cp_size=2,
+                tp_size=2,
+                seq_dim=1,
+                head_dim=2,
+                dims_str="b s(cp,zigzag) h(tp)",
+            )
+
+        args = _make_args(
+            baseline_dir / _FIXED_EXP_NAME,
+            target_dir / _FIXED_EXP_NAME,
+            diff_threshold=0.01,
+        )
+
+        records = _run_and_parse(args, capsys)
+        comp = _assert_single_comparison_passed(records)
+        assert comp.name == "hidden"
+

 # --------------------------- Assertion helpers -------------------

@@ -1011,6 +1079,65 @@ def _create_ep_cp_tp_sharded_dumps(
    return directory / _FIXED_EXP_NAME


+def _create_cp_zigzag_tp_sharded_dumps(
+    directory: Path,
+    *,
+    full_tensor: torch.Tensor,
+    name: str,
+    cp_size: int,
+    tp_size: int,
+    seq_dim: int,
+    head_dim: int,
+    dims_str: str,
+    num_steps: int = 1,
+) -> Path:
+    """Create CP-zigzag (+optional TP) sharded dump files from a full tensor."""
+    num_chunks: int = cp_size * 2
+    natural_chunks: list[torch.Tensor] = list(
+        full_tensor.chunk(num_chunks, dim=seq_dim)
+    )
+
+    zigzag_order: list[int] = []
+    for i in range(cp_size):
+        zigzag_order.append(i)
+        zigzag_order.append(num_chunks - 1 - i)
+
+    zigzagged: torch.Tensor = torch.cat(
+        [natural_chunks[idx] for idx in zigzag_order], dim=seq_dim
+    )
+
+    cp_chunks: list[torch.Tensor] = list(zigzagged.chunk(cp_size, dim=seq_dim))
+
+    rank: int = 0
+    for cp_rank in range(cp_size):
+        tp_chunks: list[torch.Tensor] = (
+            list(cp_chunks[cp_rank].chunk(tp_size, dim=head_dim))
+            if tp_size > 1
+            else [cp_chunks[cp_rank]]
+        )
+        for tp_rank in range(tp_size):
+            parallel_info: dict[str, int] = {
+                "cp_rank": cp_rank,
+                "cp_size": cp_size,
+            }
+            if tp_size > 1:
+                parallel_info["tp_rank"] = tp_rank
+                parallel_info["tp_size"] = tp_size
+
+            _create_rank_dump(
+                directory,
+                rank=rank,
+                name=name,
+                tensor=tp_chunks[tp_rank],
+                dims=dims_str,
+                parallel_info=parallel_info,
+                num_steps=num_steps,
+            )
+            rank += 1
+
+    return directory / _FIXED_EXP_NAME
+
+
 def _create_tp_sharded_dumps(
    directory: Path,
    *,