Ulysses长序列并行

问题分析

从生成式AI到科研模型，长序列训练正在变得非常重要。 在生成式AI领域，会话式AI、长文档摘要和视频生成等任务都需要在空间和时间层面对长上下文进行推理。 同样，章节和书籍级别的摘要（数万甚至数十万字）在会话式AI和摘要任务中也非常重要。现有的数据、张量和流水线等并行方法无法解决序列维度的扩展问题。

解决方案

支持 Ulysses长序列并行方案，以此解决序列维度扩展问题。

解决思路:

Ulysses 将各个样本在序列维度上分割给参与的计算设备。然后，在 attention 计算之前，它对已分割的查询(Q)、键(K)和值(V)执行 all-to-all 通信操作，以使每个计算设备接收完整的序列，但仅用于注意力头的非重叠子集。这使得参与的计算设备可以并行计算不同的注意力头。最后，Ulysses 还使用另一个 all-to-all 来在注意力头上收集结果，同时重新在序列维度上进行分区。

使用场景

num_head 要能被 tp_size*cp_size 整除。

使用方法

设置--context-parallel-size，默认为1，根据用户需求配置。 设置--context-parallel-algo，可选项为ulysses_cp_algo或者megatron_cp_algo，默认ulysses_cp_algo。

使用效果

利用多个计算设备对输入序列进行并行切分，降低单设备的内存消耗，相比不开启序列并行单步耗时增加，相比重计算计算效率提升。