入门下大模型训练中的MP、PP、TP等操作

最新推荐文章于 2025-05-23 08:36:32 发布

大模型面试

最新推荐文章于 2025-05-23 08:36:32 发布

阅读量1.3k

点赞数 13

文章标签：深度学习人工智能 AI大模型语言模型 ai LLM 大模型训练

本文链接：https://blog.csdn.net/Code1994/article/details/145209389

版权

近年来，我们在许多 NLP 基准任务上看到了更好的结果，这得益于更大规模的预训练语言模型。训练大型且深度的神经网络是一个挑战，因为它需要大量的 GPU 内存和漫长的训练时间。然而，单个 GPU 工作器的内存有限，许多大型模型的大小已经超出了单个 GPU 的容量。有几种并行范式可以实现在多个 GPU 上进行模型训练，以及各种模型架构和内存节省设计，以帮助实现训练真正大型神经网络的可能性。

一、训练并行化

训练非常大型神经网络模型的主要瓶颈是对大量 GPU 内存的强烈需求，这远远超出了单个 GPU 机器所能承载的范围。除了模型权重（例如数十亿个浮点数）之外，存储中间计算输出（如梯度和优化器状态，例如 Adam 中的动量和变化）通常成本更高。此外，训练大型模型通常伴随着大型训练语料库，因此单个进程可能需要无限长的时间。因此，需要并行化。并行化可以在不同的维度上发生，包括数据、模型架构和张量操作。

1、数据并行化

数据并行化（DP）最简单的方法是将相同的模型权重复制到多个工作器中，并将数据的一部分分配给每个工作器同时处理。

如果模型大小超过了单个 GPU 节点的内存，简单的 DP 就无法很好地工作。像 GeePS（Cui 等人，2016）这样的方法将暂时不用的参数卸载回 CPU，以便在模型太大而无法放入一台机器时与有限的 GPU 内存一起工作。数据交换传输应该在后端进行，不要干扰训练计算。

在每个小批量结束时，工作器需要同步梯度或权重，以避免过时。有两种主要的同步方法，各有优缺点。

批量同步并行（BSP）：工作器在每个小批量结束时同步数据。它防止模型权重过时，具有良好的学习效率，但每台机器都必须暂停并等待其他机器发送梯度。
异步并行（ASP）：每个 GPU 工作器异步处理数据，无需等待或停滞。然而，这很容易导致使用过时的权重，从而降低统计学习效率。尽管它增加了计算时间，但可能不会加快收敛到训练时间的速度。

介于两者之间的是每n 次迭代同步一次全局梯度。这个特性被称为“梯度累积”，自 Pytorch v1.5（Li 等人，2021）起在分布式数据并行（DDP）中实现。桶装梯度避免了立即的 AllReduce 操作，而是将多个梯度放入一个 AllReduce 中，以提高吞吐量。可以根据计算图进行计算和通信调度优化。

在这里插入图片描述

2、模型并行化

模型并行化（MP） 旨在解决模型权重无法放入单个节点的情况。计算和模型参数分布在多台机器上。与数据并行化不同，数据并行化中每个工作器托管整个模型的完整副本，MP 只在一个工作器上分配模型参数的一部分，因此内存使用和计算都减少了。

由于深度神经网络通常包含一系列垂直层，因此按层分割大型模型似乎很直接，将一小部分连续层分组为一个分区放在一个工作器上。然而，一个简单的实现，即让每个数据批次通过多个这样的工作器，存在顺序依赖，会导致大量等待时间的“气泡”和计算资源的严重未充分利用。

在这里插入图片描述

图 2. 一个简单的模型并行化设置，模型垂直分割为 4 个分区。由于顺序依赖，数据一次只能被一个工作器处理，导致大量闲置时间的“气泡”。（图片来源：Huang 等人，2019）

3、管道并行化

管道并行化（Pipeline Parallelism), 结合了模型并行化和数据并行化，以减少低效的时间“气泡”。主要思想是将一个小批量分割成多个微批次，并使每个阶段工作器同时处理一个微批次。注意，每个微批次需要两次传递，一次前向和一次反向。工作器之间的通信仅传输激活（前向）和梯度（反向）。这些传递的调度方式以及梯度的聚合方式在不同方法中有所不同。分区（工作器）的数量也称为管道深度。

在 GPipe（Huang 等人，2019）中，多个微批次的梯度在结束时同步聚合并应用。同步梯度下降保证了学习的一致性和效率，无论工作器的数量如何。如图 3 所示，仍然存在气泡，但比图 2 中的要小得多。假设微批次均匀分割，且每个微批次的前向和反向传递各需一个时间单位，气泡的比例为：

GPipe 论文观察到，如果微批次的数量超过分区数量的 4 倍（应用激活重计算时），气泡开销几乎可以忽略不计。

在这里插入图片描述

图 3. GPipe 中管道并行化的示意图，包含 4 个微批次和 4 个分区。GPipe 在每个批次结束时同步聚合和更新跨设备的梯度。（图片来源：Huang 等人，2019）

GPipe 随着设备数量的增加几乎实现了线性吞吐量加速，尽管如果模型参数在工作器之间分布不均，则不总是能保证这一点。

PipeDream（Narayanan 等人，2019）安排每个工作器交替处理前向和反向传递（1F1B）。 PipeDream 将每个模型分区称为“阶段”，每个阶段工作器可以有多个副本以运行数据并行化。在此过程中，PipeDream 使用确定性轮询负载均衡策略在阶段的多个副本之间分配工作，以确保同一小批量的前向和反向传递发生在同一副本上。

在这里插入图片描述

图 4. PipeDream 中 1F1B 微批次调度的示意图。（图片来源：Harlap 等人，2018）

由于 PipeDream 在所有工作器之间没有每个批次结束时的全局梯度同步，1F1B 的原生实现很容易导致一个微批次的前向和反向传递使用不同版本的模型权重，从而降低学习效率。PipeDream 提出了一些设计来解决这个问题：

权重存储：每个工作器跟踪几个模型版本，并确保在给定一个数据批次时，前向和反向传递使用相同版本的权重。
垂直同步（可选）：模型权重的版本与激活和梯度一起在阶段工作器之间流动。然后计算采用从上一个工作器传播的相应存储版本。这个过程保持了工作器之间的版本一致性。注意，这是异步的，与 GPipe 不同。

在训练开始时，PipeDream 首先分析模型中每一层的计算内存成本和时间，然后优化一个将层分割为阶段的解决方案，这是一个动态规划问题。

在这里插入图片描述

图 5. VGG16 在 ILSVRC12 上的结果。（顶部）准确率与时间。整数标记阶段工作器的数量。ASP = 异步并行 & BSP = 批量同步并行。（底部）不同并行化配置的训练时间加速。（图片来源：Harlap 等人，2018）

后来提出了两种 PipeDream 的变体，以减少存储模型版本所需的内存占用（Narayanan 等人，2021）。

PipeDream-flush 周期性地添加一个全局同步的管道刷新，就像 GPipe 一样。通过这种方式，它大大减少了内存占用（即只维护一个模型权重版本），以牺牲一些吞吐量为代价。

图 6. PipeDream-flush 中中管道调度的示意图。（图片来源：Narayanan 等人，2021）

PipeDream-2BW 仅维护两个模型权重版本，“2BW”是“双缓冲权重”的缩写。它每 n 个微批次生成一个新模型版本，且 n 应该大于管道深度 P。一个新更新的模型版本不能立即完全替换旧版本，因为一些剩余的反向传递仍然依赖于旧版本。总共只需要保存两个版本，因此内存成本大大降低。

在这里插入图片描述

图 7. PipeDream-2BW 中管道调度的示意图。（图片来源：Narayanan 等人，2021）

4、张量并行化

模型和管道并行化都是垂直分割模型。另一方面，我们可以水平分割一个张量操作的计算到多个设备上，称为张量并行化（TP）。

以 Transformer 为例，因其流行度而备受关注。Transformer 模型主要由 MLP 层和自注意力块组成。Megatron-LM（Shoeybi 等人，2020）采用了一种简单的方法来并行化 MLP 和自注意力的层内计算。

Transformer 中的 MLP 层包含一个 GEMM（通用矩阵乘法）后跟一个非线性 GeLU 转换。让我们按列分割权重矩阵 A得到

自注意力块根据上述分割并行运行 GEMM，使用Query、Key和Value），然后通过另一个 GEMM 结合它们以产生注意力头结果。

在这里插入图片描述

图 8. Megatron-LM 提出的 Transformer 关键组件的张量并行化示意图。（图片来源：Shoeybi 等人，2020）

Narayanan 等人（2021）结合了管道、张量和数据并行化，并采用了一种新的管道调度策略，将其方法称为 PTD-P。与其仅将一组连续的层（“模型块”）定位在设备上，每个工作器可以被分配多个较小连续子集的层块（例如，设备 1 有层 1、2、9、10；设备 2 有层 3、4、11、12；每个都有两个模型块）。一个批次中的微批次数量应能被工作器数量整除。如果有 M 个模型块每个工作器，管道气泡时间可以通过 GPipe 调度的倍数减少。

在这里插入图片描述