下载
中文
注册

Hybrid长序列并行

背景与挑战

从生成性AI到科研模型,长序列训练正在变得非常重要。 在生成性AI领域,会话式AI、长文档摘要和视频生成等任务都需要在空间和时间层面对长上下文进行推理。 同样,章节和书籍级别的摘要(数万甚至数十万字)在会话式AI和摘要任务中也非常重要。现有的数据、张量和流水线等并行方法无法在序列维度进行切分。当序列维度(S)增长时,训练内存开销会以O(S2)的速度增长。因此需要针对长序列场景进行特定的优化解决长训练场景的训练需求。

目前流行的序列并行方案,Ulysses和Ring Attention存在各自的局限性。

Ulysses需要确保num-attention-heads数可以被序列并行维度整除,在GQA(Grouped-query Attention,分组查询注意力)、MQA(Multi-query Attention,多查询注意力)场景下序列并行的大小有限制,导致序列长度的扩展有限。

Ring Attention的并行维度不受num-attention-heads数限制,因此理论上序列长度可以无限拓展。但相比于Ulysses,Ring Attention不能充分利用通信和计算带宽,在序列块大小较低时性能劣于Ulysses。

解决方案

对Ulysses和Ring Attention做融合,实现Hybrid长序列并行,以此解决两个方案互相的缺陷。

使用场景

可兼容Flash Attention,目前已默认开启Flash Attention

序列并行维度被分为Ulysses维度和Ring Attention维度,Ulysses维度和Ring Attention维度乘积即为序列并行维度。

目前仅支持单向Causal Attention。

使用方法

请参见下表。
表1 方法说明

重要参数

参数说明

--context-parallel-size [int]

设置长序列并行大小,默认为1,根据用户需求配置。

--ulysses-degree-in-cp [int]

  • 需要确保--context-parallel-size可以被该参数整除且大于1。

    例如当设置--context-parallel-size为8时,可以设置--ulysses-degree-in-cp为2或--ulysses-degree-in-cp为4。

  • 同时需要确保--ulysses-degree-in-cp可以被num-attention-heads数整除。

--context-parallel-algo hybrid_cp_algo

可选,设置长序列并行算法。

  • hybrid_cp_algo:开启Hybrid长序列并行。
  • megatron_cp_algo:开启Ring Attention长序列并行。
  • ulysses_cp_algo:开启Ulysses长序列并行,缺省值。

使用效果

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