PyTorch Lightning多GPU训练卡死问题分析与解决

问题背景

在使用PyTorch Lightning进行多GPU训练时，开发者遇到了一个棘手的问题：当使用2个NVIDIA RTX 4090 GPU进行分布式数据并行(DDP)训练时，程序会在显示"All distributed processes registered"后卡死。更严重的是，当尝试设置NCCL_P2P_DISABLE=1环境变量时，整个服务器会出现CPU软锁死(soft lockup)现象，最终导致系统崩溃需要强制重启。

环境配置

该问题出现在以下环境中：

• PyTorch 2.1.0
• CUDA 12.1
• NVIDIA驱动版本545.29.02
• 使用2块NVIDIA RTX 4090 GPU
• Python 3.10环境

问题现象分析

当运行简单的BoringModel训练脚本时，程序会在初始化分布式训练阶段卡住，具体表现为：

1. 分布式进程注册成功
2. 显示"All distributed processes registered"后无进一步输出
3. 训练流程无法继续执行

尝试设置NCCL_P2P_DISABLE=1环境变量后，问题更加严重：

1. 系统内核报出CPU软锁死错误
2. 出现"watchdog: BUG: soft lockup"错误信息
3. 最终导致服务器完全无响应

根本原因

经过分析，这个问题可能与NVIDIA驱动版本有关。545.29.02版本的驱动在与PyTorch Lightning的DDP实现交互时可能存在兼容性问题，特别是在使用NCCL通信库进行多GPU通信时。

解决方案

开发者最终通过以下方法解决了问题：

1. 将NVIDIA驱动降级到535.146.02版本
2. 保持其他软件环境不变

降级驱动后，多GPU训练能够正常进行，不再出现卡死或系统崩溃的情况。

技术建议

对于遇到类似问题的开发者，建议：

1. 首先确认NVIDIA驱动版本与PyTorch、CUDA版本的兼容性
2. 在遇到DDP初始化问题时，可以尝试不同的NCCL环境变量设置
3. 但要注意NCCL_P2P_DISABLE=1在某些情况下可能导致系统负载过高
4. 驱动版本的选择对深度学习训练稳定性至关重要
5. 新发布的驱动不一定是最稳定的选择，有时需要回退到经过验证的稳定版本

总结

PyTorch Lightning的多GPU训练功能依赖于底层的NCCL通信库和NVIDIA驱动，当这些组件之间存在兼容性问题时，可能导致训练过程中断甚至系统崩溃。通过选择合适的驱动版本，可以有效地解决这类问题，确保分布式训练的稳定性。这也提醒我们在构建深度学习环境时，需要特别注意各组件版本间的兼容性。