PyTorch 作为机器学习中广泛使用的开源框架,具有速度快、效率高等特点。而近年来广受好评的 AMD 处理器具有多核、多任务性能良好、性价比高等优势。开发者们一直希望二者联合起来,在 AMD 处理器上使用 PyTorch 进行深度学习的开发和研究。
前段时间发布的 PyTorch 1.8 新增了对 AMD ROCm 的支持,对于想在 AMD 上用 PyTorch 进行深度学习的开发者来说,这是一个好消息。
但是,对使用 AMD cpu 的开发者用 PyTorch 做 AI 开发,也许没那么顺利。
这不,我们就从 PyTorch 的 Github 上发现这么一个还未解决的 issue。
GitHub 地址:
https://github.com/pytorch/pytorch/issues/52142
有开发者表示:PyTorch 在 AMD CPU 的计算机上,用数据并行单机和多 GPU 训练 CNN 会导致死锁,而相同的代码在 Intel CPU 的计算机中就不会出现死锁。TensorFlow 也不会出现这种问题。
遇到了什么 bug
让我们来具体看一下这个死锁是怎么产生的?
在提交的 Issue 中,开发者表述在训练的 for 循环中,代码会一直运行,直到前向传播 output = model(images)。代码会一直停留在 model(images)阶段,而 gpu 的利用率为 0% (内存占用率不是 0),有 3 个 cpu 核的利用率达到 100%,而其他 cpu 核为 0%。使用 ctrl+c 和 ctrll+z 停止后,进程 PID 和 GPU 内存使用情况保持不变。kill 、pkill 和 fuser -k /dev/nvidia * 命令导致僵尸进程(也称为 Defunct 进程)。僵尸进程的父进程 pid 为 1,因此不能 kill 它。唯一的解决办法是重新启动系统。
代码在 3 台 Intel cpu 计算机上运行良好,但在 4 台使用 AMD cpu 的机器上就会出现上述问题。开发者进一步测试了 GTX 1080、Titan V、 Titan RTX、Quadro RTX 8000 和 RTX 3090,证明该 bug 与 GPU 模型无关。
此外,该项目还提到分布式数据并行(DDP)也存在类似的问题。
感兴趣的开发者可以按照如下配置复现一下这个 bug:
使用带有 AMD cpu 和多个 NVIDIA gpu 的机器;
Linux、Python3.8、cuda 11.0、pytorch 1.7.1、 torchvision 0.8.2;
编写代码在 torchvision 中训练 resnet18 模型;
可以尝试测试数据并行 (DP) 和分布式数据并行(DDP),以检查是否都会出现这种情况。
根据该项目的描述,复现之后可能会出现:
当使用 AMD cpu 时,在训练的第一个 epoch 和第一次迭代的前向传播时会发生死锁;
当使用 intel cpu 时,相同的代码运行良好。
尚未解决
对于这次发现的 bug,有网友表示也存在相同的问题。
发现该 bug 的研究者在 ImagNet、mini ImageNet、 CIFAR 10 以及其他数据集上进行了实验,由于 CIFAR 10 具有轻量级、利于调试的特点,因此开发者在 CIFAR 10 上给出了相应的代码示例:
此外,有开发者表示使用 TensorFlow 则不会出现该 bug。
提交 Issue 后,PyTorch 运维团队的成员也留言参与了讨论,在被问到「测试阶段是否包含使用 AMD CPU 的用例」时,该成员表示:「在和其他队友讨论之后,我们意识到测试中没有涉及 AMD CPU 的用例,目前我们还没有办法重现这个问题。如果今年晚些时候我们通过支持更多的硬件类型改进了测试环境,将重新讨论这个问题。」
此前有网友发现了 AMD cpu 下 PyTorch 多卡并行卡死的问题,查看日志之后找到错误原因,问题才得以解决。而这次暴露的 bug 目前仍未解决。