【论文阅读】DSP: Dynamic Sequence Parallelism for Multi-Dimensional Transformers

阅读这篇文章之前建议先阅读：序列并行
本文链接：https://link.zhihu.com/?target=https%3A//arxiv.org/pdf/2403.10266.pdf
文章源码：https://github.com/NUS-HPC-AI-Lab/OpenDiT

Abstract

• 为多维Transformer实现高效的序列并行性
• 关键思想是根据当前计算阶段动态切换并行维度，与应用传统的单维并行性相比，这种动态维度切换允许以最小的通信开销实现序列并行性。

1 Introduction

• 应用中涉及的长序列会导致大量的激活内存成本和极慢的生成速度，因此需要使用序列并行性。
• 许多应用程序跨不同维度执行注意力计算。现有的序列并行：
  • Megatron-SP
  • DeepSpeed-Ulysses
  • Ring Self-attention
• 现有的序列并行假设只有一个维度的序列，他们仍然按照传统的注意力计划将注意力分散到并行性上，并产生不必要的通信。
• Spatial-Temporal Attention时空注意力
  • 其中注意力是针对时间和空间维度分别计算的。类似的方法用于蛋白质预测、时空信息处理
• DSP通过根据当前计算阶段动态切换并行维度来利用多维Transformer的潜在特性
• 主要贡献：
  • DSP利用多维Transformer的潜在特性，并根据当前注意力维度动态切换并行维度，以最小的通信成本扩展长序列。
  • 提升了端到端的吞吐量减少通信开销
  • DSP还与ZeRO结合支持大序列长度和海量模型大小，并且还兼容各种注意力内核
  • DSP 易于使用且可移植，只需对现有框架进行最少的代码更改。

2 Preliminary and Related Work

• Transformer 的关键创新是使用自注意力机制，而不是像循环神经网络那样基于序列的建模。自注意力允许模型在进行预测时权衡输入数据的不同部分，从而更有效地捕获长期依赖性。
• 多头注意力允许模型在生成输出时关注输入中的不同单词
• 前馈层对来自注意力层的表示应用额外的变换。
• 残差连接：各层之间使用残差连接以帮助进行更深入的训练
• Layernorm：在每一层之前还应用层归一化以稳定训练过程。

2.1 多维Transformer

• 多维Transformer扩展了标准变压器的自注意力机制，能够在多个维度上运行，而不仅仅是一个序列。
• 自注意力计算不是沿着一维序列进行，而是泛化为跨越输入数据中存在的附加空间、时间或其他相关维度。
• 突出的例子：spatial-temporal transformers for video data
  • 这些transformer沿着每个视频帧的空间高度和宽度维度以及跨帧的时间维度应用自注意力。
  • 这允许捕获视频中的远程时空依赖性
  • 时空transformer块用**专门的多维前馈层（FFN）**替换标准前馈层，以处理参与的 3D 时空表示。
• 突出的例子：蛋白质预测
  • 自注意力跨越氨基酸残基的空间 3D 坐标，随着时间的推移，多维注意力涵盖特征维度和时间维度。
• 关键在于：扩展自注意力计算，以有效地对输入数据中相关附加维度的依赖关系进行建模。这使得transformer能够应用于更广泛的多维数据模式。

2.2 相关工作

• 四种并行：数据并行，张量并行，流水线并行，序列并行

3 DSP

3.1 系统设计

- Megatron-SP 和 DeepSpeed-Ulysses 都遵循传统的注意力计划，它将张量布局从序列分割转换为头部分割，使得注意力能够在本地执行。随后，张量布局被转换回序列分割或完整序列，以用于接下来的 MLP 计算。

• 多维transformer：每个维度的序列计算都是独立的

  • 计算空间变换器块时，计算与时间序列维度无关，并且数据可以跨设备分割而不影响结果
    

• 根据计算阶段动态切换序列并行度的维度，注意力模块内部就不需要复杂的通信操作；

• 在计算阶段之间转换时只需要单个 AlltoAll 操作

• DSP的优势：它能够最大限度地降低通信成本并实现可扩展的通信操作。

• 通信开销的分析：

• Megatron-SP 每个变压器块采用 4 个资源密集型集体通信操作。

  • 启动一个 AllGather 操作来聚合整个输入 x
  • 注意力模块和 MLP 模块的输出处都进行了 ReduceScatter 操作，最终通信量达到 8M

• DeepSpeed-Ulysses采用了更高效的AlltoAll方法，

  • 它将每个块的通信操作数量从 2 个增加到 4 个
  • 利用 AlltoAll 进行查询、键、值和输出布局转换。
  • 对于跨 N 个 GPU 的大小为 M 的 AlltoAll 通信，每个设备传输的通信量为 8M/N 。
  • 它显着减少了通信量，但引入了额外的通信启动开销，并且在设备数量有限的情况下，与 Megatron-SP 相比，它几乎没有表现出效率增益。

• DSP:

  • SP总共只使用两次AlltoAll操作来降低通信成本，从而显着减少通信操作的数量
  • 此外，它将通信量最小化至 2M/N ，显着优于其他序列并行技术。
  • 通信量随着节点的增加而减少——使 DSP 成为涉及长序列的大规模分布式训练和推理任务的绝佳选择

• 内存开销的分析：

• 激活的内存：

  • Megatron-SP 需要在 AllGather 操作之后保存整个激活，从而导致托管激活的内存要求更高。

• 参数的内存：

  • 遵循 DeepSpeed-Ulysses 的方法，利用 ZeRO 技术在不同设备上均匀地分片所有参数
  • Megatron-SP 也通过分割权重来对参数进行分片，但 Ring Attention 需要保存整个参数集，导致参数内存需求更高