长序列并行

问题分析

随着AI技术的飞速发展，特别是在生成式AI和科研模型领域，长序列训练的重要性日益凸显。无论是在会话式AI、长文档摘要生成，还是视频内容生成等复杂场景中，系统都需要在广阔的空间和时间维度内对长上下文进行有效推理。此外，章节和书籍级别的庞大文本摘要（涉及数万乃至数十万字的处理）同样对会话式AI和摘要任务提出了严峻的挑战。传统的数据、张量和流水线等并行处理方法在处理长序列时面临显著限制。这些方法往往无法有效扩展以适应序列维度的增加，从而影响了整体系统的性能和效率。

具体来说，传统的并行方法在处理长序列时可能遭遇以下问题：

Ulysses长序列并行

解决方案

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

解决思路

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

使用场景

num-attention-heads要能够被 tensor-model-parallel-size*context-parallel-size 整除。

使用方法

使用效果

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

Ring Attention长序列并行

解决方案

支持Ring Attention长序列并行方案，以此解决序列维度扩展问题。详细说明可参见文献Ring Attention with Blockwise Transformers for Near-Infinite Context。

在原始Ring Attention基础上设计了新的计算块切分方案，解决负载不均衡问题。

解决思路

Ring Attention借鉴了分块Softmax原理，在不需要获取整个序列的完整矩阵情况下进行分块attention计算。因此作者提出以分块方式执行自注意力和前馈网络计算，跨多个设备分布序列维度。具体地，该方法在进程之间构建注意力计算块的环状通信结构（Ring），每个进程具有一个切分后的本地QKV块。在计算完本地的attention后，通过向后发送和向前获取KV块，遍历进程设备环，以逐块的方式进行注意力和前馈网络计算。同时，本地的attention计算和KV块的通信理想情况下可以互相掩盖，从而消除了额外引入的通信开销。另外该方案在计算attention的过程中全程不需要数据拼接，支持的序列长度理论上可以无限拓展。

使用场景

当使用GPT类模型进行训练，同时数据进MoE层时实际序列长度8K以上。

不同于Ulysses方案，该方案不需要确保head_size被cp_size整除。

可兼容FlashAttention，目前已默认开启FlashAttention。

使用方法

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

重要参数

参数说明

--context-parallel-size [int]

开启CP对应的数量，默认为1，根据用户需求配置。

--seq-length [int]

输入序列的长度。

--use-cp-send-recv-overlap

建议开启，开启后支持send receive overlap功能。

--cp-attention-mask-type [full/causal]

可选，设置Mask计算类型。

  • causal：设置为causal（倒三角）Mask计算，缺省值。
  • full：设置为全量计算。

--context-parallel-algo megatron_cp_algo

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

  • megatron_cp_algo：开启Ring Attention长序列并行。
  • ulysses_cp_algo：开启Ulysses长序列并行，缺省值。

使用效果

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

注意事项

父主题： 并行算法优化