有bug!PyTorch在AMD CPU的計算機(jī)上卡死了
PyTorch 作為機(jī)器學(xué)習(xí)中廣泛使用的開源框架,具有速度快、效率高等特點(diǎn)。而近年來廣受好評的 AMD 處理器具有多核、多任務(wù)性能良好、性價比高等優(yōu)勢。開發(fā)者們一直希望二者聯(lián)合起來,在 AMD 處理器上使用 PyTorch 進(jìn)行深度學(xué)習(xí)的開發(fā)和研究。
前段時間發(fā)布的 PyTorch 1.8 新增了對 AMD ROCm 的支持,對于想在 AMD 上用 PyTorch 進(jìn)行深度學(xué)習(xí)的開發(fā)者來說,這是一個好消息。
但是,對使用 AMD cpu 的開發(fā)者用 PyTorch 做 AI 開發(fā),也許沒那么順利。
這不,我們就從 PyTorch 的 Github 上發(fā)現(xiàn)這么一個還未解決的 issue。

GitHub 地址:
https://github.com/pytorch/pytorch/issues/52142
有開發(fā)者表示:PyTorch 在 AMD CPU 的計算機(jī)上,用數(shù)據(jù)并行單機(jī)和多 GPU 訓(xùn)練 CNN 會導(dǎo)致死鎖,而相同的代碼在 Intel CPU 的計算機(jī)中就不會出現(xiàn)死鎖。TensorFlow 也不會出現(xiàn)這種問題。
遇到了什么 bug
讓我們來具體看一下這個死鎖是怎么產(chǎn)生的?
在提交的 Issue 中,開發(fā)者表述在訓(xùn)練的 for 循環(huán)中,代碼會一直運(yùn)行,直到前向傳播 output = model(images)。代碼會一直停留在 model(images)階段,而 gpu 的利用率為 0% (內(nèi)存占用率不是 0),有 3 個 cpu 核的利用率達(dá)到 100%,而其他 cpu 核為 0%。使用 ctrl+c 和 ctrll+z 停止后,進(jìn)程 PID 和 GPU 內(nèi)存使用情況保持不變。kill 、pkill 和 fuser -k /dev/nvidia * 命令導(dǎo)致僵尸進(jìn)程(也稱為 Defunct 進(jìn)程)。僵尸進(jìn)程的父進(jìn)程 pid 為 1,因此不能 kill 它。唯一的解決辦法是重新啟動系統(tǒng)。
代碼在 3 臺 Intel cpu 計算機(jī)上運(yùn)行良好,但在 4 臺使用 AMD cpu 的機(jī)器上就會出現(xiàn)上述問題。開發(fā)者進(jìn)一步測試了 GTX 1080、Titan V、 Titan RTX、Quadro RTX 8000 和 RTX 3090,證明該 bug 與 GPU 模型無關(guān)。
此外,該項目還提到分布式數(shù)據(jù)并行(DDP)也存在類似的問題。
感興趣的開發(fā)者可以按照如下配置復(fù)現(xiàn)一下這個 bug:
使用帶有 AMD cpu 和多個 NVIDIA gpu 的機(jī)器;
Linux、Python3.8、cuda 11.0、pytorch 1.7.1、 torchvision 0.8.2;
編寫代碼在 torchvision 中訓(xùn)練 resnet18 模型;
可以嘗試測試數(shù)據(jù)并行 (DP) 和分布式數(shù)據(jù)并行(DDP),以檢查是否都會出現(xiàn)這種情況。
根據(jù)該項目的描述,復(fù)現(xiàn)之后可能會出現(xiàn):
當(dāng)使用 AMD cpu 時,在訓(xùn)練的第一個 epoch 和第一次迭代的前向傳播時會發(fā)生死鎖;
當(dāng)使用 intel cpu 時,相同的代碼運(yùn)行良好。
尚未解決
對于這次發(fā)現(xiàn)的 bug,有網(wǎng)友表示也存在相同的問題。

發(fā)現(xiàn)該 bug 的研究者在 ImagNet、mini ImageNet、 CIFAR 10 以及其他數(shù)據(jù)集上進(jìn)行了實驗,由于 CIFAR 10 具有輕量級、利于調(diào)試的特點(diǎn),因此開發(fā)者在 CIFAR 10 上給出了相應(yīng)的代碼示例:

此外,有開發(fā)者表示使用 TensorFlow 則不會出現(xiàn)該 bug。
提交 Issue 后,PyTorch 運(yùn)維團(tuán)隊的成員也留言參與了討論,在被問到「測試階段是否包含使用 AMD CPU 的用例」時,該成員表示:「在和其他隊友討論之后,我們意識到測試中沒有涉及 AMD CPU 的用例,目前我們還沒有辦法重現(xiàn)這個問題。如果今年晚些時候我們通過支持更多的硬件類型改進(jìn)了測試環(huán)境,將重新討論這個問題。」

此前有網(wǎng)友發(fā)現(xiàn)了 AMD cpu 下 PyTorch 多卡并行卡死的問題,查看日志之后找到錯誤原因,問題才得以解決。而這次暴露的 bug 目前仍未解決。