Support multi axis unsharding in dump comparator (#19280)

2026-02-25 09:44:07 +08:00
parent 4bb678f28a
commit 0de1f4b07b
7 changed files with 706 additions and 96 deletions
@@ -701,6 +701,114 @@ class TestEntrypointGroupingLogical:
        assert summary.total == 2
        assert summary.passed == 2

+    def test_cp_tp_unshard(self, tmp_path, capsys):
+        """CP=2 + TP=2: multi-axis shards are unsharded before comparison."""
+        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_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) 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"
+
+    def test_cp_tp_different_sizes(self, tmp_path, capsys):
+        """Baseline CP=2+TP=2 vs target CP=1+TP=4: both sides independently unshard."""
+        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"
+
+        _create_cp_tp_sharded_dumps(
+            baseline_dir,
+            full_tensor=full_baseline,
+            name="hidden",
+            cp_size=2,
+            tp_size=2,
+            seq_dim=1,
+            head_dim=2,
+            dims_str="b s(cp) h(tp)",
+        )
+
+        _create_tp_sharded_dumps(
+            target_dir,
+            full_tensor=full_target,
+            name="hidden",
+            tp_size=4,
+            shard_dim=2,
+            dims_str="b s 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)
+        _assert_single_comparison_passed(records)
+
+    def test_ep_cp_tp_three_axis_unshard(self, tmp_path, capsys):
+        """EP=2 + CP=2 + TP=2: three-axis shards are unsharded before comparison."""
+        torch.manual_seed(42)
+        full_baseline = torch.randn(4, 8, 16, 32)
+        full_target = full_baseline + torch.randn(4, 8, 16, 32) * 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_ep_cp_tp_sharded_dumps(
+                side_dir,
+                full_tensor=full_tensor,
+                name="hidden",
+                ep_size=2,
+                cp_size=2,
+                tp_size=2,
+                expert_dim=1,
+                seq_dim=2,
+                head_dim=3,
+                dims_str="b e(ep) s(cp) 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 -------------------

@@ -825,6 +933,84 @@ def _create_rank_dump(
    return directory / _FIXED_EXP_NAME


+def _create_cp_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+TP multi-axis sharded dump files from a full tensor."""
+    cp_chunks = list(full_tensor.chunk(cp_size, dim=seq_dim))
+    rank = 0
+    for cp_rank in range(cp_size):
+        tp_chunks = list(cp_chunks[cp_rank].chunk(tp_size, dim=head_dim))
+        for tp_rank in range(tp_size):
+            _create_rank_dump(
+                directory,
+                rank=rank,
+                name=name,
+                tensor=tp_chunks[tp_rank],
+                dims=dims_str,
+                parallel_info={
+                    "cp_rank": cp_rank,
+                    "cp_size": cp_size,
+                    "tp_rank": tp_rank,
+                    "tp_size": tp_size,
+                },
+                num_steps=num_steps,
+            )
+            rank += 1
+    return directory / _FIXED_EXP_NAME
+
+
+def _create_ep_cp_tp_sharded_dumps(
+    directory: Path,
+    *,
+    full_tensor: torch.Tensor,
+    name: str,
+    ep_size: int,
+    cp_size: int,
+    tp_size: int,
+    expert_dim: int,
+    seq_dim: int,
+    head_dim: int,
+    dims_str: str,
+    num_steps: int = 1,
+) -> Path:
+    """Create EP+CP+TP three-axis sharded dump files from a full tensor."""
+    ep_chunks = list(full_tensor.chunk(ep_size, dim=expert_dim))
+    rank = 0
+    for ep_rank in range(ep_size):
+        cp_chunks = list(ep_chunks[ep_rank].chunk(cp_size, dim=seq_dim))
+        for cp_rank in range(cp_size):
+            tp_chunks = list(cp_chunks[cp_rank].chunk(tp_size, dim=head_dim))
+            for tp_rank in range(tp_size):
+                _create_rank_dump(
+                    directory,
+                    rank=rank,
+                    name=name,
+                    tensor=tp_chunks[tp_rank],
+                    dims=dims_str,
+                    parallel_info={
+                        "ep_rank": ep_rank,
+                        "ep_size": ep_size,
+                        "cp_rank": cp_rank,
+                        "cp_size": cp_size,
+                        "tp_rank": tp_rank,
+                        "tp_size": tp_size,
+                    },
+                    num_steps=num_steps,
+                )
+                rank += 1
+    return directory / _FIXED_EXP_NAME
+
+
 def _create_tp_sharded_dumps(
    directory: Path,
    *,
@@ -3,8 +3,11 @@ import sys
 import pytest
 import torch

-from sglang.srt.debug_utils.comparator.dims import parse_dims
-from sglang.srt.debug_utils.comparator.unshard.executor import execute_unshard_plan
+from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
+from sglang.srt.debug_utils.comparator.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.test.ci.ci_register import register_cpu_ci
@@ -18,37 +21,240 @@ class TestExecuteUnshardPlan:
        shards = list(full_tensor.chunk(4, dim=1))

        dim_specs = parse_dims("b h(tp) d")
-        parallel_infos = [{"tp": AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)]
-        plan = compute_unshard_plan(dim_specs, parallel_infos)
-        assert plan is not None
-        tensors_by_rank = {i: shards[i] for i in range(4)}
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
+        ]
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 1

-        result = execute_unshard_plan(plan, tensors_by_rank)
-        assert torch.allclose(result, full_tensor)
+        result = execute_unshard_plan(plans[0], shards)
+        assert len(result) == 1
+        assert torch.allclose(result[0], full_tensor)

    def test_scrambled_world_ranks_correct_result(self) -> None:
        full_tensor = torch.randn(4, 8)
        shards = list(full_tensor.chunk(4, dim=0))

        parallel_infos = [
-            {"tp": AxisInfo(axis_rank=2, axis_size=4)},
-            {"tp": AxisInfo(axis_rank=0, axis_size=4)},
-            {"tp": AxisInfo(axis_rank=3, axis_size=4)},
-            {"tp": AxisInfo(axis_rank=1, axis_size=4)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
        ]
        dim_specs = parse_dims("h(tp) d")
-        plan = compute_unshard_plan(dim_specs, parallel_infos)
-        assert plan is not None
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 1

-        tensors_by_rank = {
-            0: shards[2],
-            1: shards[0],
-            2: shards[3],
-            3: shards[1],
-        }
+        tensors_ordered_by_world_rank = [
+            shards[2],  # world_rank=0, axis_rank=2
+            shards[0],  # world_rank=1, axis_rank=0
+            shards[3],  # world_rank=2, axis_rank=3
+            shards[1],  # world_rank=3, axis_rank=1
+        ]

-        result = execute_unshard_plan(plan, tensors_by_rank)
-        assert torch.allclose(result, full_tensor)
+        result = execute_unshard_plan(plans[0], tensors_ordered_by_world_rank)
+        assert len(result) == 1
+        assert torch.allclose(result[0], full_tensor)
+
+    def test_single_step_reduces_tensor_count(self) -> None:
+        """8 tensors with 2 groups of 4 produce 2 output tensors."""
+        full_a = torch.randn(4, 8)
+        full_b = torch.randn(4, 8)
+        shards_a = list(full_a.chunk(4, dim=0))
+        shards_b = list(full_b.chunk(4, dim=0))
+
+        dim_specs = parse_dims("s(cp) h(tp)")
+        parallel_infos = []
+        for cp_rank in range(2):
+            for tp_rank in range(4):
+                parallel_infos.append(
+                    {
+                        ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                        ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=4),
+                    }
+                )
+
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 2
+
+        tensors: list[torch.Tensor] = []
+        for cp_rank in range(2):
+            source = shards_a if cp_rank == 0 else shards_b
+            for tp_rank in range(4):
+                tensors.append(source[tp_rank])
+
+        intermediate = execute_unshard_plan(plans[0], tensors)
+        assert len(intermediate) == 4
+
+        final = execute_unshard_plan(plans[1], intermediate)
+        assert len(final) == 1
+
+    def test_cp_tp_concat(self) -> None:
+        """CP=2 + TP=2: multi-step unshard reconstructs original tensor."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16)
+
+        cp_chunks = list(full_tensor.chunk(2, dim=1))
+        tensors: list[torch.Tensor] = []
+        parallel_infos = []
+        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) h(tp)")
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 2
+
+        current = tensors
+        for plan in plans:
+            current = execute_unshard_plan(plan, current)
+
+        assert len(current) == 1
+        assert torch.allclose(current[0], full_tensor)
+
+    def test_cp_tp_scrambled(self) -> None:
+        """Scrambled world_ranks for CP=2 + TP=2 still reconstruct correctly."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16)
+
+        cp_chunks = list(full_tensor.chunk(2, dim=1))
+        shard_map: dict[tuple[int, int], torch.Tensor] = {}
+        for cp_rank in range(2):
+            tp_chunks = list(cp_chunks[cp_rank].chunk(2, dim=2))
+            for tp_rank in range(2):
+                shard_map[(cp_rank, tp_rank)] = tp_chunks[tp_rank]
+
+        scrambled_assignment = [
+            (1, 1),  # world_rank=0
+            (0, 0),  # world_rank=1
+            (1, 0),  # world_rank=2
+            (0, 1),  # world_rank=3
+        ]
+
+        tensors: list[torch.Tensor] = []
+        parallel_infos = []
+        for cp_rank, tp_rank in scrambled_assignment:
+            tensors.append(shard_map[(cp_rank, 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) h(tp)")
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 2
+
+        current = tensors
+        for plan in plans:
+            current = execute_unshard_plan(plan, current)
+
+        assert len(current) == 1
+        assert torch.allclose(current[0], full_tensor)
+
+    def test_unsupported_params_type_raises(self) -> None:
+        """_apply_unshard raises ValueError for unknown params type."""
+
+        class _FakeParams:
+            pass
+
+        with pytest.raises(ValueError, match="Unsupported unshard"):
+            _apply_unshard(_FakeParams(), [torch.randn(2, 2)])
+
+    def test_cp_tp_ep_three_axis_concat(self) -> None:
+        """CP=2 + TP=2 + EP=2: three-step unshard reconstructs original tensor."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16, 32)
+
+        ep_chunks = list(full_tensor.chunk(2, dim=1))
+        shard_map: dict[tuple[int, int, int], torch.Tensor] = {}
+        for ep_rank in range(2):
+            cp_chunks = list(ep_chunks[ep_rank].chunk(2, dim=2))
+            for cp_rank in range(2):
+                tp_chunks = list(cp_chunks[cp_rank].chunk(2, dim=3))
+                for tp_rank in range(2):
+                    shard_map[(ep_rank, cp_rank, tp_rank)] = tp_chunks[tp_rank]
+
+        tensors: list[torch.Tensor] = []
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for ep_rank in range(2):
+            for cp_rank in range(2):
+                for tp_rank in range(2):
+                    tensors.append(shard_map[(ep_rank, cp_rank, tp_rank)])
+                    parallel_infos.append(
+                        {
+                            ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
+                            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 e(ep) s(cp) h(tp)")
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 3
+
+        current = tensors
+        for plan in plans:
+            current = execute_unshard_plan(plan, current)
+
+        assert len(current) == 1
+        assert torch.allclose(current[0], full_tensor)
+
+    def test_cp_tp_ep_scrambled_three_axis(self) -> None:
+        """Scrambled ranks for CP=2 + TP=2 + EP=2 still reconstruct correctly."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16, 32)
+
+        ep_chunks = list(full_tensor.chunk(2, dim=1))
+        shard_map: dict[tuple[int, int, int], torch.Tensor] = {}
+        for ep_rank in range(2):
+            cp_chunks = list(ep_chunks[ep_rank].chunk(2, dim=2))
+            for cp_rank in range(2):
+                tp_chunks = list(cp_chunks[cp_rank].chunk(2, dim=3))
+                for tp_rank in range(2):
+                    shard_map[(ep_rank, cp_rank, tp_rank)] = tp_chunks[tp_rank]
+
+        scrambled_assignment = [
+            (1, 0, 1),  # world_rank=0
+            (0, 1, 0),  # world_rank=1
+            (1, 1, 0),  # world_rank=2
+            (0, 0, 0),  # world_rank=3
+            (0, 1, 1),  # world_rank=4
+            (1, 0, 0),  # world_rank=5
+            (0, 0, 1),  # world_rank=6
+            (1, 1, 1),  # world_rank=7
+        ]
+
+        tensors: list[torch.Tensor] = []
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for ep_rank, cp_rank, tp_rank in scrambled_assignment:
+            tensors.append(shard_map[(ep_rank, cp_rank, tp_rank)])
+            parallel_infos.append(
+                {
+                    ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
+                    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 e(ep) s(cp) h(tp)")
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 3
+
+        current = tensors
+        for plan in plans:
+            current = execute_unshard_plan(plan, current)
+
+        assert len(current) == 1
+        assert torch.allclose(current[0], full_tensor)


 if __name__ == "__main__":
@@ -16,12 +16,12 @@ class TestComputeUnshardPlan:
        parallel_infos = [
            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
        ]
-        plan = compute_unshard_plan(dim_specs, parallel_infos)
+        plans = compute_unshard_plan(dim_specs, parallel_infos)

-        assert plan is not None
-        assert plan.axis == ParallelAxis.TP
-        assert plan.params.dim == 2
-        assert plan.world_ranks_by_axis_rank == [0, 1, 2, 3]
+        assert len(plans) == 1
+        assert plans[0].axis == ParallelAxis.TP
+        assert plans[0].params.dim == 2
+        assert plans[0].groups == [[0, 1, 2, 3]]

    def test_inconsistent_axis_size_raises(self) -> None:
        dim_specs = parse_dims("h(tp)")
@@ -35,7 +35,7 @@ class TestComputeUnshardPlan:
    def test_missing_axis_in_parallel_info_raises(self) -> None:
        dim_specs = parse_dims("h(tp)")
        parallel_infos = [{ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2)}]
-        with pytest.raises(ValueError, match="No parallel_info found"):
+        with pytest.raises(ValueError, match="missing parallel_info"):
            compute_unshard_plan(dim_specs, parallel_infos)

    def test_empty_parallel_infos_raises(self) -> None:
@@ -52,29 +52,193 @@ class TestComputeUnshardPlan:
            {ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
            {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
        ]
-        plan = compute_unshard_plan(dim_specs, parallel_infos)
-        assert plan is not None
-        assert plan.world_ranks_by_axis_rank == [1, 3, 0, 2]
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 1
+        assert plans[0].groups == [[1, 3, 0, 2]]

-    def test_no_sharded_axes_returns_none(self) -> None:
+    def test_no_sharded_axes_returns_empty(self) -> None:
        dim_specs = parse_dims("b s d")
        parallel_infos = [{}]
-        plan = compute_unshard_plan(dim_specs, parallel_infos)
-        assert plan is None
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert plans == []

-    def test_multi_axis_raises(self) -> None:
-        dim_specs = parse_dims("h(tp) s(cp)")
+    def test_multi_axis_plan(self) -> None:
+        """Multi-axis (TP + CP) produces a 2-step plan."""
+        dim_specs = parse_dims("s(cp) h(tp)")
        parallel_infos = [
            {
-                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
            },
            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+            {
                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
            },
        ]
-        with pytest.raises(NotImplementedError, match="Multi-axis unshard"):
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 2
+        assert plans[0].axis == ParallelAxis.CP
+        assert plans[1].axis == ParallelAxis.TP
+
+    def test_cp_tp_plan(self) -> None:
+        """CP=2 + TP=4 produces correct 2-step plan with correct groups."""
+        dim_specs = parse_dims("s(cp) h(tp)")
+        parallel_infos = []
+        for cp_rank in range(2):
+            for tp_rank in range(4):
+                parallel_infos.append(
+                    {
+                        ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                        ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=4),
+                    }
+                )
+
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 2
+
+        cp_plan = plans[0]
+        assert cp_plan.axis == ParallelAxis.CP
+        assert len(cp_plan.groups) == 4
+        for group in cp_plan.groups:
+            assert len(group) == 2
+
+        tp_plan = plans[1]
+        assert tp_plan.axis == ParallelAxis.TP
+        assert len(tp_plan.groups) == 1
+        assert len(tp_plan.groups[0]) == 4
+
+    def test_cp_tp_scrambled_ranks(self) -> None:
+        """Scrambled rank assignment still produces correct plan."""
+        dim_specs = parse_dims("s(cp) h(tp)")
+        parallel_infos = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+        ]
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 2
+
+        cp_plan = plans[0]
+        assert cp_plan.axis == ParallelAxis.CP
+        assert len(cp_plan.groups) == 2
+        for group in cp_plan.groups:
+            assert len(group) == 2
+
+        tp_plan = plans[1]
+        assert tp_plan.axis == ParallelAxis.TP
+        assert len(tp_plan.groups) == 1
+        assert len(tp_plan.groups[0]) == 2
+
+    def test_axis_rank_coverage_incomplete_raises(self) -> None:
+        """TP size=4 but only ranks 0,1,3 provided (missing rank 2)."""
+        dim_specs = parse_dims("h(tp)")
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
+        ]
+        with pytest.raises(ValueError, match="axis_rank coverage.*incomplete"):
+            compute_unshard_plan(dim_specs, parallel_infos)
+
+    def test_reduction_not_implemented_raises(self) -> None:
+        dim_specs = parse_dims("h(tp,partial)")
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
+        ]
+        with pytest.raises(NotImplementedError, match="reduction"):
+            compute_unshard_plan(dim_specs, parallel_infos)
+
+    def test_ordering_not_natural_raises(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)
+
+    def test_ordering_natural_accepted(self) -> None:
+        dim_specs = parse_dims("s(cp,natural)")
+        parallel_infos = [
+            {ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
+        ]
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+        assert len(plans) == 1
+        assert plans[0].axis == ParallelAxis.CP
+
+    def test_three_axis_plan(self) -> None:
+        """EP=2 + CP=2 + TP=2 produces a 3-step plan."""
+        dim_specs = parse_dims("b e(ep) s(cp) h(tp)")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for ep_rank in range(2):
+            for cp_rank in range(2):
+                for tp_rank in range(2):
+                    parallel_infos.append(
+                        {
+                            ParallelAxis.EP: AxisInfo(axis_rank=ep_rank, axis_size=2),
+                            ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                            ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
+                        }
+                    )
+
+        plans = compute_unshard_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 3
+        assert plans[0].axis == ParallelAxis.EP
+        assert plans[1].axis == ParallelAxis.CP
+        assert plans[2].axis == ParallelAxis.TP
+
+        # Step 0 (EP): 8 tensors → 4 (groups of 2)
+        assert len(plans[0].groups) == 4
+        for group in plans[0].groups:
+            assert len(group) == 2
+
+        # Step 1 (CP): 4 tensors → 2 (groups of 2)
+        assert len(plans[1].groups) == 2
+        for group in plans[1].groups:
+            assert len(group) == 2
+
+        # Step 2 (TP): 2 tensors → 1 (single group of 2)
+        assert len(plans[2].groups) == 1
+        assert len(plans[2].groups[0]) == 2
+
+    def test_replicated_axis_raises(self) -> None:
+        """A world_rank missing a sharded axis raises ValueError."""
+        dim_specs = parse_dims("s(cp) h(tp)")
+        parallel_infos = [
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=0, axis_size=2),
+                ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=2),
+            },
+            {
+                ParallelAxis.CP: AxisInfo(axis_rank=1, axis_size=2),
+                # missing TP — replicated
+            },
+        ]
+        with pytest.raises(ValueError, match="missing parallel_info"):
            compute_unshard_plan(dim_specs, parallel_infos)