PyTorch 作為機(jī)器學(xué)習(xí)中廣泛使用的開(kāi)源框架，具有速度快、效率高等特點(diǎn)。而近年來(lái)廣受好評(píng)的 AMD 處理器具有多核、多任務(wù)性能良好、性價(jià)比高等優(yōu)勢(shì)。開(kāi)發(fā)者們一直希望二者聯(lián)合起來(lái)，在 AMD 處理器上使用 PyTorch 進(jìn)行深度學(xué)習(xí)的開(kāi)發(fā)和研究。

前段時(shí)間發(fā)布的 PyTorch 1.8 新增了對(duì) AMD ROCm 的支持，對(duì)于想在 AMD 上用 PyTorch 進(jìn)行深度學(xué)習(xí)的開(kāi)發(fā)者來(lái)說(shuō)，這是一個(gè)好消息。

但是，對(duì)使用 AMD cpu 的開(kāi)發(fā)者用 PyTorch 做 AI 開(kāi)發(fā)，也許沒(méi)那么順利。

這不，我們就從 PyTorch 的 Github 上發(fā)現(xiàn)這么一個(gè)還未解決的 issue。

GitHub 地址：
https://github.com/pytorch/pytorch/issues/52142

有開(kāi)發(fā)者表示：PyTorch 在 AMD CPU 的計(jì)算機(jī)上，用數(shù)據(jù)并行單機(jī)和多 GPU 訓(xùn)練 CNN 會(huì)導(dǎo)致死鎖，而相同的代碼在 Intel CPU 的計(jì)算機(jī)中就不會(huì)出現(xiàn)死鎖。TensorFlow 也不會(huì)出現(xiàn)這種問(wèn)題。

遇到了什么 bug

讓我們來(lái)具體看一下這個(gè)死鎖是怎么產(chǎn)生的？

在提交的 Issue 中，開(kāi)發(fā)者表述在訓(xùn)練的 for 循環(huán)中，代碼會(huì)一直運(yùn)行，直到前向傳播 output = model(images)。代碼會(huì)一直停留在 model(images)階段，而 gpu 的利用率為 0% （內(nèi)存占用率不是 0），有 3 個(gè) 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 它。唯一的解決辦法是重新啟動(dòng)系統(tǒng)。

代碼在 3 臺(tái) Intel cpu 計(jì)算機(jī)上運(yùn)行良好，但在 4 臺(tái)使用 AMD cpu 的機(jī)器上就會(huì)出現(xiàn)上述問(wèn)題。開(kāi)發(fā)者進(jìn)一步測(cè)試了 GTX 1080、Titan V、 Titan RTX、Quadro RTX 8000 和 RTX 3090，證明該 bug 與 GPU 模型無(wú)關(guān)。

此外，該項(xiàng)目還提到分布式數(shù)據(jù)并行（DDP）也存在類似的問(wèn)題。

感興趣的開(kāi)發(fā)者可以按照如下配置復(fù)現(xiàn)一下這個(gè) bug:

使用帶有 AMD cpu 和多個(gè) NVIDIA gpu 的機(jī)器；

Linux、Python3.8、cuda 11.0、pytorch 1.7.1、 torchvision 0.8.2；

編寫(xiě)代碼在 torchvision 中訓(xùn)練 resnet18 模型；

可以嘗試測(cè)試數(shù)據(jù)并行 (DP) 和分布式數(shù)據(jù)并行(DDP)，以檢查是否都會(huì)出現(xiàn)這種情況。

根據(jù)該項(xiàng)目的描述，復(fù)現(xiàn)之后可能會(huì)出現(xiàn)：

當(dāng)使用 AMD cpu 時(shí)，在訓(xùn)練的第一個(gè) epoch 和第一次迭代的前向傳播時(shí)會(huì)發(fā)生死鎖；

當(dāng)使用 intel cpu 時(shí)，相同的代碼運(yùn)行良好。

尚未解決

對(duì)于這次發(fā)現(xiàn)的 bug，有網(wǎng)友表示也存在相同的問(wèn)題。

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

此外，有開(kāi)發(fā)者表示使用 TensorFlow 則不會(huì)出現(xiàn)該 bug。

提交 Issue 后，PyTorch 運(yùn)維團(tuán)隊(duì)的成員也留言參與了討論，在被問(wèn)到「測(cè)試階段是否包含使用 AMD CPU 的用例」時(shí)，該成員表示：「在和其他隊(duì)友討論之后，我們意識(shí)到測(cè)試中沒(méi)有涉及 AMD CPU 的用例，目前我們還沒(méi)有辦法重現(xiàn)這個(gè)問(wèn)題。如果今年晚些時(shí)候我們通過(guò)支持更多的硬件類型改進(jìn)了測(cè)試環(huán)境，將重新討論這個(gè)問(wèn)題。」

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