Fix the seq-first Ulysses all2all output layout

deepspeed/sequence/layer.py

@@ -50,7 +50,9 @@ def _generate_layout_params(scatter_idx, batch_dim_idx, seq_world_size, input):
             pre_all2all_permute_idx = None
 
             post_all2all_permute_idx = (1, 2, 0, 3, 4)
-            post_all2all_res_shape = [bs, seq_world_size * global_seq_len, num_local_head // seq_world_size, head_dim]
+            # seq-first layout: the all2all scatters the sequence and gathers the heads, so the
+            # result keeps bs in dim 1 with a local sequence and every rank's heads.
+            post_all2all_res_shape = [global_seq_len // seq_world_size, bs, seq_world_size * num_local_head, head_dim]
         else:
             local_seq_len, bs, num_total_head, head_dim = input.shape
             assert num_total_head % seq_world_size == 0, f"Number of heads ({num_total_head}) must be divisible by the sequence parallel size ({seq_world_size})!"

tests/unit/sequence_parallelism/test_ulysses.py

@@ -11,7 +11,7 @@
 import deepspeed.runtime.sequence_parallel.parallel_state_sp as sp_mpu
 from transformers import AutoModel
 from unit.common import DistributedTest
-from deepspeed.sequence.layer import _SeqAllToAll
+from deepspeed.sequence.layer import _SeqAllToAll, _generate_layout_params, post_all2all, pre_all2all_fun
 from deepspeed.sequence.fpdt_layer import _FPDTGPUOffloadingAttentionImpl_, FPDT_InputConstruct
 from unit.util import skip_on_arch
 from unit.simple_model import *
@@ -148,6 +148,63 @@ def test_alltoall_output_consistency(self, d0: int, d1: int, head_dim: int, num_
             assert torch.allclose(input_tensor, outputs[i]), f"Outputs differ for sequence dim {seq_dims[i]}"
 
 
+def _emulate_all_to_all(shards):
+    """CPU stand-in for dist.all_to_all_single: rank i sends chunk j of dim 0 to rank j."""
+    seq_world_size = len(shards)
+    return [
+        torch.cat([shards[src][dst:dst + 1] for src in range(seq_world_size)], dim=0) for dst in range(seq_world_size)
+    ]
+
+
+def _run_layout_all_to_all(scatter_idx, batch_dim_idx, seq_world_size, shards):
+    """single_all_to_all's layout math, driven by the emulated all2all above."""
+    pre_permute_idx, pre_inp_shape, post_permute_idx, post_res_shape = _generate_layout_params(
+        scatter_idx, batch_dim_idx, seq_world_size, shards[0])
+    sent = [pre_all2all_fun(pre_permute_idx, pre_inp_shape, shard) for shard in shards]
+    post_fun = post_all2all(post_permute_idx, post_res_shape)
+    return [post_fun(received) for received in _emulate_all_to_all(sent)]
+
+
+@pytest.mark.parametrize("batch_dim_idx", [0, 1])
+@pytest.mark.parametrize("seq_world_size", [2, 4])
+class TestUlyssesAll2AllLayout:
+    """_generate_layout_params is a pure function, so the shapes it hands to reshape can be
+    checked on CPU without a process group. TestUlyssesAll2All above only runs batch_dim_idx=0
+    and TestUlyssesAll2All_odd takes the uneven-head path, so the seq-first (s, b, n, h) layout
+    is otherwise never exercised."""
+
+    def _shards(self, batch_dim_idx, seq_world_size):
+        local_seq_len, bs, local_num_heads, head_dim = 3, 2, 2, 4
+        seq_len = local_seq_len * seq_world_size
+        num_heads = local_num_heads * seq_world_size
+        seq_dim = 1 if batch_dim_idx == 0 else 0
+        full = torch.arange(seq_len * bs * num_heads * head_dim, dtype=torch.float32)
+        full = full.reshape(seq_len, bs, num_heads, head_dim)
+        if batch_dim_idx == 0:
+            full = full.transpose(0, 1).contiguous()
+        # sequence parallel: every head, a slice of the sequence.
+        seq_parallel = [
+            full.narrow(seq_dim, r * local_seq_len, local_seq_len).contiguous() for r in range(seq_world_size)
+        ]
+        # head parallel: every position, a slice of the heads.
+        head_parallel = [
+            full.narrow(2, r * local_num_heads, local_num_heads).contiguous() for r in range(seq_world_size)
+        ]
+        return seq_dim, seq_parallel, head_parallel
+
+    def test_seq_to_head_parallel(self, batch_dim_idx, seq_world_size):
+        _, seq_parallel, head_parallel = self._shards(batch_dim_idx, seq_world_size)
+        got = _run_layout_all_to_all(2, batch_dim_idx, seq_world_size, seq_parallel)
+        for rank, (actual, expected) in enumerate(zip(got, head_parallel)):
+            assert torch.equal(actual, expected), f"rank {rank} got {actual.shape}, expected {expected.shape}"
+
+    def test_head_to_seq_parallel(self, batch_dim_idx, seq_world_size):
+        seq_dim, seq_parallel, head_parallel = self._shards(batch_dim_idx, seq_world_size)
+        got = _run_layout_all_to_all(seq_dim, batch_dim_idx, seq_world_size, head_parallel)
+        for rank, (actual, expected) in enumerate(zip(got, seq_parallel)):
+            assert torch.equal(actual, expected), f"rank {rank} got {actual.shape}, expected {expected.shape}"
+
+
 @pytest.mark.parametrize("d0", [2, 4])  #batch or sequence dimension
 @pytest.mark.parametrize("d1", [4, 8])  #batch or sequence dimension
 @pytest.mark.parametrize("num_heads", [3, 7])