下到 GPU · CUDA 流与缓存分配器 · PyTorch 源码漫游

张量在 GPU 上怎么活？CUDA 调用是异步入队（所以慢 Python 能喂饱快 GPU，所以计时要 synchronize）；而那个让人困惑的「PyTorch 一上来就吃光显存」，真相是它自带的缓存分配器——为了绕开又慢又同步的 cudaMalloc，自己囤内存反复用。搞懂这两点，计时和 OOM 就不再玄学。

下到 GPU · CUDA 流与缓存分配器

漫游到最底层：张量落在 GPU 上之后，是怎么运作的？这一章讲两件让无数人栽跟头的事——CUDA 的异步执行，和那个「一上来就吃光显存」的缓存分配器。搞懂它俩，你的性能计时和 OOM 调试就从玄学变成常识。

关键真相一：CUDA 调用是「异步入队」

这是最反直觉、也最重要的一点：

当你从 Python 调一个 CUDA 操作，GPU 并不会立刻算完再返回。这个操作被「入队」到一条 CUDA 流(stream)上，Python 立刻返回、往下跑；GPU 稍后按队列顺序慢慢执行。

一条 CUDA 流，就是一个有序的 GPU 工作队列：同一条流上的操作按入队顺序执行，不同流之间可以并行。

这个「异步入队」解释了一个大问题：又慢又是解释执行的 Python，凭什么能喂饱快得飞起的 GPU？ 答案就是——它不等。Python 只管飞快地把一个个操作往流里塞，GPU 在后头排着队闷头算。Python 跑在 GPU 前面，只要队列不空，GPU 就不会闲。nanoGPT 训练循环里「趁反向在 GPU 上跑、同时预取下一批数据」，能成立的底层前提正是这个异步性。

这也是给 GPU 代码计时的头号大坑。 你写 t0=time(); y=model(x); t1=time()，量到的很可能只是**「入队」的时间**，不是 GPU 真正算完的时间——因为 Python 早跑走了。要量真实耗时，必须先 torch.cuda.synchronize()（等 GPU 把队列清空）再读时间。同理，.item()、.cpu()、print(tensor) 这些「要拿到真实数值」的操作，会强制同步（必须等 GPU 算完才有值）。nanoGPT 训练循环那句注释「loss.item() is a CPU-GPU sync point」，说的就是这个：读 loss 的值，得停下来等 GPU。

顺带把 nanoGPT get_batch 那两个词也收了：pin_memory() 把主机内存锁页，这样 .to(device, non_blocking=True) 才能发起异步的主机→显存拷贝，和计算重叠。全是围绕「异步」做文章。

关键真相二：缓存分配器，和「吃光显存」的误会

第二个千古谜题：你的张量明明没多大，nvidia-smi 却显示 PyTorch 占了一大片显存。它是不是漏内存了？

不是。真相是 PyTorch 自带一个缓存分配器(caching allocator)，而它这么设计有充分理由：

向 CUDA 直接申请 / 释放显存的 cudaMalloc / cudaFree，又慢、又会强制同步（打断上面说的异步流水）。训练时每秒要创建、销毁成千上万个临时张量，每个都 cudaMalloc 一次，性能会被拖垮。
所以 PyTorch 一次向 CUDA 抓一大块，自己切成小块分给张量；张量「释放」时，内存块不还给 CUDA，而是留在自己的池子里，等下一个张量复用。

于是：

nvidia-smi 看到的，是分配器「囤」的总量(reserved)，不是你此刻真正在用的量(allocated)。 囤着是为了复用，不是漏。
两个指标要分清：torch.cuda.memory_allocated()（真正在用）vs torch.cuda.memory_reserved()（分配器囤了多少）。
torch.cuda.empty_cache() 能把池子里空闲的块还给 CUDA（让 nvidia-smi 数字降下来），但通常没必要——还回去下次还得重新抓，反而更慢。

OOM 从哪来，又怎么查

搞懂分配器，爆显存(CUDA out of memory)就好诊断了。训练时显存被这几坨吃掉：

模型参数本身；
优化器状态（如 AdamW 每个参数还要存一阶、二阶动量，约等于参数量的 2 倍）；
autograd 图留存的中间激活——训练显存的大头，随 batch 和深度涨；
分配器为复用囤着的空闲块。

还有个隐蔽杀手：碎片化。池子被切得七零八落，总空闲够、但没有一整块连续的满足一次大分配 → 照样 OOM。报错里那句 reserved but unallocated ... fragmentation 就是它。这时可以用环境变量 PYTORCH_CUDA_ALLOC_CONF（如 expandable_segments:True）调分配器策略缓解。

省显存的招数，现在都能对上号了：混合精度 / bf16 让激活和参数减半；torch.no_grad() 不建图、不留激活（第 3 坨直接归零）；梯度检查点用「反向重算」换「不留激活」；梯度累积 用小 batch 模拟大 batch、压住第 3 坨。它们各自在拧上面某一坨。你在 nanoGPT 里见过的这些技巧，拧的都是本章这几个显存旋钮。

引擎看完了，还差最后一层：编译

从 Tensor 说明书、dispatcher、ATen 内核、autograd，到 nn.Module 和本章的 GPU，你已经把 PyTorch 这台**逐个算子、即时执行(eager)**的引擎走了个遍。

但「一个算子一个算子地即时执行」有它的开销上限。最后一站，我们看 PyTorch 2.0 怎么给这台动态引擎做图编译——torch.compile 背后的 TorchDynamo 和 TorchInductor。👉 torch.compile：给动态图做编译

参考文献

[1] PyTorch 文档：“CUDA semantics.” https://pytorch.org/docs/stable/notes/cuda.html —— 异步执行、流、同步。
[2] PyTorch 文档：“Understanding CUDA Memory Usage” / “Memory management.” https://pytorch.org/docs/stable/torch_cuda_memory.html —— 缓存分配器、allocated vs reserved、碎片化、PYTORCH_CUDA_ALLOC_CONF。
[3] Paszke et al. “PyTorch: An Imperative Style, High-Performance Deep Learning Library.” NeurIPS 2019. arXiv:1912.01703. —— 缓存分配器的设计动机。
源码：c10/cuda/CUDACachingAllocator.cpp、c10/cuda/CUDAStream.h(pytorch/pytorch)。

下到 GPU · CUDA 流与缓存分配器