训练期间的前向和后向计算

在深度学习中，前向传播和后向传播是训练神经网络的两个主要步骤。在前向阶段，输入数据逐层通过网络，每层应用变换然后产生输出。在后向阶段，使用损失函数（例如交叉熵、MSE 等）计算损失，反向传播计算损失相对于权重的梯度，使用链式微分法则逐层向后传播梯度。然后使用优化算法（例如 Adam）更新权重。

在进入下一部分之前，我们应该了解一些事情：

对于一个输入样本，后向计算是在前向计算之后。
对于一个输入样本，在前向阶段，layerN 首先处理它，然后传递给 layerN+1；但在后向阶段，layerN+1 首先处理它，然后传递给 layerN。
通常我们向神经网络输入一批样本，然后更新一次网络权重。

大规模分布式训练的挑战

当模型较小并且可以装入单个 GPU 内存时，它非常简单高效：输入样本、前向计算、后向计算等，GPU 计算资源可以得到充分利用。为了加速使用更多数据训练这种小型到中型模型，可以使用数据并行（例如PyTorch Distributed — 2020 Meta AI）等技术将整个模型复制到多个计算资源（GPU），独立生成梯度，然后在每次迭代时传达这些梯度以保持模型副本的一致性。

对于大型模型，比如具有数千亿个参数的 LLM，该模型太大而无法放入单个 GPU，我们需要对模型进行分区并将其分布到具有数百或数千个 GPU 的大型集群中（例如，据报道，具有 175B 个参数的 GPT-3 模型在 10,000 个 A100 GPU 上进行训练），这种技术称为模型并行性。我们可以按层对模型进行分区（例如，GPT-3 175B 有大约 96 层），并将不同的层分布到不同的 GPU。如果单个层仍然太大而无法放入单个 GPU 内存（此外，参数，优化器状态和梯度在训练期间也需要内存），仍然可以将该单层分成几部分，并将每部分分配给专用的 GPU。还记得我在第一篇博客中介绍的专家混合吗？ Mixture-of-Experts 已经将前馈层分成了独立的分区，我们可以将这些专家分布在不同的 GPU 上，每个 Attention 层都有多个 head（GPT-3 175B 有 96 个 head），我们可以将这些 head（及其 MatMul 计算）分布在不同的 GPU 上。而且我们可以从张量级别进一步划分模型。

数据并行和模型并行并不相互排斥，它们可以一起使用来加速大规模模型训练过程。通过对模型进行分区应用模型并行后，模型的一个分区可以复制并分发到多个 GPU 并应用数据并行。

到目前为止一切看起来都很好？现在让我们来谈谈大规模分布式训练的挑战。在模型和计算分布在数百或数千个 GPU 上之后，如何充分利用这些大规模计算资源是一个巨大的挑战（还有其他挑战，例如如果一台机器在训练期间崩溃了怎么办，因为单个集群中有数千个节点，节点崩溃的概率高于小集群，如何从检查点恢复这部分，例如从空闲的 HA 池中添加新的 GPU / 节点回集群或重新平衡模型和训练任务以适应拓扑等。在这篇博客中，我们重点介绍如何充分利用资源。）。正如我在上一节中所描述的，前向计算和后向计算之间存在依赖关系，层的计算之间存在依赖关系，模型已跨不同节点进行分区。需要跨节点进行数据传输和通信。由于依赖于其他节点的计算结果，某些节点可能处于空闲状态，而计算结果仍在进行中；梯度、某一层的输出和其他数据可能需要从一个节点传输到另一个节点，这可能会达到网络带宽瓶颈并导致某些节点空闲。可能有些步骤（例如，优化器步骤）需要同步，也会导致某些节点空闲。

零气泡流水线并行性

为了在深度神经网络训练期间充分利用集群的计算资源，引入了几种管道并行技术。PipeDream ：快速高效的管道并行 DNN 训练 — 2018（微软、卡内基梅隆大学、斯坦福大学）这篇论文引入了“一前一后”（1F1B）调度策略，通过管道重叠通信和计算来提高集群的 GPU 利用率。下图中的 1、2、3、4 表示不同的训练数据小批量。

1F1B管道时间表

零气泡流水线并行性——2023 Sea AI指出，后向计算实际上包含两个部分：计算关于输入 x (B) 和层参数 W (W) 的梯度。1F1B 策略将 B 和 W 合并为 B，但不必要地增加了顺序依赖的计算。因此，零气泡流水线将 B 和 W 分成不同的阶段，以减少流水线中的气泡。并用更新后验证替换事前同步，以进一步减少优化器步骤中的气泡（下图底部）。

零泡沫管道计划

DeekSeek 中的 DualPipe

DeekSeek（从 V3 开始）引入了 DualPipe Schedule，其思想与 Zero Bubble Pipeline Schedule 类似，但进行了一些额外的更改，以进一步提高计算与通信的比率和效率：

细粒度阶段：将每个块分为 4 个部分：注意力、全对全调度（处理设备之间的通信）、MLP（多层感知器）、全对全组合（跨设备合并输出）。对于后向块，注意力和 MLP 进一步分为两个部分：后向输入（B）和后向权重（W），如 Zero Bubble。
双向流水线调度可同时从流水线两端提供微批次，并且大部分通信可以完全重叠（参见下图中的 2 个黑色箭头）。为了支持双向流水线调度，DualPipe 需要保留两份模型参数。如果我们有 8 个设备和 8 层模型，则在零气泡调度中，每个设备都有一个相应的层。但在 DualPipe 调度中，为了处理双向流水线，设备 0 应该具有模型的 layer0 和 layer7，而设备 7 应该具有模型的 layer7 和 layer0。