WSL2环境下PyTorch模型训练内存增长导致OOM问题排查求助

环境配置 #

环境1（WSL2） #

• 操作系统：Linux-5.15.146.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
• GPU：RTX4090
• Python：3.12.2
• PyTorch：2.2.2
• CUDA相关依赖：
  nvidia-cublas-cu12==12.1.3.1 nvidia-cuda-cupti-cu12==12.1.105 nvidia-cuda-nvrtc-cu12==12.1.105 nvidia-cuda-runtime-cu12==12.1.105 nvidia-cudnn-cu12==8.9.2.26 nvidia-cufft-cu12==11.0.2.54 nvidia-curand-cu12==10.3.2.106 nvidia-cusolver-cu12==11.4.5.107 nvidia-cusparse-cu12==12.1.0.106 nvidia-nccl-cu12==2.19.3 nvidia-nvjitlink-cu12==12.4.127 nvidia-nvtx-cu12==12.1.105
• 内存表现：训练过程中内存持续增长直至饱和

环境2（原生Linux） #

• 操作系统：Linux-5.15.0-1058-aws-x86_64-with-glibc2.31
• GPU：Tesla T4
• Python：3.10.14
• PyTorch：2.2.0
• CUDA版本：未知
• 内存表现：训练过程中内存保持稳定，未达饱和

核心问题 #

完全相同的训练代码在两个环境中内存表现差异极大：WSL2环境内存快速饱和，原生Linux环境内存使用稳定。

疑问 #

1. 为何相同代码在不同环境下内存使用差异如此巨大？
2. 是否存在特定PyTorch版本或Linux发行版/内核的已知问题导致该现象？
3. 如何确保不同环境下内存使用的一致性？

已排查措施 #

• 确认两个环境使用完全一致的代码
• 实时监控内存使用情况
• 对比两个环境的Python、PyTorch及CUDA依赖版本

可能原因及解决方案 #

1. 差异原因分析

• WSL2内存回收机制差异：WSL2的虚拟内存映射、显存缓存回收逻辑和原生Linux不同，torch.cuda.empty_cache()的释放效率更低，容易导致未及时回收的内存堆积；同时WSL2内核的内存管理对大张量操作的适配性不如原生Linux。
• 版本兼容性问题：PyTorch 2.2.2搭配Python 3.12的组合在WSL2环境中可能存在未覆盖的内存管理bug，而原生Linux使用的PyTorch 2.2.0+Python 3.10是更稳定的适配组合。另外WSL2中CUDA依赖版本不统一（如nvidia-nvjitlink-cu12版本高于CUDA runtime的12.1），也可能引发内存异常。
• GPU硬件差异：RTX4090与Tesla T4的显存架构、驱动适配逻辑不同，WSL2对消费级GPU的内存管理优化不如原生Linux对数据中心GPU的支持。

2. 已知问题参考

• PyTorch官方issue中存在WSL2+CUDA 12.x+PyTorch 2.2.x的内存增长报告，部分用户反馈升级WSL2内核或降级PyTorch后问题解决。
• WSL2 5.15.x系列内核存在部分内存回收bug，升级到最新内核可修复部分内存泄漏场景。

3. 解决及一致性保障方案

• 对齐环境版本：将WSL2的PyTorch降级到2.2.0，Python降级到3.10.x，与原生Linux环境版本一致，观察内存表现是否恢复稳定。
• 优化内存管理代码：在训练循环的合适节点主动调用torch.cuda.empty_cache()和gc.collect()；显式删除临时张量（如del tensor），避免循环中冗余张量堆积。
• 升级WSL2内核：执行wsl --update升级到最新WSL2内核，修复内核层面的内存管理bug。
• 统一CUDA依赖版本：将WSL2中所有CUDA相关依赖版本对齐到CUDA runtime 12.1，比如将nvidia-nvjitlink-cu12降级到12.1.x版本。
• 定位泄漏点：使用nvidia-smi监控显存变化，结合torch.cuda.memory_summary()输出，定位训练过程中内存增长的具体环节，针对性优化代码。

内容的提问来源于stack exchange，提问作者Sébastien Chapeland