小白学大模型：训练大语言模型的深度指南

在当今人工智能飞速发展的时代，大型语言模型（LLMs）正以其强大的语言理解和生成能力，改变着我们的生活和工作方式。在最近的一项研究中，科学家们为了深入了解如何高效地训练大型语言模型，进行了超过4000次的实验。这些实验动用了多达512个GPU（图形处理单元），它们协同工作，为模型训练提供了强大的计算支持。在这项研究中，研究人员特别关注了两个关键指标：吞吐量（通过标记的大小来表示）和GPU利用率（通过标记的颜色来表示）。这两个指标都根据模型的大小进行了标准化处理，以便更直观地比较不同模型在不同硬件配置下的表现。

AI模型训练的三大挑战

内存使用：有限的“容器”

想象一下，你有一个装满水的杯子，但你还需要往里面加更多的水。如果杯子满了，就再也装不下任何东西了。这就像训练AI模型时的内存问题。如果一个训练步骤所需的内存超出了GPU的容量，那么训练就无法继续。内存是一个硬性限制，我们必须在有限的空间内完成复杂的计算任务。

计算效率：让硬件“火力全开”

我们花了大价钱购买高性能的GPU，当然希望它们能时刻保持高效运转。但现实往往是，GPU在等待数据传输或者等待其他GPU完成工作时，会浪费大量时间。这就像是一个生产线上的工人，因为原材料没送到或者上一个环节还没完成，只能干等着。为了提高效率，我们需要减少这些等待时间，让硬件尽可能多地用于计算。

通信开销：减少“内耗”

在多GPU协同工作的场景中，不同GPU之间需要频繁地交换信息。但这种通信会占用大量的时间和资源，甚至会让GPU处于闲置状态。这就像是一个团队开会，如果沟通成本过高，那么真正干活的时间就会减少。因此，我们需要巧妙地利用节点内部（速度快）和节点之间（速度慢）的带宽，并尽可能让通信和计算同时进行，从而减少通信开销。

从单个GPU开始：AI模型训练的第一步在探索如何用数千个GPU训练大型AI模型之前，我们先从最基础的部分开始——在单个GPU上训练模型。这就好比在学会驾驶飞机之前，先要学会驾驶一辆汽车。别小看这一步，它是我们理解整个训练过程的关键。

单GPU训练的三个基本步骤

当你在单个GPU上训练一个模型时，整个过程通常可以分为三个步骤：

前向传播（Forward Pass）：将输入数据通过模型，得到模型的输出结果。这就好比把食材放进烤箱，等待美味的蛋糕出炉。

反向传播（Backward Pass）：计算梯度，也就是找出模型需要改进的方向。这一步就像是检查蛋糕的口感，看看哪里需要调整。

优化步骤（Optimization）：用计算出的梯度来更新模型的参数，让模型变得更好。这就好比根据反馈调整烤箱的温度和时间。

关键参数：批次大小（Batch Size）

批次大小（Batch Size）是训练过程中一个非常重要的参数。它决定了每次训练时输入模型的数据量。批次大小的选择对模型的训练效果和效率有很大影响。

小批次大小：在训练初期，小批次大小可以帮助模型快速找到一个好的学习方向，就像在迷宫中快速试探出一条路。但如果一直使用小批次，模型的梯度会比较“嘈杂”，最终可能无法达到最优性能。

• 大批量大小：大批量大小可以提供更准确的梯度估计，但同时也会让模型对每个训练样本的利用效率降低，导致训练速度变慢，甚至浪费计算资源。

批次大小与训练时间

批次大小还会影响训练一个给定数据集所需的时间。小批次大小需要更多的优化步骤来处理相同数量的数据，而每个优化步骤都需要计算时间，因此总训练时间会更长。不过，只要批次大小在最优值附近，模型的最终性能通常不会受到太大影响。

在大型语言模型（LLM）的预训练领域，批次大小通常用“token”（标记）数量而不是样本数量来表示。这样可以让训练数据与输入序列长度无关，更加通用。

挑战：内存不足怎么办？

当我们尝试将训练扩展到更大的批次大小时，第一个挑战就来了——GPU内存不足。当我们的GPU无法容纳目标批次大小时，该怎么办呢？

首先，我们需要理解为什么会出现内存不足的问题。在训练神经网络时，我们需要在内存中存储以下内容：

模型权重（Weights）

模型梯度（Gradients）

优化器状态（Optimizer States）

用于计算梯度的激活值（Activations）

训练过程中的内存动态变化

使用PyTorch的分析工具，我们可以清楚地看到，内存的使用并不是一成不变的，而是在训练过程中不断变化。例如，在一个训练步骤中，内存的使用会随着前向传播、反向传播和优化步骤而起伏。

前向传播：当我们开始训练时，模型的激活值会迅速增加，因为模型需要计算输入数据的输出结果。

反向传播：随着梯度的计算，内存中的梯度值会逐渐积累，而用于计算梯度的激活值则会被逐步清理。

• 优化步骤：最后，在更新模型参数时，我们需要使用所有梯度，并更新优化器的状态。

有趣的是，第一次训练步骤与其他步骤看起来很不一样。这是因为PyTorch的内存分配器在第一步中做了很多准备工作，为后续步骤预留内存，从而避免在后续步骤中频繁搜索空闲内存块。这种机制虽然优化了训练效率，但也可能导致一个常见问题：第一次训练步骤成功了，但后续步骤却因为内存不足而失败。这正是因为优化器状态在第一次步骤后开始占用更多内存。

如何估算模型的内存需求？

为了更好地管理内存，我们需要了解模型的各个部分分别需要多少内存。首先，我们来看看模型的权重、梯度和优化器状态。对于一个简单的Transformer模型，其参数数量可以通过以下公式估算：其中，h是隐藏层维度，v是词汇表大小，L是模型层数。从这个公式可以看出，当隐藏层维度增大时，h²项会迅速增长，成为内存占用的主要部分。

接下来，我们来计算这些参数的内存需求。在传统的全精度（FP32）训练中，每个参数和梯度需要4字节，而优化器（如Adam）还需要额外存储动量和方差，每个参数再占用8字节。

如果使用混合精度训练（如BF16），虽然计算速度更快，但内存需求会略有变化。在混合精度训练中，我们通常使用BF16（2字节）进行计算，同时保留一份FP32（4字节）的模型权重副本，以确保数值稳定性。因此，总内存需求如下：

• 模型权重（BF16）：
• 梯度（BF16）：
• 模型权重副本（FP32）：
• 优化器状态（FP32）：

模型参数规模	全精度训练（FP32）	混合精度训练（BF16 FP32副本）
10亿参数	16 GB	20 GB
70亿参数	112 GB	140 GB
700亿参数	1120 GB	1400 GB
4050亿参数	6480 GB	8100 GB

激活值内存：训练中的“内存大户”

在训练大型语言模型时，激活值内存（Activation Memory）是一个非常关键且复杂的部分。它不仅取决于模型的结构，还与输入数据的长度和批次大小密切相关。激活值内存的管理，往往决定了我们能否在有限的硬件资源上训练更大规模的模型。

激活值内存的大小并不固定，它会随着输入序列的长度和批次大小的变化而变化。经过仔细分析Transformer模型的反向传播过程，我们可以估算出激活值内存的大小。具体公式如下：

其中：L 是模型的层数；seq 是输入序列的长度；bs 是批次大小；h 是模型的隐藏层维度；n_heads 是多头注意力机制中的头数。

激活值重计算：用计算换内存的“魔法”

在传统的训练过程中，我们会存储每一步的隐藏状态（即激活值），以便在反向传播时用来计算梯度。但有了激活值重计算，我们只需要在模型的关键节点存储少量激活值，丢弃其余的激活值，并在反向传播时从最近的保存点重新计算它们。这就好比在旅行中，你只在关键地点留下标记，而不是记录整个路线的每一个细节。激活值重计算有几种不同的策略，每种策略在内存节省和计算成本之间都有不同的权衡：1. 完整重计算（Full Recomputation）这种策略会在Transformer模型的每一层之间存储激活值。因为它需要在反向传播时重新执行每一层的前向传播，所以计算成本最高，但节省的内存也最多。通常，这种策略会使计算成本和时间增加30-40%，效果非常明显。2. 选择性重计算（Selective Recomputation）

选择性重计算是一种更高效的策略。研究人员发现，注意力机制的计算通常占用大量内存，但重新计算的成本较低。因此，我们可以丢弃这些激活值，而只存储前馈网络（Feedforward）的激活值。例如，在一个1750亿参数的GPT-3模型中，这种策略可以减少70%的激活值内存占用，而计算成本仅增加2.7%。

梯度累积：用“微批次”突破内存限制

梯度累积的核心思想非常简单：将一个大批次拆分成多个小批次（微批次），然后依次对每个微批次执行前向传播和反向传播，计算梯度，并将这些梯度累加起来。最后，我们用累加后的梯度平均值来更新模型参数。这样，我们就可以在不增加内存占用的情况下，有效地模拟大批次训练的效果。

具体来说，我们把每个前向传播的批次大小称为“微批次大小”（Micro Batch Size, mbs），而把两次优化步骤之间的总批次大小称为“全局批次大小”（Global Batch Size, gbs）。如果我们每执行8次前向/反向传播后进行一次优化步骤，那么全局批次大小就是微批次大小的8倍。虽然梯度累积解决了内存问题，但它也有一个明显的缺点：每次优化步骤需要多次连续的前向/反向传播，这会增加计算开销，从而减慢训练速度。不过，如果你仔细思考，你会发现这些前向/反向传播是可以并行化的——每个微批次的计算是独立的，唯一的区别是输入样本不同。这意味着，我们可以通过多GPU并行计算来加速这一过程。数据并行：用多GPU加速模型训练在数据并行中，每个GPU处理一个独立的微批次数据，因此每个GPU上计算出的梯度是不同的。为了保持所有模型实例的一致性，我们需要在反向传播过程中，通过一个称为“全归约”（all-reduce）的操作来平均这些梯度。全归约是数据并行中的第一个“分布式通信原语”，它负责在GPU实例和节点之间同步和通信。

优化数据并行的三种策略

1. 优化一：计算与通信重叠在简单的数据并行实现中，我们通常需要等待反向传播完成后，才开始同步梯度。但其实，我们可以将通信与计算重叠起来，让它们同时进行。具体来说，当反向传播到达最后一层时，我们可以立即开始同步这些层的梯度，而不需要等待前面层的梯度计算完成。这样，大部分全归约操作可以在反向传播过程中完成，从而提高效率。在PyTorch中，我们可以通过为每个参数添加一个全归约钩子函数来实现这一点。当某个参数的梯度计算完成后，立即触发全归约操作，而其他参数的梯度计算仍在继续。这样可以显著减少等待梯度同步的时间。2. 优化二：梯度分桶（Bucketing）GPU在处理大型张量时通常比处理多个小型张量更高效，通信操作也是如此。因此，我们可以将梯度分组到“桶”中，然后对每个桶中的梯度执行一次全归约操作，而不是对每个梯度单独执行全归约。这就像打包物品时，发送几个大箱子比发送许多小箱子更高效。通过这种方式，我们可以显著减少通信开销，加快通信速度。3. 优化三：与梯度累积结合

我们之前提到，梯度累积通过多次前向和反向传播来累积梯度，然后在最后一步更新参数。当我们将梯度累积与数据并行结合时，需要注意梯度同步的时机。在简单的实现中，每次反向传播后都会触发全归约操作，但这其实是不必要的。我们可以在最后一步统一触发全归约操作，从而减少通信开销。

数据并行与梯度累积的结合

在训练大型语言模型时，全局批次大小（GBS）是一个关键参数，它直接影响模型的收敛速度和训练效率。通过引入数据并行（Data Parallelism, DP）和梯度累积（Gradient Accumulation, GA），我们可以更灵活地调整全局批次大小，同时优化训练速度和硬件利用率。全局批次大小 (GBS) = 微批次大小 (MBS) × 梯度累积步数 (GA) × 数据并行实例数 (DP)在实际应用中，人们通常会优先选择最大化数据并行的节点数量（DP），因为数据并行是并行化的，而梯度累积是顺序的。只有在数据并行无法满足目标全局批次大小时，才会增加梯度累积的步数。例如，当我们有足够的GPU时，可以通过增加数据并行的节点数量来加速训练，而不是单纯依赖梯度累积。

确定目标全局批次大小（GBS）：通过查阅文献或实验测量模型的收敛情况，确定最佳的全局批次大小（以token为单位）。

选择序列长度：根据文献或实验选择合适的训练序列长度。通常，2-8k tokens是一个可靠的选择。

确定单个GPU的最大微批次大小（MBS）：逐步增加微批次大小，直到单个GPU的内存不足。

1. 确定可用的GPU数量：根据目标数据并行实例数（DP），计算所需的梯度累积步数。全局批次大小除以数据并行实例数，即为剩余的梯度累积步数。

ZeRO：零冗余优化器，让内存管理更高效

DeepSpeed的ZeRO（Zero Redundancy Optimizer）通过将这些组件分散到不同的数据并行节点上，显著减少了内存冗余，同时仍然允许使用完整的参数集进行计算。接下来，我们将深入了解ZeRO的三个阶段：ZeRO-1、ZeRO-2和ZeRO-3。

ZeRO-1：优化器状态分区

在传统的数据并行中，所有节点在反向传播后收集相同的梯度，并同时执行相同的优化器步骤。这不仅效率低下，还浪费了大量内存。ZeRO-1通过将优化器状态分成N_d（数据并行度）等份来解决这个问题。每个模型副本只保留1/N_d的优化器状态，并在优化步骤中只更新1/N_d的浮点参数。

ZeRO-2：添加梯度分区

在ZeRO-1的基础上，ZeRO-2进一步将梯度也进行分区。由于每个副本只需要与优化器状态对应的梯度片段，因此在反向传播时，我们只需要执行“归约分散”（reduce-scatter）操作，而不是“全归约”（all-reduce）。这样，每个副本只需要保留1/N_d的梯度，从而节省更多内存。

ZeRO-3：添加参数分区

ZeRO-3（也称为FSDP，即“完全分片数据并行”）是ZeRO的最终阶段，它将模型参数也进行了分区。这意味着每个副本在需要时才会动态收集所需的参数片段，并在使用后立即释放它们。这种“按需收集”的方式进一步减少了内存占用。

张量并行：打破内存限制的“魔法”在训练大型语言模型时，激活值内存往往会成为瓶颈，尤其是在单个GPU无法容纳模型的单层时。这时，张量并行（Tensor Parallelism, TP）就派上了用场。张量并行不仅分区了模型的参数、梯度和优化器状态，还分区了激活值，且无需在计算前收集所有片段。

张量并行的基本原理

张量并行的核心在于矩阵乘法的数学性质。在神经网络中，矩阵乘法通常表示为 X × W，其中：

• X是输入或激活值；
• W是神经网络层的权重。

张量并行利用了矩阵乘法的两个基本性质：

按列分区：可以将矩阵 B 的每一列分别与矩阵 A 相乘，然后组合结果。

按行分区：可以将矩阵 A 的每一行分别与矩阵 B 相乘，然后将结果相加。

按行分区（Row-wise Sharding）

按行分区是另一种实现方式。具体步骤如下：

1. 按行分割权重矩阵：将权重矩阵 W 按行分割成多个片段，并将这些片段分配到不同的GPU上。
2. 分割输入矩阵：将输入矩阵 X 分割成多个片段（需要一个“分散”操作，scatter）。
3. 计算局部结果：每个GPU计算局部输入矩阵与局部权重矩阵的乘积。

汇总结果：通过一个“全归约”（all-reduce）操作，将所有GPU上的局部结果相加，得到最终结果。

多头注意力机制（MHA）中的张量并行

多头注意力机制是Transformer模型的核心部分，它涉及多个矩阵乘法操作（Q、K、V）。张量并行在MHA中的应用也非常直观：

按列分区（Column Parallel）：将Q、K、V矩阵按列分割，并分配到不同的GPU上。每个GPU计算一个或多个注意力头的输出。这种方法非常适合多头注意力机制，因为每个GPU可以独立计算一个或多个头的注意力结果。

1. 按行分区（Row Parallel）：对于输出投影（Output Projection），可以按行分割权重矩阵，从而减少每个GPU上的内存需求。

序列并行的基本原理

序列并行的核心思想是：在张量并行的基础上，进一步将激活值和计算分割到不同的GPU上，但这次是沿着输入序列的维度，而不是隐藏维度。这种方法特别适用于那些需要完整隐藏维度的操作，如LayerNorm和Dropout。

以LayerNorm为例，它需要完整的隐藏维度来计算均值和方差：

其中，均值和方差是在隐藏维度上计算的。尽管LayerNorm的计算成本较低，但它仍然需要大量的激活值内存，因为需要完整的隐藏维度。序列井行通过在序列维度上分割激活值，将内存负担分散到多个GPU上，从而显著减少了每个GPU的内存需求。

Diving in the GPUs – fusing, threading, mixing

Fused Kernels

Flash Attention 1-3

Mixed Precision Training