2、显存管理基础：GPU显存架构简介、PyTorch显存分配机制、TensorFlow显存分配机制

聊到Batch调参，很多人第一反应就是调大batch size能提速。但真正上手后，往往第一个拦路虎就是——显存爆了。

我见过不少同学，模型写得很漂亮，一跑就OOM（Out Of Memory）。然后开始怀疑人生，怀疑代码，甚至怀疑显卡是假的。其实问题没那么玄乎，说白了就是没搞懂显存是怎么工作的。

这一节，咱们就把显存管理这件事掰开揉碎。你理解了底层机制，调参时才能做到心中有数。

2.1 GPU显存架构：一张图看懂

先看一张我手绘的架构图。别被吓到，其实结构很清晰。

这张图我画得比较直观。从上往下看，数据流动的路径是：全局显存 → L2缓存 → 共享内存/寄存器。每一级容量都在缩小，但速度在飙升。

我个人习惯把全局显存想象成一个巨大的仓库，而寄存器就是桌面上的工具。你想想看，从仓库拿东西肯定比从桌面上拿慢得多。所以深度学习训练时，显存带宽往往比显存容量更关键。

核心要点：

全局显存：存模型参数、梯度、优化器状态、中间激活值。容量最大，但最慢。
L2缓存：所有SM共享，自动管理。你不需要手动操作它。
共享内存：每个SM独有，需要手动管理。写CUDA kernel时常用，但PyTorch/TF用户一般碰不到。
寄存器：每个线程私有，编译器自动分配。如果寄存器不够用，会 spill 到局部内存（其实还是全局显存），性能骤降。

避坑指南：

我曾经在调一个超大模型时，发现显存占用比预期高很多。排查了半天，发现是中间激活值没释放。说白了，就是反向传播时，PyTorch为了计算梯度，把前向的激活值都缓存了。这个量级，往往比模型参数本身还大。

2.2 PyTorch显存分配机制

PyTorch的显存管理，核心就一句话：预分配 + 缓存复用。它不像C++那样用多少申请多少，而是提前向驱动要一大块，然后自己内部做二次分配。

为什么会这样？因为每次向驱动申请显存，都是一次昂贵的系统调用。PyTorch为了性能，干脆一次性多要点，后面自己慢慢用。

2.2.1 缓存分配器（Caching Allocator）

PyTorch内部有一个叫 caching_allocator 的组件。它的工作流程是这样的：

申请：当你的代码需要显存时，分配器先看缓存池里有没有大小合适的空闲块。
命中：如果有，直接复用，零开销。
未命中：如果没有，向驱动申请新的显存块（通常是2MB的倍数）。
释放：当张量被销毁时，显存不还给驱动，而是放回缓存池。

嗯，这里要注意：你看到的显存占用，不等于实际使用的显存。PyTorch可能占着茅坑不拉屎——缓存池里有很多空闲块，但驱动显示已经被占用了。

常见误区：

很多人用 nvidia-smi 看到显存占用很高，就以为模型太大。其实可能是缓存碎片导致的。我建议用 torch.cuda.memory_summary() 查看详细分配情况，比nvidia-smi准得多。

2.2.2 关键API与调试

PyTorch提供了一些工具，帮你搞清楚显存到底花在哪了：

import torch

# 查看当前已分配显存（不含缓存）
print(torch.cuda.memory_allocated())

# 查看缓存池大小（已分配但未使用）
print(torch.cuda.memory_reserved())

# 清空缓存池（释放回驱动）
torch.cuda.empty_cache()

# 打印详细分配报告
print(torch.cuda.memory_summary())

我的经验：

调试显存问题时，我习惯先跑一个小batch，看 memory_allocated 和 memory_reserved 的差值。如果差值很大，说明缓存碎片严重。这时候可以尝试 torch.cuda.empty_cache()，但注意这会影响后续性能，因为下次分配又要重新申请。

2.2.3 显存峰值与Batch Size的关系

显存峰值主要来自三个部分：

组件	显存占用	与Batch Size的关系
模型参数	固定	与batch size无关
优化器状态	固定（Adam需要2倍参数显存）	与batch size无关
中间激活值	动态	与batch size成正比
梯度	固定	与batch size无关

看到了吗？中间激活值是显存峰值的主要变量。你调大batch size，激活值线性增长。这就是为什么很多人从batch size=32调到64，直接OOM。

计算公式（粗略估算）：

显存峰值 ≈ 参数显存 + 优化器显存 + 梯度显存 + batch_size × 单样本激活值显存

其中，单样本激活值显存 ≈ 模型深度 × 隐藏层维度 × 精度字节数

2.3 TensorFlow显存分配机制

TensorFlow的显存管理，和PyTorch思路完全不同。它默认是全占模式——启动时就把所有显存都申请了。

我第一次用TensorFlow时，看到nvidia-smi显示显存被占满，吓了一跳。后来才知道，这是TF的默认行为，目的是避免运行时动态分配带来的性能抖动。

2.3.1 全占模式 vs 按需增长

TensorFlow 2.x 默认使用 GPU内存增长 模式，但很多老代码还是全占模式。两者的区别：

模式	行为	优点	缺点
全占模式	启动时申请全部显存	运行时无分配开销，性能稳定	浪费显存，无法多进程共享
按需增长	用多少申请多少	节省显存，可多进程共存	可能有分配延迟

我个人建议，在调试阶段用按需增长模式，方便观察显存变化。生产环境如果独占显卡，全占模式更省心。

2.3.2 如何配置显存策略

TensorFlow 2.x 提供了 tf.config.experimental 来精细控制：

import tensorflow as tf

# 方式一：按需增长
gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

# 方式二：限制最大显存
tf.config.experimental.set_virtual_device_configuration(
    gpus[0],
    [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=4096)]  # 限制4GB
)

注意：

设置 memory_growth 必须在任何TensorFlow操作之前。我曾经在模型加载后才设置，结果没生效，排查了半天才发现是顺序问题。

2.3.3 BFC分配器

TensorFlow底层用的是 BFC（Best-Fit with Coalescing）分配器。这个名字听起来高大上，其实原理很简单：

维护一个空闲块列表，按大小排序。
申请时，找大小最接近的空闲块（Best-Fit）。
释放时，合并相邻的空闲块（Coalescing）。

说白了，就是尽量少产生碎片。我在项目中遇到过，BFC分配器在长时间训练后，碎片率会逐渐升高。这时候可以调用 tf.config.experimental.get_memory_info() 查看碎片情况。