DEEPSPEED ULYSSES: SYSTEM OPTIMIZATIONS FOR ENABLING TRAINING OF EXTREME LONG SEQUENCE TRANSFORMER MODELS

作者/机构: Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Shuaiwen Leon Song, Samyam Rajbhandari, Yuxiong He (Microsoft Inc)

A1 主要贡献

本文针对大规模语言模型（LLM）在长序列训练中面临的系统挑战，介绍了一种名为 DeepSpeed-Ulysses 的新方法。

核心问题：现有的大模型训练并行技术（如数据并行、张量并行、流水线并行）主要关注批量大小、隐藏层维度和模型深度，并未针对序列长度维度进行优化。而现有的序列并行方法存在内存和通信效率低下的问题，限制了其在长序列大模型上的扩展性。此外，这些方法通常需要侵入式且易于出错的代码重构，可用性有限。

研究目标：提出一种简单、可移植且高效的方法，以实现极长序列长度下 LLM 的高效和可扩展训练。

核心创新点 (DeepSpeed-Ulysses)：
该方法的核心思想是沿序列维度对输入数据进行分区。在进行注意力计算之前，它对分区后的查询（queries）、键（keys）和值（values）执行一次高效的 all-to-all 通信。这样，每个 GPU 会接收到完整的序列，但只负责注意力头（attention heads）的一个不重叠的子集。这使得参与的 GPU 可以并行计算不同注意力头的注意力。最后，通过另一次 all-to-all 通信来收集注意力头的计算结果，同时将结果重新沿序列维度分区。

主要贡献：
* 支持更长的序列训练：DeepSpeed-Ulysses 能够训练序列长度比现有系统长4倍的 Transformer 模型，并支持超过一百万个 token 的序列。
* 显著提升训练吞吐量：与现有系统相比，通信量减少超过10倍，吞吐量提升高达2.5倍，并能维持超过 175 TFlops/GPU 的高吞吐率（超过硬件峰值性能的54%）。
* 通用且与注意力实现无关：DeepSpeed-Ulysses 支持密集和稀疏注意力，并能与 FlashAttention v2 等高效注意力实现协同工作。
* 支持大规模模型训练：通过与 ZeRO-3 结合，不仅支持长序列，还能支持巨大规模的模型训练。
* 易于使用和移植：对现有训练框架的代码修改要求极小。

A3 背景知识与相关工作

2.1 背景知识

2.1.1 Transformer 架构

Transformer模型构建模块。一个典型的多头注意力 Transformer 架构【12, Attention is all you need, 2017, Advances in neural information processing systems】如图1所示。其构建模块包括输入序列，这些序列被投影成查询（Q）、键（K）和值（V）嵌入。QKV 通常是一个大小为 N, b, d 的3D张量，其中 N 是序列长度，b 是微批量大小，d 是隐藏维度。QKV 张量被送入注意力模块，这是 Transformer 模型的核心组件。注意力的输出是多层感知机（MLP）或逐点前馈网络的输入。注意力模块后跟 MLP 模块，这种结构被多次复制以形成编码器、解码器或编码器-解码器 Transformer 网络。

图 1: 多头注意力 Transformer

2.1.2 并行模式

数据并行及其局限性。数据并行【13, Large scale distributed deep networks, 2012, Advances in neural information processing systems】是加速神经网络训练的常用方法，它将输入数据按样本或批次维度划分到不同计算设备上，同时在每个设备上复制模型参数。当批量大小足够大以隐藏通信开销时，数据并行是有效的。然而，当模型很大，无法在每个设备上完整复制时，数据并行就受到限制。ZeRO【14, Zero: Memory optimizations toward training trillion parameter models, 2020, SC20】【15, Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning, 2021, SC ’21】通过在可用计算设备间划分模型参数来解决此问题。此外，大批量训练已知会对模型质量产生影响【16, On large-batch training for deep learning: Generalization gap and sharp minima, 2016, arXiv】。

序列并行与数据并行的互补性。本文提出的方法与数据并行和 ZeRO 是正交的，可以与两者结合使用。通过利用序列并行，我们可以在大型系统上保持全局批量大小在合理范围内，从而有效缓解大批量对模型收敛的影响。序列并行在此有两个作用：首先，对于给定的长序列长度，序列并行可以加速求解时间，即按比例增加计算资源可以减少迭代时间。其次，序列并行支持更长的序列训练或持续预训练，其中训练上下文长度随时间逐渐增加【4, Effective long-context scaling of foundation models, 2023】。例如，在1024个GPU上进行大规模训练，如果需要将序列长度从8K增加到32K，而由于对模型质量的负面影响又不能增加全局批量大小，序列并行便成为一种有效的系统优化技术，它允许将大批量跨多个GPU拆分，而无需增加全局批量大小，也避免了繁琐的超参数搜索。

张量与流水线并行。张量并行【17, Megatron-lm: Training multi-billion parameter language models using model parallelism, 2019】和流水线并行【18, Pipedream: Generalized pipeline parallelism for dnn training, 2019, SOSP 2019】【19, Gpipe: Efficient training of giant neural networks using pipeline parallelism, 2018, ArXiv】【20, Memory-efficient pipeline-parallel dnn training, 2021, ICML】是另外两种流行的大规模训练方法，统称为模型并行。与数据并行不同，当模型过大无法在数据并行各等级上完全复制时，会使用模型并行。张量并行在层内拆分计算算子（如注意力和MLP），而流水线并行则以深度（逐层）方式拆分模型。3D并行【21, DeepSpeed: Extreme-scale model training for everyone, 2020】【22, Using deepspeed and megatron to train megatron-turing nlg 530b, a large-scale generative language model, 2022, arXiv】结合了数据并行、张量并行和流水线并行，以获得比单个组件更高的吞吐量，但代价是大量的代码重写和生产力开销【23, Zero++: Extremely efficient collective communication for giant model training, 2023】。

2.2 相关工作

分布式训练方法概述。关于深度神经网络分布式训练方法的广泛概述和调查，可参见【24, Demystifying parallel and distributed deep learning: An in-depth concurrency analysis, 2019, ACM Computing Surveys (CSUR)】。这些方法大致分为数据并行和模型并行。然而，所有现有的并行方法在处理与极长序列相关的中间激活内存开销方面都存在局限性。

表 1: 本工作 (DS-Ulysses) 与其他序列并行方法的比较。

现有序列并行方法的不足。虽然最近的序列并行工作解决了内存开销问题，但它们在通信效率方面有所欠缺，因此扩展能力有限。与我们的工作类似，所有现有的序列并行工作都沿序列维度划分输入数据，但在划分哪些输入投影以及如何聚合和通信以进行注意力计算方面有所不同。

现有序列并行方法详解。文献【25, Sequence parallelism: Long sequence training from system perspective, 2022b】（后文称 ColAI-SP）引入了环形自注意力（ring self attention），这是一种环状通信集合，其中查询投影是本地的，而键和值投影则以环形方式传输以计算全局注意力，导致通信复杂度与消息大小 M 呈线性关系。Megatron-LM 的序列并行方法【26, Reducing activation recomputation in large transformer models, 2022】与 Megatron 张量并行紧密集成。它沿序列维度划分序列，并应用 all-gather 和 reduce-scatter 集合来聚合 QKV 投影以进行注意力计算。通信复杂度分析表明，与我们的方法不同，Megatron-LM 序列并行的通信量随消息大小 M 线性增加，而与计算设备数量无关。相比之下，DeepSpeed-Ulysses 通过按比例增加 GPU 数量来保持通信量恒定，具体细节见3.2节。

DeepSpeed-Ulysses的优势总结。表1总结了 DeepSpeed-Ulysses 与其他现有方法的不同之处。DeepSpeed-Ulysses 在通信效率上优于其他两种方法。它还受益于利用 ZeRO【14, Zero: Memory optimizations toward training trillion parameter models, 2020, SC20】【15, Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning, 2021, SC ’21】优化，可在序列和数据并行组之间划分模型参数。DeepSpeed-Ulysses 支持不同类型的注意力，并且易于使用。Megatron-LM 的序列并行与其张量并行紧密集成，限制了其内存效率和易用性。ColAI-SP 需要一种特定类型的注意力，不易使用，且其环形自注意力对其他注意力类型和机制的泛化能力尚不明确。

正交的相关工作。在稀疏 Transformer 领域也有相关工作，主要关注全注意力的近似，例如稀疏注意力【27, Generating long sequences with sparse transformers, 2019, CoRR】【28, Rethinking attention with performers, 2020, CoRR】【29, Big bird: Transformers for longer sequences, 2021】【1, Longformer: The long-document transformer, 2020】。最近也有关于单 GPU 内存和计算高效注意力的工作，其中一个流行的例子是 Flash attention【30, Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022】【11, Flashattention-2: Faster attention with better parallelism and work partitioning, 2023】，它利用了如切片（tiling）和重计算等已知技术来提高计算和内存效率。这些工作与我们的工作是正交的，并被相应地利用。

A2 方法细节

3.1 系统设计

DeepSpeed-Ulysses核心设计。图2展示了 DeepSpeed-Ulysses 的核心设计。与已知的 Transformer 架构一样，该设计包含在 P 个可用设备上分区的输入序列 N。每个本地 N/P 分区被投影成查询（Q）、键（K）和值（V）的嵌入。接下来，(QKV)嵌入通过参与计算设备之间高度优化的 all-to-all 集合操作聚合成全局 QKV。all-to-all 集合操作之后是每个头的注意力计算，形式如下：

注意力计算之后，另一个 all-to-all 集合操作将注意力计算的输出上下文张量转换为序列(N/P)并行形式，以供 Transformer 层块中其余模块的后续算子（如 MLP MatMul、层归一化等）使用。

图 2: DeepSpeed 序列并行 (DeepSpeed-Ulysses) 设计

3.2 通信分析

DeepSpeed-Ulysses的通信优势。DeepSpeed-Ulysses 与其他现有长序列方法的主要区别在于，我们具有更小的总通信量和更好的随序列并行度增加的可扩展性。

通信量分析。在具有节点内 NVSwitch 互连和节点间胖树 IB 拓扑的现代集群上，对于 P 个 GPU 上的总消息大小为 M 的 all-to-all 操作，每个链路传输的通信量为 M/P。对于一个隐藏大小为 h、序列长度为 N、并行度为 P 的 Transformer 模型，DS-Sequence 在注意力计算前对 QKV 投影执行 all-to-all 操作，总消息大小为 3Nh，之后对输出上下文投影执行另一次 all-to-all 操作，大小为 Nh。因此，DeepSpeed 序列并行每个链路的总通信量为 4Nh/P，其复杂度为 O(N/P)。值得注意的是，当 N 和 P 按比例增加时，此通信量是恒定的。

与现有方法的通信对比。相比之下，像 Megatron-LM 这样的现有方法，其通信量随 N 线性增加，而与 P 无关，导致通信复杂度为 O(N)。例如，Megatron-LM 在每个 Transformer 层中执行两次消息量为 Nh 的 all-gather 和两次消息量为 Nh 的 reduce-scatter。然而，当 P » 1 时，每次 all-gather 和 reduce-scatter 的成本仍然是 M，而不是 M/P。因此，Megatron-LM 序列并行的每个链路通信量为 4Nh，是 DeepSpeed 序列并行的 P 倍。这使得 DeepSpeed 序列并行能够支持极长序列的训练，同时实现比现有方法显著更高的训练效率。我们的评估结果与此分析相符。

3.3 内存效率

与ZeRO-3的集成。虽然 DeepSpeed 序列并行在训练长序列时减少了激活内存，但它不影响模型状态所消耗的内存。因此，为了支持大型语言模型进行长序列训练，DeepSpeed 序列并行与 ZeRO-3 集成。ZeRO 冗余优化器第3阶段 (ZeRO-3)【14, Zero: Memory optimizations toward training trillion parameter models, 2020, SC20】【15, Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning, 2021, SC ’21】是一种用于训练大型模型的内存优化技术。与传统的神经网络数据并行训练中模型状态在数据并行 ranks 间复制不同，ZeRO-3 通过在数据并行 ranks 间划分模型状态来优化内存使用。

扩展ZeRO-3分区。通过序列并行，训练数据可以同时在批次（样本）和序列维度上考虑，相关的并行组可以组合成一个更大的组以进行 ZeRO 并行。因此，我们将 ZeRO-3 的分区扩展到数据并行和序列并行 ranks 的组合。换句话说，在 DeepSpeed 序列并行中，ZeRO 在序列和数据并行组之间划分模型状态，并在需要时通过 all-gather 收集每个 rank 的分区。同样，梯度在数据和序列并行 ranks 之间进行 reduce 操作以更新参数。ZeRO 的支持在序列和数据维度上都带来了巨大的内存节省，不仅能扩展到长序列长度，还能扩展到大型模型。

3.4 通用和与注意力无关的解决方案

通用性与模块化设计。DeepSpeed 分布式注意力模块的实现足够通用，可以支持任何类型的注意力：例如，自注意力、交叉注意力、因果注意力，包括它们的密集和稀疏版本，以及支持本地长序列注意力的各种优化内核，如不同版本的 FlashAttention。DeepSpeed-Ulysses 的通用性源于其核心设计的模块化特性：一个以注意力为中心的序列并行设计。

以注意力为中心的设计。在注意力计算之前是 N/P 分区的序列并行，注意力计算本身是头并行，每个头进行完整的注意力计算，但头的数量较少。因此，注意力计算部分可以替换为任何类型的注意力机制，例如密集注意力或各种形式的稀疏注意力。

A4 实验

实验环境

• 模型架构：实验使用 GPT 模型【31, Language models are unsupervised multitask learners, 2019】，参数规模包括 12亿（1.2B）、70亿（7B）和 300亿（30B）。
• 硬件配置：实验在最多 256 块 A100 GPU 上进行。
• 软件配置：DeepSpeed-Ulysses (在文中也称为 DeepSpeed Sequence) 与 Megatron-LM 的序列并行实现进行了比较。
• 数据集：论文未明确指出具体使用的数据集，但实验聚焦于模型训练的系统性能（如吞吐量、扩展性）和收敛性，而非特定任务的 SOTA 结果。

实验结果

我们对 DeepSpeed-Ulysses 进行了五方面的评估：i) 序列长度的可扩展性，ii) 密集注意力的吞吐量及与现有系统的比较，iii) 稀疏注意力的吞吐量及与现有系统的比较，iv) 并行扩展性研究，以及 v) 收敛性研究。

1. 序列长度可扩展性
* 实验内容：在一个 12 亿参数的 GPT 模型上进行了序列长度的强扩展性测试，序列长度最高达到100万个 token。
* 实验结果：如图3所示，DeepSpeed 序列并行允许序列长度随 GPU 数量线性增加，并在不同序列长度和相应 GPU 数量下保持了相似的计算吞吐量，展示了线性的扩展能力。

图 3: DeepSpeed 序列并行在不同序列长度和 GPU 数量下的强扩展性评估

2. 密集注意力评估
* 实验内容：在 70 亿和 300 亿参数的 GPT 密集注意力模型上，分别使用 32 和 64 块 A100 GPU，将 DeepSpeed-Ulysses 与 Megatron-LM 的序列并行进行了比较。
* 实验结果：如图4和图5所示，DeepSpeed-Ulysses 在所有可运行的序列长度上均优于 Megatron-LM，并且能够运行更长的序列。其性能优势主要来自两方面：(1) DeepSpeed-Ulysses 结合 ZeRO-3，因内存优化可以容纳更多样本，从而获得更高吞吐量；(2) DeepSpeed-Ulysses 受益于比 all-gather 更高效的 all-to-all 通信。在 30B 模型上，DeepSpeed-Ulysses 相比 Megatron-LM 训练速度提升了 2.5 倍，并能处理 4 倍长的序列。

图 4: 在 7B 参数模型上使用密集注意力对 DeepSpeed-Ulysses 和 Megatron LM 的评估 (32 GPUs)

图 5: 在 30B 参数模型上使用密集注意力对 DeepSpeed-Ulysses 和 Megatron LM 的评估 (64 GPUs)

3. 稀疏注意力评估
* 实验内容：在 70 亿和 300 亿参数的稀疏注意力模型上，将 DeepSpeed-Ulysses 与 Megatron-LM 进行了基准测试。
* 实验结果：如图6和图7所示，稀疏注意力的实验结果与密集注意力相似。DeepSpeed-Ulysses 的吞吐量是 Megatron-LM 的2倍以上。得益于 ZeRO-3 的内存节省，DeepSpeed-Ulysses 可以扩展到比 Megatron-LM 长4倍的序列。作者指出，目前 DeepSpeed 的吞吐量受限于本地稀疏注意力的实现，随着本地稀疏注意力性能的提升，未来性能差距预计会进一步扩大。

图 6: 在 7B 参数模型上使用分块稀疏注意力对 DeepSpeed-Ulysses 和 Megatron LM 的评估 (32 GPUs)

图 7: 在 30B 参数模型上使用分块稀疏注意力对 DeepSpeed-Ulysses 和 Megatron LM 的评估 (64 GPUs)

4. 并行扩展性研究
* 实验内容：使用 GPT-7B 密集模型和全局批量大小为8，进行了两种扩展性研究：(1) 强扩展性：固定序列长度为131,072，GPU数量从64增加到256。(2) 弱扩展性（变种）：GPU数量与序列长度成比例增加。
* 实验结果：
* 强扩展性（表2）：执行时间几乎随着GPU数量的增加而线性下降。
* 弱扩展性（表3）：尽管注意力计算的复杂度随序列长度呈二次方增长，实验仍显示出良好的扩展性，在两个研究中都达到了理论峰值GPU性能的很高的百分比。吞吐量的轻微下降归因于通信开销的增加，但总体上证明了 DeepSpeed-Ulysses 良好的并行效率。

表 2: 固定序列长度的并行扩展性研究

表 3: 不同序列长度的并行扩展性研究

5. 收敛性研究
* 实验内容：在一个 13 亿参数的 GPT 模型上，使用 32K 序列长度和 8 个 A100 GPU（序列并行度为4），比较了 DeepSpeed-Ulysses（使用不同 ZeRO 阶段）和 Megatron-LM 的收敛性。
* 实验结果：如图8所示，实验验证了 DeepSpeed-Ulysses 是一种纯粹的系统优化技术，对训练模型的质量没有（负面）影响，其收敛曲线与基线一致。

图 8: DeepSpeed-Ulysses 在不同 ZeRO 内存优化阶段下的收敛性评估

A5 结论

本文提出了一种内存和通信高效的 DeepSpeed Sequence (Ulysses) 技术，为训练具有长序列的大型 Transformer 模型提供了可能。DeepSpeed Sequence 在 GPU（以及其他 AI 加速器）之间实现序列并行，将序列并行化应用于 Transformer 模型的所有组件，并无缝支持了包括 FlashAttention 在内的最先进的（密集和稀疏）注意力机制。使用 DeepSpeed Sequence 进行训练，可以使模型大小和序列长度几乎无限扩展，不受单 GPU 内存的限制，同时能达到接近峰值计算性能的高效率。