4. 深度学习框架显存分配机制
搞AI训练的人，十个有九个被OOM折磨过。但你知道吗？很多时候显存不够，不是你的模型真那么大，而是框架分配显存的方式太“野”了。
PyTorch、TensorFlow、JAX，这三个框架的显存管理思路差别很大。我之前调一个百亿参数模型，在PyTorch上跑得好好的，换到TensorFlow直接爆掉——后来才发现是TensorFlow默认申请了所有可用显存。
这一章我就带你捋清楚，它们各自怎么玩显存的，以及你怎么利用这些机制省显存。
4.1 PyTorch：懒惰分配 + 缓存复用
PyTorch默认用CUDA的cudaMalloc分配显存，但有个特点：它不会一次性把所有显存占满。PyTorch只在你要往GPU放张量的时候才去申请，而且申请了也不轻易还给系统，而是放回自己的缓存池里。
这个池子叫缓存分配器（Caching Allocator）。你释放了张量，PyTorch把这块显存标记为空闲，下次需要同样大小或更小的张量时，直接复用。这样做的好处是省掉了cudaFree的开销，但坏处是——你看着nvidia-smi显示的占用一直没降，以为泄露了。
import torch
print(torch.cuda.memory_summary())我个人习惯在debug时加一行torch.cuda.empty_cache()，虽然只清空缓存，不能真把显存还给系统，但至少能让我知道“实际占用”是多少。
4.1.1 默认行为
- 按需分配：只有张量创建或操作需要显存时才调用cudaMalloc。
- 缓存复用：释放的张量不立即归还给系统，保留在分配器缓存中。
- 自动扩大：当缓存不够时，PyTorch会额外申请一大块（默认2MB对齐），可能产生碎片。
4.1.2 关键环境变量
| 变量 | 作用 |
|---|---|
PYTORCH_CUDA_ALLOC_CONF | 可设置max_split_size_mb、roundup_power2_divisions等，控制碎片 |
CUDA_LAUNCH_BLOCKING=1 | 同步执行，方便定位显存分配位置（慎用，大幅降低性能） |
with torch.no_grad():后，显存直接砍半。PyTorch的自动求导会保留中间结果，你一定要确认哪些变量不需要梯度。
4.2 TensorFlow：霸占式分配
TensorFlow（包括Keras）的默认行为更激进——它一启动就把几乎全部可用显存都申请了。你跑个tf.config.list_physical_devices('GPU')，然后啥也没干，nvidia-smi一看，占用已经4、5个G了。
为什么这么设计？TF早期版本为了减少分配碎片，直接占山为王。但这在多卡环境或同时跑多个实验时，简直灾难。
4.2.1 设置显存按需增长
好在TF提供了配置接口，你可以让它像PyTorch一样“用多少拿多少”：
import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
try:
tf.config.set_logical_device_configuration(
gpus[0],
[tf.config.LogicalDeviceConfiguration(memory_limit=4096)])
except RuntimeError as e:
print(e)或者设置memory_growth=True：
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)你想想看，如果跑生产环境，多个模型共享一张卡，你忘了设memory_growth，卡就被第一个模型占死了。
4.2.2 图执行与显存生命周期
TF使用静态计算图，显存分配在图构建阶段就已经规划好了。和PyTorch的动态图相比，TF的显存使用更加“计划性”，但不够灵活。当你调用tf.function时，TF会把函数内的操作编译成图，然后为中间结果预先分配显存。一旦图变得复杂，可能瞬间吃掉大量显存。
我记得有一次把一个小模型加上tf.function后，显存从1.2G跳到了2.8G。排查后才发现，TF的自动微分在图中生成了大量中间节点。
- 使用
tf.profiler.experimental.Profile查看各op的显存消耗。 - 设置
TF_GPU_ALLOCATOR=cuda_malloc_async（TF 2.7+）切换到异步分配，减少碎片。 - 不要忘记
tf.keras.backend.clear_session()，释放已构建的图资源。
4.3 JAX：函数式 + 显存显式管理
JAX的设计哲学和PyTorch、TF完全不同。它基于函数式编程——没有“状态”的概念，一切变换都是纯函数。这意味着显存管理也变成了显式的。
在JAX里，每个张量（叫DeviceArray）的生存期非常明确。你不保留引用，它就立刻被回收。JAX没有缓存池，没有“显存占用一直下不去”的问题。
4.3.1 编译时分配
JAX用XLA做即时编译（JIT）。当你@jit装饰一个函数时，XLA会分析输入输出形状，提前分配好所有中间结果需要的显存。这个分配是一次性的，训练过程中不再变动。所以JAX的显存占用曲线通常是阶梯状的，不会有渐进式增长。
但是——XLA对动态形状极度不友好。如果你输入形状变动频繁，XLA会不断重新编译，每次编译都分配新显存，旧的不释放，内存会像漏了一样上涨。我曾在处理变长序列时踩过这个坑，最后靠jnp.pad把所有输入padding到固定长度才解决。
4.3.2 显存清理
# 强制释放JAX显存
from jax import lax
lax._force_dcache_cleanup()另外，JAX的pmap在多卡并行时，每张卡都有独立的显存空间。如果某个设备上pmap的函数出了异常，显存可能不被释放。这种情况下，重启JAX进程是最干净的解决办法（虽然粗暴）。
4.3.3 与PyTorch的对比
| 特性 | JAX | PyTorch |
|---|---|---|
| 分配时机 | JIT编译时一次性分配 | 运行时按需分配，缓存复用 |
| 碎片问题 | 少（形状固定时） | 较多，尤其是动态形状 |
| 调试难度 | 低（分配清晰） | 中等（缓存干扰） |
| 动态形状支持 | 差（触发重编译） | 好 |
4.4 三框架显存管理核心差异
我把关键差异整理成下面这张图，你一看就懂了：
4.5 实战：跨框架显存优化策略
不管你用哪个框架，下面几条经验都通用：
- 先确认框架版本——不同版本的默认行为可能不同。比如PyTorch 1.10+引入了
cudaMallocAsync后端，能减少碎片。 - 设置环境变量限制最大显存，防止框架“抢”光整张卡。
- 手动清理不需要的中间张量——PyTorch用
del+torch.cuda.empty_cache()，TF用keras.backend.clear_session()，JAX保持无状态即可。 - 监控碎片率：PyTorch可以用
torch.cuda.memory_stats()查看碎片比例，大于15%就要考虑优化了。
4.6 避坑指南：框架切换时的显存陷阱
如果你从PyTorch迁移到JAX，或者从TF迁移到PyTorch，下面几个坑我亲自踩过：
- PyTorch → JAX： 别把PyTorch的
cuda()写法直接搬过来。JAX的显存是纯函数式的，你创建的张量一旦离开作用域就会被回收。如果你在循环里创建张量又不显式保留引用，就会崩溃。 - TF → PyTorch： 记得关掉TF的显存增长模式，否则PyTorch分配时可能发现显存不够用。我曾经犯过傻，两个框架混用，TF占了卡，PyTorch只能分到200M。
- 多进程训练： 每个子进程都会独立申请显存。如果父进程已经占用了大部分显存，子进程会启动失败。解决方案是设置
CUDA_VISIBLE_DEVICES或者用torch.multiprocessing.set_start_method('spawn')。