返回关卡地图
闯关题：0 / 1 XP 0 / 100

Level 12 | 零基础导学关卡

Gradient Accumulation

先用小例子把必要知识学会，再用 1 道闯关题检查理解，最后回到 notebook 写真正的 PyTorch 作业。

12_Gradient_Accumulation.ipynb Gradient AccumulationMicro-batchMemory
概念输入
闯关检查
Notebook 作业
effective_batch = micro_batch × accumulation_steps

Notebook 原文导学

来自 12_Gradient_Accumulation.ipynb

12. Gradient Accumulation | 梯度累积

难度： Medium | 环境： CPU-first | 标签： 训练技巧, 显存优化, PyTorch | 目标人群： 模型微调与工程部署

在做大模型微调时，显存通常先被 batch size 吃满。梯度累积的核心思路是：把一个大 batch 拆成多个 micro-batch，分多次 backward，最后只 step 一次，从而在不增加峰值显存的前提下，模拟更大的有效 batch。可以先把它记成：把大 batch 拆成小块，分多次反传，最后一次更新。

关键词： gradient accumulation, micro-batch, effective batch

前置阅读

导语： 先把梯度累积需要的 PyTorch 训练接口和张量操作补齐，再来看小 batch 聚合更新。

  • [P0: 09. PyTorch nn.Module Basics | nn.Module 基础](../00_Prerequisites/09_PyTorch_nn_Module_Basics.ipynb)
  • [P0: 11. PyTorch Optimizers and Loss | 优化器与损失](../00_Prerequisites/11_PyTorch_Optimizers_and_Loss.ipynb)
  • [P0: 13. Simple Neural Network Training | 简单神经网络训练](../00_Prerequisites/13_Simple_Neural_Network_Training.ipynb)

相关阅读

导语： 梯度累积和训练性能、显存占用关系紧密，可结合硬件和 profiling 一起看。

  • [P1: 06. VRAM Calculation and ZeRO | 显存估算与 ZeRO](../01_Hardware_Math_and_Systems/06_VRAM_Calculation_and_ZeRO.ipynb)
  • [P1: 13. Profiling and Bottleneck Analysis | 性能分析与瓶颈定位](../01_Hardware_Math_and_Systems/13_Profiling_and_Bottleneck_Analysis.ipynb)
  • [P1: 20. NCCL and AllReduce Basics | NCCL 与 AllReduce 基础](../01_Hardware_Math_and_Systems/20_NCCL_and_AllReduce_Basics.ipynb)

Step 1: 为什么需要梯度累积

这一节先把“大 batch 更稳，但显存不够”这个矛盾说清楚。

大 batch 的好处：梯度更稳定，更新方向更平滑。

但显存不够怎么办？

- 直接增大 batch 往往会先爆显存。

- 梯度累积通过“多次反传、一次更新”保留大 batch 的优化效果。

- 只要每个 micro-batch 的 loss 按 accum_steps 做缩放，最终效果就和一次性喂入大 batch 非常接近。

Step 2: 数学等价性

这里要抓住的重点不是公式本身，而是缩放和反传的顺序。

设一个完整 batch 被切成 K 个 micro-batch。若每个 micro-batch 的损失记为 L_i，则梯度累积相当于计算：

$$

\nabla L = \frac{1}{K} \sum_{i=1}^{K} \nabla L_i

$$

工程上最关键的细节只有两个：

  • 每次 backward() 前把 loss 除以 accum_steps。
  • 只在最后一个 micro-batch 后执行 optimizer.step() 和 optimizer.zero_grad()。

这也是为什么 train_step_with_accumulation 的实现要先缩放再反传。

Step 3: 代码实现框架

这一段把完整 batch 和累积 batch 的更新路径并排对齐。

下面我们实现两个更新步骤：

  • train_step_full_batch：一次性使用完整 batch 更新。
  • train_step_with_accumulation：把 batch 拆成多个 micro-batch，累积梯度后再更新。

这两种方式在等价条件下，参数更新应该几乎一致。

参考代码与解析

代码

答案与直觉

  • 这一题要解决什么：把大 batch 的更新效果用 micro-batch 累积模拟出来。
  • 为什么这样做：显存不够时靠多次 backward、一次 step 保持等价更新。
  • 带走的直觉：梯度累积的关键不是拆 batch，而是保持梯度尺度不变并延后参数更新。

1. TODO 1 (切分当前 micro-batch)

  • 切分逻辑： 梯度累积不是一次喂完整 batch，而是先把 x / y 按 accum_steps 拆成多个 micro-batch。
  • 训练目标： 每一轮循环都只处理当前片段，这样才能模拟大 batch 的效果，同时把峰值显存压低。
  • 实现重点： 这里要先确定当前 micro-batch 的切片范围，再把输入和标签切出来。

2. TODO 2 (缩放 loss 并反传)

  • 梯度对齐： 每个 micro-batch 的 loss 必须先除以 accum_steps，再执行 backward()。
  • 等价性： 这样累积出来的总梯度才和完整 batch 的梯度一致，不会悄悄把更新幅度放大 accum_steps 倍。
  • 实现重点： 这一层的核心是“先缩放，再反传，再累加”。

3. TODO 3 (统一更新参数并返回累计 loss)

  • 先攒后更： 所有 micro-batch 都完成 backward 之后，再统一执行一次 optimizer.step() 和 optimizer.zero_grad()。
  • 闭环意义： 这样一次参数更新就等价于完整 batch 的更新，梯度累积的逻辑才真正闭环。
  • 结果记录： 最后返回累计 history 或 total_loss，方便观察训练过程中 loss 是否下降。

4. 进阶思考：为什么要做重复样本验证？

  • 一致性检查： 通过重复样本验证，可以确认梯度累积是否真的等价于完整 batch。
  • 工程价值： 只要这套链路对齐，后续再切换更复杂的数据和更大的 batch 也更稳。
  • 实践意义： 这一步把 SFT Loss、梯度累积、参数更新 连接成一个可运行的小闭环。
Mission 1 本关 TODO

Gradient Accumulation：先建立本关直觉

先补的知识

  • 累积期间不能提前 optimizer.step。
  • 每个 micro-batch 的 loss 要按累积步数缩放。
  • 最后一次统一更新，才能近似等价于大 batch。

图解原理

先不要背实现。把本关拆成输入、核心操作、输出三段，写代码时就不容易迷路。

1. 先认出本关的核心对象

把一个大 batch 拆成多个 micro-batch 累积梯度，用更多计算步骤换更低峰值显存。

Gradient AccumulationMicro-batchMemory
把名词翻译成张量、函数或训练步骤
2. 再看输入输出关系

每个 notebook 的 TODO 都在训练一种固定动作：先确认输入是什么，再确认输出应该长什么样，最后用测试检查数值或 shape。

输入核心操作输出/测试

语法热身：读 PyTorch 练习时先抓这三件事

下面是通用读法，不直接等同于本关 notebook 答案。

# 1. 先看 shape，别急着写公式
print(x.shape)

# 2. 写模块时，参数通常放在 __init__
class TinyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.proj = nn.Linear(4, 4, bias=False)

    def forward(self, x):
        return self.proj(x)

# 3. 测试通常会检查 shape 或数值接近
assert out.shape[0] == x.shape[0]
torch.allclose(out_a, out_b, atol=1e-5)

读法：先找 TODO 附近的输入变量和期望输出，再回到课程区确认这个操作在数学上想做什么。

闯关题

梯度累积时为什么常把每个 micro-batch loss 除以累积步数？

回到 notebook 的作业

这里不直接写答案。你已经拿到足够输入，最后用 notebook 的 TODO 做举一反三。

完成 1 道闯关题后，本关即算完成；作业 checklist 用来辅助你回 notebook 练习。
进度只保存在当前浏览器 localStorage，分享 HTML 不会带走你的记录。
上一关：L11 下一关：L13