一、显存基础与瓶颈分析
大家好，我是老张。今天咱们聊聊多卡训练里最让人头疼的问题——显存。说实话，我见过太多团队，模型设计得挺好，一上多卡训练就崩，十有八九是显存没管好。
1.1 GPU显存架构
先说说显存长什么样。你想想看，GPU显存本质上就是个高速仓库。以NVIDIA的Ampere架构为例，A100有40GB或80GB HBM2e显存，带宽能到2TB/s以上。但这里有个关键点——显存不是一整块，而是分成了多个内存控制器和L2缓存。
核心架构特点：
- HBM堆叠技术：通过硅中介层把多个DRAM die堆在一起，带宽高但容量受限
- L2缓存：通常几十MB，命中率直接影响性能
- 内存控制器：每个控制器管理一部分显存，跨控制器访问有额外延迟
我个人习惯把显存想象成一个大超市。L2缓存是收银台附近的货架，拿东西快但地方小。HBM是仓库深处，东西多但来回跑费时间。嗯，这个比喻虽然糙，但道理不糙。
1.2 显存带宽与容量
带宽和容量，这两个指标经常被混为一谈。我简单解释一下：
- 带宽：单位时间内能读写多少数据，好比高速公路的车流量
- 容量：能存多少数据，好比停车场的大小
我在项目中遇到过，有人买了块大显存卡，以为万事大吉。结果训练时发现带宽不够，计算单元一直在等数据，GPU利用率不到30%。说白了，容量够用只是及格线，带宽才是真正的瓶颈。
| GPU型号 | 显存容量 | 显存带宽 | 适用场景 |
|---|---|---|---|
| A100 80GB | 80GB HBM2e | 2TB/s | 大模型训练 |
| V100 32GB | 32GB HBM2 | 900GB/s | 中等规模训练 |
| RTX 4090 | 24GB GDDR6X | 1TB/s | 单卡实验 |
避坑指南：我曾经在V100上跑一个70B模型，显存容量勉强够，但带宽成了瓶颈。后来换成A100，同样的batch size，训练速度提升了近3倍。所以选卡时，别只看容量，带宽同样重要。
1.3 多卡训练中的显存瓶颈
多卡训练时，显存问题会放大。为什么？因为每张卡不仅要存自己的数据，还要处理通信和同步。我总结了几类常见瓶颈：
- 数据并行中的冗余存储：每张卡都存一份完整的模型参数，显存浪费严重
- 梯度同步开销：AllReduce操作需要额外显存来暂存梯度
- 通信缓冲区：NCCL等通信库会预留显存做缓冲区
- 激活值存储：反向传播需要前向计算的中间结果，占大头
你想想看，8张卡跑数据并行，每张卡都存一份完整的模型。如果模型有10GB参数，光参数就占了80GB显存。这还没算梯度和优化器状态。嗯，这里要注意，优化器状态（比如Adam的动量和方差）通常比参数本身还大。
警告：多卡训练时，显存占用不是简单的线性叠加。通信库、框架缓存、临时变量都会额外吃显存。我曾经见过一个案例，8卡训练时单卡显存占用比单卡训练高了40%，就是因为通信缓冲区没优化好。
1.4 显存碎片化问题
最后说说碎片化。这是个容易被忽视但很致命的问题。显存碎片化分两种：
- 内部碎片：分配的内存块比实际需要的大，比如申请3GB但给了4GB
- 外部碎片：空闲显存被分割成很多小块，无法满足大块申请
我遇到过最夸张的一次，训练一个BERT-large模型，显存显示还剩12GB空闲，但一申请4GB的tensor就报OOM。查了半天，发现是外部碎片导致的——空闲显存被分割成了几十个几百MB的小块，凑不出一个连续的4GB空间。
为什么会这样？因为PyTorch的缓存分配器默认会保留已释放的显存，避免频繁向驱动申请。但如果训练过程中tensor大小变化频繁，就容易产生碎片。
碎片化解决方案（个人经验）：
- 设置
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128，限制最大碎片大小 - 使用
torch.cuda.empty_cache()手动清理缓存（但别频繁调用，会影响性能） - 尽量复用固定大小的tensor，避免频繁申请释放
- 考虑使用
torch.cuda.memory_summary()查看碎片情况
说白了，显存管理就像整理房间。东西用完不放回原位，时间长了就乱成一团。我建议养成好习惯：训练前预估显存需求，训练中监控碎片率，训练后及时释放。
好了，这一章就到这里。显存管理是个系统工程，后面我们会一步步深入每个优化点。记住一句话：显存不是无限的，但优化空间是无限的。