题记：
2024/06/17 updates：本文之前题目是《告别DeepSpeed-Ulysses，来用Odysseus：重新审视大模型序列并行技术》，试图对Ulysses进行优化，来设计通信更优的方法。经过网友指正，我对Ulysses的通信开销估计有误，努力都是徒劳。改正后，我发现Ulysses真香。原文我也不删了，读者评论非常有意义。我对有问题的论断做了一些标注，供大家参考。
前言：
众所周知，笔者最近在研究大模型的序列并行技术。我在知乎比较过DeepSpeed-Ulysses和Ring-Attention，提出了一种融合二者的混合序列并行USP（Unified Sequence Parallelism），讨论和4D并行的最佳实践和一些实际应用前景。
USP以其易于实现和性能只增不减无副作用的优势，被很多人开始使用。大家也都发现，DeepSpeed-Ulysses在性能上往往超越Ring-Attention，因此，实践者们倾向于将Ulysses的并行度拉满，不足时再用Ring-Attention来弥补。
背景：序列并行的三家分晋到三分归晋
补充一下背景，现在市面上有三种序列并行，Ulysses、Ring和TP-SP。三者由不同机构独立提出，开始时彼此井水不犯河水，后来三者可以组成序列+张量的混合并行，归于一统。混合效果优于单独使用任意一种并行，参考我们USP论文里展示了实验结果。不过，大一统王朝背后暗流涌动，Ulysses和TP-SP略有冲突，因为二者并行度都受限于head num维度的大小。
USP: A Unified Sequence Parallelism Approach for Long Context Generative AI
之所以三者混合有收益，因为它们在内存和通信上的表现都很近似，通过混合能够构成更好的通信模式，从而适应异构互联的网络。比如下表中TP-SP，Ulysses和Ring三行。值得一提的是，Ring用P2P而非集合通信，虽然通信量低，但是它的通信拆成小块，和Attention计算重叠，计算通信的性能其实都受损，实际端到端表现往往逊于Ulysses和TP-SP。
Ulysses和TP-SP：看似相同，实则不同
在序列并行三家中，Ring的价值稳如泰山，因为它的并行度不受限于head维度大小，可以无限扩展。尽管现在开源实现性能一般，X+Ring模式仍是超大规模序列并行的刚需。而且，Ring的工程优化方式也在不断进化。
反观Ulysses和TP-SP两者，地位就有些冲突，用其中一个未尝不可舍弃另一个。更尴尬的是，Ulysses和TP-SP的算法流程其实也非常相似。它们的SDPA（Scaled Dot-Product Attention）部分计算一模一样，都是沿着head维度切分。区别在于，Ulysses将通信部分龟缩于SDPA前后，而TP-SP则主动出击在QKV三个Linear前和O Linear之后通信。不过因为Ulysses对1/N input进行All2All，通信量小于对1/N input做AllGather和ReduceScatter。
Ulysses出身名门望族，背靠3.29万亿市值巨无霸微软，可谓含着金汤匙出生，名字里带着高贵姓氏DeepSpeed，行走江湖时，更有大哥ZeRO保驾护航，一时也是风生水起。
TP-SP背景也不简单，作为Megatron-LM的核心技术，爸爸NVIDIA乃当今美股小霸王，3.24万亿市值丝毫不怵微软。TP-SP只是我们外人称呼，在Megatron-LM大家族中，人家就叫Sequence Parallel。为捧亲儿子，Megatron-LM还收了Ring-Attention做养子，改名叫Context Parallel，来辅佐TP-SP闯荡江湖。
就目前形势来说，DeepSpeed捧Ulysses+ZeRO，Megatron-LM站Ring+TP-SP，后者的工程水准和品牌效应明显更强一些。但是Megatron-LM没有ZeRO3，没法充分发挥Ring威力，有些可惜。
一些国内的框架，如华为MindSpeed，已经采用USP方式融合三者。
docs/features/hybrid-context-parallel.md · Ascend/MindSpeed - Gitee.com
2024/6/17 updates:后面的都不用看了，我之前分析Ulysses Attention通信量算错了，All2All的算法带宽是AG和1/N。所以，Ulysses是通信量最低的方案。大家序列并行以后使用USP就可以了。
赵轩磊，Lin Zhang，镜影，草木等读者都发现问题了，感谢他们的指正！
这里有一些代码还是有些用处的，大家可以用来实验不同序列并行的性能。
https://github.com/feifeibear/Odysseus-Transformer
Odysseus：和Ulysses说拜拜
为何不直接用TP-SP的Attention部分替换Ulysses的Attention部分？上表中，Ulysses和TP-SP对Activation整体通信量近似，Ulysses也还略有优势。但是，Ulysses的对Activation通信全发生在Attention内，而TP-SP则是Attention和MLP都有通信。单论Attention部分，Ulysses则远逊于TP-SP了（有误：实际上是好于)。上图Ulysses通信4xAll2All，而TP-SP只需AG+RS，通信需求直接减半。要知道，TP-SP的ReduceScatter之后，张量也是可以验证序列维度切分，这和就Ulysses的输出一模一样了。
这里提出一种名为Odysseus的序列并行方法，它的通信性质优于Ulysses（有误：实际上是弱于）。奥德修斯（Odysseus）是尤利西斯（Ulysses）的希腊名字，也寓意对序列并行技术重新考古，返璞归真。如下图所示，Odysseus对Attention部分使用TP-SP，而MLP部分使用ZeRO。它的通信和内存需求在上表面的最后一行，可见在Activation和参数/梯度的通信量都小于Ulysses，或者说它就是一个在各方面都优于Ulysses的方法。
它和Ulysses区别在于Attention部分不同，如上所述，Odysseus的Attention部分通信量更少，因此它使用场景可以覆盖Ulysses。（实际上只能覆盖一部分TP-SP使用场景，无法代替Ulysses）
它和TP-SP区别在于MLP部分不同。TP-SP还需要对MLP进行AllGather+ReduceScatter。在长序列场景下，即seqlen (L) > intermediate_size (i)，Activation通信量更大，因此Odysseus比TP-SP通信量更小。但是，短序列Odysseus就不如TP-SP。
Odysseus仓库中实现了Udysseus，TP-SP，Ulysses，Ring-Attention，并在相同任务进行比较。
实验结果如下，其实Odysseus-ZeRO3MLP是我们的最终Odysseus实现，在单机八卡上内存和速度表现远好于Ulysses，和TP-SP相比略有优势。注意，下表中64K序列长度时，Ring和Ulysses使用ZeRO3还是会OOM。
关于Odysseus代码，有一些门道和大家汇报一下：
我没用现成的大模型框架，比如Megatron-LM或者DeepSpeed，用PyTorch实现基础接口实现本项目。因为我希望控制训练流程的每个细节，保证消融分析的准确性。Megatron-LM已经是非常庞大的项目，很多未知因素都会对结果产生影响。
本项目非常灵活，更加适合测试各种创新点子，可以作为研究序列并行的试验场。把本项目中的Attention和MLP集成到其他大模型框架中也很容易。
TP-SP的实现参考了Megatron-LM的linear_with_grad_accumulation_and_async_allreduce代码。这里有一个技巧很经典，对应NVIDIA论文如下描述。（简单说，TP-SP为了避免input的AllGather结果在forward过程中一直占显存，实现了用完即释放，反向需要时再重新AllGather，这个AllGather可以和计算重叠）。通过这个项目你可以更加深入理解我之前的文章，方佳瑞：大模型训练优化：论反向传播中序列并行和张量并行的通信计算重叠方法https://zhuanlan.zhihu.com/p/700639611
From Equation 3, sequence parallelism along with tensor parallelism divides all the activations required for the backward pass along the parallel dimension except for the tensor Y that is required for the first linear operation. To alleviate this issue,we do not store the full tensor Y for the backward pass. Instead, we store only the Y s i part on the ith tensor parallel rank and perform an extra all-gather in the backward pass.To eliminate the latency introduced by this extra all-gather, we overlap this communication with the computation required to calculate gradients with respect to Y , and as a result, we reduce the overhead.
初始版本我用PyTorch FSDP来做MLP的而是ZeRO3（Odysseus性能表格中的Odysseus-FSDP），发现Odysseus的性能不如TP-SP。通过观察profiler的timeline发现是FSDP的ReduceScatter拖慢的。于是，我手动实现了ZeRO3版本的Linear，不过这个ZeRO3就是同步通信，没有和计算重叠，性能一下子变得符合预期。我觉得可能是TP-SP的linear实现用到了异步通信，和FSDP的异步通信冲突导致的。
The End
作者：方佳瑞
来源：GiantPandaCV
推荐阅读
- KV Cache优化: 层内和层间KV Cache共享
- 图解大模型计算加速系列：分离式推理架构1，从DistServe谈起
- 窥探Triton的lower(二)
- Llama也能做图像生成！港大字节推出开源自回归文生图模型，在线体验已开放
欢迎大家点赞留言，更多Arm技术文章动态请关注极术社区嵌入式AI专栏欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
