百万级超长序列大模型训练如何加速，硬核解读 MindSpeed 方案

2024-08-14
广东
本文字数：2998 字
阅读完需：约 10 分钟

摘要：针对现有长序列训练场景的痛点，MindSpeed 在并行算法、计算效率、内存占用以及通信四个维度系统性优化大模型长序列训练效率，支持大模型百万级长序列训练。

1 长序列已经成为主流大模型能力之一

23 年底 Gemini1.5 Pro 发布以来，大模型序列长度迅速增长，处理超长序列上下文（Context）已经成为了目前主流闭源和开源大模型的重要能力之一，在模型规模不变的情况下，增大序列长度可以显著提升模型效果。同时，多模态大模型由于需要处理视频、语音等超长时序信号，模型的上下文窗口需要持续增大。为了保持大模型在超长序列场景的性能，在预训练以及微调过程中，通常会加入部分长序列输入对模型进行训练。目前，主流大模型均已支持 128K 及以上的序列长度输入。

2 长序列分布式训练面临的挑战

在长序列训练场景，大模型的训练开销（内存和计算量）会以非线性的速度增长，为长序列训练带来巨大挑战。为了加速大模型长序列训练效率，业界提出了多种上下文并行（Context Parallelism，CP）加速方案，对模型的输入上下文窗口进行切分，通过分布式并行技术将计算和内存切分到多个设备，降低每个设备上的负载。目前主流大模型加速库均提供了上下文并行方案以加速长序列训练效率，例如 DeepSpeed 的 Ulysses 和 Megatron 的 Ring Attention。

然而，在实际使用场景中，两种方案存在各自的局限性。

Ulysses 会对注意力头进行切分，在 4D 并行场景中，张量并行（TP）也会对注意力头进行切分，因此上下文并行（CP）的维度会受到注意力头数（）的限制，从而影响到上下文窗口的有效扩展。

Ring Attention 不涉及对注意力头的切分，因此理论上最大序列并行维度不受限，能够实现上下文窗口的无限拓展。然而，在实际应用中，P2P 通信需在序列长度足够大的情况下才能被计算过程所掩盖。此外，与集合通信相对，P2P 通信在利用通信带宽方面存在局限性。再者，分块注意力计算对于特定位置编码、注意力掩码（attention mask）和注意力算法泛化能力较弱，因此针对不同的应用场景可能需要对分块算法进行专门的调整。

3 MindSpeed 长序列训练加速方案

针对现有长序列训练场景的痛点，MindSpeed 在并行算法、计算效率、内存占用以及通信四个维度系统性优化大模型长序列训练效率，支持大模型百万级长序列训练。

3.1 并行算法--同时支持三种上下文并行算法/兼容专家并行

MindSpeed 同时支持了 Ulysses，Ring Attention 上下文并行方案，同时还支持混合上下文并行方案[3]。该方案对 Ulysses 和 Ring Attention 进行了融合嵌套，通过混合排布 AllToAll 通信和 P2P 通信，提升了跨节点场景上下文并行的通信效率，同时缓解了 Ulysses 方案最大并行维度受到注意力头数限制的问题。

此外，MindSpeed 上下文并行兼容专家并行。专家并行是稀疏专家模型（MoE）训练场景中常用的分布式并行加速方案。在 MoE 训练场景下，MindSpeed 可以同时开启上下文并行和专家并行加速万亿 MoE 超长序列训练场景。

3.2 计算优化—支持 FlashAttention

FlashAttention(FA)[4]通过 IO 感知的分块注意力算法，大幅提升了整体注意力计算效率。MindSpeed 长序列训练加速方案使用了昇腾亲和的 FlashAttention 算子对注意力计算进行加速，大幅提升注意力计算效率。

3.3 内存优化--ALiBi 和 Reset Attention Mask 压缩

在长序列训练场景，注意力掩码带来的内存开销巨大。例如当序列长度为 256K 时，一个完整的注意力掩码的大小是，存储成本过高。

ALiBi 和 Reset Attention Mask 是长序列训练场景两种常用的技术，MindSpeed 针对这两种场景提供了 Mask 压缩方案，实现内存优化，支撑昇腾上典型长序列训练场景。

3.3.1 ALiBi 位置编码压缩

为了更好的解决位置编码的外推能力，扩展上下文长度，ALiBi（Attention with Linear Biases）[5]通过向注意力分数的每一项添加一个线性偏置项来引入位置信息，见图 3。FA 算子通过两个入参可以控制 ALiBi 生成时位置的偏移，从而在内核通过分片的方式生成小块掩码，避免传入完整的 ALiBi 偏置矩阵。

图 1: ALiBi 通过偏置矩阵添加线性偏置示意图[5]

然而，Ring Attention 方案针对负载不均衡问题对序列划分进行了调整，在该切分方式下，q、k、v 是并不连续的（如图 4 中 cp0 的 0，7 块），无法直接通过内核分片方式生成小块掩码。

图 2: Ring Attention 负载均衡场景支持 ALiBi 编码压缩计算流程图

为了保持原有负载均衡逻辑，同时不传入完整 ALiBi 偏置矩阵降低内存开销，MindSpeed 采取了将不连续的块再次拆分的策略，见图 2。具体来说，如果当前分块满足连续条件（例如 cp3 的所有块），或者对于每一行的内部连续（如 cp0 第 2 块第 2 行），则可以直接使用内核生成策略；否则进行进一步拆分，直到满足连续条件，再使用内核生成策略进行生成。如图所示，每一个黑色框内的部分单独计算（白色被掩码掉的部分不参与运算）。

在使用 ALiBI 位置编码的场景下，相比未做上述优化，完全节省掩码矩阵的内存开销，即。

3.3.2 Reset Attention Mask 压缩

为了缓解长序列训练场景同一序列内文档数量过多，文档间互相干扰的问题，Llama 3.1 在预训练过程中使用了 Reset Attention Mask。根据 Llama 3.1 技术报告显示，使用 Reset Attention Mask 在长序列训练中十分重要[6]。Reset Attention Mask 在每个文档的末尾重置注意力，仅保留文档内部的注意力。图 3 给出了一个例子，文档 1 的长度为 2，文档 2 的长度为 4，文档 3 的长度为 2，所有跨文档的注意力分数均会被掩盖（图 3 中 0 表示掩盖）。

图 3: Reset Attention Mask 示意图

在该场景下，Ring Attenton 的切分会导致部分计算步骤的掩码不规则，无法利用 FA 的稀疏模式和掩码压缩。MindSpeed 通过保留每个 Token 所属的文档 ID，在每个 Ring Attention 步骤分片计算出当前步骤的分片掩码，无需保留完整掩码。在使用 Reset Attention Mask 场景下，该优化将注意力掩码的内存开销从

降低到，其中 P 为上下文并行的维度。