【多模态大模型】系列2：如何用多GPU训练一个非常大的模型（数据/模型/流水线/张量并行、MoE、混合精度训练、压缩、激活重新计算）

1 训练并行
- 1.1 数据并行（Data Parallelism）
- 1.2 模型并行（Model Parallelism）
- 1.3 流水线并行（Pipeline Parallelism）
- 1.4 张量并行（Tensor Parallelism）
2 混合专家（Mixture-of-Experts, MoE）
3 其他节省内存的设计
- 3.1 CPU 卸载（CPU Offloading）
- 3.2 激活重新计算（Activation Recomputation）
- 3.3 混合精度训练（Mixed Precision Training）
- 3.4 压缩（Compression）
- 3.5 内存高效优化器（Memory Efficient Optimizer）
参考资料

CLIP原文中提到，CLIP模型是真的太大了，为了能让模型训起来，相似度的计算被分到了多个GPU去计算。我们趁热打铁，学习一下相关的并行技术吧。

近年来，随着更大规模的预训练语言模型的应用，我们在许多自然语言处理（NLP）基准任务上看到了更好的结果。然而，训练大型和深层神经网络是一个具有挑战性的任务，因为它需要大量的GPU内存和较长的训练时间。

单个GPU工作节点的内存是有限的，而许多大型模型的规模已经超出了单个GPU的能力范围。为了在多个GPU之间进行模型训练，出现了几种并行化范式，以及各种模型架构和节省内存的设计，帮助实现对非常大的神经网络的训练。

1 训练并行

1.1 数据并行（Data Parallelism）

数据并行（DP）最简单的方法是将相同的模型权重复制到多个工作线程中，并将全量数据分成很多部分小数据，将小部分数据分配给每个工作线程以同时处理。

如果模型大小大于单个 GPU 节点的内存，则DP无法正常工作。GeePS等方法将暂时未使用的参数卸载回 CPU，数据交换传输在后端进行，并且不会干扰训练计算。

在每个小批量数据处理结束时，工作节点需要同步梯度或权重，以避免过时问题。主要有两种同步方法，这两种方法各有明显的优缺点。

批量同步并行（Bulk Synchronous Parallels, BSP）：在每个小批量数据处理结束时，工作节点同步数据。这种方法防止了模型权重的过时，并保持了良好的学习效率，但每台机器必须暂停并等待其他机器发送梯度。
异步并行（Asynchronous Parallel, ASP）：每个GPU工作节点异步处理数据，无需等待或停止。然而，这容易导致使用过时的权重，从而降低统计学习效率。

介于两者之间，有一种方法是每x次迭代（x > 1）进行一次全局梯度同步，这种特性在分布式数据并行（Distribution Data Parallel, DDP）中被称为“梯度累积”，自PyTorch v1.5版本起被引入。通过将多个梯度打包成一个AllReduce操作来避免立即执行AllReduce操作，以此提高吞吐量。基于计算图，可以对计算和通信调度进行优化。

数据并行

1.2 模型并行（Model Parallelism）

模型并行（Model Parallelism, MP）旨在解决模型权重无法放入单个节点的情况，将计算和模型参数分配到多个机器上。与数据并行不同，数据并行中每个工作节点都保存整个模型的完整副本，模型并行则只在每个工作节点上分配模型参数的一部分，从而减少了内存使用和计算量。

由于深度神经网络通常包含一系列垂直堆叠的层，因此可以直接通过层来分割大型模型，即将一小部分连续的层组合成一个分区，并放置在一个工作节点上。但是，让每个数据批次依次通过多个具有顺序依赖关系的工作节点会导致大量的等待时间（即“气泡”），并严重浪费计算资源，如下图所示。

模型并行

1.3 流水线并行（Pipeline Parallelism）

流水线并行（Pipeline Parallelism, PP）结合了模型并行和数据并行两种方法，以减少无效的等待时间“气泡（bubbles）”。主要思想是将一个小批量数据分割成多个微批量，并允许每个阶段的工作节点同时处理一个微批量。需要注意的是，每个微批量需要两次传递：一次前向传递和一次反向传递。工作节点之间的通信仅传输激活值（前向）和梯度（反向）。这些传递如何调度以及梯度如何聚合在不同的方法中有所不同。分区（工作节点）的数量也被称为流水线深度。

在GPipe中，来自多个微批量的梯度会在最后被聚合并同步应用。这种同步梯度下降保证了学习的一致性和效率，无论工作节点的数量是多少。如下图所示，尽管仍然存在气泡，但它们比上图中的要小得多。

假设m个均匀分割的微批量和d个分区，如果每个微批量的前向和反向传递各需要一个单位时间，则气泡所占的比例为：

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

GPipe 在设备数量增加的情况下几乎实现了线性加速，尽管如果模型参数在各工作节点之间分配不均，这种线性加速并不总是能得到保证。

1.4 张量并行（Tensor Parallelism）

模型并行和流水线并行都是垂直分割模型，而另一方面，我们可以将一个张量操作的计算在多个设备上进行水平划分，这种方法被称为张量并行（Tensor Parallelism, TP）。

以目前非常流行的Transformer模型为例，Transformer模型主要由多层MLP（多层感知机）和自注意力块组成。Megatron-LM采用了一种简单的方法来并行化MLP和自注意力层内的计算。

在一个Transformer中的MLP层包含一个通用矩阵乘法（General matrix multiply, GEMM），随后是一个非线性的GeLU变换。我们可以通过列来分割权重矩阵A：

公式2

自注意力块根据上述分割方式并行运行查询（Q）、键（K）和值（V）权重的通用矩阵乘法（GEMM），然后通过另一个GEMM将它们结合起来以生成注意力头的结果。

公式3

具体示意图如下：

矩阵示意图
qkv

2 混合专家（Mixture-of-Experts, MoE）

具体介绍请查看博文：【DeepSeek背后的技术】系列一：混合专家模型（MoE）

3 其他节省内存的设计

3.1 CPU 卸载（CPU Offloading）

当 GPU 内存满时，一种选择是将暂时不用的数据卸载到 CPU，并在需要时再读回。CPU卸载的思想非常直接，但由于其带来的训练时间延迟，在近年来变得不那么受欢迎。

3.2 激活重新计算（Activation Recomputation）

激活重新计算（也称为“激活检查点”或“梯度检查点”）是一种通过增加计算时间来减少内存占用的聪明而简单的想法。它将训练一个L层深度神经网络所需的内存成本降低到O(sqrt(L))，并且每批次仅额外消耗一次前向传递的计算量。

假设我们将一个L层网络均匀地划分为d个分区，只有在分区边界处的激活值会被保存并在工作节点之间进行通信。分区内各层的中间激活值仍然需要用于计算梯度，因此它们会在反向传播过程中被重新计算。通过激活重新计算，训练M(L)的内存成本为：

内存成本

激活重新计算技巧可以实现相对于模型大小的次线性内存成本。

memory

3.3 混合精度训练（Mixed Precision Training）

Narang & Micikevicius等人（2018）提出了一种使用半精度浮点数（FP16）训练模型的方法，而不会损失模型的准确性。

mixed

避免在半精度中丢失关键信息的三种技术：

全精度权重主副本（Full-precision master copy of weights）：维护一个全精度（FP32）的模型权重副本，用于累积梯度。前向和反向传播过程中使用的数值会被舍入为半精度（FP16）。其动机是每次梯度更新（即梯度乘以学习率）可能太小，无法完全包含在FP16范围内（例如，2^-24在FP16中会变为零）。通过这种方式，可以在保持计算效率的同时避免由于数值范围限制导致的信息丢失。
损失缩放（Loss scaling）：通过放大损失值来更好地处理幅度较小的梯度（见上图）。放大梯度有助于将它们移动到可表示范围内的较大值部分，从而保留那些原本会丢失的小数值。具体来说，缩放梯度可以使其占据更大的范围，确保更多的梯度信息能够被有效利用，而不是因为数值过小而被截断或归零。
算术精度（Arithmetic precision）：对于常见的网络算术运算（如向量点积、通过求和元素进行降维），可以在FP32中累积部分结果，然后在保存到内存之前将其最终输出转换为FP16。逐点操作可以在FP16或FP32中执行。这种方法确保了在关键计算步骤中的数值稳定性，同时仍然能利用半精度带来的计算和存储优势。

histogram
在他们的实验中，损失缩放对于某些网络（如图像分类、Faster R-CNN）来说并不是必需的，但对于其他网络（如Multibox SSD、大型LSTM语言模型）则是必要的。

3.4 压缩（Compression）

中间结果通常会占用大量内存，尽管它们只在一个前向传递和一个反向传递中使用。在这两次使用之间存在明显的时差。因此，Jain等人（2018）提出了一种数据编码策略，在第一次传递后对中间结果进行压缩，然后在反向传播时解压缩回原样。

他们的系统Gist结合了两种编码方案：

层特定无损编码：专注于ReLU-Pool（“二值化”）和ReLU-Conv（“稀疏存储和密集计算”）模式。
激进有损编码：使用延迟精度降低（DPR）。他们观察到特征图的第一次立即使用应保持高精度，而第二次使用可以容忍较低的精度。

实验表明，Gist可以在5个最先进的图像分类DNN上将内存成本减少2倍，平均减少1.8倍，仅带来4%的性能开销。

3.5 内存高效优化器（Memory Efficient Optimizer）

优化器通常需要大量的内存。以流行的Adam优化器为例，它内部需要维护动量和方差，这些都与梯度和模型参数处于相同规模。突然之间，我们需要保存4倍于模型权重的内存。

已经提出了几种优化器来减少内存占用。例如，Adafactor（Shazeer等人，2018）没有像Adam那样存储完整的动量和方差，而是只跟踪每行和每列移动平均值的总和，然后基于这些总和估计第二时刻。SM3（Anil等人，2019）描述了一种不同的自适应优化方法，同样显著减少了内存需求。

ZeRO（零冗余优化器；Rajbhandari等人，2019）基于对大规模模型训练中两个主要内存消耗的观察，优化了用于训练大型模型的内存：

大部分内存被模型状态占据，包括优化器状态（如Adam动量和方差）、梯度和参数。混合精度训练需要大量内存，因为优化器需要保留FP32参数和其他优化器状态的副本，除了FP16版本之外。
其余内存被激活、临时缓冲区和无法使用的碎片内存（在论文中称为残余状态）所消耗。

ZeRO结合了两种方法：ZeRO-DP和ZeRO-R。

ZeRO-DP 是一种增强的数据并行性，避免了对模型状态的简单冗余。它通过动态通信调度将优化器状态、梯度和参数划分为多个数据并行进程，以最小化通信量。
ZeRO-R 优化了残余状态的内存消耗，使用分区激活重新计算、固定缓冲区大小和即时内存碎片整理。

参考资料

How to Train Really Large Models on Many GPUs?
https://www.bilibili.com/video/BV1nB4y1R7Yz/?spm_id_from=333.788.videopod.sections&vd_source=392e51c88665530ce8438373de9e995e