sequence parallel with communication overlap

[inkcherry]

SP is a fantastic piece of work, it is very elegant and concise， at the current stage, a transformer layer's forward and backward passes involve 8 all-to-all operations, with 5 opportunities for overlapping communication:

Forward pass: The QKV matrix operations can be pipelined alongside some of the all-to-all communications.
Backward pass: DQ, DK, DV all-to-all communications can be pipelined alongside matrix operations.
Backward pass: DO_w can be parallel with DO_input, involving matrix operations and all-to-all communications. Similar overlap-comm strategies are used in Megatron for TP/TP-sp parallelism.
I tested under conditions of 1N8C zero1, disabled activation checkpointing, ds-sp=8, and gbs=16:
1B 64K
7B 16K
They showed over 10% improvement (where I found that for mega-ds, using split QKV itself can also enhance performance due to reducing slice + cat operations in fwd/bwd), despite some TFLOPs already performing at a relatively good level.
co-work with microsoft/Megatron-DeepSpeed#415

[Edenzzzz]

overlapping only happens when computation doesn't depend on communication?

[inkcherry]

overlapping only happens when computation doesn't depend on communication?

@Edenzzzz Yes, manual sync of some dependencies is required。

[inkcherry]

we set gbs=2 ,sp=4, seq_len=16K,model size =1B, zero_stage=1, disable=activation_checkpoint, use-flash-attn-v2.

• without this patch
• with this patch, enable splitqkv+sp-overlap-comm
• with this path, disable splitqkv+sp-overlap-comm

we list the loss curve & grad norm curve and they are consistent.

[samadejacobs]

@inkcherry, many thanks for this excellent contribution to DeepSpeed codebase. To help with our review, could you please add (1) unit test(s) and (2) numbers on parallel performance improvements (throughput and latency) to the pull request? Your continuous and remarkable contributions to DeepSpeed are appreciated.

[Edenzzzz]

@inkcherry Thanks for your insight! Can I ask why we need sp_stream here, as it seems to be never used, e.g. by torch.cuda.stream(sp_stream)?

[inkcherry]

@inkcherry Thanks for your insight! Can I ask why we need sp_stream here, as it seems to be never used, e.g. by torch.cuda.stream(sp_stream)?

hi @Edenzzzz
apology for missing your comments, I noticed that DeepSpeed's sequence parallel is designed in a modular way, which means we can't freely insert communication calls (comm) anywhere we want to use async_op. When pytorch computation kernel launched before communication one, the communication one will automatically sync with the default stream, we need to use a custom stream or even an event to achieve parallelism between computation and communication. It's also crucial to maintain the dependencies between them properly. The stream setup for this is in Megatron-DeepSpeed implementation.
Here are two of the three cases mentioned in this PR that fall into this category, using an additional stream. The other case uses async_op=True with all2all.