小显存如何靠梯度累积训练大模型

小显存训练大模型的核心思路是：用时间换空间，通过梯度累积（Gradient Accumulation）把多个小 batch 的梯度“攒起来”，再一次性更新参数，从而模拟大 batch 训练。下面从原理、做法、注意事项到进阶技巧系统讲一下。

一、为什么小显存也能训练大模型？

显存主要被这几部分占用：

1. 模型参数（Parameters）
2. 前向激活值（Activations）
3. 梯度（Gradients）
4. 优化器状态（Optimizer states，如 Adam 的 m、v）

梯度累积主要解决的是：

无法一次塞进大 batch，但希望保持大 batch 的更新效果

二、梯度累积的核心原理

正常训练流程

for each batch:
forward
backward  → 计算梯度
step      → 更新参数

梯度累积训练流程

for micro_batch in range(k):
forward
backward  → 梯度累加（不更新参数）
step  → 累积 k 次梯度后，统一更新参数

数学上等价于：

大 batch = k × micro_batch_size

三、最小可用示例（PyTorch）

普通训练

for x, y in dataloader:
out = model(x)
loss = criterion(out, y)
loss.backward()
optimizer.step()
optimizer.zero_grad()

梯度累积版（显存更小）

accumulation_steps = 4  # 累积 4 次
for i, (x, y) in enumerate(dataloader):
out = model(x)
loss = criterion(out, y)
# 关键：除以累积步数
loss = loss / accumulation_steps
loss.backward()
if (i + 1) % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()

✅ 显存只取决于 micro batch size

✅ 梯度等价于大 batch

四、为什么 loss 要除以 accumulation_steps？

因为：

• loss.backward() 会累加梯度
• 累积 k 次后，梯度会变成原来的 k 倍

所以：

loss = loss / k

保证：

梯度 ≈ 大 batch 的平均梯度

五、小显存训练的完整组合拳（强烈推荐）

✅ 1. 梯度累积（核心）

• micro batch size = 1～4
• accumulation steps = 8～64

✅ 2. 混合精度（AMP）

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for x, y in dataloader:
with autocast():
out = model(x)
loss = criterion(out, y) / accumulation_steps
scaler.scale(loss).backward()
if step % k == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()

✅ 显存减少 30～50%

✅ 训练速度提升

✅ 3. ZeRO / FSDP（进阶）

• ZeRO-1/2：优化器状态、梯度分片
• FSDP：参数也分片

适合：

• 单卡 12–24GB
• 想训 7B、13B 模型

✅ 4. 梯度检查点（Gradient Checkpointing）

model.gradient_checkpointing_enable()

✅ 用 30% 时间换 50% 激活显存

✅ 5. 其他实用技巧

六、梯度累积的注意事项（很重要）

⚠️ 1. BatchNorm 会出问题

• BN 统计的是 micro batch
• 替代方案：
• 用 LayerNorm
• 或 同步 BN（多卡）

⚠️ 2. 学习率要调整

经验法则：

new_lr ≈ base_lr × (global_batch_size / base_batch_size)

⚠️ 3. 并不是“无限堆”

• 累积步数太大：
• 梯度噪声下降
• 收敛变慢
• 一般 4～32 步最常见

七、一个现实配置示例

单卡 16GB，训练 1.3B 模型
| 项目 | 设置 |

八、一句话总结

小显存训练大模型 = 小 batch + 梯度累积 + 混合精度 +（可选）ZeRO / LoRA
如果你愿意，可以告诉我：

• 显存大小（如 8G / 16G / 24G）
• 模型规模（如 7B / 13B）
• 单卡还是多卡

我可以直接给你可运行的训练配置。