一、显存基础与瓶颈分析
大家好，我是你们的老朋友。今天咱们聊聊显存——这个在深度学习里既熟悉又让人头疼的东西。
说实话，我刚开始做模型训练那会儿，对显存的理解就是「越大越好」。直到有一次，我往一张 24GB 的卡上塞了个 7B 模型，结果直接 OOM，连个 warning 都没给我留。嗯，从那以后我才开始认真研究显存到底是怎么工作的。
1.1 显存架构：HBM vs GDDR
先说说显存的两种主流架构。你可能会问：为什么有的卡用 HBM，有的用 GDDR？
HBM（High Bandwidth Memory），说白了就是堆叠起来的显存。像 HBM2e 和 HBM3，带宽能到 2TB/s 以上。我习惯把 HBM 比作「多层停车场」——占地面积小，但层数多，数据进出快。
GDDR（Graphics Double Data Rate），比如 GDDR6 和 GDDR6X，带宽一般在 500GB/s 到 1TB/s 之间。它更像「平铺的停车场」——占地大，但每层跑得快。
| 特性 | HBM（如 HBM2e） | GDDR（如 GDDR6X） |
|---|---|---|
| 带宽 | ~2 TB/s | ~1 TB/s |
| 容量 | 单颗 8-16GB | 单颗 8-32GB |
| 功耗 | 较低 | 较高 |
| 成本 | 高 | 低 |
| 典型应用 | A100, H100, MI250 | RTX 4090, A6000 |
我在项目中遇到过一个问题：用 RTX 4090 跑大模型推理，batch size 稍微大一点就 OOM。后来换成 A100，虽然单卡容量差不多，但 HBM 的带宽优势让吞吐量直接翻倍。说白了，带宽和容量你得一起看，不能只看一个。
1.2 带宽与容量的权衡
你想想看，显存带宽决定了「数据能多快被喂给计算单元」，而容量决定了「能同时放多少数据」。这两者往往是矛盾的。
- 带宽瓶颈：如果你的模型计算量不大，但数据量很大（比如视频理解），带宽就是短板。我曾经用 GDDR 卡跑一个 3D 卷积模型，计算单元大部分时间在等数据，利用率不到 30%。
- 容量瓶颈：如果你的模型参数多（比如 LLM），容量就是硬伤。我记得有一次训练 13B 模型，batch size 设为 1 都 OOM，最后不得不做梯度累积。
核心原则：带宽和容量，你得根据你的 workload 来选。计算密集型选高带宽，参数密集型选大容量。
1.3 PyTorch 显存分配机制
PyTorch 的显存分配，说白了就是「缓存池」机制。它不会每次需要显存都去问 CUDA 要，而是先预分配一大块，然后自己管理。
为什么会这样？因为 CUDA 的显存分配开销很大。我做过测试，一次 cudaMalloc 大概要几十微秒，而 PyTorch 内部的缓存分配只要几微秒。所以 PyTorch 选择「先囤着，慢慢用」。
# 查看 PyTorch 显存分配情况
import torch
# 分配一个 1GB 的张量
x = torch.randn(256, 1024, 1024, device='cuda')
# 查看当前显存使用
print(torch.cuda.memory_summary())
# 释放张量
del x
# 显存不会立即归还给系统
print(torch.cuda.memory_summary()) # 缓存还在
嗯，这里要注意：del 之后显存不会立刻释放，而是回到 PyTorch 的缓存池。如果你想强制清空缓存，可以调用 torch.cuda.empty_cache()。但我个人建议别频繁调用，因为清空缓存后下次分配又要重新申请，反而更慢。
1.4 常见 OOM 场景剖析
OOM（Out of Memory）是每个炼丹师的噩梦。我总结了几种最常见的场景：
- batch size 太大：这是最直接的。你想想看，batch size 翻倍，显存占用也翻倍。我建议从 batch size=1 开始，逐步往上加。
- 中间激活值爆炸：很多人只算了模型参数，忘了算中间激活值。比如一个 7B 模型，序列长度 2048，单层激活值可能就要 1GB 以上。我遇到过有人把序列长度设到 8192，结果显存直接爆了。
- 梯度累积没清缓存：梯度累积时，如果你不手动清空计算图，显存会越积越多。我曾经踩过这个坑，训练到一半突然 OOM，查了半天才发现是计算图没释放。
- 多进程显存泄漏：用 DataLoader 的 num_workers 时，如果子进程没有正确释放显存，会导致显存慢慢被吃光。我建议用
torch.multiprocessing的set_start_method('spawn')来避免。
避坑指南：我曾经在训练一个 30B 模型时，发现显存占用一直在涨，但明明没有新的张量分配。后来发现是 DataLoader 的 worker 进程里有个 torch.no_grad() 没加，导致计算图被保留。加了之后，显存就稳定了。
1.5 知识体系总览
下面这张图是我自己画的，把本章的核心逻辑串起来了。你可以把它当作一个「显存优化地图」。
个人经验：我建议你在开始训练之前，先用 torch.cuda.memory_summary() 看一下显存分配情况。这样能提前发现潜在问题，而不是等到 OOM 了才去排查。
好了，这一章就到这里。显存这东西，说白了就是「知己知彼，百战不殆」。你理解了它的架构、分配机制和常见坑，后面优化起来就顺手多了。