4、显存复用技术：Tensor In-Place操作、梯度累积、激活值重计算

显存，说白了就是钱。做多卡训练时，你经常会发现——卡还没跑起来，显存先爆了。我早期做百亿参数模型时，最头疼的就是这个。后来慢慢摸清了几个套路，今天跟你聊聊。

显存复用技术，核心就三招：Tensor In-Place操作、梯度累积、激活值重计算（Activation Checkpointing）。这三招用好了，显存占用能砍掉一半以上。

4.1 Tensor In-Place操作：原地修改，省一份拷贝

先讲最简单的。In-Place操作，就是直接在原tensor上改，不新分配内存。

举个例子，PyTorch里 a.add_(b) 和 a + b 的区别。前者直接在 a 的内存上累加，后者会新创建一个tensor。

我习惯在以下场景用In-Place：

• 激活函数：比如 torch.relu_(x)，省掉一份中间结果
• 梯度裁剪：直接在梯度tensor上做clamp
• 权重更新：optimizer.step() 内部其实都是In-Place

但要注意——反向传播会受影响。你想想看，如果前向时把输入改了，反向计算梯度时拿到的就是错误的值。所以In-Place操作只适合那些不需要保留原始值的场景。

4.2 梯度累积：用时间换空间

梯度累积，说白了就是——攒几个batch的梯度，再统一更新参数。

为什么能省显存？因为每个batch的中间结果（激活值）用完就释放了，只保留梯度。你想想看，原本batch size=32才能训，现在batch size=8，累积4步，效果差不多，但显存占用直接降到1/4。

实现上很简单：

# 梯度累积示例
accumulation_steps = 4
optimizer.zero_grad()

for i, (inputs, labels) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps  # 平均一下
    loss.backward()
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

嗯，这里要注意几个细节：

• Loss要除以累积步数，否则梯度会放大
• BatchNorm层会受影响——它统计的是每个micro-batch的均值和方差，不是全局的。我建议用SyncBN或者干脆冻结BN
• 多卡场景下，梯度累积配合分布式训练，效果更好。每个卡算自己的梯度，累积到一定步数再all-reduce

4.3 激活值重计算（Activation Checkpointing）

这是最狠的一招。说白了就是——前向时不保存中间激活值，反向时重新算一遍。

你想想看，一个100层的Transformer，前向时每层都要保存激活值。如果每层保存10MB，100层就是1GB。用了Checkpointing，只保存输入和输出，中间的全扔掉。反向时从输入重新跑一遍前向，拿到中间值再算梯度。

代价是什么？多了一次前向计算。时间换空间，大概多花20-30%的计算时间，但显存能省50-70%。

PyTorch里用起来很简单：

from torch.utils.checkpoint import checkpoint

# 原始前向
def forward(self, x):
    x = self.layer1(x)
    x = self.layer2(x)
    return x

# 使用checkpoint
def forward(self, x):
    x = checkpoint(self.layer1, x)
    x = checkpoint(self.layer2, x)
    return x

但要注意——不是所有层都适合用Checkpointing。我建议只对计算量大的层用，比如Transformer的Attention和FFN。像Embedding、LayerNorm这种轻量层，省不了多少显存，反而增加调度开销。

4.4 三种技术的对比与选择

这三种技术各有适用场景。我整理了一个表格，方便你对比：

技术	省显存效果	额外计算开销	适用场景
In-Place操作	低（5-15%）	无	激活函数、梯度裁剪
梯度累积	中（30-50%）	无（但收敛变慢）	大batch训练、显存不足
激活值重计算	高（50-70%）	20-30%	深层模型、Transformer

我个人习惯是：先上In-Place操作，零成本，能省一点是一点。如果还不够，加梯度累积。最后实在不行，上Checkpointing。三种技术可以叠加使用，效果更好。

4.5 多卡场景下的显存复用

多卡训练时，显存复用有个额外的技巧——通信和计算重叠。

你想想看，All-Reduce通信时，GPU在等数据，这时候显存是空闲的。我们可以利用这个时间做点别的事，比如释放一些中间tensor，或者提前计算下一层的梯度。

具体做法：

• 梯度异步通信：用 torch.distributed.all_reduce 的async版本，不阻塞计算
• 分层释放：每算完一层的梯度，立即释放该层的激活值，不等整个模型算完
• 显存池化：预分配一块大显存，重复利用，避免频繁malloc/free

4.6 核心逻辑图

下面这张图展示了三种显存复用技术的核心逻辑和关系：

从这张图可以看出，三种技术是递进关系。In-Place操作最简单，零成本；梯度累积中等复杂度；Checkpointing效果最好但代价也最大。多卡场景下，还要考虑通信和计算的叠加优化。