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

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

- 作者：David（道雾轩）
- 专栏：PyTorch 源码漫游（https://daiw.org/manual/pytorch-internals.md）
- 最后更新：2026-07-13
- 原文：https://daiw.org/manual/pytorch-internals/cuda-and-memory
- 转载与引用：请注明出处并附原文链接（https://daiw.org/about/copyright）

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

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

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

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

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

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

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

<Callout type="warn">
  **这也是给 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」](https://daiw.org/manual/nanogpt/training-loop)，说的就是这个：读 loss 的值，得停下来等 GPU。
</Callout>

顺带把 [nanoGPT get_batch 那两个词](https://daiw.org/manual/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)就好诊断了。训练时显存被这几坨吃掉：

1. **模型参数**本身；
2. **优化器状态**（如 [AdamW](https://daiw.org/manual/nanogpt/optimizer) 每个参数还要存一阶、二阶动量，约等于参数量的 2 倍）；
3. **[autograd 图留存的中间激活](https://daiw.org/manual/pytorch-internals/autograd-graph)**——训练显存的大头，随 batch 和深度涨；
4. 分配器为复用**囤着的空闲块**。

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

<Callout type="info">
  省显存的招数，现在都能对上号了：**[混合精度 / bf16](https://daiw.org/manual/nanogpt/scaling-and-speed)** 让激活和参数减半；**[`torch.no_grad()`](https://daiw.org/manual/pytorch-internals/autograd-graph)** 不建图、不留激活（第 3 坨直接归零）；**梯度检查点**用「反向重算」换「不留激活」；**[梯度累积](https://daiw.org/manual/nanogpt/training-loop)** 用小 batch 模拟大 batch、压住第 3 坨。它们各自在拧上面某一坨。你在 nanoGPT 里见过的这些技巧，拧的都是本章这几个显存旋钮。
</Callout>

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

从 [Tensor 说明书](https://daiw.org/manual/pytorch-internals/tensor-and-storage)、[dispatcher](https://daiw.org/manual/pytorch-internals/dispatch-journey)、[ATen 内核](https://daiw.org/manual/pytorch-internals/operators-and-kernels)、[autograd](https://daiw.org/manual/pytorch-internals/autograd-graph)，到 [nn.Module](https://daiw.org/manual/pytorch-internals/nn-module) 和本章的 GPU，你已经把 PyTorch 这台**逐个算子、即时执行(eager)**的引擎走了个遍。

但「一个算子一个算子地即时执行」有它的开销上限。最后一站，我们看 PyTorch 2.0 怎么给这台**动态引擎**做**图编译**——`torch.compile` 背后的 TorchDynamo 和 TorchInductor。👉 [torch.compile：给动态图做编译](https://daiw.org/manual/pytorch-internals/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)。