DualPipe 详解：一份即使没有分布式训练背景也能理解的 DualPipe 综合指南

社区文章 2025年2月28日发布

上周，DeepSeek 在社交媒体上宣布了“开源周”活动，连续五天发布了一系列开源软件库。在前几天，他们推出了 FlashMLA（高效 Hopper GPU MLA 解码内核）、DeepEP（用于 MoE 的专家并行通信库）和 DeepGEMM（支持 FP8 的 GEMM 库）。第四天，他们一口气开源了三个主要组件：DualPipe、EPLB 和 profile-data。其中，DualPipe 因其“双向流水线并行”的核心思想，引发了广泛讨论。

本篇博文将重点介绍 DualPipe 的核心概念：如何在大型模型训练中完全重叠前向和后向传播，从而大大减少流水线执行中的“气泡时间”。为了帮助您理解这些概念，我们将从一个简单的类比开始——**“机械车间中的工艺优化”**——首先通过机械加工的视角介绍每个概念，然后将其映射回深度学习中的并行训练。阅读完本文后，您将对这些概念的工作原理有一个清晰的认识。我们还将深入探讨 DualPipe 的源代码级别细节，探索它如何进一步减少流水线气泡，重叠前向和后向传播，最小化通信压力，并集成到更复杂的混合并行场景中。

1. 引言

随着 GPT-3、PaLM 和 LLama 等大型语言模型 (LLM) 的盛行，分布式训练已成为突破单 GPU 限制并成功训练超大型模型的必要技术。数据并行、模型并行和流水线并行等术语经常出现，但对于初学者来说，理解它们的区别和联系可能具有挑战性。当您遇到 DeepSeek-V3 中的 DualPipe 等高级功能时，可能会感到更加望而生畏。

另一方面，在工业环境中，优化生产过程通常需要无数次试验，而在 AI 中，训练大型语言模型也同样需要调整大量参数。这两项活动看似无关，但却有着惊人的相似之处。让我们通过托尼的机械车间的故事，看看多台机床的生产线如何帮助我们理解大型模型训练中的四种主要并行类型。

1.1 单 GPU 时代：小型手工作坊

在苏州工业园区，托尼拥有一家专注于优化铸造温度、淬火时间、切削角度等制造工艺的中型机械公司。当新订单到来时，托尼首先设计一套初始制造参数（“工艺手册”），然后进行加工，接着检查最终零件并反向传播调整：如果零件有空洞缺陷，他会调高铸造温度，依此类推。

托尼的工艺与大型模型训练惊人地相似：“工艺手册”就像模型的参数，正在加工的零件类似于训练数据，每个独立的工艺（铸造、热处理等）对应于神经网络中的不同层，机床则类似于 GPU。最后，质量检查类似于计算损失函数，根据检查调整参数就像反向传播。

托尼刚开始创业时，接的订单相对简单，比如加工螺丝。对于这类零件，只需要一台多功能机床就能处理两个步骤：切割和抛光。如果零件抛光不均匀，就调整抛光步骤；如果切割不准确，就调整切割角度。一切都在一台机器上完成——就像一个小型手工作坊。这对于基本零件来说已经足够，但没有可扩展性。

这对应于单 GPU 训练。所有模型层（所有“工艺步骤”）都在同一个 GPU（同一个机器）上运行。前向传播（加工）和反向传播（调整参数）都在单个设备上进行。它简单可靠，但一旦任务变得更加复杂，单个设备就会成为瓶颈。

1.2 模型并行：拆分工艺手册的艺术

有一天，托尼接到一个更大的订单，需要优化发动机曲轴的制造工艺。他很快意识到一台机器无法处理所有必需的步骤。因此，他将工艺（铸造、热处理、精密加工）分摊到三台专用机器上，每台机器都有自己的操作说明。他还必须跟踪调整铸造参数可能对后续工艺产生的影响。这引入了一个新问题——机器空闲时间。当第一台机器忙于铸造时，其他机器可能正在等待。此外，将物品从一台机器移动到下一台机器需要额外的时间。如果规划不仔细，这些机器转换可能会导致额外的空闲期。

在大型语言模型中，这就是模型并行。当一个模型太大，无法容纳在单个 GPU 的内存中时，您需要将其拆分到多个 GPU 上（例如，每个 GPU 负责不同的层或模块）。在这个类比中，铸造就像输入层，热处理是中间层，精密加工是输出层。在训练过程中，每个 GPU 负责模型的一部分，并且必须将中间输出传输给其他 GPU。这不可避免地导致 GPU 之间存在空闲时间以及频繁的跨设备数据传输。托尼面临的问题类似于 GPU 之间的调度和通信挑战。

1.3 数据并行：克隆车间的计划

为了进一步加快工艺参数优化，托尼（现在资金更充足了）在隔壁建造了三个完全相同的车间。每个车间都拥有相同的完整流水线，只是处理不同批次的涡轮盘（数据分片）。一天结束时，四个车间经理会聚在一起交流经验，统一工艺标准。以前需要一个月才能完成的 10000 个原始零件的订单，现在只需要大约两周。托尼纳闷：“为什么我的车间产能翻了两番，速度却只翻了一番？”与经理们交谈后，他发现每批原材料都可能遇到独特的问题，导致一些车间完成得更晚。当它们完成时，它们必须等待最慢的那个，才能总结当天的工作成果。

这就是数据并行。每个车间（GPU）都拥有工艺手册（模型参数）的完整副本，但处理数据（不同的微批次）的不同部分。一旦每个车间计算出其局部梯度（检查零件并找出参数调整），它们必须收集并平均这些梯度（通过 All-Reduce）以形成一组统一的参数。瓶颈在于最慢的车间（落后者 GPU）会拖累所有人，并且随着并行车间数量的增加，通信开销（All-Reduce 带宽）会急剧增加。

1.4 张量并行：协作完成超大型单一组件

有一天，托尼升级了，接到了一个巨大的项目，需要优化飞机零部件的工艺——例如，一个飞机机翼。即使仅仅是机翼本身也大到需要多台相同机器协同完成同一个子步骤。换句话说，尽管这个子步骤属于一个特定的工艺阶段，但它仍然超出了单台机器的能力。因此，托尼让多台相同机器协作完成一个巨大的单一零件。

这类似于大型模型训练中的张量并行。即使将模型拆分到不同的层或模块后，您可能仍然发现单个模块对于一个 GPU 来说太大了。在这种情况下，您需要拆分模块的张量本身——就像将一个大矩阵拆分到多个 GPU 上以执行并行矩阵乘法一样。一旦完成每次部分乘法，您就合并结果。这种方法将单个非常大的层的工作负载分配到多个 GPU 上。

至此，托尼意识到调度已成为提高效率的关键。机器必须密切协作，因为它们中的许多只处理零件的一小部分。该零件可能会在不同机器之间多次移动；后续步骤通常无法在之前的步骤完成之前开始，并且在反馈循环中，早期步骤无法在后期步骤完成自己的参数调整之前最终确定参数调整。所有这些都会导致到处都是空闲期：有些机器等待它们的“兄弟机器”完成，下游进程等待上游结果，上游进程等待反馈。托尼觉得他需要做些什么来进一步提高效率。

这就是张量并行实际的工作方式。当即使单个过程（模型中的单个层）超过单个 GPU 的容量时，您将大张量分成较小的块，以便 GPU 组可以处理。例如，抛光飞机机翼可能需要四台机器在不同的机翼区域工作，然后将它们重新组合在一起。这种级别的协作引入了通信开销（合并部分输出）、同步开销（一台慢机器会拖累其余机器）以及额外的梯度同步“反馈循环”。总的来说，这些都会引入新的空闲时间——**“协作气泡”**。

1.5 流水线并行：让不同机器同时工作

为了应对不同工艺之间的协作挑战，托尼设计了一个巧妙的流水线系统。如果原始加工路线是铸造 → 锻造 → 热处理 → 抛光，那么铸造机一完成第一批次，该批次就立即移交给锻造机；铸造机随即开始处理第二批次。当第一批次从锻造机进入热处理时，第二批次到达锻造机，第三批次则可以送往铸造机——就像多米诺骨牌一样。

不使用流水线并行时的数据流概述：当数据传递到模型的不同部分（在不同的 GPU 上）时，一次只有一个 GPU 工作（无论是前向还是后向）。这使得 GPU 效率非常低，因此不推荐使用。

在实施这个流水线系统之前，每个车间的工作流程如下图所示。在第一批次的生命周期（T1-T4）中，一次只有一台机器工作；当进行检查和参数反馈（T5-T12）时，同样只有一台机器处于活跃状态，而其他机器则空闲（图中的灰色区域）。托尼很快发现了这种低效率：一旦第一批次进入第二阶段，第一阶段就可以处理第二批次，依此类推。因此，流水线变成了

1F1B 流水线并行策略概述。

上图展示了著名的 1F1B（一次前向，一次后向）流水线并行方案。其原则是：当 GPU（或机器）准备好对最近的微批次执行后向传播时，它将优先执行后向传播。例如，在时间 T5，设备 4 必须选择是为第二个批次执行前向传播还是为第一个批次执行后向传播，它优先选择后向传播。同时，所有数据都以微批次顺序进行反向传播；例如，第二个微批次只有在第一个微批次的后向传播开始后才开始其后向传播。此外，每个设备都保留有限数量的前向激活（图中为 4 个），以避免存储过多的中间结果用于后向计算。回到我们的类比，这些中间“激活数据”就像车间存储的制造日志，用于辅助最终的质量评估和参数更新。如您所见，1F1B 流水线中仍然存在空闲时间。在大型模型训练中，这种空闲时间被称为气泡——GPU 等待而不是工作的时间。我们的目标是尽可能减少这些气泡。

经过仔细观察，托尼发现流水线“气泡”的一个关键来源是参数反馈过程所需的时间——几乎是每个批次实际处理时间的两倍。这意味着一旦第一阶段处理完第四批次，它必须等待很长时间才能收到第一批次的反馈更新。为了解决这个问题，他想出了一个新颖的想法：既然反馈消耗了这么多时间，那么将其拆分成两个独立的部分，这样它们就可以解耦。例如，每个过程可能涉及夹具设计和制造设计。如果夹具设计可以根据质量检查报告立即更新，并且独立于制造设计更新，我们就可以消除一些空闲时间。基于这个想法，托尼设计了改进后的流水线

ZB1P 流水线并行策略概述。

在这里，每个批次的反馈被分为两个阶段：夹具设计调整（浅蓝色）和制造设计调整（深蓝色）。在同一个工艺阶段内，相同的反馈阶段之间存在依赖关系（例如，锻造的夹具设计必须在铸造更新其夹具设计之前完成），但不同阶段之间不相互依赖。这种解耦允许每个阶段的夹具设计更新更早完成，而制造设计更新可以延迟，从而与原始的 1F1B 方法相比减少了空闲气泡。

上述方案是 DeepSeek-V3 技术报告中提到的 ZB1P（Zero Bubble 1F1B）基线。它将深度学习中的反向传播分为两个子步骤

输入梯度计算：将梯度从当前层传递回上一层。

参数梯度计算：计算当前层的参数梯度，以便进行更新。

对于一个线性层 $\mathbf{y} = \mathbf{W}\mathbf{x}$ 和损失 $L$ ，一旦我们收到后续层传来的 $\frac{\partial L}{\partial \mathbf{y}}$ ，我们需要计算

$\frac{\partial L}{\partial \mathbf{x}}$ ：将梯度反向传播到前一层；

$\frac{\partial L}{\partial \mathbf{W}}$ ：用于层自身的参数更新。

有趣的是，这两个计算之间没有严格的顺序依赖关系。即使一个层只完成了计算 (1) 但没有计算 (2)，梯度仍然可以传播到前一层。利用这一事实，ZB1P 解耦了 (1) 和 (2)，以便可以尽早执行输入梯度 (1)，而参数梯度 (2) 则推迟，从而大大增加了流水线调度的灵活性和效率。

至此，您已经了解了 DeepSeek-V3 报告中与 DualPipe 进行比较的两种基线流水线方案。下面是它们的效率对比总结表

方法	气泡	激活
1F1B	$(P P - 1) (F + B)$	$1 \times PP$
ZB1P	$(P P - 1) (F + B - 2 W)$	$1 \times PP$

关键参数:

$P P$ ：流水线深度，即并行处理的工艺阶段数量。
$F$ ：前向传播所需时间（例如，每个车间的初始加工）。
$B$ ：后向传播所需时间（例如，每个车间的反馈调整）。
$W$ ：激活数据累积的窗口大小——即反向传播存储的中间激活上限。

从这个表格中可以看出，与 1F1B 相比，ZB1P 在激活内存使用量相同的情况下，大幅减少了气泡。通过解耦后向计算，ZB1P 实现了前向和后向之间更大的重叠，从而减少了空闲时间。更高级的调度策略（如 DualPipe）进一步推动了这一思想，旨在最大化大型模型训练中的资源利用率和并行性能。

简而言之，托尼不断发展的流水线方案——从最初的 1F1B 到 ZB1P，再到 DualPipe——反映了大型语言模型训练的进步。每项新创新都减少了气泡（空闲时间），并推动系统达到更高的性能和更好的资源利用率。

1.6 当流水线仍有“盲点”：ZB1P 的局限性

尽管 ZB1P 的解耦方法显著缩短了空闲时间，但托尼的车间仍然存在一些遗留的间隙。想象一下一个铸造 → 锻造 → 热处理 → 抛光的流水线：一旦铸造机完成其“夹具设计更新”，它就会将该信息传递给锻造机，但更深阶段的“制造设计更新”可能仍然需要等待所有之前的更新完成。由于零件需要经过多个紧密耦合的阶段，即使一个阶段的小延迟也可能波及并产生新的气泡。

在大型模型训练中，ZB1P 确实提高了前向和后向之间的重叠，但它无法实现真正同时的前向和后向传播。为什么？

流水线气泡仍然存在
传统流水线实现通常严格分离前向和后向：首先以向前模式处理所有微批次，然后再进行向后处理。这导致在每个阶段 GPU 空闲。
手动调度复杂
传统流水线方案需要编写大量逻辑：何时发送激活张量？何时接收梯度？如果安排不当，您将面临额外的等待或通信瓶颈。
前向和后向可能互相干扰
随着越来越多的微批次在前向传播中累积，后向传播最终可能会耗尽前向传播的 GPU 资源（例如，内存或带宽）。不良的资源管理可能导致额外的停滞。
前后重叠不足
理想情况下，我们希望其他微批次的前向传播与不同 GPU 上的后向传播同时运行。然而，这样做需要复杂的调度设计。

回到托尼的车间例子：如果在某个时刻铸造机和锻造机都得到了充分利用，但热处理或抛光机由于某种原因降低了产量（类似于 GPU 负载不平衡），您可能很快就会回到“前等后、后等前”的场景。ZB1P 已经提高了调度灵活性，但仍然存在空闲时间的“盲点”。

为了寻求更好的方法，托尼旨在进一步分解前向和后向之间的依赖关系，使它们可以在相邻的流水线节点中完全交错。换句话说，如果一台机器正在处理后向传播，它仍然可以处理下一个微批次的前向传播。这种方法极大地提高了流水线并发性并最小化了停机时间。托尼为他的车间设计的新方案与一种新的模型训练调度策略——DualPipe——本文的主要焦点——并行。

2. DualPipe：双管齐下的流水线策略

2.1 双向调度：将“前后”同时推上生产线

在 传统（单向）流水线（如 1F1B 或 ZB1P）中，一台机器要么执行前向处理，要么等待反馈信号进行后向调整。这两种模式通常是互斥的。然而，在 DualPipe 中，托尼为每台机器配备了“分时模式”和灵活的前后运输系统，允许同一台机器同时处理前向和后向任务。机器可以从“前端”接收新原材料，同时也可以从“后端”收到反馈报告。

通过引入这种双向运输系统，机器可以持续接收新的零件进行处理（前向传播），同时处理来自下游的后向传播更新。在大型语言模型中，这意味着前向和后向传播可以真正并行发生，显著提高 GPU 利用率。

与单向流水线不同，DualPipe 从流水线的两端注入微批次。如果你将流水线想象成一条直线传送带，以前我们只从一侧输入数据，并一直传送到末端。然而，DualPipe 也开启了另一端，因此一台机器可以为来自左侧（上游）的任务进行前向处理，同时从右侧（下游）接收后向梯度。这种方法极大地提高了整体流水线效率并减少了空闲时间。GPU 可以同时处理来自“前端”的前向任务和来自“后端”的后向任务，从而避免了早期方法中导致等待时间的严格单向流。

下表比较了 1F1B、ZB1P 和 DualPipe 在流水线气泡和激活使用方面的表现，说明了每种方法如何解决延迟和调度问题

方法	气泡	激活
1F1B	$(P P - 1) (F + B)$	$1 \times PP$
ZB1P	$(P P - 1) (F + B - 2 W)$	$1 \times PP$
DualPipe	$\left(\frac{PP}{2} - 1\right) (F\&B + B - 3W)$	$2 \times PP + 1$

其中：

$P P$ 是流水线深度（流水线阶段数），
$F$ 和 $B$ 是前向和后向时间，
$W$ 是激活窗口大小。

通过比较三者

1F1B
最简单的单向流水线。气泡时间为 $(P P - 1) (F + B)$ 。每个阶段必须以大部分串行的方式等待。激活存储为 $1 \times PP$ 。
ZB1P
通过解耦后向步骤（拆分输入梯度和参数梯度），气泡时间缩小到 $(P P - 1) (F - 2 W) (PP - 1)(F + B - 2W)$ 。激活使用量仍为 $1 \times PP$ 。
DualPipe
使用双向流水线加上计算和通信的重叠，大大减少了空闲时间。气泡时间进一步降至 $\left(\frac{PP}{2} - 1\right) (F\&B + B - 3W)$ 。然而，它需要增加激活存储，最高可达 $2 \times PP + 1$ ，因为每个 GPU 都持有额外的激活数据以支持同时的前向和后向传播。

简而言之，DualPipe 牺牲了少量内存使用，以换取前向和后向传播之间显著更高的重叠，从而在大型分布式训练中实现更高的吞吐量——就像托尼的车间通过让机器同时向两个方向工作来获得更高的利用率一样。

DualPipe 流水线并行概述。

2.2 GPU 内部的“计算-通信”重叠：数据移动时进行计算

DualPipe 效率的另一个关键是“基于块”或“分段”传输。在大型分布式场景中，计算和通信都可能是主要的耗时瓶颈。如果它们严格地一个接一个地发生（计算 → 通信 → 计算 → 通信），我们将不可避免地在通信期间面临 GPU 空闲，在计算期间面临网络带宽空闲。

DualPipe 将数据传输拆分为更小的块——“微批次流传输”——并将其与部分计算任务交错，因此 GPU 可以在部分数据到达后立即开始计算，而不是等待整个数据集传输完毕。以下是基于块传输至关重要的原因

2.2.1 为什么“分块传输”能提高效率

想象一下托尼的铸造机有 1000 个零件要发送给锻造机。如果采用一次性传输，锻造机必须等到所有 1000 个零件都交付完毕才能开始工作。在此期间，锻造机是空闲的。交付后，锻造机可能很快完成一些任务，然后又开始等待。这会导致“你等我，我等你”的局面。

但是，如果我们将 1000 个零件分成几个小批次（例如 4 或 10 块），锻造机可以在第二批次运输途中立即开始处理第一批次，从而消除空闲时间，确保任务的稳定流动。

在 GPU 中，这类似于将大型张量传输或全对全通信拆分成更小的片段，从而在剩余数据仍在传输时计算部分结果。

DualPipe 中的通信设计概述。

2.2.2 “计算-通信”如何在 GPU 上重叠

资源分区
现代 GPU 拥有多个 SM（流式多处理器）。我们可以分配一些用于处理通信（发送/接收数据包），而另一些则用于处理计算内核。如果通信被分块成更小的消息，计算 SMs 就可以在通信 SMs 处理数据传输的同时保持忙碌。
张量的细粒度拆分
我们不再等待整个模型层的数据传输完成，而是将每层的输出或梯度拆分成更小的片段。GPU 可以在下一个片段仍在传输时对第一个片段进行计算。
并行调度
像 PyTorch 这样的框架支持多流异步操作。您在一个流中发出通信调用（例如，cudaMemcpyAsync 或集合操作），同时在另一个流中继续计算。只有当您真正需要结果时才进行同步。
流水线机制
正如 DualPipe 引入双向流一样，每个 GPU 都既从上一阶段接收数据，又向下一阶段发送更新。通过仔细调度更小的块，从“前向”到“后向”以及从“后向”到“前向”的通信可以与本地计算重叠，从而在微批次和 GPU 之间形成紧密的流水线。

所有这些都实现了我们所说的计算和通信之间的时间重叠——这是 GPU 不会因为等待大型传输完成而空闲的根本原因。

在“单块”场景中，您可能会看到

Time: ======================> Comm: [send all data] -> (wait) -> Compute:[start only after entire data is available]

导致通信期间 GPU 周期浪费，计算期间网络带宽浪费。相比之下，分块

Time:     ==================================>
Info:      [comm chunk1]  [comm chunk2]   ...
Compute:         [compute chunk1]   [compute chunk2]  ...

它们像相互啮合的齿轮一样重叠，大大减少了空闲时间，确保了工作流的持续性。这就是 DualPipe 即使模型大小和并行度增加也能保持高吞吐量的原因。

就像托尼的多阶段车间一样——零件的部分批次来回运输并进行处理，而无需等待整个批次传输——DualPipe 利用基于块的数据流来使 GPU 保持接近完全利用，有效地将通信开销隐藏在并发计算之后。

3. 源代码如何实现“两路并进”：DualPipe 的核心逻辑

在上一节中，我们通过一个制造类比来阐述 DualPipe 如何使前向和后向传播真正在同一个流水线中“同时”运行，从而最大限度地提高 GPU 利用率。在这里，我们将探讨 DualPipe 的 核心源代码 如何实现这些思想。具体来说，我们将了解它如何处理批次拆分、管理通信以及巧妙地交织前向-后向任务，以实现深度重叠和最小的流水线气泡。

请注意，我们不会逐行解析代码。相反，我们将重点关注关键功能和工作流程，并将它们与“机械车间”类比联系起来，以帮助您更直观地理解它们。

3.1 基本类结构和初始化

class DualPipe(nn.Module):
    def __init__(
        self,
        modules: Tuple[nn.Module, nn.Module],
        ...
    ) -> None:
        super().__init__()
        
        ...
        self.module = nn.ModuleList(modules)
        self.overlaped_forward_backward = ...
        ...

modules：此参数是两个 nn.Module 实例的元组，通常表示流水线的“前半部分”和“后半部分”。在我们的机械类比中，这可以对应于一组组合的“前向处理”机器（例如，铸造 + 锻造）和另一组组合（例如，质量检查 + 参数调整）。
overlaped_forward_backward：这是一个专门的函数，用于检查两个 Module 对象是否支持前向-后向重叠。只有当两个 Module 都实现 overlaped_forward_backward 方法（并且类型相同）时，随后的工作流程才会对同一批次应用真正的前向-后向交错。

此外，代码还设置了与分布式训练相关的参数和标志，例如：

self.group 和 self.rank：这些与 torch.distributed 相关联，用于管理当前进程的流水线等级（即阶段）。
self.is_first_rank, self.is_last_rank, self.is_in_second_half：标记此节点（机器）是否处于流水线的“左端”或“右端”，即它是否属于前半部分或后半部分的标志。

这些标志和映射类似于 Tony 可能附加到车间中每台机器的标签，例如“铸造线”、“抛光线”或“倒数第二步”。

3.2 状态管理和重置

下面是 _reset_states 方法：

def _reset_states(self) -> None:
    WeightGradStore.clear()

    self.input_chunks = ([], [])
    self.output_chunks = ([], [])
    self.input_grad_chunks = ([], [])
    self.output_grad_chunks = ([], [])
    self.labels = None
    self.loss_chunks = []
    self.criterion = None

    ...
    # Various counters for tracking which chunk is being processed
    self.current_f_chunk_id = [0, 0]
    self.current_b_chunk_id = [0, 0]
    ...
    self.comm_ops = []
    self.to_free = []

_reset_states 类似于每当 Tony 收到新订单时，清理车间并重新布置工具和记录本。
- WeightGradStore.clear() 移除任何存储的“参数梯度”回调，为零气泡或部分重叠策略做准备。
- input_chunks, output_chunks, input_grad_chunks, output_grad_chunks：这些就像“正在进行中的部件”及其“梯度信息”。将它们初始化为空列表，以便我们可以用每个微批次填充它们并使其移动。
- 各种计数器跟踪前向块、后向块、发送/接收事件等的数量。

3.3 前向和后向计算：在同一设备上实现重叠

在大型模型训练中，“前向计算”和“后向计算”本质上涉及在张量上运行前向或后向传播，然后将梯度传递给前一层。在机械类比中，前向是“加工零件”，而后向是“检查缺陷并调整生产参数”。DualPipe 将这些过程封装在几个方法中：

3.3.1 `_forward_compute_chunk(self, phase)`

def _forward_compute_chunk(self, phase: int) -> None:
    phase ^= self.is_in_second_half  # dynamically correct the phase
    chunk_id = self.current_f_chunk_id[phase]
    self.current_f_chunk_id[phase] += 1
    inputs = self.input_chunks[phase][chunk_id]
    ...
    outputs = self.module[phase](*inputs)
    ...

在这里，phase 被调整以选择实际使用哪个子模块（流水线的“前”或“后”部分）。然后，我们从 input_chunks[phase] 中获取批处理数据并对其运行前向计算。
结果 outputs 将存储在 self.output_chunks[phase] 中。
如果这是最后阶段（is_last_stage）并且定义了 criterion，则“损失”将放置在 self.loss_chunks 中。可以将其视为最终车间输出“缺陷分数”以供检查。

3.3.2 `_backward_compute_chunk(self, phase, enable_zb: bool = False)`

def _backward_compute_chunk(self, phase: int, enable_zb: bool = False) -> None:
    if self.forward_only:
        return
    phase ^= self.is_in_second_half
    chunk_id = self.current_b_chunk_id[phase]
    ...
    if is_last_stage:
        # at the last stage, directly call backward on the loss
        loss = self.loss_chunks[chunk_id]
        loss.backward()
        loss.detach_()
    else:
        # run_backward on outputs + output_grads
        outputs = self.output_chunks[phase][chunk_id]
        ...
        run_backward(outputs, output_grads)
    ...
    if enable_zb:
        WeightGradStore.flush()

    # update input_grads
    ...

对于后向传播，主要逻辑是：如果处于最终阶段，则对 loss 调用 backward()；否则，对中间输出调用 run_backward() 以将梯度传递给上游。
enable_zb 激活零气泡（例如，ZB1P）方法，其中一些参数-梯度计算通过将它们放入 WeightGradStore 并在适当的时候刷新它们来延迟。这与我们之前对解耦“输入梯度计算”和“参数梯度计算”的解释相符。
后向传播完成后，上游阶段的“梯度”将放置在 self.input_grad_chunks[phase] 中，类似于 Tony 将一些缺陷报告返回给前一台机器。

3.3.3 `_forward_backward_compute_chunk(self, phase0, phase1)`

def _forward_backward_compute_chunk(self, phase0: int, phase1: int) -> None:
    if self.forward_only:
        self._forward_compute_chunk(phase0)
        return

    if not self.overlaped_forward_backward:
        self._forward_compute_chunk(phase0)
        self._backward_compute_chunk(phase1)
        return

    # 1) pre-forward
    # 2) pre-backward
    # 3) overlaped_forward_backward(...)
    # 4) post-forward
    # 5) post-backward

DualPipe 的真正核心：如果 overlaped_forward_backward 为 True，则同一 GPU 可以融合某些批次的前向和后向传播（就像“两只手协同工作”）。

该函数首先从 input_chunks 中获取前向数据，从 output_chunks 中获取后向数据，然后调用 module0.overlaped_forward_backward(...)。
- 在车间类比中，操作员收集要加工的新一批零件，以及来自先前批次的缺陷报告和调整，然后使用单个“组合机器操作”来处理前向和后向任务。
它最终将新输出和梯度存储到 output_chunks 和 input_grad_chunks 中。这意味着部分前向-后向步骤可以在同一阶段内发生，而不是单独调用它们或等待另一个阶段（如 ZB1P）。

3.4 通信管理：将“传输时间”隐藏在计算之后

在大型模型流水线并行中，每个阶段的 GPU 经常需要与上游/下游 GPU 通信（例如，将半成品从铸造运到锻造）。DualPipe 使用几个函数来分解和调度通信，以便计算和通信广泛重叠：

_recv_forward(self, phase) / _send_forward(self, phase)
用于接收/发送“前向输出”或“前向输入”。在车间类比中，将部分加工的零件交给下一台机器。
_recv_backward(self, phase) / _send_backward(self, phase)
用于接收/发送“后向梯度”或“后向输入”，类似于在机器之间传递检查报告。
_commit_and_wait_comm(self)
内部调用类似 dist.batch_isend_irecv(self.comm_ops) 的函数来启动所有非阻塞通信，然后 wait()。这使得计算可以与通信并行进行，因此“运输时间”被“隐藏”在机器空闲或可以将资源用于传输的时间段内。

3.5. WeightGradStore：延迟梯度更新设计

在后向传播过程中，相同的权重可能会多次累积梯度。WeightGradStore 使用静态队列来保存这些更新，然后在适当的时候集体执行它们。优点：

更少的同步和内存写入：您可以在每次微批处理时批量更新参数，而不是对每个微批处理进行参数更新，可以在适当的时候批量更新它们。
流水线集成：避免中断流水线并发，并在与其他任务重叠的时间进行同步。

    class WeightGradStore:
        enabled: bool = False
        cache: List[Callable] = []
        funcs_queue = queue.Queue()

        @classmethod
        def put(cls, func: Callable) -> None:
            cls.cache.append(func)

        @classmethod
        def flush(cls) -> None:
            cls.funcs_queue.put(cls.cache)
            cls.cache = []

        @classmethod
        def pop(cls) -> None:
            funcs = cls.funcs_queue.get()
            for func in funcs:
                func()

注意：phase ^= self.is_in_second_half 技巧是一种巧妙地根据流水线等级是否在后半部分“翻转”阶段的方法，允许同一函数处理从左到右和从右到左的数据传输。

3.6 整体调度：`step(...)` 方法的 8 个阶段

核心逻辑存在于 step(...) 中。这个函数就像 Tony 的中央命令，它协调所有机器。为了实现 DualPipe 的双向流水线，它会经过这些阶段（简化视图，详情请参阅代码注释）：

第一步：nF0
在流水线启动时，让一端首先处理一定数量的前向批次。就像在右侧空闲时，左侧部分“预处理”原材料。
第二步：nF0F1
逐渐让另一端开始前向操作，以便两个方向都激活——交替调用 _forward_chunk(0) 和 _forward_chunk(1)。
第三步：nB1W1F1
当某些批次出现后向传播时，将后向传播（_backward_chunk(1)）与前向传播（_forward_chunk(1)）混合，再加上 _weight_chunk() 用于延迟权重更新（与 ZeroBubble 相关）。
第四步：nF0B1F1B0（主循环）
DualPipe 的主要“交错”步骤：_forward_backward_chunk(0,1) / _forward_backward_chunk(1,0) 用于同时调度前向和后向。
第五步：nB1F1B0
继续进行后向+前向交错。
第六步：nB1B0
更侧重于后向，让 WeightGradStore 刷新部分梯度数据。
第七步：nWB0
更多的“权重更新+后向”循环。
第八步：nW
最终的权重更新，确保所有任务都干净利落地完成。

在车间类比中，这是一个包含多个时间段的“调度时间表”：首先让左侧机器进行一些初始运行，然后到某个点后，右侧也输入原材料；两侧不断传递半成品零件和缺陷报告，直到所有工作完成，最后进行统一和完成步骤。

多个步骤可能看起来很复杂，但这是因为双向流水线需要仔细管理不同阶段的微批处理以保持气泡最小。

3.7 总结

从 step(...) 编排到 _forward_backward_compute_chunk(...) 的“前向和后向融合”，您可以看到 DualPipe 如何利用细粒度分区和双向注入。在代码方面，它极大地隐藏了通信开销，并使单个 GPU 能够在适当的时机运行前向和后向。

回到车间类比：在同一个时间段内，机器 A 可能正在处理新的原材料（前向），而机器 B 正在处理来自前一批次的缺陷反馈（后向），并且“运输卡车”（通信）主要在那些空闲或可以分配资源进行传输的时候快速移动——减少了每台机器的空闲时间。

其优点显而易见：

最小的流水线气泡：DualPipe 真正实现了“前后并发”，降低了串行等待时间。
广泛的通信重叠：通过 _commit_and_wait_comm() 和异步发送/接收，大部分跨节点的全对全或流水线通信在 GPU 忙于计算时完成。
可扩展性：即使集群增长（更多 GPU，更多节点），只要计算-通信比管理得当，这种重叠仍然有效。
内存优化：将最浅层和最深层放置在与零气泡相同的流水线等级上，有助于减少中间层存储开销。

这种“两路并进”的解决方案显著加快了需要大量跨节点专家并行通信的大型模型训练（例如 MoE）。它允许您在有限的 GPU 资源（如 H800）上训练千亿参数的网络，而不会陷入停滞。

4. 结论与展望

从机械车间类比到大型语言模型并行训练，我们涵盖了无并行、模型并行、数据并行和流水线并行。在流水线并行中，最困难的部分是最小化“流水线气泡”、减少通信等待时间以及重叠前向-后向执行。DualPipe 是一种高级解决方案，它巧妙地安排微批处理和前向-后向时序，以达到“零气泡”理想（或接近理想），极大地提高了流水线训练速度和利用率。

简而言之，DualPipe 不仅在概念上重叠了前向和后向传播，而且还封装了通信和调度细节，让用户更容易地进行复杂的流水线并行训练。对于那些希望深入研究大型模型分布式训练的人来说，研究其源代码（加上实践测试）对于开发更高效的并行解决方案非常有益。

总而言之：就像一个精心编排的装配车间一样，DualPipe 同步前向和后向处理以协同进行，极大地提高了多 GPU 效率，并为大型模型时代奠定了重要的技术基础。

社区

通过拖放到文本输入框、粘贴或点击此处上传图片、音频和视频。

点击或粘贴此处以上传图片

· 注册或登录以评论