Skip to content

[Ulysses] Carry the KV head count per DistributedAttention #8316

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
delock merged 6 commits into from
Sep 3, 2026
Merged

[Ulysses] Carry the KV head count per DistributedAttention #8316

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
cfbeedf
[Ulysses] Carry the KV head count per call instead of a module global
alanhuangyoo Aug 26, 2026
f9d0757
Cover batch_dim_idx 1 in the GQA case
alanhuangyoo Aug 26, 2026
74ad9c7
Move tests/unit/sequence_parallelism to tests/unit/v1
alanhuangyoo Aug 26, 2026
fa7caac
Merge branch 'master' into fix/ulysses-per-model-head-count
delock Sep 1, 2026
0a12a59
Merge branch 'master' into fix/ulysses-per-model-head-count
delock Sep 2, 2026
c4e46ac
Merge branch 'master' into fix/ulysses-per-model-head-count
delock Sep 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .github/workflows/nv-flash-attn.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@ on:
pull_request:
paths:
- 'deepspeed/sequence/**'
- 'tests/unit/sequence_parallelism/**'
- 'tests/unit/v1/sequence_parallelism/**'
- '.github/workflows/nv-flash-attn.yml'
schedule:
- cron: "0 0 * * *"
Expand Down Expand Up @@ -57,7 +57,7 @@ jobs:
run: |
unset TORCH_CUDA_ARCH_LIST # only jit compile for current arch
cd tests
python -m pytest --color=yes --durations=0 --verbose -rF unit/sequence_parallelism/test_ulysses.py --torch_ver="2.6" --cuda_ver="12"
python -m pytest --color=yes --durations=0 --verbose -rF unit/v1/sequence_parallelism/test_ulysses.py --torch_ver="2.6" --cuda_ver="12"
- name: Open GitHub issue if nightly CI fails
if: ${{ failure() && (github.event_name == 'schedule') }}
uses: JasonEtco/create-an-issue@v2
Expand Down
2 changes: 1 addition & 1 deletion deepspeed/runtime/engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1923,7 +1923,7 @@ def _configure_distributed_model(self, model):
"DeepSpeed Sequence Parallelism (Ulysses) with PyTorch < 2.3 may encounter "
"rank indexing errors during backward pass when sp_size < world_size. "
"Please use the weighted all-reduce workaround shown in the regression test "
"(https://github.com/deepspeedai/DeepSpeed/blob/master/tests/unit/sequence_parallelism/test_ulysses.py) "
"(https://github.com/deepspeedai/DeepSpeed/blob/master/tests/unit/v1/sequence_parallelism/test_ulysses.py) "
"or upgrade to PyTorch 2.3+.")
self.communication_data_type = self._config.seq_parallel_communication_data_type
self.seq_parallel_group = groups._get_sequence_parallel_group()
Expand Down
90 changes: 71 additions & 19 deletions deepspeed/sequence/fpdt_layer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,9 @@ def forward(ctx: Any,
cpu_offloading=True):

do_save = layernorm_output.requires_grad
# The all-to-all partitions by KV group, and only the caller knows that count: the
# gather-direction and backward calls below see an already-sharded head dim.
num_kv_heads = kv_projection_size // hidden_size_per_attention_head

if rotary_pos_emb is not None:
pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3)
Expand All @@ -180,6 +183,7 @@ def forward(ctx: Any,
ctx.dtype = layernorm_output.dtype
ctx.projection_size = projection_size
ctx.kv_projection_size = kv_projection_size
ctx.num_kv_heads = num_kv_heads

global_q = []
global_k = []
Expand All @@ -206,7 +210,7 @@ def forward(ctx: Any,
q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape(
qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd
q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg)
q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)
global_q_chunk_len = q_chunk.shape[1]
if rotary_pos_emb is not None:
q_chunk = apply_rotary_pos_emb(q_chunk,
Expand All @@ -217,7 +221,7 @@ def forward(ctx: Any,
k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape(
qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd
k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg)
k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)
if rotary_pos_emb is not None:
k_chunk = apply_rotary_pos_emb(k_chunk,
pos_emb_cos[:, global_q_chunk_len * i:global_q_chunk_len * (i + 1)],
Expand All @@ -227,7 +231,7 @@ def forward(ctx: Any,
v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape(
qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd
v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg)
v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)
global_v.append(v_chunk)

for k_i in range(len(global_k)):
Expand Down Expand Up @@ -276,7 +280,12 @@ def forward(ctx: Any,

for i in range(num_chunks):
global_lse[i] = global_lse[i][:, :, :, 0].permute(0, 2, 1).contiguous()
output[i] = single_all_to_all(global_o[i].to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg)
output[i] = single_all_to_all(global_o[i].to(ctx.dtype).contiguous(),
gather_idx,
scatter_idx,
0,
spg,
num_kv_heads=num_kv_heads)
output = torch.cat(output, dim=1)

head_dim = output.shape[-1]
Expand Down Expand Up @@ -312,6 +321,7 @@ def backward(ctx, grad_output):

projection_size = ctx.projection_size
kv_projection_size = ctx.kv_projection_size
num_kv_heads = ctx.num_kv_heads

layernorm_output = ctx.saved_tensors[0]

Expand All @@ -338,7 +348,12 @@ def backward(ctx, grad_output):
st = chunk_size * i
ed = st + chunk_size
grad_global_attn_output.append(
single_all_to_all(grad_output[:, st:ed].contiguous(), scatter_idx, gather_idx, 0, spg))
single_all_to_all(grad_output[:, st:ed].contiguous(),
scatter_idx,
gather_idx,
0,
spg,
num_kv_heads=num_kv_heads))

del grad_output

Expand Down Expand Up @@ -440,8 +455,8 @@ def backward(ctx, grad_output):
else:
dk[i] = dk[i].to(dtype)
dv[i] = dv[i].to(dtype)
dk[i] = single_all_to_all(dk[i].contiguous(), gather_idx, scatter_idx, 0, spg)
dv[i] = single_all_to_all(dv[i].contiguous(), gather_idx, scatter_idx, 0, spg)
dk[i] = single_all_to_all(dk[i].contiguous(), gather_idx, scatter_idx, 0, spg, num_kv_heads=num_kv_heads)
dv[i] = single_all_to_all(dv[i].contiguous(), gather_idx, scatter_idx, 0, spg, num_kv_heads=num_kv_heads)

input_st = i * input_chunk_size
input_ed = input_st + input_chunk_size
Expand Down Expand Up @@ -474,7 +489,12 @@ def backward(ctx, grad_output):
ctx.pos_emb_sin[:, dq_seq_len * i:dq_seq_len * (i + 1)])
else:
dq[i] = dq[i].to(dtype)
dq[i] = single_all_to_all(dq[i].to(dtype).contiguous(), gather_idx, scatter_idx, 0, spg)
dq[i] = single_all_to_all(dq[i].to(dtype).contiguous(),
gather_idx,
scatter_idx,
0,
spg,
num_kv_heads=num_kv_heads)

input_chunk = layernorm_output[:input_chunk_size].reshape(-1, layernorm_output.shape[-1])
layernorm_output = layernorm_output[input_chunk_size:]
Expand Down Expand Up @@ -565,6 +585,9 @@ def forward(ctx: Any,
cpu_offloading=True):

do_save = layernorm_output.requires_grad
# The all-to-all partitions by KV group, and only the caller knows that count: the
# gather-direction and backward calls below see an already-sharded head dim.
num_kv_heads = kv_projection_size // hidden_size_per_attention_head

if rotary_pos_emb is not None:
pos_emb_cos, pos_emb_sin = rotary_pos_emb[0].permute(1, 0, 2, 3), rotary_pos_emb[1].permute(1, 0, 2, 3)
Expand All @@ -590,6 +613,7 @@ def forward(ctx: Any,
ctx.dtype = layernorm_output.dtype
ctx.projection_size = projection_size
ctx.kv_projection_size = kv_projection_size
ctx.num_kv_heads = num_kv_heads

global_q = []
global_k = []
Expand Down Expand Up @@ -628,18 +652,18 @@ def forward(ctx: Any,
q_chunk = qkv_chunk[:, :, :projection_size].contiguous().reshape(
qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd
q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg)
q_chunk = single_all_to_all(q_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)
global_q_chunk_len = q_chunk.shape[1]

k_chunk = qkv_chunk[:, :, projection_size:projection_size + kv_projection_size].contiguous().reshape(
qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd
k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg)
k_chunk = single_all_to_all(k_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)

v_chunk = qkv_chunk[:, :, projection_size + kv_projection_size:].contiguous().reshape(
qkv_chunk.shape[0], qkv_chunk.shape[1], -1,
hidden_size_per_attention_head).permute(1, 0, 2, 3).contiguous() # b, l, nh, hd
v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg)
v_chunk = single_all_to_all(v_chunk, scatter_idx, gather_idx, 0, spg, num_kv_heads=num_kv_heads)

dist.barrier()

Expand Down Expand Up @@ -739,8 +763,12 @@ def forward(ctx: Any,
global_q[q_compute_chunk_idx].offload()
q_compute_chunk_idx += 1

all2all_output = single_all_to_all(
cur_attn_output.to(ctx.dtype).contiguous(), gather_idx, scatter_idx, 0, spg)
all2all_output = single_all_to_all(cur_attn_output.to(ctx.dtype).contiguous(),
gather_idx,
scatter_idx,
0,
spg,
num_kv_heads=num_kv_heads)
final_output.append(all2all_output)
with get_accelerator().stream(general_offload_stream):
global_o.append(SequenceChunk(cur_attn_output.to(ctx.dtype)))
Expand Down Expand Up @@ -783,6 +811,7 @@ def backward(ctx, grad_output):

projection_size = ctx.projection_size
kv_projection_size = ctx.kv_projection_size
num_kv_heads = ctx.num_kv_heads

layernorm_output = ctx.layernorm_output

Expand Down Expand Up @@ -821,8 +850,12 @@ def backward(ctx, grad_output):
device=qkv_linear_weight.device,
dtype=torch.float)

grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(), scatter_idx,
gather_idx, 0, spg)
grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(),
scatter_idx,
gather_idx,
0,
spg,
num_kv_heads=num_kv_heads)
get_accelerator().synchronize()
grad_output = grad_output[:, chunk_size:]

Expand Down Expand Up @@ -931,7 +964,11 @@ def backward(ctx, grad_output):

if grad_global_attn_output[next_q_compute_chunk_idx] is None:
grad_global_attn_output_chunk = single_all_to_all(grad_output[:, :chunk_size].contiguous(),
scatter_idx, gather_idx, 0, spg)
scatter_idx,
gather_idx,
0,
spg,
num_kv_heads=num_kv_heads)
dist.barrier()
grad_output = grad_output[:, chunk_size:]
grad_global_attn_output[next_q_compute_chunk_idx] = SequenceChunk(
Expand Down Expand Up @@ -976,9 +1013,24 @@ def backward(ctx, grad_output):
dk_accum = dk_accum.to(dtype)
dv_accum = dv_accum.to(dtype)

dq_accum = single_all_to_all(dq_accum.contiguous(), gather_idx, scatter_idx, 0, spg)
dk_accum = single_all_to_all(dk_accum.contiguous(), gather_idx, scatter_idx, 0, spg)
dv_accum = single_all_to_all(dv_accum.contiguous(), gather_idx, scatter_idx, 0, spg)
dq_accum = single_all_to_all(dq_accum.contiguous(),
gather_idx,
scatter_idx,
0,
spg,
num_kv_heads=num_kv_heads)
dk_accum = single_all_to_all(dk_accum.contiguous(),
gather_idx,
scatter_idx,
0,
spg,
num_kv_heads=num_kv_heads)
dv_accum = single_all_to_all(dv_accum.contiguous(),
gather_idx,
scatter_idx,
0,
spg,
num_kv_heads=num_kv_heads)

general_offload_stream.synchronize()
compute_stream.wait_stream(general_offload_stream)
Expand Down
Loading
Loading