xDiT项目中SP+PP混合并行模式的最优通信方案设计

在分布式深度学习训练中，序列并行(SP)与管道并行(PP)的混合使用是提升大模型训练效率的重要手段。xDiT项目团队近期发现，在现有实现中SP+PP混合并行的通信模式存在优化空间，特别是在键值(KV)缓存的处理方面。

现有通信模式的问题

当前实现中，当采用SP+PP混合并行时，输入SP注意力模块的KV数据量仅为完整数据的1/sp_degree。例如当PP并行度为4、SP并行度为2时，输入数据量为完整KV的1/2，而理论上最优情况应该是1/(4*2)=1/8。这种非最优状态会导致额外的通信开销和计算资源浪费。

问题根源分析

造成这一现象的根本原因是：

1. 当前实现中，各PP并行rank之间无法共享KV缓存
2. 标准的SP注意力实现会丢弃中间计算结果
3. 为了弥补KV数据不足，现有方案会拼接pp_degree份KV碎片作为SP注意力输入

优化方案设计

项目团队提出了一种创新的解决方案：

1. KV缓存保留机制：修改yunchang库实现，保留SP注意力计算中的中间结果
2. 全量KV获取：在完成一个diffusion step后，每个设备将拥有该step的全部KV数据
3. 具体实现修改：
   • 在ulysses all2all操作后缓存中间结果
   • 在ring p2p通信后保存数据到KV缓冲区

技术优势

这种优化方案具有以下显著优势：

1. 通信效率提升：避免了不必要的KV数据重复传输
2. 计算资源节约：减少了冗余计算量
3. 内存使用优化：通过智能缓存策略平衡了内存占用
4. 扩展性增强：方案在不同并行度配置下都能保持高效

实现示意图

优化后的数据流可以描述为：

1. 初始阶段各设备只拥有部分KV
2. 通过SP内部计算和通信，逐步获取完整KV
3. 最终每个设备都拥有完整的KV缓存
4. 后续计算可以直接利用本地完整KV数据

预期效果

以PP=8、SP=2的配置为例：

• 原方案：每个SP注意力输入为1/2完整KV
• 优化后：每个SP注意力输入为1/16完整KV
• 通信量减少为原来的1/8

这种优化对于大规模模型训练尤为重要，可以显著降低通信开销，提升整体训练效率。xDiT项目的这一改进展示了混合并行模式下通信优化的重要性和可行性，为后续更大规模模型训练提供了宝贵经验。