Sequence Parallelism 相关

初版

SP架构

初版符号表

环形计算，devices按行切分$QKV$矩阵，得到$Q^{n}, K^{n}, V^{n} \in \mathbb{R}^{L/N\times A}$，之后执行$N-1$轮环形通信，计算$S^{n} = \mathrm{concat}_{i=1}^{N}\bigl(Q^{n}(K^{i})^{T}\bigr)\in \mathrm{R}^{L/N\times L}$，之后再进行$N-1$轮环形通信，计算attention输出$O^{n} = S^{n}V = \mathrm{concat}_{i=1}^{N}\bigl(S^{n}_{i}V^{i}\bigr)$，其中$S^{n}_{i}$代表$S^{n}$按列切分为$N$段

性能比较

SP与TP性能比较1

SP与TP性能比较2

因此若$SP > TP$则：

• $BL > 32H$
• $BL > 16 AZ$

通信比较

二者相同，都是$8(N-1)BZ(L/N)A$，详细推导省略

对于TP来说的优势

在与PP结合的时候可以减少一次all-gather，这是因为TP会在进入下一阶段的时候对输入也做split，在计算结束之后在all-gather，但是由于SP在初始阶段就会做split，所以就可以失去这个操作

Ulysses

Ulysses架构

沿着head cnt做切分，这样一张卡上是若干个完整的头，可以直接计算Attention，并且可以用一些常规的方法加速，例如FlashAttention等
注意在上图中$d = hc\times hs$，是将多头的直接concat起来得到的向量

Ring-Attention

Ring Attention架构

Flash Attention的并行化版本，利用异步P2P通信，在一个设备上计算局部attention的同时传递K-V块