1. 显存去哪了

大家好，我是你们的老朋友。今天咱们聊聊一个让所有炼丹师都头疼的问题——显存。

说实话，我刚开始做深度学习那会儿，最怕的就是 CUDA out of memory。明明模型不大，怎么一跑就爆显存？后来我仔细算了一笔账，才发现显存这玩意儿，真不是只给模型参数住的。

一、显存的四大“住户”

咱们把显存想象成一套房子。模型训练时，里面住了四户人家：

• 模型参数：就是你的权重矩阵，最显眼的住户
• 梯度：反向传播时产生的中间产物
• 优化器状态：比如Adam里的动量、方差，这玩意儿特别能吃显存
• 中间激活值：前向传播时每层的输出，这个才是真正的“显存杀手”

我见过太多新手，以为显存只存参数。你想想看，一个1亿参数的模型，FP32下也就400MB。但实际训练时，显存占用轻松上10GB。为什么？因为中间激活值才是大头！

二、逐项拆解：每样东西到底吃多少显存？

咱们拿一个具体的例子说话。假设你有一个BERT-base模型，1.1亿参数，用FP32训练，batch size=32，序列长度=128。

项目	计算公式	显存占用	占比
模型参数	1.1亿 × 4字节	~420 MB	3%
梯度	同参数大小	~420 MB	3%
优化器状态(Adam)	参数 × 2 × 4字节	~840 MB	6%
中间激活值	每层输出 × 层数 × batch	~12 GB	88%

看到没？中间激活值占了将近9成！我当年第一次算这个数的时候，整个人都傻了。原来我一直在跟空气斗智斗勇，优化了半天参数，结果大头在激活值上。

三、混合精度训练：从源头砍掉一半显存

好了，问题找到了。怎么解决？混合精度训练（AMP）就是干这个的。

它的核心思路很简单：能用半精度（FP16）的地方，绝不用全精度（FP32）。

具体来说：

• 前向传播：用FP16计算，中间激活值直接减半
• 反向传播：梯度也用FP16存储
• 权重更新：保留一份FP32的权重副本（主权重），保证精度

你想想看，中间激活值从FP32变成FP16，直接砍掉一半。12GB变6GB，这效果立竿见影。

四、AMP的显存节省明细

咱们还是拿BERT-base举例，看看AMP到底省在哪：

项目	FP32训练	AMP训练	节省
模型参数	420 MB	420 MB (主权重) + 210 MB (FP16副本)	无节省（反而多了一点）
梯度	420 MB	210 MB (FP16)	50%
优化器状态	840 MB	840 MB (仍用FP32)	无节省
中间激活值	12 GB	6 GB (FP16)	50%
总计	~13.7 GB	~7.5 GB	~45%

注意看，参数和优化器状态其实没怎么省。但中间激活值砍掉6GB，梯度砍掉200MB，加起来就是实打实的6.2GB。这就是AMP的威力所在。

五、一张图看懂AMP的显存流向

下面这张SVG图，是我自己画的AMP显存流向图。你可以看到数据是怎么流动的，以及每个阶段显存是怎么变化的。

从这张图可以看得很清楚：前向传播和反向传播都在FP16下进行，中间激活值和梯度直接减半。只有优化器更新这一步，需要FP32的主权重来保证精度。这就是“混合精度”这个名字的由来——FP16和FP32混着用。

六、实际代码：怎么在PyTorch里开AMP？

说了这么多理论，咱们来点实际的。PyTorch从1.6开始就原生支持AMP了，用起来特别简单。

import torch
from torch.cuda.amp import autocast, GradScaler

# 初始化
model = MyModel().cuda()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
scaler = GradScaler()  # 这个很重要，后面会讲

for data, target in dataloader:
    data, target = data.cuda(), target.cuda()
    
    # 前向传播用autocast自动管理精度
    with autocast():
        output = model(data)
        loss = loss_fn(output, target)
    
    # 反向传播用scaler缩放梯度
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

你看，就多了三行代码：autocast、GradScaler、scaler.scale。但显存直接砍半。我刚开始用的时候，心里还犯嘀咕：这么简单？结果一试，真香。

七、AMP的代价：精度和速度的权衡

当然，天下没有免费的午餐。AMP也有它的代价：

• 精度损失：FP16的数值范围只有FP32的1/16，某些小梯度可能会被截断
• 额外开销：FP32和FP16之间的转换需要时间，不过GPU有专门的硬件单元，基本可以忽略
• 调试难度：如果loss爆炸，你得先排查是不是AMP的问题

我个人习惯是：先跑一个epoch的FP32作为baseline，再开AMP对比。如果loss曲线基本重合，那就放心用。如果出现差异，再考虑调整loss scaling策略。

好了，这一讲就到这里。显存去哪了这个问题，说白了就是四个住户在抢房子。AMP做的事情，就是让其中两个住户（中间激活值和梯度）住小一点的房间。效果嘛，你试试就知道了。