?作者 | 偉大是熬出來的

單位 | 同濟大學

研究方向 | 機器閱讀理解

前言

因為課題組發的卡還沒有下來，先向導師問了實驗室的兩張卡借用。之前都是單卡訓練模型，正好在這個機會實踐以下單機多卡訓練模型的方法。關于 DDP 網上有很多資料，但都比較零碎（有些博客的代碼甚至沒辦法 run），Pytorch 給出的官方文檔看起來也比較吃力。因此這篇文章的主要目的是梳理一下筆者學習過程中認為比較好的資料，用通俗的語言介紹一下 DDP 的原理，最后給出使用 DDP 的模板以及一份詳細的運行案例。

當代研究生應當掌握的并行訓練方法（單機多卡）:

https://zhuanlan.zhihu.com/p/98535650

這篇文章中介紹了目前常用的并行訓練方法。其中，nn.Dataparallel?的使用最簡單，只需要使用 Dataparallel?包裝模型，再設置一些參數就可以實現。參數中需要指定參與訓練的 GPU，device_ids=gpus；匯總梯度的 GPU，output_device=gpus[0]。

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

nn.Dataparallel?方法實際上是使用單進程將模型和數據加載到多個 GPU 上，控制數據在 GPU 之間流動，協同不同的 GPU 上的模型進行并行訓練。這篇文章 [8] 中提到 nn.Dataparallel?方法的弊端：在訓練的過程中，每個 batch 的模型權重是在一個進程上計算出來之后，再分發到每個 GPU 上。

這會導致負載不均衡的問題，可能第一個 GPU（12GB）占用了 10GB，剩余 GPU 卻只使用了 4GB。因為在數據并行的時候，loss 會在第一個 GPU 上相加計算，更新好以后把權重分發到其余卡。這就造成了第一個 GPU 的負載遠大于其他顯卡。

nn.DistributedDataParallel 原理

與 nn.Dataparallel?使用單進程控制多個 GPU 不同， nn.DistributedDataParallel 為每個 GPU 都創建一個進程。這些 GPU 可以位于同一個結點上（單機多卡），也可以分布在多個節點上（多機多卡）。每個進程都執行相同的任務，每個進程都與其他進程進行通信。

另外一點不同是，只有梯度會在進程（GPU）之間傳播。以單機多卡舉例，假設我們有三張卡并行訓練，那么在每個 epoch 中，數據集會被劃分成三份給三個 GPU，每個 GPU 使用自己的 minibatch 數據做自己的前向計算，然后梯度在 GPU 之間全部約簡。在反向傳播結束的時候，每個 GPU 都有平均的梯度，確保模型權值保持同步（synchronized）。

運行模板

這一小節以官方文檔給出的 demo 作為例子，介紹 DDP 的使用模板以及運行流程：

https://pytorch.org/tutorials/intermediate/ddp_tutorial.html

首先我們導入庫文件，其中 torch.multiprocessing?用于創建進程，后面會詳細介紹。

import?os import?torch import?torch.distributed?as?dist import?torch.nn?as?nn import?torch.optim?as?optim import?torch.multiprocessing?as?mpfrom?torch.nn.parallel?import?DistributedDataParallel?as?DDP

前文提到，DDP 模型對于每個 GPU 都會創建一個單獨的進程管理。在程序并發執行的過程中，進程之間需要同步和通信。因此我們需要一個方法管理進程組，這個方法需要知道如何找到進程 0。也要知道進程組中同步了多少個進程。init_process_group 方法能夠實現上述功能，其中參數的含義解釋如下：

• backend：使用的后端。包括 mpi, gloo, 和 nccl。根據官方文檔的介紹，nccl 是運行速度最快的，因此大多設置為這個

• rank：當前進程的等級。在 DDP 管理的進程組中，每個獨立的進程需要知道自己在所有進程中的階序，我們稱為 rank

• world_size：在 DDP 管理的進程組中，每個獨立的進程還需要知道進程組中管理進程的數量，我們稱為 world_size

下面 setup?以及?cleanup?分別實現了進程組的設置以及銷毀。

ef?setup(rank,?world_size):os.environ['MASTER_ADDR']?=?'localhost'os.environ['MASTER_PORT']?=?'12355'#?initialize?the?process?groupdist.init_process_group("nccl",?rank=rank,?world_size=world_size)def?cleanup():dist.destroy_process_group()

本例中我們訓練一個簡單的網絡結構?ToyModel

class?ToyModel(nn.Module):def?__init__(self):super(ToyModel,?self).__init__()self.net1?=?nn.Linear(10,?10)self.relu?=?nn.ReLU()self.net2?=?nn.Linear(10,?5)def?forward(self,?x):return?self.net2(self.relu(self.net1(x)))

下面是模型訓練部分的模板。數據和模型加載到當前進程使用的 GPU 中，正常進行正反向傳播，需要注意以下幾點：

• 每個進程都需要復制一份模型以及數據。我們需要根據前文提到的?rank?和?world_size?兩個參數初始化進程組。這樣進程之間才能相互通信。使用我們前文定義的?setup()?方法實現；

• model = ToyModel().to(rank)?這條語句將我們的模型移動到對應的 GPU中，?rank?參數作為進程之間的階序，可以理解為當前進程 index。由于每個進程都管理自己的 GPU，因此通過階序可以索引到對應的 GPU；

• ddp_model = DDP（model, device_ids=[rank]）這條語句包裝了我們的模型；

• 其他與 pytorch 中訓練模型的模板相同，最后一點需要注意的是，在我們將 tensor 移動到 GPU 的時候，同樣需要使用 rank?索引，代碼中體現在第 14 行。

def?demo_basic(rank,?world_size):print(f"Running?basic?DDP?example?on?rank?{rank}.")setup(rank,?world_size)#?create?model?and?move?it?to?GPU?with?id?rankmodel?=?ToyModel().to(rank)ddp_model?=?DDP(model,?device_ids=[rank])loss_fn?=?nn.MSELoss()optimizer?=?optim.SGD(ddp_model.parameters(),?lr=0.001)optimizer.zero_grad()outputs?=?ddp_model(torch.randn(20,?10))labels?=?torch.randn(20,?5).to(rank)loss_fn(outputs,?labels).backward()optimizer.step()cleanup()

最后是啟動器的介紹。DDP 的啟動有兩種方式，分別對應不同的代碼。

• torch.distributed.launch 啟動器，用于在命令行分布式地執行 python 文件。在執行過程中，啟動器會將當前進程的（其實就是 GPU 的）index 通過參數傳遞給 python。在使用的時候執行語句CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py?調用啟動器 torch.distributed.launch。
  

• 本例中使用的是第二種方法 torch.multiprocessing.spawn。使用時，只需要調用 torch.multiprocessing.spawn，torch.multiprocessing 就會幫助我們自動創建進程。
  

如下面的代碼所示，spawn 開啟了 world_size 個進程，每個進程執行 demo_fn 并向其中傳入 local_rank（當前進程 index）作為參數。這里需要結合前文 demo_basic 的定義來看。args 中的 world_size 對應 demo_basic?的 world_size 參數；mp.spawn 中 nprocs 則是創建進程的數量；至于demo_basic?中的 rank 參數，應當是框架內部實現了索引機制因此不需要我們顯示對應（筆者自己的理解）。

def?run_demo(demo_fn,?world_size):mp.spawn(demo_fn,args=(world_size,),nprocs=world_size,join=True)

完整代碼如下：

import?os import?sys import?tempfile import?torch import?torch.distributed?as?dist import?torch.nn?as?nn import?torch.optim?as?optim import?torch.multiprocessing?as?mpfrom?torch.nn.parallel?import?DistributedDataParallel?as?DDP#?On?Windows?platform,?the?torch.distributed?package?only #?supports?Gloo?backend,?FileStore?and?TcpStore. #?For?FileStore,?set?init_method?parameter?in?init_process_group #?to?a?local?file.?Example?as?follow: #?init_method="file:///f:/libtmp/some_file" #?dist.init_process_group( #????"gloo", #????rank=rank, #????init_method=init_method, #????world_size=world_size) #?For?TcpStore,?same?way?as?on?Linux.def?setup(rank,?world_size):os.environ['MASTER_ADDR']?=?'localhost'os.environ['MASTER_PORT']?=?'12355'#?initialize?the?process?groupdist.init_process_group("nccl",?rank=rank,?world_size=world_size)def?cleanup():dist.destroy_process_group()class?ToyModel(nn.Module):def?__init__(self):super(ToyModel,?self).__init__()self.net1?=?nn.Linear(10,?10)self.relu?=?nn.ReLU()self.net2?=?nn.Linear(10,?5)def?forward(self,?x):return?self.net2(self.relu(self.net1(x)))def?demo_basic(rank,?world_size):print(f"Running?basic?DDP?example?on?rank?{rank}.")setup(rank,?world_size)#?create?model?and?move?it?to?GPU?with?id?rankmodel?=?ToyModel().to(rank)ddp_model?=?DDP(model,?device_ids=[rank])loss_fn?=?nn.MSELoss()optimizer?=?optim.SGD(ddp_model.parameters(),?lr=0.001)optimizer.zero_grad()outputs?=?ddp_model(torch.randn(20,?10))labels?=?torch.randn(20,?5).to(rank)loss_fn(outputs,?labels).backward()optimizer.step()cleanup()def?run_demo(demo_fn,?world_size):mp.spawn(demo_fn,args=(world_size,),nprocs=world_size,join=True)if?__name__?==?"__main__":n_gpus?=?torch.cuda.device_count()assert?n_gpus?>=?2,?f"Requires?at?least?2?GPUs?to?run,?but?got?{n_gpus}"world_size?=?n_gpusrun_demo(demo_basic,?world_size)

這份代碼可以直接運行，輸入結果如下,筆者使用兩張卡，因此對應的 rank 分別是 0 和 1：

Running?basic?DDP??example?on?Rank?1 Running?basic?DDP??example?on?Rank?0

在官網給的這份 demo 之外，其實還有一點需要注意。我們使用 pytorch 處理數據集創建 dataloader 的過程中，需要使用 DistributedSampler 采樣器。我們已經知道，每個進程都會拷貝一份模型和數據的副本，但是在并行計算的過程中，單個進程只會處理自己的 minibatch 的數據。

假設我們使用五個 GPU 并行，其對應的五個進程都有模型和數據的副本，我們假設訓練集有一萬條數據，那么在單個 epoch 中每個進程實際上只需要使用兩千條數據訓練，之后進行梯度整合。那么進程如何知道自己需要處理哪些數據呢？這是 DistributedSampler 的功能。

關于 DistributedSampler 具體做了什么，可以參考文獻 7。

在此基礎上，筆者根據官網的 demo 總結了一份使用 DDP 進行多卡并行加速模型的模板，讀者在使用過程中根據需要進行簡單更改即可使用：

import?os import?torch import?torch.distributed?as?dist import?torch.nn?as?nn import?torch.optim?as?optim import?torch.multiprocessing?as?mpfrom?torch.nn.parallel?import?DistributedDataParallel?as?DDPdef?setup(rank,?world_size):os.environ['MASTER_ADDR']?=?'localhost'os.environ['MASTER_PORT']?=?'12355'#?initialize?the?process?groupdist.init_process_group("nccl",?rank=rank,?world_size=world_size)def?run(demo_fn,?world_size):setup(rank,?world_size)torch.manual_seed(18)torch.cuda.manual_seed_all(18)torch.backends.cudnn.deterministic?=?Truetorch.cuda.set_device(rank)?#?這里設置?device?，后面可以直接使用?data.cuda(),否則需要指定?ranktrain_dataset?=?...train_sampler?=?torch.utils.data.distributed.DistributedSampler(train_dataset)train_loader?=?torch.utils.data.DataLoader(train_dataset,?batch_size=...,?sampler=train_sampler)model?=?...model?=?torch.nn.parallel.DistributedDataParallel(model,?device_ids=[rank])optimizer?=?optim.SGD(model.parameters())for?epoch?in?range(100):train_sampler.set_epoch(epoch)for?batch_idx,?(data,?target)?in?enumerate(train_loader):data?=?data.cuda()target?=?target.cuda()...output?=?model(images)loss?=?criterion(output,?target)...optimizer.zero_grad()loss.backward()optimizer.step()if?__name__?==?"__main__":n_gpus?=?torch.cuda.device_count()assert?n_gpus?>=?2,?f"Requires?at?least?2?GPUs?to?run,?but?got?{n_gpus}"world_size?=?n_gpusmp.spawn(run,args=(world_size,),nprocs=world_size,join=True)

總結一下，使用 DDP 進行多卡并行加速模型的重點：

• init_process_group 函數管理進程組
  

• 在創建 Dataloader 的過程中，需要使用 DistributedSampler 采樣器
  

• 正反向傳播之前需要將數據以及模型移動到對應 GPU，通過參數?rank?進行索引，還要將模型使用 DistributedDataParallel 進行包裝
  

• 在每個 epoch 開始之前，需要使用 train_sampler.set_epoch（epoch）為 train_sampler 指定 epoch，這樣做可以使每個 epoch 劃分給不同進程的? minibatch 不同，從而在整個訓練過程中，不同的進程有機會接觸到更多的訓練數據
  

• 使用啟動器進行啟動。不同啟動器對應不同的代碼。torch.distributed.launch 通過命令行的方法執行，torch.multiprocessing.spawn 則可以直接運行程序。
  

最后，筆者使用上述模板實現了一份基于 Roberta 的文本分類任務，使用 DDP 進行單機雙卡并行加速，運行的結果如下，有感興趣的讀者再放代碼吧，這里展示一下運行結果：

參考文獻

注：下述文獻大多使用 torch.distributed.launch 啟動器執行程序

[1] 當代研究生應當掌握的并行訓練方法（單機多卡）：

https://zhuanlan.zhihu.com/p/98535650

[2] PyTorch Parallel Training（單機多卡并行、混合精度、同步BN訓練指南文檔）：https://zhuanlan.zhihu.com/p/145427849

[3] pytorch多卡分布式訓練簡要分析：https://zhuanlan.zhihu.com/p/159404316

[4] Pytorch中的Distributed Data Parallel與混合精度訓練（Apex）：https://zhuanlan.zhihu.com/p/105755472

[5] PyTorch分布式訓練基礎--DDP使用：https://zhuanlan.zhihu.com/p/358974461

[6] 使用PyTorch編寫分布式應用程序：https://www.jianshu.com/p/be9f8b90a1b8?utm_campaign=hugo&utm_medium=reader_share&utm_content=note&utm_source=weixin-friends

[7] DistributedSampler 具體做了什么：https://blog.csdn.net/searobbers_duck/article/details/115299691?utm_medium=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7Edefault-1.no_search_link&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7Edefault-1.no_search_link

[8]?Pytorch的nn.DataParallel：https://zhuanlan.zhihu.com/p/102697821

更多閱讀

#投 稿?通 道#

?讓你的文字被更多人看到?

如何才能讓更多的優質內容以更短路徑到達讀者群體，縮短讀者尋找優質內容的成本呢？答案就是：你不認識的人。

總有一些你不認識的人，知道你想知道的東西。PaperWeekly 或許可以成為一座橋梁，促使不同背景、不同方向的學者和學術靈感相互碰撞，迸發出更多的可能性。?

PaperWeekly 鼓勵高校實驗室或個人，在我們的平臺上分享各類優質內容，可以是最新論文解讀，也可以是學術熱點剖析、科研心得或競賽經驗講解等。我們的目的只有一個，讓知識真正流動起來。

📝?稿件基本要求：

? 文章確系個人原創作品，未曾在公開渠道發表，如為其他平臺已發表或待發表的文章，請明確標注?

? 稿件建議以?markdown?格式撰寫，文中配圖以附件形式發送，要求圖片清晰，無版權問題

? PaperWeekly 尊重原作者署名權，并將為每篇被采納的原創首發稿件，提供業內具有競爭力稿酬，具體依據文章閱讀量和文章質量階梯制結算

📬?投稿通道：

? 投稿郵箱：hr@paperweekly.site?

? 來稿請備注即時聯系方式（微信），以便我們在稿件選用的第一時間聯系作者

? 您也可以直接添加小編微信（pwbot02）快速投稿，備注：姓名-投稿

△長按添加PaperWeekly小編

🔍

現在，在「知乎」也能找到我們了

進入知乎首頁搜索「PaperWeekly」

點擊「關注」訂閱我們的專欄吧

·

總結

以上是生活随笔為你收集整理的PyTorch多卡分布式训练：DistributedDataParallel (DDP) 简要分析的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。