从零开始构建 LLM（第32k部分）：本地训练与梯度累积优化Writing an LLM from scratch, part 32k -- Interventions: training a better model locally with gradient accumulation

归档

分类目录

友情链接

我一直在基于 Sebastian Raschka 的书籍《Build a Large Language Model (from Scratch)》开发一个类似 GPT-2-small 的小型语言模型。我在云端训练了它的多个版本，以确定对模型和训练代码进行哪些干预能最有效地降低在特定测试数据集上的损失值。现在，我想在本地运行一次训练，复现其中效果最佳的那一套方案。为此，我需要让本地的批处理大小与云端训练时保持一致。

在我刚开始学习这些内容的时候，批处理（batching）看起来只是一个性能优化手段——对于像 GPU 这样高度并行的系统来说，通常可以将（比如）两个输入一起送入模型处理，所用时间不到单独处理一个的两倍，因此将多个样本打包成一批处理是有意义的。

在推理（inference）阶段确实能获得这种优势，但在训练过程中，我越来越清楚地认识到，批处理还能提升模型的质量。我最直观的理解是：如果每次只输入一个样本就立即调整参数，模型很容易“矫枉过正”，反复修正反而偏离了正确方向。而使用批次时，所有样本的梯度会被平均化，从而平滑更新过程，使训练更加稳定。

当然，过度使用也会适得其反。举个极端的例子，假设你能把整个训练集塞进一个批次里——那你只需要前向传播一次、反向传播一次，然后一次性更新所有参数。显然，这种做法效果会很差，因为初始参数几乎是随机的，仅凭一次更新很难学到有效信息。

在我的本地机器上，RTX 3090 最多只能容纳包含六个序列的批次。当我迁移到云主机后，发现更大的批处理显著改善了模型在测试中的损失表现。通过粗略拟合曲线，我估计该次训练下最优批处理大小约为 97。巧合的是，这接近于一台配备 8 块 A100 40 GiB/GPU 的云服务器所能支持的最大容量，所以我采用 96 的批处理大小来测试各种干预措施的效果。

最终，当我把所有有助于训练的干预组合在一起时，结果（多少让我有些意外）显示，它们的综合效果——损失降低了 0.113765——还不到增大批处理带来的损失改善幅度（0.252474）的一半。

这一切都说明：如果我希望本地训练能达到云端训练模型的质量，不仅需要应用那些经过详细测试的干预措施，还必须匹配云端的批处理规模。而这，就需要我了解梯度累积（gradient accumulation）技术了。

梯度累积基础

梯度累积顾名思义：它不是常规的“前向传播→计算损失→反向传播求梯度→用优化器更新参数”的单步流程，而是执行多次前向-反向传播循环，让梯度逐步累积，最后再统一进行一次优化器更新。

这样做可以在不牺牲性能的前提下，获得大批量训练带来的稳定性优势。听起来很简单，理论上确实如此，但实际实现起来却要复杂一些。

让我们一步步来看。

首先，想象一下你有一个非常简单的训练循环：

for inputs, targets in batched_dataset:
    # Other stuff

    logits = model(inputs)
    loss = calculate_loss(logits, targets)
    loss.backward()

    # Other stuff

    optimizer.step()
    optimizer.zero_grad()

    # Other stuff

向其中添加梯度累积非常简单！

假设 batched_dataset 的长度可以被 gradient_accumulation_steps（我们希望在更新优化器之前执行的步数）整除。作为一个初步（但不完全准确）的近似，你可以直接这样做：

for step, (inputs, targets) in enumerate(batched_dataset):
    # Other stuff

    logits = model(inputs)
    loss = calculate_loss(logits, targets)
    loss.backward()

    # Other stuff

    if (step + 1) % gradient_accumulation_steps == 0 or step == len(batched_dataset) - 1:
        optimizer.step()
        optimizer.zero_grad()

    # Other stuff

可以看到，我们每经过 gradient_accumulation_steps 步才更新一次优化器。另一种实现方式是使用一个内层循环：

dataset_ix = 0
while dataset_ix < len(batched_dataset):
    # Other stuff

    for ii in range(gradient_accumulation_steps):
        if dataset_ix >= len(batched_dataset):
            break
        inputs, targets = batched_dataset[dataset_ix]
        logits = model(inputs)
        loss = calculate_loss(logits, targets)
        loss.backward()
        dataset_ix += 1

    # Other stuff

    optimizer.step()
    optimizer.zero_grad()

    # Other stuff

哪种方法更好取决于训练循环的具体细节——一般来说，如果你希望“其他操作”在每个训练批次执行一次，那么第一种方式更合适；而如果希望这些操作在每个优化器步骤执行一次，则第二种方式更容易实现。正如我稍后会展示的，我的代码采用了第二种方案。

然而，我们需要对上述任一方法做一个小修正才能使其正确工作。请记住，当你在一个批次上计算损失时——例如像这样的交叉熵损失：

torch.nn.functional.cross_entropy(logits.flatten(0, 1), targets.flatten())

...你得到的是该批次中所有样本的平均损失，因此在反向传播时，你获得的是平均梯度。相比之下，在上面的代码中，我们在每一步都对完整的损失进行反向传播，这意味着每次反向传播产生的梯度会被累加起来——最终得到的是所有梯度的总和而非平均值。因此，优化器应用时的梯度会是实际应有的 gradient_accumulation_steps 倍大——相当于我们将学习率乘以了这个数值！

但这很容易修复。若干步的平均梯度等于这些步的梯度之和除以步数，我们可以提前通过将损失按比例缩小来实现这一除法。将其加入上面的第一个示例中：

for step, (inputs, targets) in enumerate(batched_dataset):
    logits = model(inputs)
    loss = calculate_loss(logits, targets)
    (loss / gradient_accumulation_steps).backward()

    if (step + 1) % gradient_accumulation_steps == 0 or step == len(batched_dataset) - 1:
        optimizer.step()
        optimizer.zero_grad()

基本上就是这样了；经过这些修改后，原始的简单训练循环就变成了支持梯度累积的版本。有效批量大小等于真实批量大小乘以梯度累积步数。

梯度累积的实际应用

不过，我用于这些实验的真实训练循环比那个简单示例要复杂得多。它包含了检查点保存、AMP（自动混合精度），最重要的是——还能处理使用 DistributedDataParallel 的多 GPU 训练。这使得问题变得稍微复杂了一些。

首先要解决的是如何选取用于训练的数据。我的数据集已经是按批次组织的，但需要将这些批次拆分到各个 GPU 上。代码中的解决方案是计算出共有多少个全局步（global steps）——每个全局步代表一个批次在所有机器上的 GPU 上完成一次前向和反向传播：

total_global_steps = len(train_ds) // world_size

world_size 如果你还记得 DDP 文章的话，它表示在一次多 GPU 训练中运行进程的数量——每个 GPU 对应一个进程。

接下来，在训练循环中，我对全局步进行迭代：

progress_bar = tqdm(
    range(start_global_step, total_global_steps),
    disable=(rank != 0)
)
for global_step in progress_bar:
    # Get the data
    # Forward and backward pass
    # Step the optimiser.

...对于每一个全局步，获取当前运行的特定 GPU 所需的数据批次：

    inputs, targets = train_ds[global_step * world_size + rank]

rank 是一个从零开始编号的数字，每个 GPU 进程都有唯一的 rank。因此，这实际上是将 train_ds 分割成若干长度为 world_size 的块，然后每个 GPU 接收其对应 rank 偏移位置处的批次数据。

我希望保持训练流程的一致性，这样在本地进行梯度累积时，其行为能尽可能接近云端使用每 GPU 微批次（per-GPU batching）的运行方式。具体来说：我在云端训练时使用八块 GPU，每块 GPU 的微批次大小为 12，因此总批次大小为 96；而本地我只能在一块 GPU 上容纳 6 的批次大小，所以我需要通过梯度累积执行 96 / 6 = 16 步。

为了尽可能保持一致，我决定让“全局步数”（global step）的概念在两种运行方式之间对齐。也就是说，它不再仅仅表示“每 GPU 一次批次”，而是扩展为“每 GPU 一次优化器步”。因此，每次进入 global_step 循环时，我们会执行多次前向-后向传播，然后才进行一次优化器更新。这意味着实现方式应更接近上述两段示例代码中的后者——即带有内层循环而非取模运算的那个版本。

也许用代码更容易说明这一点：

total_global_steps = (len(train_ds) // world_size) // gradient_accumulation_steps

...

for global_step in progress_bar:
    ...
    for accumulation_step in range(gradient_accumulation_steps):
        # Forward and backward passes
        ...

    # Step the optimiser

这需要对数据查找逻辑进行修改：我决定将 train_ds 分割成大小为 gradient_accumulation_steps * world_size 的块，然后再将每个块按 world_size 进一步切分。这样一来，获取给定循环中对应批次的方式就变成了如下形式：

inputs, targets = train_ds[((global_step * gradient_accumulation_steps) + accumulation_step) * world_size + rank]

这要求在 load_dataset 中也做出相应调整，以确保 train_ds 能被 world size、每 GPU 批次（“微批次”）大小以及梯度累积步数整除，但这并不复杂：

one_full_batch_tokens = world_size * microbatch_size * seq_length

...变成了如下形式：

one_full_batch_tokens = world_size * microbatch_size * gradient_accumulation_steps * seq_length

这些改动足以实现梯度累积！接下来，我需要修改反向传播部分的代码，将损失值缩放，以便获得平均而非求和的梯度。由于我们可能在使用 AMP 配合 scaler，因此 loss.backward 并非简单的调用：

if scaler is not None:
    scaler.scale(train_loss).backward()
else:
    train_loss.backward()

...但修改思路是显而易见的：

if scaler is not None:
    scaler.scale(train_loss / gradient_accumulation_steps).backward()
else:
    (train_loss / gradient_accumulation_steps).backward()

所有这些改动，再加上一些代码整理，已经足够构建一个正确的梯度累积训练循环！不过我还需要做一个小调整。

当使用 DDP（DistributedDataParallel）时，不同 GPU 进程之间的梯度必须同步。回顾一下过程如下：

在我最初实现的梯度累积版本中，实际发生的情况如下：

虽然这在逻辑上是正确的，但效率不高——我们在每次累积步骤中都发送并平均梯度。然而，由于每个 GPU 进程都在维护自己的“局部”平均梯度（通过累积已缩放的梯度），我们实际上只需要在最后一步之前，将所有局部平均值汇总并计算全局平均即可。这样做可以节省大量通信开销。

避免频繁同步的关键在于利用 DistributedDataParallel 类提供的 no_sync 方法。我们的模型正是被这个类包装。我们的目标是抑制除最后一次累积步骤外的所有步骤中的梯度同步。

判断是否处于最后一次梯度累积步骤很容易：

is_last = accumulation_step == gradient_accumulation_steps - 1

现在，我们需要做的是将以下内容封装起来：

if scaler is not None:
    scaler.scale(train_loss / gradient_accumulation_steps).backward()
else:
    (train_loss / gradient_accumulation_steps).backward()

...用 model.no_sync 包装，但仅当 is_last 为 false 时。

条件语句有时会显得繁琐，不过 Python 的 contextlib 模块提供了一个“什么都不做”的 nullcontext 上下文管理器，也就是说，

with nullcontext():
    do_something()

...它与下面这段代码完全等价：

do_something()

因此我们可以这样结合三元运算符使用：

    with model.no_sync() if not is_last else nullcontext():
        if scaler is not None:
            scaler.scale(train_loss / gradient_accumulation_steps).backward()
        else:
            (train_loss / gradient_accumulation_steps).backward()

...这正好实现了我们想要的效果。

经过这次改动，我对此感到满意；你可以在这里查看差异。

是时候开始一次训练运行了！

新的本地基线

我原本计划直接基于上次在云端运行的配置进行本地训练，启用所有我认为值得使用的干预措施，并结合梯度累积。

然而，我决定先尝试一个新的“基线”训练。此前我在本地做过一些训练，随后通过在云端使用完全相同配置、在 8 块 A100 40 GiB 显卡上以整体批量大小 96 进行训练，建立了一个基准模型。现在我可以在本地复现这一过程，采用梯度累积方式，这将验证两件事（或者更准确地说，是同一现象的不同侧面）：

这有助于确认我的理解：在云端取得良好效果的关键因素是大批量训练本身，而非其他架构差异；同时也能作为梯度累积代码正确性的有力测试。

这是本次训练运行的配置。我已启动任务：

giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-baseline datasets/
Fetching 4 files: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 810.65it/s]
Starting rank 0 training at global step 0
  0%|                                                                                                                  | 0/33165 [00:04<?, ?it/s, loss=10.991, tps=20,451]


Checkpoint

Continuing training
  0%|                                                                                                       | 23/33165 [01:50<43:29:50,  4.72s/it, loss=7.625, tps=20,549]

全局步数看起来是正确的，与我之前在云端训练时观察到的数值一致。整个训练耗时约 44 小时，这个时间也合理：最初的本地训练用了 48 小时，但那段时间我花了不少精力做验证，而当前代码并未包含这部分开销。

不到两天后：

Training complete in 155,402.289 seconds
Tokens seen: 3,260,252,160
Throughput: 20,979 tokens/second
Final train loss: 3.738

一切进展顺利。损失曲线如下所示：

为了对比，这里是使用相同配置但在云端训练时的损失曲线（不过使用的是更大的实际批量而非梯度累积）：

可以看到两者非常相似，但不完全一样——这完全可以预料！因为两次训练分别运行在不同硬件上——RTX 3090 与 A100，所以 CUDA 内核可能存在差异，此外 PyTorch 的 AMP（自动混合精度）也可能做出不同的决策（例如在合适情况下使用 16 位浮点数而非 32 位）。我认为如果我们在单张 A100 机器上运行，梯度累积的结果会与大批量训练的结果更加接近（甚至可能完全相同），尤其是在关闭 AMP 的情况下。

我将模型上传至 Hugging Face，接下来进入评估阶段。首先是冒烟测试：

Every effort moves you to a more realistic approach and is always welcome, especially for a younger child!
I’

如往常一样，输出内容基本连贯。但真正重要的是它在测试集上的损失值：

Loss against our test dataset: 3.683835

结果相当不错！云端训练的基准模型损失为 3.691526，而这个本地模型的损失略低，仅为 3.683835，差距仅为 0.007691。如此接近正是我们所期望看到的 :-)

现在该看看加入各项干预措施会带来什么影响了。

带干预措施的本地训练

回顾一下，本次训练对配置所做的改动如下：

未包含 QKV bias

这是配置。 我启动了它，然后：

giles@perry:~/Dev/ddp-base-model-from-scratch (main)$ uv run torchrun --nproc_per_node=1 ddp_train.py 1xrtx3090-stacked-interventions datasets
Fetching 4 files: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:00<00:00, 749.32it/s]
Starting rank 0 training at global step 0
  0%|                                                                                                                  | 0/33165 [00:04<?, ?it/s, loss=10.994, tps=21,589]


Checkpoint

Continuing training
  0%|                                                                                                      | 19/33165 [01:25<40:52:06,  4.44s/it, loss=10.318, tps=21,867]

看起来需要运行40小时；这与云端的执行情况一致，因为移除 dropout 会显著加快训练速度。

仅仅不到两天后：

Training complete in 146,299.816 seconds
Tokens seen: 3,260,252,160
Throughput: 22,285 tokens/second
Final train loss: 3.519

训练过程中的损失曲线如下所示：

非常平滑，没有出现损失突增。作为对比，这是我们之前在云端进行相同训练时的曲线；可以看到，其波动比本地版本更剧烈一些。

梯度范数（gradient norm）图表也很有意思：

与下方云端训练的梯度范数图相比可以看出，本地的梯度噪声实际上更大——云端训练在初期有几个明显的梯度尖峰，但从大约全局步数6,000之后趋于平稳；而本地训练则一直持续到约3,000步都较为震荡，之后才稳定下来，但在约10,000步时又出现了一个巨大的峰值。

学习率我们无需比较，但值得检查一下以确保我们的训练方式确实正确：

总体来看一切正常。这次本地训练与云端版本确实存在差异，但（如前几次基线训练所示）两者足够相似。A100 和本地 RTX 3090 之间的架构差异可能是造成这种影响的原因。

评估阶段

我将模型上传到了 Hugging Face，现在该进行模型评估了。首先是冒烟测试：

Every effort moves you to give your customers the opportunity you give them in the event of a failure.<|endoftext|>A couple friends

输出还算连贯——而且我认为这是我第一次在冒烟测试中看到 这个 token！不过关键指标仍然是损失值，结果如下：

Loss against our test dataset: 3.538161

让我们将本次本地实验以及之前的基线结果添加到所有干预措施的结果表中：

这太奇怪了！使用干预措施的本地训练版本（1xrtx3090-stacked-interventions），其损失值比云端相同训练配置的版本（8xa100m40-stacked-interventions-1）低了0.039600，表现更好。虽然较低的损失总是好事，但这令人困惑——这个提升幅度甚至超过了某些干预措施本身的效果。

理论上讲，本次本地训练与云端版本之间唯一的变化就是硬件架构。我预期架构差异会带来一定影响，但认为其影响应该很小——正如基线训练 8xa100m40-baseline 和 1xrtx3090-baseline 所显示的那样，两者的损失差异仅为0.007691，大约是前者效果的约五分之一。

然而，在研究噪声对训练损失的影响时，我发现一个现象：改变权重初始化所使用的随机种子（但训练过程本身仍从相同的随机种子开始）对最终模型质量的影响，远比保持权重完全相同但改变训练起始种子的影响要大得多。即：不同初始化但相同训练流程产生的模型标准差，约为相同初始化但不同训练流程产生模型标准差的近两倍。

这个发现很有趣，尽管不能直接类比——那些实验是在固定架构下进行的，每次都在同一台8卡A100 40GB机器上运行。

不过，我觉得至少应该验证一下本地和云端训练是否真的使用了相同的初始权重。我的猜测是它们很可能是一样的，因为权重初始化使用的是确定性的非GPU代码，只要种子相同，无论在哪台机器上都应生成相同的权重。此外，本地与云端基线训练的损失结果高度相似，也支持这一推测。

但值得一试。我创建了一个训练代码的临时分支，该分支在——创建模型之后——只是将模型权重转储到文件中，然后退出。我使用 1xrtx3090-stacked-interventions 配置在本地运行它，然后在 Lambda 上启动了另一台 8x A100 40 GiB 的机器，在那里运行相同的代码，这次使用的是 8xa100m40-stacked-interventions-1 配置，然后将权重 scped 下来。

giles@perry:~/Dev/ddp-base-model-from-scratch (dump-weights)$ diff cloud-init-weights.safetensors local-init-weights.safetensors
giles@perry:~/Dev/ddp-base-model-from-scratch (dump-weights)$

结果一致。这让人放心！

我曾考虑对此进行更多分析；例如，在我对噪声的研究中，我发现保持相同的权重，但改变其余训练运行的随机种子，得到的结果标准差为 0.008672——比带干预措施的本机和云训练之间的差异小四倍以上。我能否用这个数字进行某种比较？

然而，我认为这实际上并不具可比性。那个数字是通过改变随机种子得出的，但保持了相同的架构。并没有充分的理由相信保持种子不变但改变架构会产生相同程度的差异。它们可能更相似，也可能差异更大。

我想我们真正能说的就是，机器的变化以某种方式改变了训练动态的某些方面，碰巧让我们获得了更低的损失。我很容易想象，如果我稍微做点不同的事——比如使用本地的 RTX 4090——结果也可能朝相反的方向发展。

至少让人感到欣慰的是，这种改进比之前最让我信服的那些干预措施要小；唯一比它小的改进是全精度 float32、梯度裁剪和 QKV 偏置——这些我已经认为可能只是因为噪声才带来好处。最重要的是，它与最初从本地训练迁移到更大批次的云训练时看到的 0.252474 的改进相比，小了几个数量级。

结论

所以，我想我已经完成了这一系列训练实验。我们从本地训练的模型开始，它在测试集上的损失为 3.943522，与原始的 GPT-2 small 模型（损失为 3.499677）相比。

我尝试了一系列干预措施，试图让我的模型更接近目标，最终成功将其损失降至 3.538161，几乎达到了目标。这真的很令人满意！

我认为在完全结束这个“干预”小系列并回到从头开始构建主流 LLM 的工作之前，还有两件事要做。

首先，我应该重新审视指令微调测试，我在进行这些训练运行时暂停了这项工作。这将为我们提供一些线索，判断损失改进是仅仅让一个数值变小，还是真正提升了模型的实用性。

其次，我认为我真的需要写一篇总结。我从去年十二月开始断断续续地研究这些内容，我觉得对所做的工作做一个总结会很不错！

我会很快发布；别换台 :-)

这是本系列的下一篇文章的链接。