4. 梯度检查点与显存复用

说实话，做过多卡训练的人，十有八九都被显存OOM折磨过。我印象最深的一次，是在训练一个百亿参数模型时，明明每张卡都有80GB显存，结果batch size开到2就炸了。那时候我就在想——显存都去哪了？

后来我发现，大部分显存都被中间激活值吃掉了。你想想看，前向传播时每一层都要保存激活值，反向传播才能算梯度。模型越深，激活值越恐怖。那怎么办？梯度检查点（Checkpointing）就是干这个的。

Checkpointing原理：用计算换显存

核心思想其实很简单——不保存所有中间激活值，只保存关键节点的激活值。反向传播时，需要哪些中间值，就重新算一遍。

举个例子，一个50层的Transformer，正常训练需要保存50份激活值。用了Checkpointing后，我只在第10、20、30、40层保存激活值。反向传播时，比如要算第25层的梯度，我就从第20层的激活值开始，重新前向计算到第25层。

我个人习惯把Checkpointing分成两种：

• 手动Checkpointing：自己选哪些层做检查点，灵活性高但麻烦
• 自动Checkpointing：框架帮你决定，比如PyTorch的torch.utils.checkpoint

代码实现其实不复杂：

import torch
from torch.utils.checkpoint import checkpoint

class TransformerLayer(nn.Module):
    def forward(self, x):
        # 用checkpoint包装前向计算
        return checkpoint(self._forward_impl, x, use_reentrant=False)
    
    def _forward_impl(self, x):
        # 实际的层计算
        attn_out = self.self_attention(x)
        ff_out = self.ffn(attn_out)
        return ff_out

我在项目中遇到过一个问题——use_reentrant参数。早期版本默认是True，但容易导致梯度计算错误。我建议新项目直接用use_reentrant=False，虽然多占一点显存，但稳定得多。

梯度累积策略：小batch也能跑大模型

梯度累积，说白了就是攒够梯度再更新。你想想看，如果显存只够跑batch size=4，但模型需要batch size=32才能收敛，怎么办？

那就跑8次前向+反向，每次梯度都累加起来，第8次再更新参数。这样等效于batch size=32，但显存占用只有batch size=4的水平。

实现上要注意几个坑：

accumulation_steps = 8
optimizer.zero_grad()

for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    
    # 除以累积步数，保持loss尺度一致
    loss = loss / accumulation_steps
    loss.backward()
    
    # 每accumulation_steps步更新一次
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

我曾经踩过一个坑——忘记把loss除以累积步数。结果梯度被放大了8倍，模型直接炸了。嗯，这种低级错误犯过一次就不会再犯了。

显存池复用：别让显存碎片化

多卡训练时，显存分配和释放非常频繁。如果不加管理，显存会变得支离破碎——明明总显存还有20GB，但就是分配不出一块连续的10GB空间。

显存池复用的思路是：预先申请一大块显存，训练过程中重复使用。就像你租了一个仓库，里面放了很多箱子，用完的箱子不扔掉，而是清空后给下一个用。

PyTorch的缓存分配器其实已经做了这件事：

# 查看显存使用情况
print(torch.cuda.memory_summary())

# 清空缓存（不推荐频繁使用）
torch.cuda.empty_cache()

# 设置显存分配策略
torch.cuda.set_per_process_memory_fraction(0.9)  # 限制使用90%显存

我自己在项目中常用的技巧是：

• 用torch.cuda.max_memory_allocated()监控峰值显存
• 用torch.cuda.reset_peak_memory_stats()重置统计
• 在关键节点打印显存使用，定位OOM源头

Inplace操作技巧：省显存但危险

Inplace操作，就是直接在原变量上修改，不创建新变量。比如a += 1就比a = a + 1省显存。

但inplace操作在深度学习里是个双刃剑。好处是省显存，坏处是容易破坏计算图，导致反向传播出错。

我总结了几条安全使用原则：

1. 只在不需要梯度的张量上用inplace，比如输入数据预处理
2. ReLU的inplace版本是安全的，因为它的梯度计算不依赖输入值
3. 避免对共享张量做inplace操作，比如多个层共用一个权重

# 安全的inplace用法
x = torch.randn(100, 100, requires_grad=True)
x.relu_()  # 不推荐，会破坏计算图

# 推荐做法
x = x.relu()  # 创建新张量，保留计算图

# 安全的场景：数据预处理
data = data.cuda().half()  # 直接覆盖原变量，没问题

我曾经在项目里用nn.ReLU(inplace=True)，结果模型训练到一半梯度变成NaN。排查了半天才发现是inplace操作导致的计算图异常。从那以后，我所有inplace操作都加了注释，标注为什么这里可以安全使用。

知识体系总览

下面这张图是我自己整理的显存优化知识体系，你可以对照着看：

说实话，这四种技术单独用效果有限，但组合起来威力巨大。我最近一个项目，用Checkpointing省了60%显存，配合梯度累积把batch size从4提到了32，再加上显存池复用避免了碎片化，最终训练速度反而比原来还快了10%。

嗯，显存优化这块内容就这些。记住一个原则——没有银弹。每种技术都有代价，关键是在你的场景下找到最优平衡点。