先進的深度學習模型參數正以指數級速度增長:去年的GPT-2有大約7.5億個參數，今年的GPT-3有1750億個參數。雖然GPT是一個比較極端的例子但是各種SOTA模型正在推動越來越大的模型進入生產應用程式，這裡的最大挑戰是使用GPU卡在合理的時間內完成模型訓練工作的能力。
為了解決這些問題，從業者越來越多地轉向分布式訓練。 分布式訓練是使用多個GPU和/或多個機器訓練深度學習模型的技術。 分布式訓練作業使您能夠克服單GPU內存瓶頸，通過同時利用多個GPU來開發更大，功能更強大的模型。
這篇文章是使用torch.nn.parallel.DistributedDataParallel API在純PyTorch中進行分布式訓練的簡介。 我們會：
討論一般的分布式訓練方式，尤其是數據並行化涵蓋torch.dist和DistributedDataParallel的相關功能，並舉例說明如何使用它們測試真實的訓練腳本，以節省時間什麼是分布式訓練?
在研究分布式和數據並行之前，我們需要先了解一些關於分布式訓練的背景知識。
目前普遍使用的分布式訓練基本上有兩種不同形式:數據並行化和模型並行化。
在數據並行化中，模型訓練作業是在數據上進行分割的。作業中的每個GPU接收到自己獨立的數據批處理切片。每個GPU使用這些數據來獨立計算梯度更新。例如，如果你要使用兩個GPU和32的批處理大小，一個GPU將處理前16條記錄的向前和向後傳播，第二個處理後16條記錄的向後和向前傳播。這些梯度更新然後在gpu之間同步，一起平均，最後應用到模型。
(同步步驟在技術上是可選的，但理論上更快的異步更新策略仍是一個活躍的研究領域)
在模型並行化中，模型訓練作業是在模型上進行分割的。工作中的每個GPU接收模型的一個切片，例如它的層的一個子集。例如，一個GPU負責它的輸出頭，另一個負責輸入層，另一個負責中間的隱藏層。
雖然這兩種技術各有優缺點，但數據並行化在這兩種技術中更容易實現(它不需要了解底層網絡架構)，因此通常首先嘗試這種策略。
(也可以結合使用這些技術，例如同時使用模型和數據並行化，但這是一個高級主題，我們不在這裡介紹)
因為這篇文章是對DistributedDataParallel並行API的介紹，所以我們不會再進一步討論模型並行化的細節——但請關注以後關於這個主題的文章!
數據並行是如何工作的
在前一節中，我給出了數據並行化的概述。在這一節中，我們將深入研究細節。
第一個被廣泛採用的數據並行技術是TensorFlow中的參數伺服器策略。這個功能實際上早於TensorFlow的第一個版本，早在2012年google內部的前身DistBelief中就已經實現了。這一策略在下圖中得到了很好的說明:
在參數伺服器策略中，worker和parameter進程的數量是可變的，每個worker進程在GPU內存中維護自己的模型獨立副本。梯度更新計算如下:
在接收到開始信號後，每個工作進程為其特定的批處理片積累梯度。這些工人以扇出的方式將更新發送到參數伺服器。參數伺服器會一直等待，直到它們擁有所有worker更新，然後對它們負責的梯度更新參數空間的那部分梯度求平均。梯度更新被分散到worker上，然後將它們加起來，應用到內存中模型權重的副本上(從而保持worker模型同步)。一旦每個worker都應用了更新，新的一批訓練就可以開始了。雖然很容易實現，但是這個策略有一些主要的限制。其中最重要的一點是，每個附加的參數伺服器在每個同步步驟中都需要n_workers額外的網絡調用——一個O(n)複雜度代價。計算的總體速度取決於最慢的連接，因此基於大參數伺服器的模型訓練作業在實踐中效率非常低，將網絡GPU利用率推到50%或以下。
更現代的分布式培訓策略廢除了參數伺服器，在DistributedDataParallel 並行策略中，每個進程都是一個工作進程。每個進程仍然在內存中維護模型權重的完整副本，但是批處理片梯度更新現在是同步的，並且直接在工作進程本身上平均。這是使用從高性能計算領域借來的技術來實現的：全歸約算法（all-reduce algorithm）
該圖顯示了全歸約算法的一種特定實現方式，即循環全歸約。 該算法提供了一種優雅的方式來同步一組進程之間的一組變量（在本例中為張量）的狀態。 向量在直接的worker到worker連接的序列中直接傳遞。 這消除了worker與參數伺服器之間的連接所造成的網絡瓶頸，從而大大提高了性能。
在該方案中，梯度更新計算如下:
每個worker維護它自己的模型權重副本和它自己的數據集副本。在接收到開始信號後，每個工作進程從數據集中提取一個分離的批處理，並為該批處理計算一個梯度。worker使用all-reduce算法來同步他們各自的梯度，本地計算所有節點上相同的平均梯度。每個worker都將梯度更新應用到它的本地模型副本上。下一批訓練開始。在2017年百度的一篇論文《Bringing HPC Techniques to Deep Learning》中，這種全精簡策略被提上了日程。它的偉大之處在於它基於眾所周知的HPC技術以及長期存在的開源實現。All-reduce包含在消息傳遞接口(MPI)的標準中，這就是為什麼PyTorch不少於三個不同的後端實現:Open MPI、NVIDIA NCCL和Facebook Gloo（一般情況下建議使用NVIDIA NCCL）
數據分發，第1部分:流程初始化
不幸的是，將訓練腳本修改為使用DistributedDataParallel 並行策略並不是簡單的一行更改。
為了演示API是如何工作的，我們將構建一個完整的分布式訓練腳本(在本文後面的基準測試中，我們將繼續討論這個腳本)。
您需要處理的第一個也是最複雜的新事情是進程初始化。普通的PyTorch訓練腳本在單個進程中執行其代碼的單一副本。使用數據並行模型，情況就更加複雜了:現在訓練腳本的同步副本與訓練集群中的gpu數量一樣多，每個gpu運行在不同的進程中。
考慮以下最小的例子:
# multi_init.pyimport torchimport torch.distributed as distimport torch.multiprocessing as mpdef init_process(rank, size, backend='gloo'):""" Initialize the distributed environment. """ os.environ['MASTER_ADDR'] = '127.0.0.1' os.environ['MASTER_PORT'] = '29500' dist.init_process_group(backend, rank=rank, world_size=size)def train(rank, num_epochs, world_size): init_process(rank, world_size) print( f"Rank {rank + 1}/{world_size} process initialized.\n" ) # rest of the training script goes here!WORLD_SIZE = torch.cuda.device_count()if __name__=="__main__": mp.spawn( train, args=(NUM_EPOCHS, WORLD_SIZE), nprocs=WORLD_SIZE, join=True )
在MPI的世界中，WORLDSIZE是編排的進程數量，(全局)rank是當前進程在該MPI中的位置。例如，如果這個腳本要在一個有4個gpu的強大機器上執行，WORLDSIZE應該是4(因為torch.cuda.device_count() == 4)，所以是mp.spawn會產生4個不同的進程，它們的等級 分別是0、1、2或3。等級為0的進程被賦予一些額外的職責，因此被稱為主進程。
當前進程的等級將作為派生入口點（在本例中為訓練方法）作為其第一個參數傳遞。 在訓練時可以執行任何工作之前，它需要首先建立與對等點對點的連接。 這是dist.initprocessgroup的工作。 在主進程中運行時，此方法在MASTERADDR：MASTERPORT上設置套接字偵聽器，並開始處理來自其他進程的連接。 一旦所有進程都已連接，此方法將處理建立對等連接，以允許進程進行通信。
請注意，此代碼僅適用於在一臺多GPU機器上進行訓練！ 同一臺機器用於啟動作業中的每個流程，因此訓練只能利用連接到該特定機器的GPU。 這使事情變得容易：設置IPC就像在localhost上找到一個空閒埠一樣容易，該埠對於該計算機上的所有進程都是立即可見的。 跨計算機的IPC更為複雜，因為它需要配置一個對所有計算機可見的外部IP位址。
在本入門教程中，我們將特別關注單機訓練（也稱為垂直擴展）。 即使在單主機，垂直擴展也是一個非常強大的工具。 如果在雲端，垂直擴展可讓您將深度學習訓練工作一直擴展到8xV100實例（例如AWS上的p3.16xlarge）。
我們將在以後的博客文章中討論水平擴展和數據並行化。
數據分發，第2部分:流程同步
現在我們了解了初始化過程，我們可以開始完成所有工作的train方法的主體。
回想一下我們到目前為止：
def train(rank, num_epochs, world_size):init_process(rank, world_size) print( f"{rank + 1}/{world_size} process initialized.\n" ) # rest of the training script goes here!
我們的四個訓練過程中的每一個都會運行此函數直到完成，然後在完成時退出。 如果我們現在（通過python multi_init.py）運行此代碼，我們將在控制臺上看到類似以下內容：
$ python multi_init.py1/4 process initialized.3/4 process initialized.2/4 process initialized.4/4 process initialized.
這些過程是獨立執行的，並且不能保證訓練循環中任一點處於什麼狀態。 所以這裡需要對初始化過程進行一些仔細的更改。
（1）任何下載數據的方法都應隔離到主進程中。
否則，將在所有過程之間複製數據下載過程，從而導致四個過程同時寫入同一文件，這是造成數據損壞的原因。
修改這很容易做到：
# import torch.distributed as distif rank == 0:downloading_dataset() downloading_model_weights()dist.barrier()print( f"Rank {rank + 1}/{world_size} training process passed data download barrier.\n")
此代碼示例中的dist.barrier將阻塞調用，直到完成主進程（rank == 0）downloadingdataset和downloadingmodel_weights為止。 這樣可以將所有網絡I / O隔離到一個進程中，並防止其他進程繼續前進。
（2）數據加載器需要使用DistributedSampler。
代碼示例：
def get_dataloader(rank, world_size):dataset = PascalVOCSegmentationDataset() sampler = DistributedSampler( dataset, rank=rank, num_replicas=world_size, shuffle=True ) dataloader = DataLoader( dataset, batch_size=8, sampler=sampler )
DistributedSampler使用rank和worldsize找出如何將整個過程中的數據集拆分為不重疊的批次。 工作進程的每個訓練步驟都從其本地數據集副本中檢索batchsize觀測值。 在四個GPU的示例情況下，這意味著有效批大小為8 * 4 = 32。
（3）在正確的設備中加載張量。
為此，請使用該進程正在管理的設備的rank來參數化.cuda（）調用：
batch = batch.cuda(rank)segmap = segmap.cuda(rank)model = model.cuda(rank)
（4）必須禁用模型初始化中的任何隨機性。
在整個訓練過程中，模型必須啟動並保持同步，這一點非常重要。 否則，您將獲得不正確的漸變，並且模型將無法收斂。
可以使用以下代碼使torch.nn.init.kaimingnormal之類的隨機初始化方法具有確定性：
torch.manual_seed(0)torch.backends.cudnn.deterministic = Truetorch.backends.cudnn.benchmark = Falsenp.random.seed(0)
PyTorch文檔有一整個頁面專門討論此主題：randomness.html
（5）任何執行文件I / O的方法都應與主進程隔離。
這與隔離網絡I / O的原因相同，是必要的：由於並發寫入同一文件而導致的效率低下和潛在的數據損壞。 同樣，使用簡單的條件邏輯很容易做到這一點：
if rank == 0:if not os.path.exists('/spell/checkpoints/'): os.mkdir('/spell/checkpoints/') torch.save( model.state_dict(), f'/spell/checkpoints/model_{epoch}.pth' )
順便說一句，請注意，要記錄的所有全局損失值或統計信息都需要您自己同步數據。 可以使用torch.distributed中的其他MPI原語來完成此操作，本教程未對此進行深入介紹。 可以參閱Distributed Communication Package PyTorch文檔頁面以獲取詳細的API參考。
（6）將模型包裝在DistributedDataParallel中。
假設您已正確完成所有其他操作，這就是神奇的地方。
model = DistributedDataParallel(model, device_ids=[rank])
這一步時必須的也是最後一步，如果你做完了恭喜你，你的模型現在可以在分布式數據並行模式下訓練！
那DataParallel呢？
熟悉PyTorch API的讀者可能知道PyTorch中還有另一種數據並行化策略，即torch.nn.DataParallel。 該API易於使用。 您要做的就是包裝模型初始化，如下所示：
model = nn.DataParallel(model)
所有的改動只有一行！ 為什麼不使用它呢？
在後臺，DataParallel使用多線程而不是多處理來管理其GPU工作器。 這極大地簡化了實現：由於工作進程是同一進程的所有不同線程，因此它們都可以訪問相同的共享狀態，而無需任何其他同步步驟。
但是，由於存在全局解釋器鎖，在Python中將多線程用於計算作業的效果很差。 如下一節中的基準測試所示，使用DataParallel並行化的模型比使用DistributedDataParallel並行化的模型要慢得多。
儘管如此，如果你不想花費額外的時間和精力郵箱使用多GPU訓練，DataParallel實可以考慮的。
基準測試
為了對分布式模型訓練性能進行基準測試，我在PASCAL VOC 2012數據集（來自torchvision數據集）上訓練了20個輪次的DeepLabV3-ResNet 101模型（通過Torch Hub）。 我啟動了五個不同版本的模型巡訓練工作：一次在單個V100上（在AWS上為p3.2xlarge），一次在V100x4（p3.8xlarge）和V100x8（p3.16xlarge）上使用 DistributedDataParallel和DataParallel。 該基準測試不包括運行開始時花在下載數據上的時間-僅模型訓練和節省時間計數。
DistributedDataParallel的效率明顯高於DataParallel，但還遠遠不夠完美。 從V100x1切換到V100x4是原始GPU功耗的4倍，但模型訓練速度僅為3倍。 通過升級到V100x8使計算進一步加倍，只會使訓練速度提高約30％,但是DataParallel的效率幾乎達到了DistributedDataParallel的水平。
結論
在本文中，我們討論了分布式訓練和數據並行化，了解了DistributedDataParallel和DataParallel API，並將其應用於實際模型並進行了一個簡單的基準測試。
分布式計算的領域還有很多可以改進，PyTorch團隊剛剛在本月獲得了新的PR，該PR承諾將對DistributedDataParallel的性能進行重大改進。 希望這些時間在將來的版本中降下來！
我認為討論不多的事情是分布式訓練對開發人員生產力的影響。 從「需要三個小時的訓練」到「需要一個小時的訓練」，即使採用中等大小的模型，也可以極大地增加您可以在一天之內和使用該模型進行的實驗的數量，這對開發人員而言是一個巨大的進步。
作者：Aleksey Bilogur
github/spellml/deeplab-voc-2012
deephub翻譯組
註：本文是我們和pytorch-handbook的作者合作翻譯的第一篇文章，後續還會有更深入的合作，另外介紹一下pytorch-handbook .這本pytorch的中文手冊已經在github上獲取了12000+的star是一本非常詳細的pytorch入門教程和查詢手冊，如果是想深入的學習，趕緊關注這個項目吧。
github/zergtant/pytorch-handbook