Add the missing view operations from sequence parallel(async). by inkcherry · Pull Request #6750 · deepspeedai/DeepSpeed

inkcherry · 2024-11-14T13:57:34Z

a view operation was missing in some updates compared to the original version https://github.com/microsoft/DeepSpeed/blob/17ed7c77c58611a923a6c8d2a3d21d359cd046e8/deepspeed/sequence/layer.py#L56

add missing view operation.
The shape required for the view cannot be easily obtained in the current function， so refactor layout params code.

inkcherry · 2024-12-02T05:12:14Z

@loadams The CI encountered a "no space left" issue, which doesn't seem to be caused by this patch. Could you please retrigger it? thanks！

loadams · 2024-12-05T17:15:01Z

@loadams The CI encountered a "no space left" issue, which doesn't seem to be caused by this patch. Could you please retrigger it? thanks！

Hi @inkcherry - yes, sorry that does seem to be an intermittent failure, I'll re-trigger them, but it looks like there are merged conflicts now, could you take a look?

inkcherry · 2024-12-06T03:29:04Z

@loadams Thank you for the reminder, I have resolved and verified the merge conflicts(both ds+megads) : )

Copilot

Copilot reviewed 1 out of 1 changed files in this pull request and generated no comments.

Comments suppressed due to low confidence (1)

deepspeed/sequence/layer.py:67

The docstring for the post_all2all function should use the word 'view' instead of 'reshape' to match the actual operation being performed.

output = input.reshape(res_shape).contiguous()

@loadams

…peedai#6750) FYI @loadams a view operation was missing in some updates compared to the original version https://github.com/microsoft/DeepSpeed/blob/17ed7c77c58611a923a6c8d2a3d21d359cd046e8/deepspeed/sequence/layer.py#L56 add missing view operation. The shape required for the view cannot be easily obtained in the current function， so refactor layout params code. --------- Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Signed-off-by: siqi <siqi@tecorigin.com>

@loadams

…peedai#6750) FYI @loadams a view operation was missing in some updates compared to the original version https://github.com/microsoft/DeepSpeed/blob/17ed7c77c58611a923a6c8d2a3d21d359cd046e8/deepspeed/sequence/layer.py#L56 add missing view operation. The shape required for the view cannot be easily obtained in the current function， so refactor layout params code. --------- Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>

@loadams

…peedai#6750) FYI @loadams a view operation was missing in some updates compared to the original version https://github.com/microsoft/DeepSpeed/blob/17ed7c77c58611a923a6c8d2a3d21d359cd046e8/deepspeed/sequence/layer.py#L56 add missing view operation. The shape required for the view cannot be easily obtained in the current function， so refactor layout params code. --------- Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Signed-off-by: gyou2021 <ganmei.you@intel.com>

@loadams

…peedai#6750) FYI @loadams a view operation was missing in some updates compared to the original version https://github.com/microsoft/DeepSpeed/blob/17ed7c77c58611a923a6c8d2a3d21d359cd046e8/deepspeed/sequence/layer.py#L56 add missing view operation. The shape required for the view cannot be easily obtained in the current function， so refactor layout params code. --------- Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Signed-off-by: gyou2021 <ganmei.you@intel.com>

@loadams

…peedai#6750) FYI @loadams a view operation was missing in some updates compared to the original version https://github.com/microsoft/DeepSpeed/blob/17ed7c77c58611a923a6c8d2a3d21d359cd046e8/deepspeed/sequence/layer.py#L56 add missing view operation. The shape required for the view cannot be easily obtained in the current function， so refactor layout params code. --------- Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Signed-off-by: gyou2021 <ganmei.you@intel.com>

@loadams

…peedai#6750) FYI @loadams a view operation was missing in some updates compared to the original version https://github.com/microsoft/DeepSpeed/blob/17ed7c77c58611a923a6c8d2a3d21d359cd046e8/deepspeed/sequence/layer.py#L56 add missing view operation. The shape required for the view cannot be easily obtained in the current function， so refactor layout params code. --------- Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com> Signed-off-by: yisheng <yi.sheng@intel.com>

@loadams

…peedai#6750) FYI @loadams a view operation was missing in some updates compared to the original version https://github.com/microsoft/DeepSpeed/blob/17ed7c77c58611a923a6c8d2a3d21d359cd046e8/deepspeed/sequence/layer.py#L56 add missing view operation. The shape required for the view cannot be easily obtained in the current function， so refactor layout params code. --------- Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>

_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 deepspeedai#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 <vineethsai4444@gmail.com>

inkcherry added 2 commits November 14, 2024 13:33

Add the missing view operations and refactor the code.

87679c4

rename function

6ce27a9

inkcherry requested review from loadams and tjruwase as code owners November 14, 2024 13:57

inkcherry mentioned this pull request Nov 14, 2024

support split qkv linear and sp overlap comm deepspeedai/Megatron-DeepSpeed#415

Open

loadams and others added 4 commits November 14, 2024 13:03

Merge branch 'master' into ds_overlap_fix

a947ee4

Merge branch 'master' into ds_overlap_fix

572edcb

Merge branch 'master' into ds_overlap_fix

df36800

Merge branch 'master' into ds_overlap_fix

4368928

Merge remote-tracking branch 'upstream/master' into ds_overlap_fix

a6b64c4

inkcherry and others added 4 commits December 8, 2024 16:39

Merge branch 'master' into ds_overlap_fix

147ad8d

Merge branch 'master' into ds_overlap_fix

6369b0b

Merge branch 'master' into ds_overlap_fix

47f573b

Merge branch 'master' into ds_overlap_fix

61ea87f

loadams requested a review from tohtana as a code owner December 18, 2024 18:59

loadams and others added 4 commits December 19, 2024 09:39

Merge branch 'master' into ds_overlap_fix

9c131fe

Merge branch 'master' into ds_overlap_fix

40d0ada

Merge branch 'master' into ds_overlap_fix

44c3ec0

Merge branch 'master' into ds_overlap_fix

c388565

loadams requested a review from Copilot January 13, 2025 19:04

Copilot AI reviewed Jan 13, 2025

View reviewed changes

loadams self-assigned this Jan 13, 2025

tohtana approved these changes Jan 17, 2025

View reviewed changes

Merge branch 'master' into ds_overlap_fix

9486eaf

loadams added this pull request to the merge queue Jan 21, 2025

Merged via the queue into deepspeedai:master with commit bc76b04 Jan 21, 2025

vineethsaivs mentioned this pull request Aug 25, 2026

Fix the seq-first Ulysses all2all output layout #8317

Open

Navigation Menu

Uh oh!

Add the missing view operations from sequence parallel(async). - #6750

Add the missing view operations from sequence parallel(async).#6750
loadams merged 16 commits into
deepspeedai:masterfrom
inkcherry:ds_overlap_fix

inkcherry commented Nov 14, 2024 •

edited

Loading

Uh oh!

inkcherry commented Dec 2, 2024

Uh oh!

loadams commented Dec 5, 2024

Uh oh!

inkcherry commented Dec 6, 2024

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

inkcherry commented Nov 14, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

inkcherry commented Dec 2, 2024

Uh oh!

loadams commented Dec 5, 2024

Uh oh!

inkcherry commented Dec 6, 2024

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

inkcherry commented Nov 14, 2024 •

edited

Loading