显存管理基础：GPU显存架构、显存分配与释放机制、显存碎片化问题

聊到Batch队列优化，显存管理是绕不开的坎。我见过太多同学模型写得飞起，结果一跑就OOM（显存溢出），然后一脸懵地来问我：「明明显存还有好几个G，怎么就爆了？」

嗯，这背后其实藏着三个核心问题：显存长什么样、系统怎么给你分显存、以及为什么分着分着就碎了。今天咱们就把这三件事掰开揉碎讲清楚。

先看一张图，这是我个人习惯用来理解显存层次的方式：

这张图我想强调一个关键点：显存不是一块「大平地」。它是有层次的，从最慢但最大的HBM显存，到极快但极小的寄存器，每一层都在跟延迟做斗争。

我在项目中遇到过最典型的误解：有人觉得显存就是一块连续的大内存，随便怎么分配都行。其实不然。GPU的显存控制器（Memory Controller）管理着多个显存颗粒，数据在颗粒间是交错存储的。你分配一块显存，系统可能把它打散到多个颗粒上——这本身是为了带宽，但也带来了碎片化的隐患。

  核心要点：显存层次决定了数据访问模式。频繁在HBM和L2之间搬运数据，会吃掉大量带宽。Batch队列优化的本质，就是尽量减少这种「跨层搬运」。

好，现在咱们聊聊分配和释放。很多同学写PyTorch代码时，习惯这样：

import torch

# 分配一块显存
x = torch.randn(1024, 1024).cuda()
# 做一些计算
y = x * 2
# 释放
del x
del y
torch.cuda.empty_cache()

然后跑一下`nvidia-smi`，发现显存占用没变。为什么？

原因很简单：CUDA的显存分配器是「懒惰」的。它从驱动那里申请显存时，一次会拿一大块（比如2GB），然后自己内部切成小块分给各个tensor。当你`del`掉tensor时，这块显存只是回到了CUDA分配器的「空闲池」里，并没有真正还给操作系统。

说白了，你看到的`torch.cuda.empty_cache()`，只是让CUDA分配器把空闲池里的显存归还给驱动，但驱动也不一定立刻还给OS。这就像你去图书馆还书，管理员把书放回书架，但书架上的空间并没有立刻消失。

我的建议：不要频繁调用`empty_cache()`。它本身有开销，而且会打断CUDA的分配策略。我一般只在训练完一个epoch、准备做验证时调用一次，确保验证时显存干净。

那什么时候显存会真正归还？嗯，只有进程退出时。或者你显式调用了`cuda.cudart().cudaDeviceReset()`——但这会重置整个设备，代价很大。

这是最坑的地方。显存碎片化分两种：

我举个例子你就明白了。假设显存里有这么几块：

|--- 已用 100MB ---|-- 空闲 50MB --|--- 已用 200MB ---|-- 空闲 30MB --|--- 已用 150MB ---|-- 空闲 20MB --|

总空闲是100MB，但最大连续块只有50MB。这时候你申请一个80MB的tensor，就会报OOM。明明显存还有100MB空闲，但就是没法用。

我曾经在一个NLP项目里遇到过这个问题。模型训练到一半，突然OOM。我查了显存占用，发现还有1.2GB空闲，但就是分配不出一个512MB的tensor。后来用`torch.cuda.memory_summary()`一看，好家伙，空闲块有300多个，最大的才30MB。这就是典型的外部碎片。

避坑指南：我曾经因为动态创建和销毁大量小tensor，导致显存碎片化严重。后来改成预先分配大块显存（比如用torch.empty()预分配），然后手动管理子区域，碎片问题大幅缓解。

说了这么多，怎么知道自己有没有碎片问题？我一般用这几个方法：

方法	命令/代码	能看出什么
查看总占用	`nvidia-smi`	显存总量、已用、进程占用
查看分配器状态	`torch.cuda.memory_summary()`	空闲块数量、最大空闲块大小、碎片率
跟踪分配历史	`torch.cuda.memory._record_memory_history()`	每个tensor的分配时间、大小、调用栈
查看CUDA事件	`nsys profile`	显存分配/释放的时间线、峰值

我个人习惯先用`memory_summary()`快速扫一眼。如果看到「Number of allocated blocks」超过1000，或者「Max free block size」远小于总空闲，那基本可以断定有碎片问题。

  一句话总结：显存碎片化是Batch队列优化的隐形杀手。你调了半天的batch size、梯度累积，结果被碎片问题卡住，那就太冤了。下一节我们会讲如何通过「显存池化」和「预分配」来对抗碎片化。

好了，显存管理的基础就聊到这儿。记住三个关键词：层次架构、懒惰分配、碎片陷阱。搞懂这些，后面讲Batch队列优化时你才能理解为什么有些「骚操作」能省显存。