InternEvo：通过混合并行和冗余分片实现高效长序列大语言模型训练

24年1月来自上海AI实验室等的论文“InternEvo: Efficient Long-Sequence Large Language Model Training via Hybrid Parallelism and Redundant Sharding”。

具有长序列的大语言模型 (LLM) 为越来越多的新应用程序提供动力。现有的长序列 LLM 训练方法既不高效，也不与常用的训练算法（如 Flash Attention）兼容。InternEvo 则是解决这些问题。InternEvo 将所有分片维度解耦为一个新的分层空间，并系统地分析 LLM 训练的内存和通信成本。然后，它生成一种有效的混合并行策略。一种选择性重叠机制，可以减轻混合并行引入的通信开销。采取内存管理技术来减少 GPU 内存碎片。评估结果表明，InternEvo 生成的并行化策略在模型 FLOPS 利用率方面与现有方法相当或优于现有方法。

背景

训练具有长序列的 LLM 需要大量的内存资源。为了适应 GPU 的内存限制，直接应用 3D 并行性 [21, 22] 或现有的自动并行化框架 [9, 23–26] 会降低训练速度 [27]，因为这些系统要求整个序列保存在一个 GPU 中，这限制了输入序列的长度。已经有几种方法来训练具有长序列的 LLM，同时又不会显著牺牲训练速度 [28–31]。然而，它们存在两个局限性。首先，可用性不令人满意。这些并行化策略和算法 [29–31] 在计算自注意方面和FlashAttention [32] 冲突。FlashAttention 以其速度和内存效率而闻名，被广泛应用于许多 LLM 训练 [33, 34]。其次，具有良好可用性的解决方案在训练性能方面仍然不够高效。先前的方法（例如 DeepSpeed Ulysses）引入了大量的节点间通信开销 [28]；其他方法（例如 Megatron-LM）则大大增加了通信需求 [27]。此外，这些解决方案采用了严格的并行化策略。

为了量化这一点，如图展示了不同系统在最佳性能下模型状态和激活的实际内存使用量以及相应的通信成本。在128个GPU，训练7B的模型。为了缓解序列长度为 32K 的内存压力，DeepSpeed 和 Megatron-LM 都只是在同一级对模型状态和激活进行分区（图 （a））。不幸的是，这种简单性会导致相应的高通信开销（图（b））。它们没有考虑使用各种配置（例如模型大小、批次大小、序列长度）进行训练时内存需求和通信模式的差异，使得它们在不同情况下并不总是最佳的。

并行技术包括：数据（TP）、张量（TP）、流水线（PP）和序列（SP）。

处理长序列对于支持重要应用（例如处理聊天应用程序中的较长历史记录）至关重要。为此，序列并行（SP）已成为一种旨在减轻 Transformer 训练期间激活内存占用的技术。在 SP 中，每个 Transformer 层的输入张量沿序列维度划分，从而允许跨多个 GPU 并行计算。这种分割与激活重计算相结合，可将激活内存需求大幅减少 ssp 倍。

现有的 SP 方法分为 3 个阶段，分别对应于将序列切分为 Norm 和 Dropout 模块、线性模块和 MHA 模块。SP的三个阶段如下依次介绍。
如图是长序列LLM训练的背景示意图，图（a）是Transformer的层架构，具有一个带有 D 个注意头和一个 MLP 层的注意块。输入/输出维度为 B × S × H（B：微批次大小，S：序列长度，H：隐维度）。图中展示以 16 位浮点格式存储每个元素的激活 (ACT) 所需的内存。FlashAttention [32] 的融合核用于高效的多头注意 (MHA) 计算。

SP-1：Norm 和 Dropout 模块。如图 (b) 所示，Megatron-LM 利用 张量并行（TP） 来并行化线性层和 MHA，这是训练过程中最耗时的组件。同时，它在 Norm 和 Dropout 模块上使用 SP，有效地将这些层的激活内存减少了 ssp 倍，其中 ssp =stp。为了保持计算结果的一致性，它集成了必要的通信，包括 all-gather 和 reduce-scatter，在前向和后向传递中传输激活。当激活大小随着序列长度的增加而增加时，这种通信激活的方式将产生较高的开销。

SP-2：添加线性模块。DeepSpeed Ulysses 在线性、Norm 和 Dropout 层上利用了序列并行（SP），如图 © 所示。对形状为 B × 3 × S/ssp × H 的 Linear/qkv 的输出应用all-to-all通信。这使每个 GPU 能够接收注意头的一个子集中 Q、K 和 V 的完整序列（每个序列的形状为 B×S×H/ssp）。随后，用 FlashAttention 进行高效的 MHA 计算，但限制注意头的数量 D 不应超过 ssp。然后用另一个all-to-all通信来收集 MHA 层结果，并沿序列维重新划分结果。虽然与 Megatron-LM 相比，DeepSpeed Ulysses 将激活内存使用量和每个 GPU 的通信大小减少了 ssp 倍，但它需要每个 GPU 上所有 Linear 层的完整权重矩阵。因此，DeepSpeed Ulysses 应用 ZeRO 等其他技术来减少模型状态的内存开销。

SP-3：添加 MHA 模块。现有系统（如 [29–31]）在包括 MHA 在内的所有模块上都应用序列并行（SP）。在 MHA 中，这些方法沿序列维划分 Q、K 和 V 与一系列对等通信操作相结合，以促进协作注意计算。这允许比注意头数量更高的并行维度（ssp > D），并能够训练具有相当长序列的模型，最多可达 100 万个 tokens。然而，MHA 中的划分与 FlashAttention 等内存高效机制不兼容，后者可以将内存复杂度从 O(S^2) 降低到 O(S)。

搜索空间分析

InternEvo 旨在高效处理模型状态和激活的内存需求，同时减少通信开销并提高分布式设置中的训练性能。

基于三个准则：

1. 准则 1：将模型状态和激活的内存管理策略解耦。
2. 准则 2：将模型状态的内存管理解耦为 P、G 和 OS。
3. 准则 3：通过部分冗余的交易来优化通信。

基于前两个准则，将 GPU 内存管理维度分为两大类：模型状态和激活，如图所示。模型状态内存由流水线并行 (PP) 策略和三种分片策略控制：参数分片 (PS)、梯度分片 (GS) 和优化器状态分片 (OSS)。另一方面，激活内存由数据并行 (DP)、张量并行 (TP) 和序列并行（包括 SP-1 和 SP-2）管理。

总体而言，InternEvo 定义4 个并行维度和 3 个分片维度，如图中的蓝色和橙色块所示，表示为 S = [b, s, a, spp, sdp, stp, ssp, sps, sgs, soss]。根据准则 3，不同的 si 值对通信成本有显著影响。通过优化 si 实现部分冗余存储，可以减少某些节点 i 的昂贵节点间通信。这种策略调整使得节点 i 的通信能够合并在带宽更高的同一节点内。

InternEvo 将激活存储在 Transformer 层边界并在后向传递中重新计算剩余的激活来优化内存使用。当 a = 1 时，InternEvo 需要存储 2bSH 字节用于 Transformer 层的激活，适用于没有流水线并行（PP）或序列并行（SP）的场景。相反，如果 a = 1，InternEvo 使用 34bSH 字节进行激活存储。此外，为了提高多头注意计算的效率，InternEvo 默认使用 FlashAttention 的融合核，该核已经结合选择性激活重计算技术，最大限度地减少注意计算期间的内存使用量 [33]。

InternEvo 使用多种并行策略在 GPU 之间高效分配 Transformer 层和模块。参数、梯度和优化器状态可以在 GPU 之间复制或分片。
如图所示，InternEvo 采用三种分片策略，影响内存占用和通信成本。分片因子 F 表示张量分片的 GPU 数量。对于具有 E 个元素和 R 个 GPU 的张量，三种策略可以概括如下：

• 完全复制 (F = 1)：具有 E 个元素张量的完整副本存储在R 个GPU的每一个上。
• 完全分片 (F = R)：每个 GPU 仅存储 E/R 个元素，可以在所有 G 个 GPU 上执行all-gather来重建完整张量。
• 混合分片 (1 < F < R)：此策略结合分片和复制。每个 GPU 都保存 E/F 个元素，每个元素都在 R/F GPU 上复制。可以通过在 F 个 GPU 上进行all-together来重建完整张量。

系统设计

系统架构和工作流

如图所示，InternEvo由两个关键组件组成：策略规划器和作业执行器。（1）策略规划器用于确定并行模型训练的最佳策略S。这涉及两个核心模块：首先是作业分析器，它对提交的工作负载进行细致的分析以提取关键的执行信息；其次是执行模拟器，它负责模拟S内的各种策略以找出最有效的解决方案。（2）作业执行器负责在GPU节点上实际执行具有给定训练策略S的训练作业。该组件集成了两个基本模块来增强训练性能。通信调度器通过其在细粒度上重叠计算和通信的能力，显着减轻通信开销。同时，碎片内存管理器确保统一的GPU内存管理，有效减少由于GPU内存碎片而导致的内存不足（OOM）错误发生，特别是在长序列训练场景中。

工作流程。 1. InternEvo 首先让用户指定 Transformer 模型架构（如注意头数和层数）、超参（如序列长度和全局批次大小）以及训练集群设置（如总 GPU 数量和内存容量）。 2. 作业分析器分析使用的集体通信操作有效带宽、模型内模块的计算时间以及识别通信-计算重叠的实例来进行全面分析。获得的分析结果驱动执行模拟器估算并行和分片策略的各种组合内存成本和步进时间。利用优化问题求解器，InternEvo 确定前 10 种策略，并采用实际系统分析来确定最佳策略。 3. 通过这种探索，可以确定具有最短步进时间的并行化策略，从而确保不会出现内存不足错误。 4. 然后，作业执行器利用相应的策略运行训练作业，训练速度快，内存利用率高。

作业分析器

作业分析器由三部分组成：（1）VPro 离线分析 NCCL 通信的有效带宽。（2）UPro 在线对 Transformer 层的单次前向传递进行详细分析。（3）OPro 离线分析计算和通信之间的重叠。

通信调度器

InternEvo 中的通信调度程序模块策略性地协调通信和计算，提高整体系统性能。

具体而言，该方法侧重于通过实施改进的层内策略来优化反向传递过程中的重叠。这涉及在模块粒度（例如线性模块）上利用参数预取（all-gather）和梯度同步（reduce-scatter）。在前向传递中，采用层间重叠策略来减轻延迟，方法是预取即将到来的层的整个参数集，同时计算当前层。当 sps > 1 时，这尤其有益。迭代选择性重叠方法巧妙地管理通信开销和计算执行时间，从而显著提高整体系统性能。

执行模拟器

执行模拟器使用给定的执行规划 S 估计训练性能，并通过解决整数规划问题获得最佳解决方案候选。

InternEvo 制定了一个优化问题，通过最小化前向-后向传递的执行时间与参数更新时间总和来寻找最佳执行规划，但要受到各种约束。
对于 S 中的每个 s，在设备网格上分配特定张量会引入多种可能的配置，从而影响通信延迟。利用实际经验来更有效地简化分配策略的探索，而不需要详尽地迭代 S 所有潜在 GPU 分配。

对于给定的执行规划 S，可以估计每个并行策略和分片方法的通信进程组分配。随后，能够使用估计的设备网格生成可执行文件，以便使用 UPro 进行精确分析。

评估实验设置

实现。InternEvo 是使用 Python 实现的，包含大约 35,000 行代码 (LOC)。代码库包括策略规划器的 5,000行代码（LOC） 和作业执行器的 30,000 行代码（LOC）。

测试平台。InternEvo 在使用 LLaMA 架构训练基于 Transformer 的模型方面进行了全面评估。这些模型包含的参数范围从 70 亿到 650 亿。训练过程发生在由 16 台 GPU 服务器组成的物理集群上。每台服务器配备 8 个 GPU 和 128 个 CPU 核心，总共 128 个 NVIDIA Ampere GPU。每个 GPU 都拥有 80GB 内存。GPU 通过 NVLink 和 NVSwitch 互连，而节点间通信则由四个 NVIDIA Mellanox 200Gbps HDR InfiniBand 促进。

基线和指标。将 InternEvo 与两个最先进长序列 LLM 训练系统进行基准测试，即 DeepSpeed Ulysses [28] 和 Megatron-LM [27]。评估指标是模型 FLOPS 利用率 (MFU) [47]。还在训练期间评估每 GPU 每秒的tokens数 (TGS)。虽然 InternEvo 会自动搜索并行化策略，但会手动设置 DeepSpeed Ulysses 和 Megatron-LM 的分区数。为了公平比较，这些基线选择的策略是最快的，可以避免触发内存不足错误。值得注意的是，在实践中，开发人员为 DeepSpeed Ulysses 和 Megatron-LM 选择的并行化策略可能与实验中选择的策略相同或更慢。

实例。如图展示了 InternEvo 在 128 个 GPU 上训练具有 256k 序列长度的 7B 模型时获得的执行规划 S 可视化。InternEvo 选择在训练期间不使用流水线并行和激活重计算机制。InternEvo 将 ssp 设置为 32，以适应设备上的 256k 序列。它还将 soss 设置为 2，将优化器状态 (Adam) 的内存使用量从 42 GB 减少到 5 GB。鉴于参数和梯度的模型状态在 7B 模型中占整体内存使用量的一小部分，InternEvo 将 sps 设置为 4 进一步限制通信成本，从而限制通信参与者的数量。