序列并行DeepSpeed-FPDT

序列并行DeepSpeed-FPDT

原创 手抓饼熊​  2024年11月05日 13:20 四川

原文：https://zhuanlan.zhihu.com/p/720387198

1 背景

1.1 序列并行回顾

本文主要介绍DeepSpeed-FPDT，在此之前，笔者已经分享了很多序列并行相关的内容，主要是序列并行云台28将系列。

在序列并行上篇，主要通过十几篇论文，详细介绍了业界序列并行的做法，最终总结出了三个方向（Ulysess、RingAttention、AllGatherKV）。

在序列并行中篇，主要详细分析Ulysess和RingAttention，并且分析USP（集Ulysess和RingAttention之大成者）及基于USP的一些新的文章。

在序列并行下篇，主要详细分析推理场景的序列并行，其中序列流水并行（序列并行+流水并行）被广泛使用。序列并行上、中、下三篇基本上已经分析了业界所有的序列并行方案。这些序列并行策略，尽管已被证明能够训练具有长上下文的LLM，但需要大量的GPU。

例如，使用Megatron-SP结果训练一个具有256K标记上下文窗口的7B参数模型需要超过32个A100 80G GPU。同样，使用DeepSpeed Ulysses训练一个小的1.2B参数GPT模型，其上下文长度为1M标记，需要64个A100 GPU。上述的序列并行方案和激活重计算可以一起结合使用，进一步扩展序列长度，然而激活重计算会带来额外的计算。本文思路是Ulysess + 激活offload + 分块的思想。想更好的理解本文可以详细阅读第8节和USP内容。

1.2 DeepSpeed-FPDT创新点

本文提出一种新的全流水线分布式Transformer（FPDT）来解决这些挑战，以训练长上下文LLM。该方法利用现代GPU集群中的多个内存层次结构，增强其硬件效率和成本效益，同时达到极高的MFU。

本文主要的创新点如下：

本文设计了一个基于DeepSpeed Ulysses的完全流水线式分布式Transformer，专为具有数百万令牌序列长度的LLM定制，利用GPU和主机CPU内存以及预取功能，实现几乎零开销的训练流程。

本文显著减少了LLM训练期间激活的GPU内存占用，同时利用专用的双缓冲设计，几乎将所有预取与计算重叠。

2 基础理论

2.1 transformer显存占用

随着Transformer成为当今生成模型中最流行的架构，其整体流程也变得同质化。如下表格展示transformer块前向和后向每个步骤所需的GPU内存。

在这个表格中值得注意的是，为了获取查询、键和值，内存占用量直接增加了三倍，这可能会导致内存溢出问题，特别是当序列本身过长而无法适应GPU内存时。此外，由于大多数Alltoall实现不支持原位操作，故需要额外创建一个接收缓冲区。如果启用了异步通信，则需要一次性为查询、键和值准备接收缓冲区，导致6Nd的内存占用量。

FlashAttention被引入以将内存消耗降至O(N)，然而在实践中，它也可能带来巨大的内存占用量。例如，在后向传播中，FlashAttention需要以下输入：q、k、v、of、og、dq、dk、dv，其中of表示前向传递的注意力输出，og表示后向传递中输出的梯度。由于这些张量需要同时存在于GPU内存中，这直接导致8Nd的内存占用量。目前没有任何现有解决方案专门针对这些内存负担，然而，在长上下文LLM训练中变得明显。

2.2 Ulysess和ZeRO3回顾

在序列并行策略中，DeepSpeed-Ulysses以其高效的通信模式脱颖而出，并且与诸如DeepSpeed ZeRO等最先进的基于模型的训练方案相辅相成。本文首先回顾一下DeepSpeed Ulysses的关键特性，以及它如何与ZeRO-3配合工作。

DeepSpeed-Ulysses

图2（a）显示了DeepSpeed Ulysses序列并行通信模式。每个GPU最初都有整个序列的一部分，并且拥有完整的头部，在进行注意力操作时，序列片段被收集到每个GPU，而头部则被分散到不同的GPU。

ZeRO3

ZeRO3将所有参数、梯度和优化器状态沿着数据并行的GPU组进行分区。由于序列由token组成，它们可以很容易地被视为数据的批次，在这里可以直接应用数据并行。如图3所示，在这种组合训练方案中，序列并行用于减少激活内存，而ZeRO-3则减少了与模型相关的内存。

3. DeepSpeed-FPDT详解

3.1 流水线和调度实现

由于QKV projection，Alltoall通信，attention和FFN等操作会创建多个中间缓冲区，导致内存波动严重，特别是在反向传播过程中，为了使Transformer块中的序列计算完全流水线化和内存高效，本文的分块和卸载设计将从初始输入张量（即隐藏状态）开始。首先，为了方便解释，在整个论文中使用以下符号。除分布式注意力之外的操作，每个GPU持有和处理一个[b，

，d]张量，其中

表示本地序列长度，

表示总头数。对于分布式注意力，每个GPU将处理整个序列，但具有特定的头部，即一个[b，

，d]张量。

对于第一个QKV projection，由于令牌是逐个处理的，直接将本地序列张量[b，

，d]切片成u个块，每个块都是一个[b，

，d]张量。将这些块表示为Ti，其中i∈{0, 1, . . . , u − 1}。如图4所示，ci被投影为查询qi，键ki和值vi。然后在序列并行组之间执行Alltoall通信。回想一下，在表2中由于Alltoall通信无法原地执行，需要预先分配额外的接收缓冲区。在本文的块设计中，由于将序列长度大大减少了u倍，即使仍需要qˆi，kˆi，vˆi的接收缓冲区，与具有完整上下文的原始序列张量相比，它们变得微不足道。

使用Alltoall来分散头部和聚集序列后，每个块qˆi、kˆi、vˆi都是一个[b,

, d]张量。请注意，对于生成式LLMs，由于因果掩码，在计算后，qˆi将不再在前向传递中使用，而kˆi、vˆi需要参与接下来的qˆj计算，其中i < j。因此可以将kˆi、vˆi缓存到主机内存中。具体来说，对于qˆ0、kˆ0、vˆ0，可以直接获得块T0的最终输出，因为qˆ0将不会关注剩余序列。对于块Ti，i ≥ 0，处理每个块只会给出中间结果，这些结果将在下一个块的计算中重新缩放。由于在线注意力广泛使用，在调度注意力计算时采用类似的策略。

图5给出了如何执行块Tm计算的示例。在Alltoall操作之后，GPUj接收qˆm，kˆm和vˆm。然后从主机内存逐块地提取之前的序列块到GPUj，并使用当前的qˆm执行在线注意力，并相应地更新输出块。请注意，在严格的情况下，任何给定时间，只有一组块kˆi，vˆi被放置在GPU的HBM上，将内存占用量减少到1u，与非卸载版本相比。

与推断不同，推断中，激活和中间结果可以在当前Transformer块的前向传播完成后立即释放，训练则需要保存必要的中间输出，因为它在反向传播中被重复使用。因此，一旦它们完成前向计算，可以将qˆi，kˆi，vˆi卸载到主机内存。

3.2 数据重排布

另一个需要注意的是，这种方法需要根据分块的Alltoall序列收集来重新调整序列。在原始序列并行方案中，输入序列被切片并分发到每个GPU，根据它们的排名。如图6所示，如果直接在每个GPU上对第i块执行Alltoall操作，那么收集的序列就不符合因果掩码。

例如，当收集第二块时，每个GPU将获得一个序列，其中包含特定排名的T1、T5、T9、T13块。然而，由于每个GPU在先前的块计算中已经收到了T0、T4、T8、T12，因此对于qˆ1、kˆ0的因果掩码需要专门设计，因为T5只能参与T0和T4等令牌。另一个解决方法是只让第i个GPU将其序列分散到所有GPU，这不需要重新排序序列，但会导致NVLINK负载不均衡。在训练过程中，由于每个GPU计算其所持有序列的损失，所以还相应地重新排序标签，以使损失仍然匹配。请注意，在数据加载器中我们对输入的令牌ID和标签进行了混洗；因此，在重新排序序列时不会增加额外开销。

与环形注意力不同的是，在本文的设计中，每个GPU始终在任何给定时间处理相同的序列片段，唯一的区别是特定排名的头部。因此，在计算注意力时，GPU始终是负载平衡的。这种统一的工作流程还减少了协调并行组序列所需的同步和障碍。

4 双buffer机制优化

4.1 chunk大小对计算影响

chunk太小问题

尽管利用主机内存保存未使用的序列的想法很直观，但不匹配的硬件传输带宽在充分利用计算能力方面提出了重要挑战。对于典型的HPC节点，GPU通过高带宽的NVLink连接，可以达到超过100 GB/s的点对点带宽。然而，通用的PCIe Gen-4连接仅提供16条通道的32 GB/s单向带宽，这也要求主机内存和GPU在同一NUMA域中。

图8显示了一个GPU饥饿的情况，其中序列块太短，导致注意力计算的延迟小于数据获取延迟。在这种情况下，训练将受到PCIe带宽的限制，导致模型FLOPs利用率低。

chunk太大问题

图9显示了一个相反的情况，即每个序列块都太长了。因为传递给内核（即注意力前向/后向内核）的张量需要作为一个整体驻留在HBM上，这将占用太多的GPU内存，这是不必要的，因为一旦前面的标记计算完成，标记就可以从主机内存中获取。

此外，由于GPU计算吞吐量和PCIe链路带宽之间的不匹配，将offload方法整合到常规计算管道中会带来了严重困难。图10显示了在GPU节点上执行的不同操作的延迟（有关详细信息，请参见5.1）。请注意，这里比较了transformer块中的实际张量操作，即[b，s/p，h，d]张量上的Alltoall操作，[b，s，h/p，d]张量上的注意力操作，以及[3，b，s，h/p，d]张量上的提取操作代替查询、键和值。由于注意力计算的复杂度呈二次增长，如果让每个序列块足够大（例如512k），它很容易主导总体延迟，此时通信和提取的延迟变得微不足道。然而，在实践中，大块大小也会增加内存压力，这与最初的目标相违背。找到充分利用计算核心并隐藏卸载/提取延迟的甜点非常关键，以实现最佳的训练效率。正如图10所示，Alltoall的速度要快得多，因为这只是使用NVLink进行节点内通信。

因此，本文更关心注意力计算的延迟何时超过主机到设备提取的延迟。

chunk的获取策略

与GPU之间的连接不同，主机内存和GPU之间的数据传输是不均匀的。在本文的情况下，所有GPU将共享PCIe带宽，特别是当每个GPU同时发出主机到设备内存复制时。因此，本文调查了从主机内存到相应GPU获取数据的两种不同策略。

第一种策略：最直观的方法是让每个GPU发出自己的HtoD传输，这样可以利用所有可用的DMA引擎，然而，这可能会导致资源争用，因为每个GPU都会尝试消耗一些PCIe带宽以及PCIe通道。

第二种策略：让每个节点中的一个GPU提取所有相关张量，然后使用散列操作将相应的块发送到其他设备。

在本文的分析中，发现尽管在较小的数据大小时多GPU HtoD策略表现更差，由于通道争用的开销，但两种方法的延迟在32k至64k左右被注意力计算超越，因此，它们的差异在序列长度增大时变得微不足道。由于从一个GPU获取并分散到其他GPU需要额外的同步和障碍，本文选择第一种策略，即允许每个GPU获取其数据而不需要任何额外的同步。

4.2 双buffer机制

双buffer机制在CUDA计算优化中比较的流行，如矩阵乘法优化，本文也使用了双buffer机制进一步提高性能。图7说明了buffer工作流程。请注意，这里展示了FPDT块的反向传递，因为正向工作流程更简单，也可以通过检查反向传递来反映。

如上图所示，本文假设总共有4个块，对上图说明如下：

由于在前向传递期间全局序列块qˆi，kˆi，vˆi已被缓存，因此可以直接从主机内存中获取它们，而无需引入额外的Alltoall。

可以使用嵌套循环来计算和更新查询dq ˆi，键dk ˆi和值dv ˆi的梯度。外部循环在键和值上，而内部循环在查询上。这种设计确保了内部循环中的注意力计算只需处理获取下一个查询的延迟，否则，它需要同时处理键和值的预取。

在内部循环中，由于kˆ0和vˆ0被qˆ0，qˆ1，qˆ2，qˆ3使用，因此每次更新dk ˆ0，dv ˆ0。对于dq ˆ0，在第一个内部循环后得到其最终结果，这是由于LLM的自回归特性，即如果i < j，则qi永远不会关注kj。

类似地，在第一个外部循环后得到dk ˆ0和dv ˆ0的最终结果。启动Alltoall将全局序列块分散到其原始GPU，其中dq0，dk0，dv0用于计算输入隐藏状态dc0的梯度。

请注意，这里将Alltoall和投影反向与qˆ1，kˆ1和vˆ1的预取重叠，这将在下一个外部循环中使用。

在本文的实现中，部署了三个CUDA流，因为输入隐藏状态h0的预取将只在投影反向中同步。对于每个内部循环后将不再使用的块缓冲区，可以将其标记为自由内存，可以分配给剩下的块。

以上内容即是DeepSpeed-FPDT整体技术。

参考文档

Training Ultra Long Context Language Model with Fully Pipelined Distributed Transformer

（https://arxiv.org/abs/2408.16978）

大模型产品、算法、应用

58篇原创内容

公众号