DeepSeek开源周 Day04：从DualPipe聊聊大模型分布式训练的并行策略

DualPipe简介

今天是DeepSeek开源周的第四天，官方开源了一种新型并行计算优化策略——DualPipe。 其实大家阅读过Deepseek-V3技术报告的同学，对这个技术并不陌生。

开源地址：https://github.com/deepseek-ai/DualPipe

核心亮点

• DualPipe：双向流水线并行算法
  DualPipe通过创新的流水线设计，让前向传播和反向传播的计算任务可以在独立的管道中并行执行，大幅提高计算效率。

• 完全计算-通信重叠
  DualPipe不仅能并行计算，还能实现计算和通信阶段的完美重叠，极大减少了GPU的空闲时间。

• 高效的内存使用与低气泡时间
  与传统方法相比，DualPipe显著减少了内存占用，并优化了流水线气泡问题，使得大规模分布式训练更加高效。

在今天本次分享中，我们将重点介绍如何通过双向流水线并行算法提升V3/R1训练中的计算-通信重叠效率，彻底改变分布式训练的速度和资源利用率。

在今天本次分享中，常见流水线并行技术，可能实际大模型训练过程中直接调包用就行了，毕竟开源的太硬核，不接触Infra底层优化使用场景比较少了。

首先，我们先了解训练巨大模型的显存和计算效率的挑战，为什么要有流水行并行。

显存与计算效率挑战

在训练巨大模型时，显存和计算效率面临着巨大的挑战，特别是在大规模模型（例如万亿参数模型）的训练过程中。以下是这两个方面的具体问题：

显存效率

训练万亿参数的模型对显存的需求极为庞大，远超过单个GPU的承载能力。例如，使用Adam优化器进行混合精度训练时，仅仅保存模型状态（包括参数、梯度和优化器状态）就需要大约12TB的显存空间。以英伟达A100为例，其显存容量为40GB，这意味着仅为存储模型状态就需要多达400张A100 GPU。

此外，在前向传播过程中，模型的中间激活值需要被保存，直到反向传播完成并计算出损失函数后才会释放。即使将批量大小（batch size）设为1，训练一个万亿参数的模型也会产生超过400GB的激活显存需求。虽然可以通过Checkpoint技术（以计算时间换取显存空间）来处理部分激活显存问题，但整体显存需求依然巨大。

为避免显存溢出，必须通过分布式训练策略将模型状态和激活显存有效地分散到多个GPU设备上，这是实现大规模模型训练的关键。

计算效率

训练一个万亿参数的模型估计需要消耗约5000 Zflops的计算量。这一庞大的计算需求意味着，即使使用4000张A100 GPU并以50%的计算效率运行，也需要约100天才能完成训练。

尽管大型GPU集群可能配备超过4000个GPU，但由于批量大小的限制，要在如此规模上实现高效计算仍然面临挑战。计算效率与计算时间和通信时间的比率相关，而这个比率又受到批量大小的直接影响。更大的批量有助于提高计算效率，但若批量大小过大，则可能影响模型的收敛性。以GPT-3为例，其训练批量大小约为1500，而若将批量大小提高至4000，单张GPU的批量大小仍然只有1，这将严重限制训练的扩展性。

因此，如何在保证模型收敛的前提下，优化批量大小和计算效率，是大规模训练中的核心问题。

流水线并行

流水线并行性使得训练无法放入单个 GPU 内存的大型模型成为可能。

例如：Huggingface 的BLOOM模型是一个 175B 参数的 Transformer 模型。将权重存储为 bfloat16 需要 350GB，但他们用于训练 BLOOM 的 GPU 内存只有 80GB，而训练所需的内存远不止加载模型权重。因此，他们的最终训练分布在 384 个 GPU 上。

这是通过将模型的不同层分配给不同的 GPU 来实现的，这一过程称为模型分区。如果模型分区的实现方式简单，则会导致 GPU 利用率低。

下面我们介绍一些流水行并行方式。

朴素模型并行-Naive PP

朴素模型并行（Naive Model Parallelism） 是指将模型层组分布在多个 GPU 上。每当数据进出时，会切换到与模型层相同的设备，其余部分保持不变。

例如，下图显示一个 8 层模型：

我们将模型垂直切分为两部分，0-3 层放在 GPU0 上，4-7 层放在 GPU1 上。在 0 到 3 层的数据传输是常规的计算过程，但当数据需要从第 3 层传输到第 4 层时，就涉及到 GPU0 到 GPU1 的跨设备通信，这会产生通信开销。如果 GPU 位于同一计算节点（例如同一台物理机器），则通信速度较快，但如果 GPU 位于不同的计算节点（如多台机器），则通信开销会显著增加。接下来，4 到 7 层的计算过程与正常模型一致，在第 7 层完成后，我们需要将数据发送回标签所在的 0 层（或将标签发送到最后一层），然后计算损失并开始优化。

下面我们再看一个4卡并行的时序图

模型通过层将其垂直分成4个部分。Worker 1负责网络的第一层（离输入最近），而Worker 4负责第4层（离输出最近）。图中的“F”、“B”和“U”分别表示前向传播、反向传播和更新操作。下标表示操作在哪个Worker上运行。由于数据处理需要一个一个地按顺序进行，这会导致每个Worker之间出现很大的“空闲时间”。

同步流水线并行-GPipe

GPipe（Easy Scaling with Micro-Batch Pipeline Parallelism），由谷歌提出的一种流水线并行方案。Gpipe 流水线并行主要用来解决这两个问题：

第一，提高模型训练的并行度。 Gpipe 在朴素流水线并行的基础上，利用数据并行的思想，将 mini-batch 细分为多个更小的 micro-batch，送入GPU进行训练，来提高并行程度。

第二，通过重计算（Re-materialization）降低显存消耗。 在模型训练过程中的前向传播时，会记录每一个算子的计算结果，用于反向传播时的梯度计算。

来自 GPipe论文的插图上半部分显示了Naive MP，下半部分显示了 PP：

从下半部分中可以很容易地看出 PP 的死区较少 - 死区指GPU处于空闲状态，空闲部分被称为“bubble”（气泡）。

图的两部分都显示了 4 级并行性。也就是说，有 4 个 GPU 参与流水线。因此，有 4 个管道阶段 F0、F1、F2 和 F3 的前向路径，然后有 B3、B2、B1 和 B0 的后向路径。

PP 引入了一个新的超参数来调整，即块chunks，它定义了通过同一管道阶段按顺序发送多少个数据块。例如，上图下半部分中，你可以看到chunks = 4。GPU0 对块 0、1、2 和 3（F0,0、F0,1、F0,2、F0,3）执行相同的前向路径，然后等待其他 GPU 完成其工作，并且只有当它们的工作开始完成时，GPU0 才会再次开始工作，对块 3、2、1 和 0（B0,3、B0,2、B0,1、B0,0）执行后向路径。

使用 chunks=1 时，你最终会得到Navie MP，这是非常低效的。使用非常大的 chunks 值时，你最终会得到非常小的微批次大小，这可能也不是非常高效。因此，必须进行实验才能找到让 GPU 达到最高利用率的值。

简而言之，GPipe 通过纵向对模型进行切分解决了单个设备无法训练大模型的问题；同时，又通过微批量流水线增加了多设备上的并行程度，除此之外，还使用re-materialization降低了单设备上的显存峰值。

F-then-B 策略

F-then-B 模式，先进行前向计算，再进行反向计算。F-then-B 模式由于缓存了多个 micro-batch 的中间变量和梯度，显存的实际利用率并不高。

来源：https://juejin.cn/post/7262274383287484476

1F1B 策略

1F1B （在流水线并行中，pipeline stage 前向计算和反向计算交叉进行的方式）流水线并行方式解决了这个问题。在 1F1B 模式下，前向计算和反向计算交叉进行，可以及时释放不必要的中间变量。

来源：https://juejin.cn/post/7262274383287484476

异步并行流水线并行-PipeDream

尽管 PipeDream 与 GPipe 同期发布，但其并行思想与 GPipe 截然不同。PipeDream 将前向传递的执行流水线化，并将其与后向传递穿插在一起，以期最大限度地提高硬件利用率和吞吐量。它将小批量连续插入到流水线中，并在后向传递后异步更新参数。

PipeDream 和 GPipe 之间的区别很明显：PipeDream 应用异步后向更新，而 GPipe 应用同步后向更新。

如上图所示，PipeDream为了确保任何时候没有GPU处于空闲状态，这个框架会将多个小批次数据依次注入到流水线中。在完成一个小批次的前向传播后，每个阶段会异步地将输出的激活值传递给下一个阶段，同时开始处理下一个小批次的数据。 类似地，在完成一个小批次的反向传播后，每个阶段会异步地将输出的梯度传递给前一个阶段，同时开始计算下一个小批次的数据。这样，不同的GPU可以同时处理不同的小批次，避免了空闲等待的情况。

1F1B (上) vs 1F1B 交错模式 (下)

交错模式为了进一步减少流水线空闲时间，每个设备可以计算多个不连续的层集合，称为模型块，而不是一个连续的层集合。例如，设备1不再处理第1到第4层，设备2也不再处理第5到第8层，而是将每个设备分配两个模型块，每个模型块包含两层。这样，设备1处理的层是1、2、9、10，设备2处理的层是3、4、11、12。这种方案将多个流水线阶段分配给每个设备，每个阶段处理的层数比单个连续集合少，从而进一步减少了流水线空闲时间的开销。

总结来说，PipeDream中使用了1F1B（1 Forward 1 Backward）设计，表示一个模块交替执行前向传递和反向传递，从图中不难看出，1F1B的气泡占比不变，但同时因为更早的进行了Backward，显著降低了activation的显存占用。

张量并行-TP

张量并行训练是将一个张量沿特定维度分成 N 块，每个设备只持有整个张量的 1/N，同时不影响计算图的正确性。这需要额外的通信来确保结果的正确性。

在张量并行中，每个 GPU 仅处理张量的一部分，并且仅为需要整个张量的操作聚合整个张量。

按照 Megatron 的论文符号，我们可以将其点积部分写为Y = GeLU(XA)，其中X和Y分别是输入和输出向量，A是权重矩阵。

如果我们以矩阵形式查看计算，很容易看出矩阵乘法如何在多个 GPU 之间分配：

如果我们将权重矩阵A按列拆分到N各个 GPU 并并行执行矩阵乘法XA_1，XA_n那么我们最终会得到可以独立输入的N输出向量： Y_1, Y_2, …, Y_nGeLU

$$[Y_1,Y_2]=[\text{GeLU}(X{A}_1),\text{GeLU}(X{A}_2)]$$

利用这一原理，我们可以更新任意深度的 MLP，而无需 GPU 之间进行任何同步，直到最后，我们才需要从碎片中重建输出向量。Megatron-LM 论文作者为此提供了一个有用的例子：

并行化多头注意力层更加简单，因为它们具有多个独立的头，本质上已经是并行的！

注意：TP 需要非常快的网络，因此不建议在多个节点上进行 TP。

数据并行-DP

数据集分为n块，每块随机分配到m个设备(worker)中，相当于m个batch并行训练n/m轮，模型也被复制为n块，每块模型均在每块数据上进行训练，各自完成前向和后向的计算得到梯度，对梯度进行更新，更新后，再传回各个worker。以确保每个worker具有相同的模型参数。

3D并行-DP+PP+TP

3D并行是由数据并行(DP)、张量并行(TP)和流水线并行(PP)组成。将这三者结合起来，将得到一个3D网格，其中每个参数、优化器状态等都映射到每个GPU上。这就是3D并行。

参考资料

• 一文解析 DeepSeek 大模型高效训练背后的极限 AI 工程优化，看完这一篇你就懂了！！
• DeepSeek V3 精读(2)-DualPipe
• DeepSeek-V3 关键点解读-Infra&硬件篇
• Techniques for training large neural networks
• Pipeline-Parallelism: Distributed Training via Model Partitioning
• 图解大模型训练之：流水线并行（Pipeline Parallelism），以Gpipe为例
• 大模型训练 Pipeline Parallel 流水并行性能分析
• 大模型分布式训练并行技术（三）-流水线并行
• 并行技术