[论文解读] Sequence Parallelism: Making 4D Parallelism Possible

Shenggui Li, Fuzhao Xue+2 more|arXiv (Cornell University)|May 26, 2021

Advanced Neural Network Applications参考文献 23被引用 12

一句话总结

本文提出序列并行性（sequence parallelism），一种内存高效的训练方法，通过将长序列拆分到多个GPU上，利用环形风格的 all-reduce 通信机制传递注意力嵌入，实现更长序列的建模。在 64 块 P100 GPU 上，该方法相比张量并行性实现了 13.7 倍更大的批量大小和 3.0 倍更长的序列长度，且与 PyTorch 中的数据并行性和流水线并行性完全兼容。

ABSTRACT

Within Transformer, self-attention is the key module to learn powerful context-aware representations. However, self-attention suffers from quadratic memory requirements with respect to the sequence length, which limits us to process longer sequence on GPU. In this work, we propose sequence parallelism, a memory efficient parallelism method to help us break input sequence length limitation and train with longer sequence on GPUs. Compared with existing parallelism, our approach no longer requires a single device to hold the whole sequence. Specifically, we split the input sequence into multiple chunks and feed each chunk into its corresponding device (i.e. GPU). To compute the attention output, we communicate attention embeddings among GPUs. Inspired by ring all-reduce, we integrated ring-style communication with self-attention calculation and proposed Ring Self-Attention (RSA). Our implementation is fully based on PyTorch. Without extra compiler or library changes, our approach is compatible with data parallelism and pipeline parallelism. Experiments show that sequence parallelism performs well when scaling with batch size and sequence length. Compared with tensor parallelism, our approach achieved $13.7 imes$ and $3.0 imes$ maximum batch size and sequence length respectively when scaling up to 64 NVIDIA P100 GPUs. We plan to integrate our sequence parallelism with data, pipeline and tensor parallelism to further train large-scale models with 4D parallelism in our future work.

研究动机与目标

• 解决 Transformer 中自注意力机制导致的二次方内存增长问题，从而缓解 GPU 上序列长度受限的问题。
• 突破传统方法中必须由单个设备存储完整输入序列才能进行注意力计算的瓶颈。
• 通过将序列块分发到多个 GPU 上，实现无需修改模型或硬件即可高效训练长序列。
• 设计一种通信高效的注意力机制，与现有并行策略（如数据并行和流水线并行）兼容。
• 通过将序列并行性整合到四维并行框架中，实现批量大小和序列长度的可扩展训练。

提出的方法

• 将输入序列划分为多个块，并将每个块分配给不同的 GPU，从而消除单个设备必须存储完整序列的需求。
• 在自注意力计算过程中，使用环形风格 all-reduce 通信在 GPU 之间交换注意力嵌入，实现分布式注意力计算。
• 提出环形自注意力（Ring Self-Attention, RSA），一种新型注意力机制，将环形通信直接集成到注意力计算流水线中。
• 原生在 PyTorch 中实现该方法，无需编译器或库的修改，确保与现有训练框架的广泛兼容性。
• 将序列并行性与数据并行性和流水线并行性结合，实现大规模模型训练的四维并行。
• 通过在反向传播过程中使用环形 all-reduce 同步各设备上的注意力输出，保持梯度一致性和模型精度。

实验结果

研究问题

• RQ1 序列并行性是否能通过将序列块分发到多个 GPU 上，突破 Transformer 中的序列长度限制？
• RQ2 在最大批量大小和序列长度扩展方面，序列并行性与张量并行性相比表现如何？
• RQ3 当与数据并行性和流水线并行性结合时，序列并行性在批量大小和序列长度上的扩展能力如何？
• RQ4 是否可能在不进行架构或编译器级别修改的前提下，将序列并行性与现有并行策略（如数据并行和流水线并行）无缝集成？
• RQ5 与标准 all-reduce 或参数共享相比，序列并行性中环形通信机制在注意力计算中的效率如何？

主要发现

• 在扩展到 64 块 NVIDIA P100 GPU 时，序列并行性相比张量并行性实现了 13.7 倍更大的批量大小。
• 在相同硬件配置下，该方法支持的输入序列长度比张量并行性长 3.0 倍。
• 该方法在批量大小和序列长度同时增加时表现出高度可扩展性，展现出优异的线性扩展效率。
• 实现与数据并行性和流水线并行性的完全兼容，可无缝集成到现有训练流水线中。
• 环形自注意力（RSA）通过利用环形 all-reduce 有效降低了通信开销，在训练过程中保持了高吞吐量。
• 该方法无需对 PyTorch 或外部库进行任何修改，可轻松部署到现有深度学习工作流中。