Support multi sharding group on the same dimension in dump comparator (#19601)

2026-03-01 10:36:48 +08:00
parent 46960e65cf
commit ea6ff7b01f
14 changed files with 469 additions and 151 deletions
@@ -654,7 +654,7 @@ class TestReduceSum:
        part_a = full_tensor * 0.6
        part_b = full_tensor * 0.4

-        dim_specs = parse_dims("h(tp,partial) d")
+        dim_specs = parse_dims("h(tp:partial) d")
        parallel_infos = [
            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
        ]
@@ -676,7 +676,7 @@ class TestReduceSum:
        full_tensor = torch.randn(4, 8)
        parts: list[torch.Tensor] = [full_tensor * 0.25 for _ in range(4)]

-        dim_specs = parse_dims("h(tp,partial) d")
+        dim_specs = parse_dims("h(tp:partial) d")
        parallel_infos = [
            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
        ]
@@ -710,7 +710,7 @@ class TestReduceSum:
                    }
                )

-        dim_specs = parse_dims("b s(cp) h(tp,partial)")
+        dim_specs = parse_dims("b s(cp) h(tp:partial)")
        plans = compute_unsharder_plan(dim_specs, parallel_infos)
        assert len(plans) == 2

@@ -739,7 +739,7 @@ class TestReduceSum:
            {ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
            {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
        ]
-        dim_specs = parse_dims("h(tp,partial) d")
+        dim_specs = parse_dims("h(tp:partial) d")
        plans = compute_unsharder_plan(dim_specs, parallel_infos)

        named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
@@ -752,7 +752,7 @@ class TestReduceSum:

    def test_reduce_preserves_named_dims(self) -> None:
        """Named tensor dimensions are preserved through reduce_sum."""
-        dim_specs = parse_dims("h(tp,partial) d")
+        dim_specs = parse_dims("h(tp:partial) d")
        part_a = torch.randn(4, 8).refine_names("h", "d")
        part_b = torch.randn(4, 8).refine_names("h", "d")

@@ -175,7 +175,7 @@ class TestComputeUnsharderPlan:
            compute_unsharder_plan(dim_specs, parallel_infos)

    def test_reduction_partial_returns_reduce_sum(self) -> None:
-        dim_specs = parse_dims("h(tp,partial)")
+        dim_specs = parse_dims("h(tp:partial)")
        parallel_infos = [
            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
        ]
@@ -188,7 +188,7 @@ class TestComputeUnsharderPlan:

    def test_reduction_partial_tp4(self) -> None:
        """TP=4 with partial reduction produces a single ReduceSumParams step."""
-        dim_specs = parse_dims("h(tp,partial)")
+        dim_specs = parse_dims("h(tp:partial)")
        parallel_infos = [
            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
        ]
@@ -200,7 +200,7 @@ class TestComputeUnsharderPlan:

    def test_multi_axis_with_reduction_on_one(self) -> None:
        """CP concat + TP reduce produces a 2-step plan."""
-        dim_specs = parse_dims("s(cp) h(tp,partial)")
+        dim_specs = parse_dims("s(cp) h(tp:partial)")
        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
        for cp_rank in range(2):
            for tp_rank in range(2):
@@ -221,7 +221,7 @@ class TestComputeUnsharderPlan:

    def test_reduction_scrambled_ranks(self) -> None:
        """Scrambled world_rank order with partial reduction."""
-        dim_specs = parse_dims("h(tp,partial)")
+        dim_specs = parse_dims("h(tp:partial)")
        parallel_infos = [
            {ParallelAxis.TP: AxisInfo(axis_rank=2, axis_size=4)},
            {ParallelAxis.TP: AxisInfo(axis_rank=0, axis_size=4)},
@@ -235,7 +235,7 @@ class TestComputeUnsharderPlan:
        assert plans[0].groups == [[1, 3, 0, 2]]

    def test_ordering_zigzag_accepted(self) -> None:
-        dim_specs = parse_dims("s(cp,zigzag)")
+        dim_specs = parse_dims("s(cp:zigzag)")
        parallel_infos = [
            {ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
        ]
@@ -244,7 +244,7 @@ class TestComputeUnsharderPlan:
        assert plans[0].axis == ParallelAxis.CP

    def test_ordering_natural_accepted(self) -> None:
-        dim_specs = parse_dims("s(cp,natural)")
+        dim_specs = parse_dims("s(cp:natural)")
        parallel_infos = [
            {ParallelAxis.CP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
        ]
@@ -288,6 +288,77 @@ class TestComputeUnsharderPlan:
        assert len(plans[2].groups) == 1
        assert len(plans[2].groups[0]) == 2

+    def test_same_dim_cp_sp_plan(self) -> None:
+        """t(cp:zigzag,sp) with CP=2 SP=2: SP unshards first (inner), then CP."""
+        dim_specs = parse_dims("t(cp:zigzag,sp) 1 h")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for cp_rank in range(2):
+            for sp_rank in range(2):
+                parallel_infos.append(
+                    {
+                        ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                        ParallelAxis.SP: AxisInfo(axis_rank=sp_rank, axis_size=2),
+                    }
+                )
+
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 2
+
+        # SP unshards first (rightmost modifier = innermost shard)
+        sp_plan = plans[0]
+        assert sp_plan.axis == ParallelAxis.SP
+        assert isinstance(sp_plan.params, ConcatParams)
+        assert sp_plan.params.dim_name == "t"
+        assert len(sp_plan.groups) == 2
+        for group in sp_plan.groups:
+            assert len(group) == 2
+
+        # CP unshards second (leftmost modifier = outermost shard)
+        cp_plan = plans[1]
+        assert cp_plan.axis == ParallelAxis.CP
+        assert isinstance(cp_plan.params, ConcatParams)
+        assert cp_plan.params.dim_name == "t"
+        assert len(cp_plan.groups) == 1
+        assert len(cp_plan.groups[0]) == 2
+
+    def test_same_dim_cp_sp_with_thd(self) -> None:
+        """t(cp:zigzag,sp) with THD: SP → ConcatParams, CP → CpThdConcatParams."""
+        from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
+            CpThdConcatParams,
+        )
+
+        dim_specs = parse_dims("t(cp:zigzag,sp) h")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for cp_rank in range(2):
+            for sp_rank in range(2):
+                parallel_infos.append(
+                    {
+                        ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                        ParallelAxis.SP: AxisInfo(axis_rank=sp_rank, axis_size=2),
+                    }
+                )
+
+        thd_global_seq_lens: list[int] = [100, 64]
+        plans = compute_unsharder_plan(
+            dim_specs, parallel_infos, thd_global_seq_lens=thd_global_seq_lens
+        )
+
+        assert len(plans) == 2
+
+        # SP unshards first: plain concat (SP is not CP, no THD special handling)
+        sp_plan = plans[0]
+        assert sp_plan.axis == ParallelAxis.SP
+        assert isinstance(sp_plan.params, ConcatParams)
+        assert sp_plan.params.dim_name == "t"
+
+        # CP unshards second: THD concat because dim is 't' + axis is CP + thd_global_seq_lens provided
+        cp_plan = plans[1]
+        assert cp_plan.axis == ParallelAxis.CP
+        assert isinstance(cp_plan.params, CpThdConcatParams)
+        assert cp_plan.params.dim_name == "t"
+        assert cp_plan.params.seq_lens_per_rank == [50, 32]
+
    def test_sp_in_dims_but_not_in_parallel_info(self) -> None:
        """s(sp) in dims but SP absent from parallel_info (SP disabled), should auto-skip."""
        dim_specs = parse_dims("s(sp) b h(tp)")