一、显存复用概述：为什么需要显存复用？
大家好，我是你们的老朋友。今天咱们聊聊显存复用这个话题。
说实话，做深度学习这几年，我见过太多人卡在显存OOM上。模型跑着跑着，突然蹦出个"CUDA out of memory"，那种感觉，就像你正写到高潮的代码，啪，电脑蓝屏了。嗯，我当年第一次遇到这问题，还以为是自己显卡坏了。
1.1 显存瓶颈：到底有多痛？
先看一组我实际项目中的数据。假设你训练一个BERT-Large模型：
| 组件 | 显存占用（FP32） | 说明 |
|---|---|---|
| 模型参数 | 1.3 GB | 3.4亿参数 × 4字节 |
| 优化器状态 | 5.2 GB | Adam：参数×4（m+v+grad+param） |
| 中间激活值 | 8-12 GB | 取决于batch size和序列长度 |
| 总占用 | 14.5-18.5 GB | 一张V100（32GB）勉强够用 |
你看，光一个BERT就吃掉大半张卡。要是换成GPT-3那种1750亿参数的怪兽，单卡根本别想跑。这就是为什么我们需要显存复用——说白了，就是让同一块显存空间，在不同时间干不同的事。
1.2 OOM的真相：不是显存不够，是管理太差
我在项目中遇到过这样一个案例：某团队训练一个图像分割模型，batch size设了32，结果OOM了。他们第一反应是降batch size。但我一看代码，发现他们在每个iter都创建了新的tensor，旧的也没释放。
为什么会这样？因为PyTorch的显存管理是懒惰的。你创建tensor时，它不一定立即分配显存；但当你不再引用它时，它也不会立即释放。这导致大量"僵尸显存"占着茅坑不拉屎。
核心观点：显存复用不是魔法，而是把显存的生命周期管理做到极致。说白了，就是让每字节显存都物尽其用。
1.3 显存复用的三种基本模式
我个人习惯把显存复用分成三个层次：
- 时间复用：同一块显存，不同时间存不同数据。比如前向传播的激活值，反向传播完就可以回收。
- 空间复用：不同算子共享同一块显存。比如两个不会同时执行的算子，可以用同一块buffer。
- 结构复用：模型结构中的重复模块共享参数。比如Transformer的12层，每层结构一样，但参数不同——这里其实也可以做文章。
你想想看，这三种模式如果能组合使用，显存节省效果是乘法级别的。我见过最极致的案例，有人把32GB的模型压到了8GB跑，靠的就是这三板斧。
1.4 避坑指南：显存复用的代价
我曾经踩过一个坑：为了省显存，我让两个大tensor共享同一块内存。结果反向传播时，梯度更新把另一个tensor的数据给污染了。模型直接炸了，loss变成NaN。
教训：显存复用一定要考虑生命周期和写保护。不是所有tensor都能随便复用的。
显存复用不是免费的午餐。它带来的代价包括：
- 计算开销：内存拷贝、指针重定向需要额外时间
- 调试难度：共享内存的bug极难复现，我排查过三天才找到问题
- 代码复杂度：手动管理显存生命周期，代码可读性下降
1.5 知识体系总览
下面这张图是我自己画的，把显存复用的核心逻辑串起来了：
我的建议：刚开始接触显存复用，别想着一步到位。先从时间复用入手，比如用PyTorch的checkpoint功能。等熟悉了生命周期管理，再尝试空间复用。结构复用需要改模型架构，建议放在最后。
1.6 一个简单的例子：手动管理显存
来看一段代码，感受一下显存复用的基本思路：
# 不好的做法：每次创建新tensor
def bad_forward(x):
a = x @ W1 # 分配显存
b = a @ W2 # 又分配
c = b @ W3 # 再分配
return c
# 好的做法：复用buffer
def good_forward(x, buffer):
# buffer是一块预先分配好的显存
buffer.copy_(x @ W1) # 复用buffer
buffer.copy_(buffer @ W2) # 原地操作
buffer.copy_(buffer @ W3)
return buffer你看，第二种做法只用了1块显存，而不是3块。虽然代码丑了点，但在显存紧张时，这就是救命稻草。
注意：上面的例子为了演示做了简化。实际项目中，in-place操作要非常小心，尤其是涉及自动求导时。我建议用torch.no_grad()包裹，或者用.detach()切断梯度流。
好了，这一章就到这里。显存复用不是银弹，但它是每个深度学习工程师必须掌握的技能。下一章我们会深入时间复用，聊聊checkpointing的具体实现和trade-off。