2. 单卡显存管理

聊到GPU集群，很多人第一反应是网络、是分布式通信。但说实话，我这些年踩过最深的坑，往往出在单卡显存管理上。你想想看，集群里几百张卡，如果每张卡的显存都管不好，那上层再好的调度策略也是白搭。今天我们就来扒一扒PyTorch的显存分配机制、CUDA缓存分配器，还有那个让人头疼的显存碎片化问题。

PyTorch的显存分配机制：它到底在干什么？

先问个问题：你在PyTorch里创建一个Tensor，显存是立刻分配的吗？

答案是：不一定。PyTorch的显存分配，其实分了两层。上层是PyTorch自己的缓存分配器（Caching Allocator），下层是CUDA驱动提供的cudaMalloc。我刚开始做模型训练时，以为每次创建Tensor都会调用cudaMalloc，后来发现根本不是这么回事。

PyTorch的默认分配器，我习惯叫它“小气鬼”。它从CUDA那里申请一大块显存，然后自己切成小块分给Tensor。用完不还，留着下次用。这样做的好处很明显——减少cudaMalloc的调用次数，提升速度。坏处呢？嗯，就是容易造成碎片化。

我个人习惯用torch.cuda.memory_summary()来查看当前显存的使用情况。这个命令会告诉你分配器内部的状态，包括缓存了多少、用了多少、碎片有多少。我在项目中遇到过好几次，明明显存还有剩余，但就是分配不出新Tensor，一查，全是碎片。

CUDA缓存分配器：底层的那点事

CUDA本身也有一套内存管理机制。cudaMalloc分配的是物理显存，cudaFree释放回系统。但频繁调用这两个API，性能会大打折扣。所以PyTorch在中间加了一层缓存。

说白了，PyTorch的缓存分配器就是一个“二房东”。它从CUDA那里整租一大片显存，然后分租给各个Tensor。Tensor用完退租，它也不还给CUDA，而是留着给下一个Tensor用。

这个机制有几个关键参数：

参数	作用	我常用的值
max_split_size_mb	控制缓存块的最大分割粒度	128 或 256
roundup_power2_divisions	分配时是否对齐到2的幂次	开启（默认）
caching_allocator_config	分配器的详细配置	根据模型调整

这里有个坑。我曾经调大max_split_size_mb，想着减少碎片。结果呢？大块是少了，但小Tensor分配时浪费的空间更多了。所以这个值不是越大越好，得根据你的模型Tensor大小分布来定。

显存碎片化问题：为什么显存还有，却分配不出新Tensor？

这个问题，我估计做深度学习的老铁都遇到过。显存占用显示80%，但一创建新Tensor就OOM。为什么会这样？

原因就是碎片化。想象一下，你有一块完整的显存，被切成了很多小块。有些块被占用了，有些块空闲着。但空闲块都不连续，最大的连续空闲块可能只有100MB。这时候你要分配一个200MB的Tensor，对不起，OOM。

碎片化主要有两种：

• 内部碎片：分配器给你一个块，但你只用了一部分，剩下的浪费了。比如你申请100MB，分配器给了你128MB的块，那28MB就是内部碎片。
• 外部碎片：空闲块分散在各处，无法合并成一个大块。这是最头疼的，因为显存总量够，但连续空间不够。

我在项目中遇到过最夸张的一次，显存占用只有60%，但就是跑不动一个batch。用torch.cuda.memory_snapshot()一看，好家伙，空闲块有几百个，最大的才几十MB。这就是典型的外部碎片。

如何诊断和缓解碎片化？

诊断碎片化，我一般用这几个工具：

• torch.cuda.memory_summary()：查看整体分配情况
• torch.cuda.memory_snapshot()：获取详细的分配记录
• NVIDIA的nvidia-smi：看物理显存使用量

缓解碎片化，我有几个实战技巧：

1. 预分配策略：在训练开始前，先创建好所有固定大小的Tensor，避免训练过程中频繁分配和释放。
2. 使用Tensor的.contiguous()方法：确保Tensor在显存中是连续存储的，减少碎片。
3. 调整缓存分配器参数：通过环境变量PYTORCH_CUDA_ALLOC_CONF调整max_split_size_mb和roundup_power2_divisions。
4. 使用torch.cuda.set_per_process_memory_fraction：限制每个进程的显存使用比例，避免一个进程占满所有显存导致碎片。

嗯，这里要注意一点。很多人觉得碎片化是PyTorch的问题，其实不然。CUDA驱动本身也有内存管理，PyTorch只是在上层做了一层缓存。真正的问题在于，深度学习模型的Tensor大小变化太大，从几KB的标量到几GB的权重，分配器很难做到完美适配。

一张图看懂显存分配流程

下面这张图，是我自己总结的PyTorch显存分配流程。从Tensor创建到物理显存分配，每一步都标出来了。

从这张图可以看得很清楚。Tensor创建时，先走缓存分配器。缓存里有合适的块，直接拿来用。没有，才去CUDA那里要。这个流程本身没问题，但问题出在“合适的块”这个判断上。如果缓存里的块大小和请求的大小不匹配，就会产生碎片。

我个人建议，在训练大规模模型之前，先用小批量跑一次，看看显存的分配模式。用torch.cuda.memory_snapshot()记录下分配记录，分析一下Tensor的大小分布。然后根据这个分布去调整缓存分配器的参数。这样做，能有效减少碎片化。

好了，单卡显存管理这块，核心就是理解PyTorch的缓存分配器怎么工作，以及如何应对碎片化。记住一点：显存不是无限的，但碎片化是可以管理的。用好工具，调好参数，你的单卡就能发挥出最大性能。