From 2caa549c82c8af79e5691241744b1131a5c3d211 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Tue, 25 Aug 2026 11:39:59 -0700 Subject: [PATCH] Fix the seq-first Ulysses all2all output layout _generate_layout_params builds the reshape target for every all2all in DistributedAttention. For batch_dim_idx=1 (s, b, n, h) with scatter_idx < 2 it returns [bs, seq_world_size * global_seq_len, num_local_head // seq_world_size, head_dim], which is the batch_dim_idx=0 / scatter_idx >= 2 shape: it puts the batch first, multiplies the sequence and divides the heads, when this direction scatters the sequence and gathers the heads. Before #6750 extracted this function, post_all2all computed [seq_len // seq_world_size, bs, seq_world_size * num_head, head_dim] for that case, so the refactor copied the wrong sibling branch. Restore that shape. The element count still matches whenever num_local_head is divisible by seq_world_size, so the reshape succeeds and silently returns a transposed, mis-strided tensor; when it is not divisible, the floor division makes a dimension 0 and the reshape raises. Both are reachable from DistributedAttention, whose default gather_idx is 0: the output projection all2all and the backward of the q/k/v all2alls both run scatter_idx < 2. The existing coverage misses it. TestUlyssesAll2All only runs batch_dim_idx=0, and TestUlyssesAll2All_odd sets num_kv_heads on its first call so every later call takes uneven_heads_all2all instead of _generate_layout_params. _generate_layout_params is pure, so add TestUlyssesAll2AllLayout, which drives it with an emulated all_to_all_single and checks that both directions land the right (sequence, head) shard of a known tensor. It needs no process group and no accelerator, so it runs in the CPU CI. Against the current code the two batch_dim_idx=1 head-to-sequence cases fail (2 failed, 6 passed: one shape assertion, one reshape RuntimeError) and all 8 pass with the fix. Signed-off-by: Vineeth Sai --- deepspeed/sequence/layer.py | 4 +- .../unit/sequence_parallelism/test_ulysses.py | 59 ++++++++++++++++++- 2 files changed, 61 insertions(+), 2 deletions(-) diff --git a/deepspeed/sequence/layer.py b/deepspeed/sequence/layer.py index 1ae18c14ced8..a6c926d14f6f 100644 --- a/deepspeed/sequence/layer.py +++ b/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})!" diff --git a/tests/unit/sequence_parallelism/test_ulysses.py b/tests/unit/sequence_parallelism/test_ulysses.py index abb9899892a8..336b7b2a6385 100644 --- a/tests/unit/sequence_parallelism/test_ulysses.py +++ b/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])