在分布式大模型训练中，GPU内存泄漏是一个常见但棘手的问题，尤其在长时间训练过程中容易导致训练中断。本文分享一次典型的GPU内存泄漏排查经验。
问题现象：使用PyTorch DDP训练一个7B参数的Transformer模型时，在训练到第1500个step后，GPU显存从18GB逐渐增长至24GB，最终导致OOM（Out of Memory）中断。单卡训练时无此问题，但多卡分布式训练中频繁出现。
排查步骤：
- 使用
nvidia-smi监控显存变化，发现显存持续增长而非稳定； - 添加内存监控代码：
import torch
print(f'Allocated: {torch.cuda.memory_allocated() / 1024**3:.2f} GB')
print(f'Reserved: {torch.cuda.memory_reserved() / 1024**3:.2f} GB')
- 发现
reserved内存持续增长，而allocated保持稳定。
关键发现：通过设置环境变量PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512后问题缓解，但根本原因在于模型并行时的梯度同步机制未正确释放中间缓存。在optimizer.step()后增加torch.cuda.empty_cache()可临时解决，但影响训练性能。
解决方案：
- 使用
torch.compile（PyTorch 2.0+）优化内存分配； - 设置
torch.backends.cudnn.benchmark = False避免内存碎片； - 在分布式训练中，合理设置
gradient_checkpointing参数，减少中间激活存储。
最终通过调整上述配置，成功将训练稳定运行至预期epoch数。经验表明，分布式训练的内存泄漏问题通常与缓存管理机制相关，需结合具体框架版本进行针对性优化。
讨论