该篇摘自The Ultra-Scale Playbook: Training LLMs on GPU Clusters.

在单个GPU上训练

在单个GPU上训练，通常包括三个步骤： 1. forward pass：将输入传入模型，产生输出； 2. backward pass：计算梯度； 3. optimization：使用梯度更新参数。

batch size的影响

超参数batch size：小的batch size在训练初期有助于快速完成训练过程，达到一个较优learning point；但在训练后期，小的batch size导致梯度噪声增大，模型难以收敛至最优性能点；大的batch size虽然能给出精确的梯度估计，但会降低每个训练样本的利用效率，从而导致收敛变慢，并可能浪费计算资源。

batch size影响在给定dataset上的训练时间：小的batch size在相同数量样本上，需要更多的优化步骤（优化步骤是计算密集型的，导致训练时间比大的batch size更长）。但是，batch size大小通常可以在最优值附近大幅调整，而不会对模型的最终性能产生重大影响（前提是在最优值附近）。

在LLM的预训练中，batch size通常定义为：token的数量（bst：Batch Size Tokens），使得训练次数和训练中使用的输入序列长度基本独立。在单个机器上训练，bs（样本计数）和bst（token计数）可由下计算： \[ bst=bs * seq \] 其中，seq为输入序列长度。

近期 LLM 训练的理想批量大小通常在每批次 400 万到 6000 万个 token 之间。批量大小和训练语料库的规模近年来一直在稳步增加：Llama 1 的训练使用了大约 400 万个 token 的批量大小，训练了 1.4 万亿个 tokens，而 DeepSeek 则使用了大约 6000 万个 token 的批量大小，训练了 14 万亿个 tokens。

然而，一个挑战是在将模型训练扩展到大的batch size时，将遇到显存不足的问题：当 GPU 的显存不足以容纳目标batch size的完整批次时，该怎么办？

Transformer上的内存使用

当训练一个神经网络时，将以下内容存储在内存中：模型权重、模型梯度、优化器状态、（用于计算梯度的）激活值。

以上内容作为tensor（张量）存储在内存中，分别对应不同的shapes和precisions。

训练基本步骤：前向传播时，激活值迅速增加；反向传播时，梯度逐渐积累，且计算梯度的激活值会逐步被清除；最后，执行优化步骤，此时需要所有的梯度，并更新优化器状态；然后才开始下一次的前向传播。

第一步和后续步骤明显不同的原因：激活值快速增加，再保持一段时间的平稳。在第一步中，torch 的缓存分配器进行大量准备工作，预先分配内存；后续步骤不再需要寻找空闲内存块，从而加速）

weights/grads/optimizer state的内存

对于一个简单的transformer LLM，参数数量如下： \[ N=h*v+L*(12*h^2+13*h)+2*h \] \(h\)是隐藏层维度，\(v\)是词汇大小，\(L\)是模型的层数；可以看到，当隐藏层维度较大时，主导项是\(h^2\)项。

**内存需求：参数数量*每个参数的字节数**

传统FP32训练中，参数、梯度均需4字节，优化器（例如Adam）需要存储动量和方差，为每个参数增加另外两个4字节。

\(m_{params}=4*N\)

\(m_{grad}=4*N\)

\(m_{opt}=(4+4)*N\)

若使用高低混合精度训练，当前默认做法是：使用BF16进行大部分计算（每个参数、梯度分别需要2字节），额外复制一份模型权重和梯度为 FP32，因此每个参数总共需要 12 字节。即：

\(m_{params}=2*N\)

\(m_{grad}=2*N\)

\(m_{params_{fp32}}=4*N\)

\(m_{opt}=(4+4)*N\)

混合精度本身并不会节省整体内存，它只是将内存在三个组件之间重新分配。在前向和反向传播中使用半精度计算可以： 1. 在 GPU 上使用经过优化的低精度操作，这些操作更快； 2. 减少前向传播过程中的激活内存需求，而激活内存占用了大量内存。

若使用 FP8 训练代替 BF16，内存使用量会进一步减少（但它的稳定性较差）。

模型参数数量 FP32 或 BF16（不使用 FP32 梯度累积） BF16（使用 FP32 梯度累积）
1B 16 GB 20 GB
7B 112 GB 140 GB
70B 1120 GB 1400 GB
405B 6480 GB 8100 GB

可以观察到，一旦达到 7B 参数，权重和优化器的内存需求就会显著增加，并超过典型 GPU 内存的大小。

activations的内存

依赖于模型的输入。总内存如下： \[ m_{act}=L*seq*bs*h*(34+\frac{5*n_{heads}*seq}{h}) \] 其中，\(L\)是层数，\(seq\)是序列长度，\(bs\)是batch size，\(h\)是模型的隐藏维度，\(n_{heads}\)是注意力头的数量。

可以观察到，内存使用量会随着批量大小线性增长，并随着序列长度的平方增长，那么：激活内存是最容易“膨胀”的部分。

对于短序列（或者小批量大小），激活几乎可以忽略不计；但从大约 2-4k 个 token 开始，它们就会占用大量内存，而参数、梯度和优化器状态的使用，则基本上与序列长度和批量大小无关。

控制activation增长的策略

activation重计算（gradient checkpoints）

也叫做：梯度检查点，重物化。在前向传播时，抛弃一些activations；在后向传播时，实时重新计算activations。

  • Full（全量重计算）：在Transformer的每层transition point上，设置activations checkpoints：要求每层进行一次前向传播，即在反向传播过程中增加一次完整的前向传播。
    • 可以节省最多内存，但在计算上最昂贵。
  • Selective（选择性重计算）：注意力激活值增长较多且在FLOP上计算便宜，因此抛弃他们。
    • 对于一个GPT-3（175B）模型，可以减少70%的激活内存，而计算成本仅为2.7%；DeepSeek V3使用“多头潜在注意力”（MLA）来优化激活内存。

当前大多数框架使用Flash Attention，在其优化策略中，原生集成了activation重计算：在反向传播中，计算（而非存储）注意力分数和矩阵。

activation重计算略微增加FLOPs的数量；但显著减少内存开销。该策略对具备小型高速内存的硬件尤为有利（比如GPU）。

梯度累积（gradient accumulation）

梯度累积将batch拆分为若干个小的micro-batch；依次在每个micro-batch上进行前向、反向传播，计算梯度，在执行优化步骤前，将所有micro-batch梯度相加。实际上，优化步骤是基于梯度的平均值（而非总和）进行的，因此结果与梯度累积步骤的数量无关。有： \[ bs=gbs=mbs*grad_{acc} \] 其中：每次前向传播的batch size为\(mbs\)；每两个优化步骤之间的batch size为\(gbs\)。假设在每进行8次前向/反向传播后执行一次优化步骤，则\(gbs\)将是\(mbs\)的8倍。

梯度累积的一个缺点：在每个优化步骤中，需要执行多个连续的前向/反向传播，从而增加计算开销，减慢计算速度。

然而，每个micro-batch的前向/反向传播可以并行运行。前向/反向传播是相互独立的，唯一的区别是输入样本。因此可以将训练扩展至多个GPU！

并行策略

数据并行（Data Parallelism）

思想：将模型复制到多个GPU上；在每个GPU上，对不同的micro batches执行前向/反向传播。

在每个GPU上使用不同的micro batch，那么每个GPU上的梯度不同；为了保持不同GPU上的模型实例同步，使用all-reduce对模型实例的梯度进行平均，该过程在优化之前的反向传播中执行。

  • all-reduce原语：处理 GPU 实例和节点之间的同步和通信。

一个朴素的实现方式：等待反向传播完成所有的梯度计算；触发all-reduce操作，进行通信以同步这些梯度。然而，这会导致通信时GPU空闲，而我们希望通信和计算能并行。有哪些方法呢？

优化策略

优化一：梯度同步（通信）与反向传播（计算）并行

一旦最后一层的反向传播计算完成，这些梯度可以立即被收集、求和；而反向传播计算会继续向左传播，计算更早层的梯度。

1
2
3
4
5
6
7
8
def register_backward_hook(self, hook):
    """
    Registers a backward hook for all parameters of the model that 
    require gradients.
    """
    for p in self.module.parameters():
        if p.requires_grad is True:
            p.register_post_accumulate_grad_hook(hook)
优化二：梯度分桶

GPU 操作通常在大tensor上执行时效率更高；通信操作亦然。因此，可以通过将梯度分组到多个桶中，并为每个桶内的所有梯度启动一个单独的 all-reduce 操作，（而不是为每个梯度执行独立的 all-reduce 操作）。显著减少通信开销，加速通信操作。

优化三：配合梯度累积

何时同步梯度？

在一个简单版本中，每次反向传播后，自动触发一个 all-reduce 操作，这样效率较低：在最终步骤之后执行一次 reduce 操作能达到相同效果，同时减少开销。

在 PyTorch 中，通常在不需要进行梯度同步的反向传播上添加 model.no_sync() 装饰器，来解决这个问题。

加入DP和梯度累积参数后，global batch size更新如下： \[ bs=gbs=mbs*grad_{acc}*dp \] 其中，\(grad_{acc}\)是梯度累积的步数，\(dp\)是DP中并行实例的数量。

实际上，一般倾向于最大化DP中并行节点的数量：因为DP是并行的，梯度累积是顺序的。在数据并行扩展不足时，再加上梯度累积，以达到目标的global batch size。

DP步骤

总结一下采用DP进行训练的配置步骤：

  1. 确定最佳的global batch size(in tokens)；
  2. 选择训练的序列长度（2~8k个tokens当前结果不错）；
  3. 寻找单个GPU上最大的local batch size(mbs)（不断增加，直到耗尽内存）；
  4. 确定DP使用的GPU数量：GBS 与 DP 的比值决定所需的梯度累积步数。

例子： 假设要训练一个global batch size=4M的模型，序列长度为4k；则批量大小为1024个样本。

假设观察到单个 GPU 只能容纳 MBS=2 的内存，并且有 128 个 GPU 可供训练。那么：通过4步梯度累积，将实现每个训练步骤 1024 个样本或 4M tokens 的目标。

如果突然有 512 个 GPU 可用，仍然可以保持 MBS=2，并将梯度累积步数设置为 1，从而实现更快的训练！

注意：在使用 512+ 个 GPU 的规模时，取决于所使用的网络，通信操作将开始受到环延迟的限制，这会降低计算效率，并影响吞吐量。

虽然DP将梯度同步的 all-reduce 操作与反向传播计算重叠以节省时间，但这种好处在大规模下开始失效。为什么？因为随着添加更多的 GPU（成百上千个），它们之间的协调开销会显著增加，导致网络需求变得过大，抵消了带来的好处。随着每个新 GPU 加入，设置DP的效率将越来越低。

DeepSpeed ZeRO（零冗余优化器）

在每个DP rank上对优化状态、梯度、参数进行赋值，将导致大量内存冗余。ZeRO通过在数据并行维度上，对优化器状态、梯度和参数进行分区来消除内存冗余，同时仍然允许使用完整的参数集进行计算。

activations不参与分区：每个DP replica接收不同的micro-batch，因此每个DP节点上的activations也不同，不参与复制。

考虑如下场景：使用混合精度训练和Adam优化器时，假设模型参数量为 \(\psi\)​，那么每张GPU中的显存内容分为两类： 1. 模型状态： * 模型参数（半精度，bf16/fp16）：\(2\psi\) * 模型梯度（半精度，bf16/fp16）：\(2\psi\) * Adam优化器状态（FP32格式的模型参数备份、FP32的momentum和FP32的variance）：\(4\psi+4\psi+4\psi\) Adam状态占比75%。

  1. 剩余状态： 除了模型状态之外的显存占用，包括activation、各种buffer以及无法使用的显存碎片（fragmentation）。

混合精度训练：同时存在fp16和fp32两种格式的数值，其中模型参数、模型梯度都是fp16，此外还有fp32的模型参数，如果优化器是Adam，则还有fp32的momentum和variance。、

假设显卡数量为\(N\)，提出以下三种ZeRO算法： * ZeRO-1：只对优化器状态进行分片，每张卡保存\(\frac{1}{N}\)​的状态量。此时，每张卡所需显存是\(4\psi+\frac{12\psi}{N}\)字节，当\(N\)较大时，趋向于\(4\psi\)，记为\(P_{os}\)； * ZeRO-2：对优化器状态和梯度进行分片，此时，每张卡所需显存是\(2\psi+\frac{2\psi+12\psi}{N}\)字节，当\(N\)较大时，趋向于\(2\psi\)，记为\(P_{os+g}\)； * ZeRO-3：将模型参数、梯度、优化器状态三者都进行分片，此时，每张卡所需显存是\(\frac{16\psi}{N}\)字节，当\(N\)较大时，趋向于\(0\)，记为\(P_{os+g+p}\)； * ZeRO-3对应Pytorch FSDP

ZeRO-1，ZeRO-2，ZeRO-3通信量分析

集群通信：

reduce-scatter: all-gather: Ring all-reduce:由reduce-scatter，all-gather两个步骤组成：

传统的DP在每一步计算梯度后，需要一次all-reduce操作计算梯度均值，当前常用Ring all-reduce，分为reduce-scatter和all-gather两步。

  • ZeRO-1，ZeRO-2将all-reduce梯度通信改为：reduce-scatter操作，并在优化器步骤之后，增加了对所有参数的all-scatter操作。

    • \(P_{os}\)，\(P_{os+g}\)和传统DP的通信量相同

  • ZeRO-3：

    • 前向传播：依次通过各个layer，按需获取必要的参数；在参数不再需要时，立即从显存中清除。
    • 反向传播：生成梯度分片。

    需要在前向/反向传播中，持续执行all-gathers操作，那么与ZeRO-2相比，需要额外执行\(2*numLayers-1\)次all-gather，每次操作都会带来一个小的基础延迟开销。

    在前向传播时，需要参数时执行all-gather操作，产生一个\(\psi\)的通信开销，立即清除不需要的参数，因此反向传播时还需要一次all-gather操作；最后，与ZeRO-2相同，进行reduce-scatter操作处理梯度，产生\(\psi\)的通信开销。总通信开销为：\(3\psi\)（ZeRO-2的通信开销为\(2\psi\)）

ZeRO-R

在进行tensor并行时，前向传播中的activations会在各个GPU中重复存储，因此：ZeRO-R将所有的中间activations分片存储，即只对activation checkpoints分片（其他activations已被抛弃）

见：重计算

正常情况：保存前向传播中，每一个activations，用于反向传播时计算梯度；每一个前向中的activation，到计算完对应梯度节点后，才能释放。

  • 缺点：需要保存大量的中间激活值，导致占用了大量显存，并且所需的显存是随着层数n线性增长的。

优化一：将所有的中间激活值全部丢弃，反向传播需要时，再重新计算；

  • 缺点：训练速度慢，每个前向节点原本只需要计算一次，现在最多需要计算n次！

折中做法：选取一些前向节点作为checkpoint，训练时，这些checkpoint节点的激活值会一直保存在显存中，而其他节点的激活值会被丢弃。 * 优点：计算反向梯度节点时，只需要从离它最近的checkpoint节点开始计算，而不用把每个节点都重新计算一遍。

ZeRO-Offload

GPU显存不够用，则：将一部分计算和存储下放到CPU和内存，并且不让CPU和GPU之间的通信成为瓶颈，也不让CPU参与过多计算，避免CPU计算成为瓶颈。

Adma优化器中，每一层迭代如下：

将数据流图切分成CPU和GPU两部分。ZeRO-Offload策略如下：它将计算复杂度较高的前向FWD和反向BWD放在GPU上；而参数更新和float2half这两个计算操作放在CPU上。因此，优化器状态也放在内存中，

述方法仅仅针对单卡场景。在多卡场景下，ZeRO-Offload利用ZeRO-2方法。ZeRO-2将优化器状态和梯度分片，每张卡只存储\(\frac{1}{N}\)​，而ZeRO-Offload将这\(\frac{1}{N}\)个​优化器状态和梯度都下放到内存，只在CPU上进行参数更新.

更多内容参考：大模型并行训练技术（一）—— ZeRO系列