前阵子搭一套多卡训练流程，被RuntimeError: CUDA out of memory折腾了一整天。在PyTorch里看到这个报错，绝大多数人的第一反应都是"显存不够了"，然后顺手把batch size砍半。但实际调试下来，这个报错背后至少藏着三种完全不同的病因：物理显存真的不足、PyTorch缓存分配器没释放干净、以及显存碎片化严重。每一种的处理方式完全不同，用错方案只会越改越乱。这篇内容我就把PyTorch里CUDA out of memory的完整解决思路梳理一遍，从读懂报错、常规调参、混合精度、梯度检查点到碎片化配置，一次性讲透这个报错背后到底发生了什么，以及不同场景下该优先动哪个参数。适合正在被显存卡脖子的训练场景，也适合想系统地给模型"减负"的人。
1. 先别急着改代码：把报错信息逐句拆开再动手
1.1 报错里的三个数字分别代表什么
一个典型的报错长这样：
RuntimeError: CUDA out of memory. Tried to allocate 512.00 MiB (GPU 0; 8.00 GiB total capacity; 6.20 GiB already allocated; 980.00 MiB free; 6.50 GiB reserved in total by PyTorch)很多人只扫一眼开头的"out of memory"就开始翻代码，完全忽略后面那串数字。其实这串信息才是诊断的核心。逐项解释一下：
Tried to allocate：这次前向或反向传播中，PyTorch想再从显存里拿多大一块空间。already allocated：当前实际被张量占用的显存大小，这部分是真实在用的。free：驱动层报告的空闲显存。注意这个数字包含了PyTorch已经预留但还没用掉的部分，所以它"看起来还有"不代表真的能分配出来。reserved in total by PyTorch：PyTorch的缓存分配器从驱动层预占的总量。它通常会大于already allocated，多出来的部分就是缓存池里闲置、等待复用的块。
这里有个关键认知：PyTorch不会在你释放张量后立刻把显存还给系统。训练过程中del掉一个tensor，显存大概率留在PyTorch自己的缓存池里，为的是下次分配时能快速复用。所以free显示很小，不一定意味着显存真的不够，很可能只是缓存没有还给驱动。反过来说，如果报错时already allocated已经逼近物理上限，同时reserved高得离谱，那就要从算法层面减负，而不是指望清理缓存能解决问题。
1.2 用一行代码判断是"真不够"还是"假报警"
实操时，我会先打开nvidia-smi看一眼GPU占用，再在代码里打印几组数值：
import torch print("当前已分配:", torch.cuda.memory_allocated() / 1024**3, "GiB") print("当前已预留:", torch.cuda.memory_reserved() / 1024**3, "GiB") print("历史峰值分配:", torch.cuda.max_memory_allocated() / 1024**3, "GiB") # 更完整的分配报告 print(torch.cuda.memory_summary())判断逻辑很简单：如果memory_allocated数值很低，但nvidia-smi显示显存占用接近100%，那多半是缓存没释放或碎片化严重，优先尝试torch.cuda.empty_cache()和调整分配器配置。如果memory_allocated本身就接近物理显存大小，说明模型、batch、序列长度带来的真实需求已经超标，这时候再怎么清缓存都没用，得动模型结构或者训练策略。
max_memory_allocated这个值特别值得记录，它保存的是程序启动以来的峰值占用量。如果你当前占用不高，但峰值顶到了物理上限，说明某个环节的瞬时需求超标。我会在报错堆栈的上下文中看这个峰值出现在哪个阶段，从而精准定位到是前向、反向还是优化器更新阶段出的问题。
2. 底层的节流方案：batch size、梯度累积和缓存释放的正确用法
2.1 为什么"batch size直接减半"经常没用
很多人的第一反正是把batch size从16降到8，期望显存需求也减半。但实际经常发现减半后依然报错，或者报错的位置变了。原因是显存占用主要分两部分：一部分是模型参数加优化器状态，这些跟batch size没有关系；另一部分是前向传播保存的中间激活值，这才跟batch size直接相关。
如果模型本身就很大，比如在做大模型微调，参数和优化器状态可能已经吃掉一大半显存，你把batch从16减到8，节省的只是激活值那一小部分，峰值拉不回安全线。另外，大幅改变batch size会影响BatchNorm的统计量，可能导致收敛行为变化。很多模型减半batch size后虽然不报错了，但loss曲线变得毛躁，训练效果打了折扣。所以"减半"这个操作远没有想象中那么万能。
2.2 梯度累积：保持等效batch size的正确减负方式
与其硬调batch size，我更推荐用梯度累积来模拟大batch。思路是：每个小batch照常前向、反向算梯度，但先不更新参数，等累积若干个batch的梯度后，再统一执行一次optimizer.step()。这样等效于更大的batch，而显存峰值始终只有一个小batch的量。
标准写法：
accumulation_steps = 4 model.train() optimizer.zero_grad() for batch_idx, (inputs, labels) in enumerate(train_loader): outputs = model(inputs) loss = criterion(outputs, labels) / accumulation_steps loss.backward() if (batch_idx + 1) % accumulation_steps == 0: optimizer.step() optimizer.zero_grad()这里有两个我踩过的坑要提醒。
第一，loss必须除以accumulation_steps再做backward。梯度累积的本质是多个batch的梯度相加，如果不做除法，等效batch变大了，但学习率对应的梯度尺度对不上，训练很容易震荡。这个细节网上很多简化示例都会漏掉，照抄会吃大亏。
第二，BatchNorm在这种模式下和真实大batch不一致。BN看到的是一个个小batch的统计量，累积训练并不会自动修正这一点。如果模型对BN统计量敏感，可以考虑改用SyncBN，或者换用GroupNorm这类和batch大小无关的归一化层。我在一个分割模型上遇到过类似问题，切换归一化方式之后训练稳定性明显好转。
2.3 empty_cache和del的正确使用时机
网上流传很多"训练循环里加一句torch.cuda.empty_cache()"的土方子。坦率说，这个调用是有代价的：它会把缓存池里的块归还给驱动，下一次分配又要重新走CUDA malloc流程，频繁调用不仅拖慢训练，反而可能因为频繁分配制造更多碎片。
我实际使用中，empty_cache只用在几个特定场景：训练和验证切换的间隙、数据加载阶段产生了一批一次性大张量之后、以及一个训练进程结束而下一个进程还没启动的时候。训练主循环内部我基本不用它。
至于手动del变量和gc.collect()，只对那些确实不再需要、又占用巨大空间的中间张量有意义。比如推理时保存下来用于可视化的某一层特征图。普通训练里，autograd本来就会在反向传播过程中逐步释放不再需要的激活值，手动del更多是用来排查某个变量是否在无意中被引用、攥着显存不放。
3. 显存效率翻倍的正规军：混合精度和梯度检查点
3.1 AMP混合精度到底省在哪里、怎么省
PyTorch的自动混合精度（AMP）是我最推荐的显存优化手段，因为它改动极小、收益极大。原理说白了：前向和反向的计算张量用FP16存储和运算，主权重保持FP32，优化器更新也在FP32下完成。中间激活值全用半精度保存，显存占用接近减半；同时靠GradScaler动态缩放loss，避免小梯度在FP16下被下溢成0。
标准用法：
scaler = torch.cuda.amp.GradScaler() for inputs, labels in dataloader: optimizer.zero_grad() with torch.cuda.amp.autocast(): outputs = model(inputs) loss = criterion(outputs, labels) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()新版PyTorch里也可以直接用torch.autocast(device_type='cuda')，本质是同一个机制。
不过要提醒一句：不是所有模型都适合AMP。某些数值动态范围极大的算子，或者自定义op，在FP16下可能出现NaN或loss不降。判断方法很简单：开启AMP后连续跑几个step，盯住scaler.get_scale()的变化和loss数值。如果loss突然变成inf/nan，或者scale值持续异常下降，说明模型对梯度动态范围敏感，那就需要把特定层排除在FP16之外，或者干脆放弃AMP只做梯度检查点。
3.2 梯度检查点：拿时间换显存的性价比之王
梯度检查点（activation checkpointing）的思路是：前向传播时不再把每一层的激活值全部保存下来，只保留少数关键节点；反向传播需要某段中间结果时，现场重新做一次前向把它算出来。典型能省50%到70%的显存，代价是训练时间增加20%到30%。当显存卡脖子、又没法换卡的时候，这是非常划算的交换。
PyTorch里的基本用法：
import torch.utils.checkpoint as cp # 针对单个模块：把forward改成checkpoint包裹 class MyBlock(torch.nn.Module): def forward(self, x): return cp.checkpoint(self._forward_impl, x, use_reentrant=True) # 针对nn.Sequential：按段切分，每段单独重算 from torch.utils.checkpoint import checkpoint_sequential output = checkpoint_sequential(model, segments=4, input=x)使用中有几个细节要注意。被checkpoint包裹的模块，输入必须是requires_grad=True的，如果传入的是某个预处理后的固定特征，需要手动设置x.requires_grad_()，否则会直接报错。另一个是in-place操作的问题：如果checkpoint内部对输入张量做了in-place修改，反向重算时输入可能已经被改过了，计算图就乱了。所以被包裹的模块里，尽量避免写x += something这类原地修改。
3.3 AMP和梯度检查点叠加时的坑
两个技巧一起上，显存能压得很狠，但坑也变多了。梯度检查点重算的那段前向，必须落在autocast的上下文范围内，否则会以FP32重算，重算结果和原始前向不一致，精度会悄悄下降。我的做法是：checkpoint包裹的目标函数内部不写死精度，统一依赖外层autocast上下文，并且把整个step的代码都用autocast包起来。
另外，GradScaler初始化时默认scale是2的16次方。如果模型在反向时梯度波动很大，scaler会自动把scale调小。这时候如果发现开AMP后效果不进反退，去看日志里scale值是不是一直在降，就能判断出问题所在。对我来说，这类模型宁可只做梯度检查点，也不硬开AMP。
4. 被低估的隐形杀手：显存碎片化与PYTORCH_CUDA_ALLOC_CONF
4.1 为什么显存明明还剩好几个G，却连512MiB都分配不出来
真实训练里经常出现一种令人抓狂的现象：GPU卡上剩余显存显示还有2到3GB，但程序报"Tried to allocate 512.00 MiB"直接失败。原因在PyTorch缓存分配器的块管理策略上。显存被切成大大小小的块，当模型里张量尺寸变化剧烈时（比如NLP里每个batch的序列长度不同，或者检测模型里目标数量不同），空闲块会被切得七零八落。新请求要找一个连续的空闲块，分配器找不到足够大的连续块，即使总空闲容量够，也只能报OOM。
这个场景最典型的特征是：报错里的free数字不小，reserved也充足，但Tried to allocate的请求尺寸并不大，却死活分配不出来。
4.2 max_split_size_mb的调参逻辑与实测经验
应对碎片化，官方提供了环境变量配置：
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:32这个参数的含义是：当一个空闲块比请求尺寸大，且超过max_split_size_mb这个阈值时，分配器才允许对它进行切分。把值调小，意味着小块更容易被切出来满足小请求，碎片化会减轻；但切分行为变多，显存管理开销也会上升。调大则相反，更倾向于保留大块满足大请求。
实操经验是：先看报错里Tried to allocate的尺寸。如果只有几十到几百MiB，而free显示还有好几个GiB，那优先把max_split_size_mb设成32或64试一次。如果请求是好几个GiB的大张量，多半不是碎片化，而是真实峰值需求超限，调这个参数帮不上忙。我自己的方法是先跑一次默认配置，然后打印memory_reserved和memory_allocated，两者长期差距超过60%，基本可以判定缓存池利用率很低，碎片化严重，这时候再上max_split_size_mb才有效。
4.3 动态shape是碎片化的温床
把碎片化的根源展开说：凡是batch内尺寸不固定的模型，出问题的概率都很高。NLP里同一个batch的句子padding不统一，检测模型里不同图的物体数量不同，都会让前向过程产生大量尺寸跨度很大的临时张量，缓存池被一点点切碎。针对这类模型，统一尺寸、引入合理的padding策略、或者限制batch内样本的尺寸差异，比调分配器参数更治本。尺寸稳定之后，GPU利用率也会同步提升，属于一举两得。
5. 从源头省显存：输入尺寸、生命周期和模型结构的选择
5.1 输入形状直接决定激活值数量级
对视觉模型来说，显存里的大头往往不是权重，而是中间激活值。举个例子，batch为16、输入224x224、输出通道512的特征图，一张图的特征图占用约512乘224乘224乘2字节（半精度），CNN层数一深，累计的激活值体积就很可观。所以输入尺寸每降一档，激活值显存可能指数级下降。很多OOM问题靠"把图片resize小一点"就能直接解决，代价是精度可能下降。
平衡方案是先小分辨率预训练、再大分辨率微调，很多视觉任务吃这一套，效果损失很小。Transformer对序列长度更敏感，attention的计算量和显存对长度是平方关系，长度从512减到256，激活值往往能直接省一半以上。如果你的模型是NLP方向的，先检查序列长度是不是被无脑设置成了最大值。
5.2 变量生命周期与no_grad的边界
另一类显存问题的根源是变量没被及时释放。常见场面包括：验证循环里忘了写with torch.no_grad()，导致推理过程也建了一整条计算图；或者在循环里用一个list累积保存每个样本的输出，美其名曰"后面分析用"，结果list越长显存越紧。前者的修复是推理、评估阶段强制包no_grad并确保model.eval()；后者的修复是改变保存策略，把分析逻辑放到训练结束后，或者先把结果转移到CPU上再存。
还有一个容易忽略的点：自定义层里大量使用in-place操作，除了可能中断梯度流，还会让某些张量不能被分配器及时回收。我的原则是，新写的层默认返回新tensor，除非非常明确自己为什么需要原地修改，否则不要为了省一次拷贝而埋下隐患。
5.3 收尾手段：CPU offload与模型裁剪的取舍
如果上面所有招数都试过还是不够，那就剩两条路。
第一，把一部分计算挪到CPU。比如embedding这类参数大、计算简单的层，放在CPU上，前后做好.to(device)和.to('cpu')的搬运。Tensor在CPU和GPU之间移动成本不低，所以搬运频率要控制住，否则同步开销会吃掉所有收益。这个方案适合那些"GPU上放不下但CPU内存很充裕"的场景。
第二，模型结构级裁剪。减少通道数、降低注意力头数、换更小的backbone、或者做量化。这是最彻底的改动，也意味着要重新验证精度。我的建议是先定量分析：用torch.profiler或者memory_summary看哪一层占显存最多，针对最大的那个头做结构优化，比盲目砍模型高效得多。有一次我只把某一层的通道数从1024降到768，显存峰值立刻降了20%还多，精度几乎没动。
6. 一套可以直接照抄的OOM排查流程
6.1 固定步骤：从报错信息到定位根因
我现在遇到CUDA out of memory，基本按下列顺序走一遍，不再瞎试：
- 看报错最后几行的堆栈，找到具体是哪一行代码触发的分配请求。
- 看报错里的三个数字：
Tried to allocate、already allocated、free。 nvidia-smi瞥一眼卡上占用，快速判断是物理不够还是缓存/碎片问题。- 打印
memory_allocated、memory_reserved、max_memory_allocated，确认峰值出现在哪个阶段。 - 如果报错发生在
backward阶段，大概率是激活值峰值问题；如果发生在optimizer.step()，大概率是优化器状态或梯度缓冲问题。 - 按优先级依次尝试：换小batch加梯度累积、在合适时机做
empty_cache、开AMP、上梯度检查点、调碎片化配置、最后做结构裁剪。
这套流程帮我快速区分了"显存物理不足"和"分配器效率问题"两个完全不同的方向，避免在一个错误的思路上反复浪费时间。
6.2 几条被反复验证过的实战心得
分享几个实际项目中验证过多次的经验。
第一，在训练脚本里提前加好OOM异常捕获和显存快照，比事后分析强很多。当OOM抛出时自动打印memory_summary，能立刻看到reserved和active的分布，省掉重复复现的时间。
第二，多卡场景下，显存问题经常是单卡分配不均导致的。DataLoader的shuffle、某些采样器可能让一个batch里恰好集中了长序列样本，导致某一张卡突然爆掉。打开分布式框架的日志看每张卡的峰值，找到那个"尖峰卡"，再针对性调整采样策略。
第三，PyTorch版本之间显存分配策略有差异，同样的代码在新旧版本上可能表现完全不同。遇到诡异的OOM，且代码逻辑看不出问题时，可以考虑升级或回退版本。如果你在维护一个老项目，这个问题特别常见。有一次我就是排查了三天，最后发现是框架版本和CUDA驱动版本不匹配导致分配器行为异常，升级之后问题消失。
我个人在训练脚本里最终沉淀下来的习惯是：把显存监控和峰值记录做成常规日志的一部分，每次训练跑完都看一眼。显存瓶颈很少是一瞬间突然出现的，它通常随着训练步数慢慢爬升。提前看到趋势，比等到某一步报错再抢救，要从容得多。
