Support partial tensors waiting for reduction and pipeline parallel in dump comparator (#19595)

2026-03-01 10:33:39 +08:00
parent 67810828cf
commit 003ad6daaa
8 changed files with 484 additions and 108 deletions
@@ -15,6 +15,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
    AxisInfo,
    CpThdConcatParams,
    PickParams,
+    ReduceSumParams,
    UnsharderPlan,
 )
 from sglang.srt.debug_utils.comparator.dims import (
@@ -639,118 +640,125 @@ class TestThdCpConcat:
        )


-class TestThdCpConcat:
-    def test_single_seq(self) -> None:
-        """Single seq THD unshard: 2 ranks → per-seq concat."""
-        rank0 = torch.tensor([1, 2, 3]).refine_names("t")
-        rank1 = torch.tensor([4, 5, 6]).refine_names("t")
-
-        plan = UnsharderPlan(
-            axis=ParallelAxis.CP,
-            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3]),
-            groups=[[0, 1]],
-        )
-        with warning_sink.context():
-            result = execute_unsharder_plan(plan, [rank0, rank1])
-
-        assert len(result) == 1
-        expected = torch.tensor([1, 2, 3, 4, 5, 6])
-        assert torch.equal(result[0].rename(None), expected)
-
-    def test_multi_seq(self) -> None:
-        """Multi-seq THD unshard: 2 ranks, seq_lens=[50, 32, 46]."""
-        # rank0: [seqA_r0(50) | seqB_r0(32) | pad_r0(46)]
-        # rank1: [seqA_r1(50) | seqB_r1(32) | pad_r1(46)]
-        seq_a_r0 = torch.arange(0, 50)
-        seq_b_r0 = torch.arange(100, 132)
-        pad_r0 = torch.full((46,), -1)
-        rank0 = torch.cat([seq_a_r0, seq_b_r0, pad_r0]).refine_names("t")
-
-        seq_a_r1 = torch.arange(50, 100)
-        seq_b_r1 = torch.arange(132, 164)
-        pad_r1 = torch.full((46,), -2)
-        rank1 = torch.cat([seq_a_r1, seq_b_r1, pad_r1]).refine_names("t")
-
-        plan = UnsharderPlan(
-            axis=ParallelAxis.CP,
-            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[50, 32, 46]),
-            groups=[[0, 1]],
-        )
-        with warning_sink.context():
-            result = execute_unsharder_plan(plan, [rank0, rank1])
-
-        assert len(result) == 1
-        unsharded: torch.Tensor = result[0].rename(None)
-
-        # seqA: r0(50) + r1(50) = 100 tokens, values 0..99
-        assert torch.equal(unsharded[:100], torch.cat([seq_a_r0, seq_a_r1]))
-        # seqB: r0(32) + r1(32) = 64 tokens
-        assert torch.equal(unsharded[100:164], torch.cat([seq_b_r0, seq_b_r1]))
-        # pad: r0(46) + r1(46) = 92 tokens
-        assert torch.equal(unsharded[164:256], torch.cat([pad_r0, pad_r1]))
-
-    def test_with_hidden_dim(self) -> None:
-        """THD unshard with trailing hidden dim: shape [T, H]."""
+class TestReduceSum:
+    def test_basic_tp2_reduce(self) -> None:
+        """2 partial tensors sum to full tensor."""
        torch.manual_seed(42)
-        hidden: int = 4
-        # rank0: [seqA_r0(3, 4) | seqB_r0(2, 4)]
-        # rank1: [seqA_r1(3, 4) | seqB_r1(2, 4)]
-        seq_a_r0 = torch.randn(3, hidden)
-        seq_b_r0 = torch.randn(2, hidden)
-        rank0 = torch.cat([seq_a_r0, seq_b_r0]).refine_names("t", "h")
+        full_tensor = torch.randn(4, 8)
+        part_a = full_tensor * 0.6
+        part_b = full_tensor * 0.4

-        seq_a_r1 = torch.randn(3, hidden)
-        seq_b_r1 = torch.randn(2, hidden)
-        rank1 = torch.cat([seq_a_r1, seq_b_r1]).refine_names("t", "h")
+        dim_specs = parse_dims("h(tp,partial) d")
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=2)} for i in range(2)
+        ]
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        assert len(plans) == 1
+        assert isinstance(plans[0].params, ReduceSumParams)

-        plan = UnsharderPlan(
-            axis=ParallelAxis.CP,
-            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
-            groups=[[0, 1]],
-        )
+        named_parts: list[torch.Tensor] = _name_tensors([part_a, part_b], dim_specs)
        with warning_sink.context():
-            result = execute_unsharder_plan(plan, [rank0, rank1])
+            result = execute_unsharder_plan(plans[0], named_parts)

        assert len(result) == 1
-        unsharded: torch.Tensor = result[0].rename(None)
+        assert torch.allclose(result[0].rename(None), full_tensor)

-        assert unsharded.shape == (10, hidden)
-        assert torch.equal(unsharded[:6], torch.cat([seq_a_r0, seq_a_r1]))
-        assert torch.equal(unsharded[6:10], torch.cat([seq_b_r0, seq_b_r1]))
-
-    def test_with_leading_batch_dim(self) -> None:
-        """THD unshard with leading batch dim: shape [B, T, H], t is dim=1."""
+    def test_tp4_reduce(self) -> None:
+        """4 partial tensors sum to full tensor."""
        torch.manual_seed(42)
-        batch: int = 2
-        hidden: int = 4
-        # rank0: [seqA_r0(3) | seqB_r0(2)] per batch item
-        # rank1: [seqA_r1(3) | seqB_r1(2)] per batch item
-        seq_a_r0 = torch.randn(batch, 3, hidden)
-        seq_b_r0 = torch.randn(batch, 2, hidden)
-        rank0 = torch.cat([seq_a_r0, seq_b_r0], dim=1).refine_names("b", "t", "h")
+        full_tensor = torch.randn(4, 8)
+        parts: list[torch.Tensor] = [full_tensor * 0.25 for _ in range(4)]

-        seq_a_r1 = torch.randn(batch, 3, hidden)
-        seq_b_r1 = torch.randn(batch, 2, hidden)
-        rank1 = torch.cat([seq_a_r1, seq_b_r1], dim=1).refine_names("b", "t", "h")
+        dim_specs = parse_dims("h(tp,partial) d")
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
+        ]
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        assert len(plans) == 1
+
+        named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
+        with warning_sink.context():
+            result = execute_unsharder_plan(plans[0], named_parts)
+
+        assert len(result) == 1
+        assert torch.allclose(result[0].rename(None), full_tensor)
+
+    def test_multi_axis_concat_then_reduce(self) -> None:
+        """CP concat + TP reduce end-to-end."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8, 16)
+
+        cp_chunks = list(full_tensor.chunk(2, dim=1))
+        # Each CP chunk is held as partial sums across TP ranks
+        tensors: list[torch.Tensor] = []
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for cp_rank in range(2):
+            for tp_rank in range(2):
+                tensors.append(cp_chunks[cp_rank] * 0.5)
+                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,partial)")
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+        assert len(plans) == 2
+
+        current: list[torch.Tensor] = _name_tensors(tensors, dim_specs)
+        with warning_sink.context():
+            for plan in plans:
+                current = execute_unsharder_plan(plan, current)
+
+        assert len(current) == 1
+        assert torch.allclose(current[0].rename(None), full_tensor)
+
+    def test_reduce_scrambled_ranks(self) -> None:
+        """Scrambled rank order — sum is commutative so result is the same."""
+        torch.manual_seed(42)
+        full_tensor = torch.randn(4, 8)
+        parts: list[torch.Tensor] = [
+            full_tensor * 0.1,
+            full_tensor * 0.2,
+            full_tensor * 0.3,
+            full_tensor * 0.4,
+        ]
+
+        parallel_infos = [
+            {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,partial) d")
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+
+        named_parts: list[torch.Tensor] = _name_tensors(parts, dim_specs)
+        with warning_sink.context():
+            result = execute_unsharder_plan(plans[0], named_parts)
+
+        assert len(result) == 1
+        assert torch.allclose(result[0].rename(None), full_tensor)
+
+    def test_reduce_preserves_named_dims(self) -> None:
+        """Named tensor dimensions are preserved through reduce_sum."""
+        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")

        plan = UnsharderPlan(
-            axis=ParallelAxis.CP,
-            params=CpThdConcatParams(dim_name="t", seq_lens_per_rank=[3, 2]),
+            axis=ParallelAxis.TP,
+            params=ReduceSumParams(),
            groups=[[0, 1]],
        )
        with warning_sink.context():
-            result = execute_unsharder_plan(plan, [rank0, rank1])
+            result = execute_unsharder_plan(plan, [part_a, part_b])

        assert len(result) == 1
-        unsharded: torch.Tensor = result[0].rename(None)
-
-        assert unsharded.shape == (batch, 10, hidden)
-        # seqA: r0(3) + r1(3) = 6 tokens per batch
-        assert torch.equal(unsharded[:, :6, :], torch.cat([seq_a_r0, seq_a_r1], dim=1))
-        # seqB: r0(2) + r1(2) = 4 tokens per batch
-        assert torch.equal(
-            unsharded[:, 6:10, :], torch.cat([seq_b_r0, seq_b_r1], dim=1)
-        )
+        assert result[0].names == ("h", "d")
+        expected = (part_a.rename(None) + part_b.rename(None)).refine_names("h", "d")
+        assert torch.allclose(result[0].rename(None), expected.rename(None))


 if __name__ == "__main__":
@@ -9,6 +9,7 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
    AxisInfo,
    ConcatParams,
    PickParams,
+    ReduceSumParams,
 )
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis, parse_dims
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -173,13 +174,65 @@ class TestComputeUnsharderPlan:
        with pytest.raises(ValueError, match="axis_rank coverage.*incomplete"):
            compute_unsharder_plan(dim_specs, parallel_infos)

-    def test_reduction_not_implemented_raises(self) -> None:
+    def test_reduction_partial_returns_reduce_sum(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_unsharder_plan(dim_specs, parallel_infos)
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 1
+        assert plans[0].axis == ParallelAxis.TP
+        assert isinstance(plans[0].params, ReduceSumParams)
+        assert plans[0].groups == [[0, 1]]
+
+    def test_reduction_partial_tp4(self) -> None:
+        """TP=4 with partial reduction produces a single ReduceSumParams step."""
+        dim_specs = parse_dims("h(tp,partial)")
+        parallel_infos = [
+            {ParallelAxis.TP: AxisInfo(axis_rank=i, axis_size=4)} for i in range(4)
+        ]
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 1
+        assert isinstance(plans[0].params, ReduceSumParams)
+        assert plans[0].groups == [[0, 1, 2, 3]]
+
+    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)")
+        parallel_infos: list[dict[ParallelAxis, AxisInfo]] = []
+        for cp_rank in range(2):
+            for tp_rank in range(2):
+                parallel_infos.append(
+                    {
+                        ParallelAxis.CP: AxisInfo(axis_rank=cp_rank, axis_size=2),
+                        ParallelAxis.TP: AxisInfo(axis_rank=tp_rank, axis_size=2),
+                    }
+                )
+
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 2
+        assert plans[0].axis == ParallelAxis.CP
+        assert isinstance(plans[0].params, ConcatParams)
+        assert plans[1].axis == ParallelAxis.TP
+        assert isinstance(plans[1].params, ReduceSumParams)
+
+    def test_reduction_scrambled_ranks(self) -> None:
+        """Scrambled world_rank order with partial reduction."""
+        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)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=3, axis_size=4)},
+            {ParallelAxis.TP: AxisInfo(axis_rank=1, axis_size=4)},
+        ]
+        plans = compute_unsharder_plan(dim_specs, parallel_infos)
+
+        assert len(plans) == 1
+        assert isinstance(plans[0].params, ReduceSumParams)
+        assert plans[0].groups == [[1, 3, 0, 2]]

    def test_ordering_zigzag_accepted(self) -> None:
        dim_specs = parse_dims("s(cp,zigzag)")