PyTorch训练时显存突然爆掉，但之前同样的代码没事，怎么回事？

SGD动量会多存一份buffer，显存未必比AdamW省，先试试torch.cuda.empty_cache()再看碎片。

换SGD+momentum这个点挺关键的，虽然参数更新本身占的显存少了，但momentum buffer是按参数全量存的，1.2亿参数用SGD的momentum大概多占0.5G左右，再加上如果之前AdamW的exp_avg和exp_avg_sq是fp32，换SGD后反而可能因为state初始化时机不同导致峰值变化。不过更可疑的是你说前两个epoch正常、第三个epoch才炸，这种时间点上的OOM通常不是静态占用问题，而是碎片化累积或者某个临时张量在特定batch上变大了。

gradient checkpointing只在forward开的话，backward重计算时中间激活会重新生成，如果某些层的输出依赖动态shape（比如attention mask或者变长序列有padding变化），重计算阶段的峰值可能比第一次forward高不少。建议你用torch.cuda.memory_summary()在OOM前打印一下，重点看reserved和allocated的差值，如果reserved远大于allocated那就是碎片化没跑。

另外PyTorch 2.x之后默认的cudaMallocAsync分配器跟之前版本行为差别挺大的，如果你最近升级过torch，可以试试设PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync或者反过来用native，有时候能直接绕过去。还有个小坑是数据加载器如果用了pin_memory加上num_workers较多，CPU侧锁页内存暴涨也可能间接影响CUDA上下文，不过你这个报错明确是CUDA OOM，概率小一点。

我遇到过类似情况，最后发现是验证集评估时没加torch.no_grad()，虽然训练循环里没事，但epoch末尾的eval把计算图带进了下一步，第三个epoch刚好累积到临界点。你可以检查下有没有在epoch边界做logging或者保存checkpoint时不小心保留了计算图引用。实在不行就先torch.cuda.empty_cache()在每个epoch后手动清一下，虽然治标不治本，但能帮你确认是不是碎片问题。

第3个epoch才OOM，而且前两个epoch正常，感觉更像是内存碎片累积导致的。你换SGD+momentum后，优化器状态确实省了，但SGD的更新频率和AdamW不同，可能导致某些中间激活的生命周期变长，碎片更容易堆起来。可以试试在训练循环里每隔几个epoch调一下torch.cuda.empty_cache()，或者把PYTORCH_CUDA_ALLOC_CONF设成expandable_segments:True，对碎片化挺管用的。另外gradient checkpointing只在forward开，backward时如果没配合好，反而可能让某些中间变量保留更久，值得检查下。

SGD动量会多存一份动量buffer，显存反而涨了，你查下optimizer state。

这情况挺像典型的显存碎片问题，尤其你换了SGD+momentum之后，优化器状态的内存分配模式变了，虽然理论占用更小，但分配时机不同很容易让缓存碎片化。gradient checkpointing只在forward开的话，backward阶段重计算时还是会额外占显存，如果恰好赶上碎片多的时刻就容易炸。可以试试在训练循环里定期调torch.cuda.empty_cache()，或者设PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True看能不能缓解。另外确认下第3个epoch是不是数据分布或长度有变化，有时候个别超长样本会偷偷把峰值拉高。

SGD动量会多存一份buffer，碎片化加上checkpointing的临时张量，epoch3才炸很正常。试试torch.cuda.empty_cache()或者调小batch。

SGD+momentum其实不一定更省，AdamW的exp_avg状态占两份参数显存，但SGD momentum只占一份，理论上确实少点，不过这点差异不至于让18G直接爆掉。我倒觉得更像是内存碎片的问题，尤其是gradient checkpointing开了之后中间激活反复申请释放，跑到第三个epoch碎片攒够了就OOM。可以试试torch.cuda.empty_cache()或者设PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True，新版pytorch对这个支持挺好的。另外确认下是不是换了SGD后没加weight decay，导致某些层梯度爆炸让activation变大了？

SGD+momentum确实比AdamW省显存，但它对显存碎片的敏感度反而可能更高，因为动量缓冲区虽然小，但分配时机和AdamW不一样，容易在几个epoch后把碎片堆到临界点。你试试在训练循环里每个epoch结束调一下torch.cuda.empty_cache()，或者设个PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True，我这边类似情况基本能压住。另外确认下gradient checkpointing有没有和SGD的no_sync配合出问题，有时候反向图的重建会多占一块。