PyTorch 單機多卡操作總結(jié)：分布式DataParallel，混合精度，Horovod)

↑ 點擊藍字?關注極市平臺

作者丨科技猛獸@知乎

來源丨h(huán)ttps://zhuanlan.zhihu.com/p/158375055

編輯丨極市平臺

極市導讀

?

本文介紹了數(shù)種實現(xiàn)單機多卡操作的方法，含有大量代碼，并給出了實踐中作者踩過的坑及其解決方案。>>加入極市CV技術交流群，走在計算機視覺的最前沿

在上一篇文章中（https://zhuanlan.zhihu.com/p/158375254）我們看到了多GPU訓練，也就是最簡單的單機多卡操作nn.DataParallel。但是很遺憾這種操作還不夠優(yōu)秀，于是就有了今天這篇文章~

寫這篇文章的時候看了很多的tutorials，附在文末了，在此先向文末的每位作者致敬，感謝大佬們！

其實單機多卡的辦法還有很多(如下)，而且上篇的方法是相對較慢的。

1、nn.DataParallel 簡單方便的 nn.DataParallel

2、torch.distributed 使用 torch.distributed 加速并行訓練

3、apex 使用 apex 再加速。

這里，記錄了使用 4 塊 Tesla V100-PICE 在 ImageNet 進行了運行時間的測試，測試結(jié)果發(fā)現(xiàn) Apex 的加速效果最好，但與 Horovod/Distributed 差別不大，平時可以直接使用內(nèi)置的 Distributed。Dataparallel 較慢，不推薦使用。那好像上一篇白寫了~

看到這里你可能已經(jīng)懵逼了，莫慌，下面會分別進行介紹(這里先附上一篇教程，看不懂的話直接當做沒看見即可)：

https://yangkky.github.io/2019/07/08/distributed-pytorch-tutorial.html

1. 先問兩個問題

問1：為啥非要單機多卡？

答1：加速神經(jīng)網(wǎng)絡訓練最簡單的辦法就是上GPU，如果一塊GPU還是不夠，就多上幾塊。

事實上，比如BERT和GPT-2這樣的大型語言模型甚至是在上百塊GPU上訓練的。

為了實現(xiàn)多GPU訓練，我們必須想一個辦法在多個GPU上分發(fā)數(shù)據(jù)和模型，并且協(xié)調(diào)訓練過程。

問2：上一篇講得單機多卡操作nn.DataParallel，哪里不好？

答2：要回答這個問題我們得先簡單回顧一下nn.DataParallel，要使用這玩意，我們將模型和數(shù)據(jù)加載到多個 GPU 中，控制數(shù)據(jù)在 GPU 之間的流動，協(xié)同不同 GPU 上的模型進行并行訓練。具體怎么操作？

我們只需要用 DataParallel 包裝模型，再設置一些參數(shù)即可。需要定義的參數(shù)包括：

• 參與訓練的 GPU 有哪些，device_ids=gpus。
• 用于匯總梯度的 GPU 是哪個，output_device=gpus[0] 。

DataParallel 會自動幫我們將數(shù)據(jù)切分 load 到相應 GPU，將模型復制到相應 GPU，進行正向傳播計算梯度并匯總：

model = nn.DataParallel(model.cuda(), device_ids=gpus, output_device=gpus[0])

值得注意的是，模型和數(shù)據(jù)都需要先 load 進 GPU 中，DataParallel 的 module 才能對其進行處理，否則會報錯：

# main.py
import torch
import torch.distributed as dist

gpus = [0, 1, 2, 3]
torch.cuda.set_device('cuda:{}'.format(gpus[0]))

train_dataset = ...

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=...)

model = ...
model = nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0])

optimizer = optim.SGD(model.parameters())

for epoch in range(100):
   for batch_idx, (data, target) in enumerate(train_loader):
      images = images.cuda(non_blocking=True)
      target = target.cuda(non_blocking=True)
      ...
      output = model(images)
      loss = criterion(output, target)
      ...
      optimizer.zero_grad()
      loss.backward()
      optimizer.step()

稍微解釋幾句：model.to(device)將模型遷移到GPU里面，images.cuda，target.cuda把數(shù)據(jù)遷移到GPU里面。
nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0])包裝模型。

缺點：

• 在每個訓練批次（batch）中，因為模型的權重都是在一個進程上先算出來，然后再把他們分發(fā)到每個GPU上，所以網(wǎng)絡通信就成為了一個瓶頸，而GPU使用率也通常很低。
• 除此之外，nn.DataParallel 需要所有的GPU都在一個節(jié)點（一臺機器）上，且并不支持 Apex 的 混合精度訓練。

一句話，一個進程算權重使通信成為瓶頸，nn.DataParallel慢而且不支持混合精度訓練。

2. 使用 torch.distributed 加速并行訓練：

DataParallel：單進程控制多 GPU。

DistributedDataParallel：多進程控制多 GPU，一起訓練模型。

2.1 介紹

在 1.0 之后，官方終于對分布式的常用方法進行了封裝，支持 all-reduce，broadcast，send 和 receive 等等。通過 MPI 實現(xiàn) CPU 通信，通過 NCCL 實現(xiàn) GPU 通信。官方也曾經(jīng)提到用 DistributedDataParallel 解決 DataParallel 速度慢，GPU 負載不均衡的問題，目前已經(jīng)很成熟了。

與 DataParallel 的單進程控制多 GPU 不同，在 distributed 的幫助下，我們只需要編寫一份代碼，torch 就會自動將其分配給n個進程，分別在n個 GPU 上運行。

和單進程訓練不同的是，多進程訓練需要注意以下事項：

• 在喂數(shù)據(jù)的時候，一個batch被分到了好幾個進程，每個進程在取數(shù)據(jù)的時候要確保拿到的是不同的數(shù)據(jù)（DistributedSampler）；
• 要告訴每個進程自己是誰，使用哪塊GPU（args.local_rank）；
• 在做BatchNormalization的時候要注意同步數(shù)據(jù)。
  
  

2.2 使用方式

2.2.1 啟動方式的改變

在多進程的啟動方面，我們不用自己手寫 multiprocess 進行一系列復雜的C