[源碼解析] PyTorch 分布式(8) -------- DistributedDataParallel之論文篇
0x00 摘要
工欲善其事，必先利其器，為了更好的分析代碼，我們先來學習一下相關論文。
PyTorch 開發者在實現的同時，發布了一篇論文：[ PyTorch Distributed: Experiences on Accelerating Data Parallel Training ] Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, Soumith Chintal。
其地址為：http://www.vldb.org/pvldb/vol13/p3005-li.pdf
因為論文較長，所以本文翻譯其思路和實現之中的部分內容，在后文之中將以這篇論文為基礎，結合源碼來進行分析。本文不完全按照原論文的順序進行翻譯，筆者會對其重點做標注，也會按照自己的理解進行調整，另外，原文是基於 PyTorch 1.5，與最新 PyTorch 有部分出入。
本系列其他文章如下：
[源碼解析]PyTorch如何實現前向傳播(1) --- 基礎類(上)
[源碼解析]PyTorch如何實現前向傳播(2) --- 基礎類(下)
[源碼解析] PyTorch如何實現前向傳播(3) --- 具體實現
[源碼解析] Pytorch 如何實現后向傳播 (1)---- 調用引擎
[源碼解析] Pytorch 如何實現后向傳播 (2)---- 引擎靜態結構
[源碼解析] Pytorch 如何實現后向傳播 (3)---- 引擎動態邏輯
[源碼解析] PyTorch 如何實現后向傳播 (4)---- 具體算法
[源碼解析] PyTorch 分布式(1)------歷史和概述
[源碼解析] PyTorch 分布式(2) ----- DataParallel(上)
[源碼解析] PyTorch 分布式(3) ----- DataParallel(下)
[源碼解析] PyTorch 分布式(4)------分布式應用基礎概念
[源碼解析] PyTorch分布式(5) ------ DistributedDataParallel 總述&如何使用
[源碼解析] PyTorch分布式(6) -------- DistributedDataParallel -- init_method & store
[源碼解析] PyTorch 分布式(7) ----- DistributedDataParallel 之進程組
0x01 原文摘要
深度學習的最新進展證明了大型數據集和大型模型的價值，這就需要將模型訓練擴展到更多計算資源之上。由於其簡單的原理和廣泛的適用性，數據並行已成為分布式培訓的一種流行解決方案。通常，分布式數據並行技術在每個計算源上復制模型以在每個worker之上獨立地生成梯度，然后在每次迭代中通信這些梯度以保持模型副本的一致性。盡管該技術概念簡單，但計算和通信之間的微妙依賴性使得優化分布式訓練效率非常重要。從1.5版開始，Pytorch 提供了幾種加速分布式數據並行的技術，包括bucketing梯度、通信重疊計算和跳過梯度同步。評估表明，當適當配置時，Pyrotch分布式數據並行模塊可使用256 GPU實現近似線性的可擴展性。
0x02 引論
訓練DNN模型通常重復執行以下三個步驟：
- 向前傳遞以計算損失。
- 向后傳播以計算梯度。
- 以及優化器步驟以更新參數。
數據並行性的概念普遍適用於此類框架：應用程序可以創建一個模型的多個副本，每個模型副本處理一部分訓練數據，並獨立執行向前和向后傳播。之后，模型副本可以根據算法同步其梯度或更新的參數。
2.1 挑戰
看起來，完全在應用程序端構建數據並行的工作版本是可能的，因為它只需要在每次迭代中插入適當的通信。然而，擠出最后一點性能需要在設計和調整方面付出巨大的努力。在平台端提供本機分布式數據並行API將幫助應用程序開發人員專注於優化其模型，而平台開發團隊可以持續透明地提高訓練速度。
要提供一個通用的分布式數據並行包，有三個方面的挑戰。
- 數學等價：數據並行的目的是加速對大型數據集的訓練。應用程序希望獲得相同的結果模型，就好像所有培訓都是在本地進行，沒有模型復制一樣。這就要求盡管它是分布式訓練，但是應該數學等價於本地訓練。
- 非侵入式和攔截式API：應用程序開發通常從本地模型開始，然后在必要時擴展。所以需要有一個從本地模型開始，修改代碼以適應分布式的過程。
- 為了避免這個從本地模型到分布式模型的過渡期間太過麻煩，API在應用程序代碼中必須是非侵入性的。
- 另一方面，API也需要允許一個內部實現來及時截獲各種信號，以便執行通信和系統優化。
- 高性能：數據並行訓練受制於計算和通信之間微妙的依賴關系。設計和實現必須探索解決方案空間，以有效地將更多資源轉換為更高的訓練吞吐量。
2.2 實現和評估
PyTorch以nn.Module類的形式提供分布式數據並行，其中應用程序在構建時以子模塊的形式提供其模型。為了保證數學等效性，所有副本都從相同的模型參數初始值開始，並同步梯度，以便在整個訓練迭代中保持參數一致。為了最大限度地降低集成度，該實現（分布式數據並行模型）暴露了與用戶模型相同的forward API，這允許應用程序無縫地用分布式數據並行模型對象替換之前出現的用戶模型，而無需額外的代碼更改。設計中集成了多種技術，以提供高性能培訓，包括bucketing gradients，與計算的重疊通信和跳過同步。
評估是在一個專用的32 GPU集群和一個更大的共享權限中的256 GPU上進行的。我們開發了基准程序來評估不同規模的分布式包，以深入了解不同優化技術和配置的性能影響。實驗還包括NCCL和Gloo通信庫之間的比較。結果表明：
- 通信是影響訓練延遲的主要因素，其影響隨模型尺寸的增大而增大；
- 存儲桶大小對通信效率有很大影響，如果配置正確，可能會導致2倍以上的加速；
- 適當跳過同步將顯著減少分攤的通信開銷，而不會顯著降低收斂速度。
0x03 背景
3.1 PyTorch
PyTorch將值組織成張量，張量是具有豐富數據操作集的通用n維數組。模塊定義了從輸入值到輸出值的轉換，其正向傳遞期間的行為由其 forward 成員函數指定。模塊可以包含張量作為參數。例如，線性模塊包含權重參數和偏差參數，其正向函數通過將輸入乘以權重並添加偏差來生成輸出。
應用程序通過將本機模塊（如線性、卷積等）和自定義forward函數中的Function（如relu、pool等）粘合在一起，構成自己的模塊。典型的訓練迭代包括使用輸入和標簽生成損失的前向傳遞，計算參數梯度的后向傳遞，以及使用梯度更新參數的優化器步驟。更具體地說，在向前傳播過程中，PyTorch構建了一個autograd圖來記錄所執行的動作。然后，在后向過程中，使用autograd圖進行反向傳播以生成梯度。最后，優化器應用梯度來更新參數。訓練過程重復這三個步驟，直到模型收斂。
3.2 數據並行
PyTorch 提供了多種工具來促進分布式訓練，包括：
-
DataParallel，用於在同一台機器上使用多個GPU的單進程多線程進行數據並行訓練。
-
DistributedDataParallel，用於跨GPU和機器的多進程數據並行訓練。
-
RPC，用於一般分布式模型並行訓練（例如，參數服務器）。
論文的其余部分主要關注分布式數據並行。數據並行通過在操作優化步驟之前進行梯度通信來實現分布式訓練，這樣可以確保使用完全相同的梯度集來更新所有模型副本的參數，因此模型副本可以在迭代中保持一致。
參數平均是擴展模型訓練的另一種流行技術。類似地，它可以跨多台機器啟動多個過程，但不是同步梯度，而是直接計算所有模型參數的平均值。這發生在本地優化器步驟之后，這意味着參數平均可以完全作為一個輔助步驟實現，完全不需要與本地訓練步驟交互，這很有吸引力，因為它可以輕松、干凈地解耦分布式訓練和本地迭代的代碼。但是參數平均有幾個注意事項。
-
與局部訓練相比，參數平均可產生截然不同的結果，這有時會對模型精度造成不利影響。根本原因是，參數平均在數學上並不等同於本地處理所有輸入數據，尤其是當優化器依賴於過去的本地梯度值（如動量）時。由於不同的模型副本可能會看到不同的梯度，因此optimizers中的狀態可能會逐漸發散，從而導致梯度下降方向沖突。當從局部優化模型切換到大規模部署模型時，這可能會導致性能上莫名其妙的差異。
-
參數平均的結構將計算（即反向傳遞）和通信（即計算平均值）協調到非重疊階段，使用optimizer step() 函數作為硬分離點。無論我們如何大力優化計算或通信，一種類型的資源在任何給定時間都將處於空閑狀態，從而放棄大量性能優化機會。
鑒於上述基本缺陷，我們決定使用數據並行性來同步梯度而不是參數來實施分布式訓練。請注意，應用程序仍然可以使用PyTorch輕松構建參數平均值。事實上，后文中描述的集合通信特性是該用例的合適解決方案。應用程序只需要顯式地啟動AllReduce操作來相應地計算平均參數。
3.3 AllReduce
AllReduce是一個基礎通信API，其被 DistributedDataParallel 用於計算所有進程的梯度求和。
多個通信庫都提供了AllReduce ，包括NCCL、Gloo和MPI。AllReduce操作要求每個參與進程都提供一個大小相等的張量，然后將給定的算術運算（如sum、prod、min、max）應用於所有進程的輸入張量，並向每個參與者返回相同的結果張量。
一個 AllReduce 簡單的實現可以簡單地讓每個進程向所有對等進程廣播其輸入張量，然后獨立地應用算術運算。然而，由於AllReduce對分布式訓練速度有顯著影響，通信庫實現了更復雜、更高效的算法，如基於環的AllReduce和基於樹的AllReduce。由於一個AllReduce操作在所有進程加入之前無法啟動，因此它被認為是一種同步通信，而不是參數服務器中使用的P2P通信。
0x04 系統設計
PyTorch 提供了分布式數據並行（DDP）模塊，這有助於輕松地跨多個進程和機器來進行並行化訓練。在分布式培訓期間，每個流程都有自己的本地模型副本和本地優化器。就正確性而言，分布式數據並行訓練和本地訓練必須在數學上等價。DDP可以通過如下來確保訓練正確性：
-
所有模型副本從完全相同的模型狀態開始，並在每次向后傳播之后，得到相同的參數梯度。
-
因此，即使來自不同流程的優化器都是獨立的，它們也應該能夠在每次迭代結束時將其本地模型副本置於相同的狀態
下圖示出了DDP的構建塊，它包含Python API前端、C++梯度歸並核心算法，並使用 c10d 集合通信庫。以下部分按此堆棧圖的自上而下順序顯示。
第4.1節介紹了推動DDP API設計的一般原則。第4.2節介紹Pyrotch分布式數據並行包中使用的擴展梯度歸並技術。最后，第4.3節討論了DDP的集合通信后端選項。
4.1 API
在設計API時，我們定義了兩個設計目標來實現必要的功能。
- 非侵入性：API必須對應用程序是非侵入的。應用程序開發人員通常從編寫本地培訓腳本開始，並在單個計算機上達到資源限制時擴展。在這一點上，要求開發人員重寫整個應用程序以支持分布式數據並行訓練是不可接受的。相反，開發人員應該能夠通過最少的修改來重用本地訓練腳本。
- 攔截：API需要允許實現攔截各種信號以便及時觸發適當的算法。分布式數據並行旨在通過使用更多的計算資源來加速訓練。這一過程需要在計算和通信方面進行微妙的優化，以實現最佳性能。因此，API必須對內部實現提供盡可能多的優化機會。
鑒於上述要求，我們將分布式數據並行實現為一個nn 模塊，該模塊將本地模型作為構造函數參數，並透明地同步反向過程中的數據。下面的代碼片段顯示了使用DDP模塊的示例。
- 本例使用nn.Linear層在第10行創建局部模型。
- 然后，它在第11行將本地模型轉換為分布式訓練模型，並在第12行設置優化器。
- 第14行到第23行是典型的前向傳播、后向傳播和優化器步驟實現。
在這個玩具分布式培訓示例中，第11行是將本地培訓應用程序轉換為分布式應用程序的唯一區別，它滿足了非侵入性需求，還滿足交互要求。構造器允許DDP檢查模型結構和參數。構造完成后，本地模型將被分布式模型替換，然后分布式模型可以很容易地攔截forward()調用以執行相應的必要操作。對於向后傳播，DDP依靠向后鈎子觸發梯度規約，即，在損失張量上執行backward()時，autograd引擎將執行梯度規約。
4.2 梯度規約
DDP中的梯度規約算法在過去的版本中有所發展。為了介紹當前實現的結構，讓我們從一個簡單的解決方案開始，逐步引入更多的復雜性，並在PyTorch v1.5.0中使用當前版本。這還將解釋第 4.1節中描述的相同簡單API如何允許我們安裝各種性能優化算法。
4.2.1 A Naive Solution
如第4節開頭所述，DDP通過讓所有訓練過程（1）從相同的模型狀態開始，以及（2）在每次迭代中使用相同的梯度，來保證正確性。前者可以通過在DDP構建時將模型狀態從一個進程廣播到所有其他進程來實現。為了實現后者，一個簡單的解決方案是：可以在本地向后傳播之后和更新本地參數之前插入梯度同步階段。但是，第4.1節中顯示的API沒有為此階段提供明確的入口點，因為backward()和step()之間沒有任何內容。幸運的是，PyTorch autograd引擎接受定制的后向hook。DDP可以注冊autograd鈎子，以在每次向后傳播后觸發計算。鈎子函數被激發時，每個鈎子掃描所有局部模型參數，並從每個參數檢索梯度張量。然后，它使用AllReduce 集合通信操作來計算所有進程中每個參數的平均梯度，並將結果寫回梯度張量。
Naive Solution 工作正常，但存在兩個性能問題：
-
集合通信在小張量上表現不佳，這在具有大量小參數的大型模型上尤為突出。
-
將梯度計算和同步分離會因為兩者之間的硬邊界而喪失計算與通信重疊的機會。
4.2.2 Gradient Bucketing
梯度bucketing的思想是基於這樣一個觀察，即集合通信在大張量上更有效。下圖（a）和（b）提供了定量視圖，顯示了AllReduce 60M torch.float32的總執行時間。每個AllReduce的參數數量不同。
為了最大限度地提高帶寬利用率，所有reduce操作都是異步啟動的，並同時阻塞等待所有操作，以便模仿DDP的梯度歸並算法。實驗在一台支持NVLink[3]的服務器上進行，該服務器帶有兩個NVIDIA Quadro GP100 GPU。NCCL AllReduce直接在CUDA輸入張量上運行，而Gloo AllReduce則在CPU輸入張量上運行，以便消除在使用Gloo后端時將CUDA內存復制到CPU內存的開銷。對於NCCL和Gloo，當使用較大的輸入張量時，總通信時間明顯減少。Gloo在每個輸入張量約500K參數時達到最高速度，而NVLink上的NCCL甚至沒有20M參數GPU張量的明顯飽和信號。
這些實驗表明，如果DDP在短時間內等待並將多個梯度存儲到一個AllReduce操作中，它可以實現更高的吞吐量和更低的延遲，而不是在每個梯度存儲可用時立即啟動專用的AllReduce。這對於具有許多小參數的模型尤其有用。但是，DDP不應在一個AllReduce中傳輸所有數據，否則，在計算結束之前無法啟動任何通信。上圖（c）和（d）顯示了包含大約60M參數的ResNet152 的GPU和CPU反向計算時間。X軸是准備好的梯度數量，Y軸是自向后傳播開始以來經過的時間。GPU上的后向傳播大約需要250毫秒才能完成，這與NVLink上的NCCL的數量級相同。這一結論也適用於Gloo和CPU后向傳播。這些測量預示着，對於相對較小的bucket大小，DDP可以在向后傳播的同時啟動AllReduce操作，以使通信與計算重疊，這將改變每次迭代的延遲。
4.2.3 Overlap Computation with Communication
在梯度上的AllReduce操作可以在本地向后傳播完成之前開始。使用bucketing，DDP需要等待同一個bucket中的所有內容，然后開始啟動通信。
在這種設置下，只是在向后傳播結束時觸發AllReduce不再足夠。我們需要對更頻繁的信號作出反應，並更迅速地啟動 AllReduce。因此，DDP為每個梯度累加器注冊一個autograd hook。Hook 在其相應的累加器更新梯度之后被觸發，並將檢查其所屬的bucket。如果相同桶中所有梯度的鈎子都已觸發，則最后一個鈎子將觸發該桶上的異步AllReduce。
有兩點需要注意。
- 首先，所有進程的歸並順序必須相同，否則，AllReduce內容可能不匹配，導致不正確的歸並結果或程序崩潰。然而，PyTorch在每次向前傳播時都會動態地構建autograd圖，不同的過程可能在梯度就緒順序上不一致。下圖（a）給出了一個示例，其中兩個垂直軸表示時間，虛線表示梯度准備就緒的時間。在過程1中，四個梯度按順序計算，但梯度g2在過程2的g3和g4之后計算。在這種情況下，如果所有進程都在准備就緒后立即AllReduce bucket，則AllReduce內容將不匹配。因此，所有流程必須使用相同的bucketing順序，並且沒有流程可以在裝載bucket i之前 就在bucket i+1上啟動AllReduce。如果bucket 0是最后一個准備就緒的bucket，那么通信就不可能與計算重疊【筆者：因為 bucket 0 是最后就緒的，所以其他bucket在這之前都不會被執行計算，就不能與通信重疊了】。PyTorch v1.5.0通過使用model.parameters()的相反順序作為bucketing順序來解決此問題，我們做如下假設：層（layers）可能按照正向過程中調用的相同順序進行注冊。因此，其反向順序就是反向過程中的梯度計算順序的近似表示。誠然，這並不是一個完美的解決方案，但它是一個近似方案，我們可以用最少的工程開銷來實現它。
- 其次，一次訓練迭代可能只涉及模型中的一個子圖，並且子圖在每次迭代中可能不同，這意味着在某些迭代中可能會跳過某些Gradient。然而，由於 gradient-to-bucket 的映射是在構建時確定的，這些缺少的梯度將使一些bucket 永遠看不到最終的自動裝載hook，並且無法將bucket標記為就緒。因此，向后傳播可能會暫停。下圖（b）示出了一個示例，其中在一次迭代中跳過了與梯度g3相對應的參數，導致g3缺少就緒信號。為了解決這個問題，DDP從前向傳播的輸出張量遍歷autograd圖，以找到所有參與的參數。這些參與張量的准備就緒是結束向后傳播完成的有效信號。因此，DDP可以通過在向前傳播結束時主動標記剩余的參數梯度來避免等待。請注意，此更改並不妨礙我們開發非侵入式API，因為應用程序可以直接調用DDP上的forward函數，並且DDP可以輕松地將此步驟插入其成員函數中。
下面算法給出了DDP的偽碼。
Constructor包含兩個主要步驟，廣播模型狀態和安裝autograd掛鈎。DDP的 forwad 函數是本地模型 forwad 函數的簡單包裝器。它遍歷autograd圖以相應地標記未使用的參數。autograd鈎子將內部參數索引作為輸入，這有助於找到參數張量及其所屬范圍。它將局部梯度寫入bucket中的正確偏移量，然后啟動異步AllReduce操作。偽代碼中省略了一個附加的結束步驟，它等待AllReduce操作，並在反向過程結束時將值寫回梯度。
下圖闡明了DDP在向前和向后傳播期間如何與局部模型交互。
上述解決方案適用於大多數用例。但是，由於DDP總是計算所有梯度的平均值，並將它們寫回parameter.grad字段，因此優化器無法區分梯度是否參與了最后一次向后傳播。由於DDP和優化器的解耦設計，DDP沒有旁側通道向優化器暗示該信息。如果沒有這些信息，訓練過程可能會受到模型精度回歸的影響，例如，當優化器使用梯度感知信息跳過動量值更新時。為了解決這個問題，DDP應該只接觸哪些確實涉及向后傳播的梯度。然而，由於在對等DDP過程中，前向/后向過程可能仍然涉及到局部缺失梯度，因此無法僅從局部autograd圖中提取該信息。因此，DDP使用位圖跟蹤本地參數參與者，並啟動另外一個AllReduce來收集全局未使用的參數。不幸的是，由於元素類型可能不匹配，DDP無法將此位圖合並到其他梯度AllReduce操作中。只有當應用程序顯式地告訴DDP查找未使用的參數時，這種額外的開銷才會出現，因此只有在必要時才會支付代價。
4.2.4 Gradient Accumulation
加速分布式數據並行訓練的一種常用技術是降低梯度同步頻率。在全局同步梯度之前，應用程序可以執行n次局部訓練迭代，而不是在每次迭代中啟動AllReduce。如果輸入批次太大而無法裝入設備，這也很有幫助，因為應用程序可以將一個輸入批次拆分為多個微批次，在每個微批次上運行局部向前和向后傳播，並且僅在大批次的邊界處啟動梯度同步。理論上，這應該產生與大批量數據一次性處理相同的結果，因為梯度將簡單地累積到相同的張量。然而，這在某種程度上與第 4.2.3節中討論的梯度歸並算法相沖突。該算法將在每次向前傳遞結束時將未使用的參數標記為就緒，而一次迭代中未使用的參數仍可以參與后續迭代。此外，DDP無法區分應用程序是否應該在向后或通過多次迭代累積梯度后立即調用optimizer.step()。因此，我們需要為這個用例引入一個額外的接口（即，no_sync ）。
在內部，no_sync 的實現非常簡單。上下文管理器只是在進入和退出上下文時切換一個標志，該標志在DDP的forward 功能中使用。在 no_sync 。全局未使用參數的信息也會累積在位圖中，並在下次通信發生時使用。下面是一個示例代碼段。
4.3 Collective Communication
分布式數據並行訓練使用一種特殊的通信模式：每個參與者提供一個相同尺寸的張量，並收集所有參與者的全局和（global sum）。這可以通過如下方式來實現：首先是一個gather操作，然后使用點對點通信對每個參與者進行局部歸並（local reductions），但這將喪失性能優化的機會。
DDP構建在集合通信庫之上，包括三個選項：NCCL、Gloo和MPI。DDPs從三個庫中獲取API，並將它們包裝到同一個ProcessGroup API中。該名稱預示着ProcessGroup希望多個進程作為一個組一起工作。
所有ProcessGroup實例通過使用集合服務（rendezvous service）同時構造，其中第一個實例將進行阻塞，一直等待，直到最后一個實例加入。對於NCCL后端，ProcessGroup為通信維護一組專用的CUDA流，以便通信不會阻止默認流中的計算。由於所有通信都是集合操作，因此所有ProcessGroup實例上的后續操作必須在大小和類型上匹配，並遵循相同的順序。對所有庫使用相同的ProcessGroup API允許我們使用相同的DDP實現來試驗不同的通信算法。例如，PyTorch v1.5提供了一個round-robin ProcessGroup實現，它獲取ProcessGroup實例列表，並以循環方式向這些ProcessGroup實例發送集合通信。通過使用round-robin ProcessGroup，在單個NCCL、Gloo或MPI處理組無法飽和鏈路容量的情況下，DDP可以獲得更高的帶寬利用率。
0x05 實施
在過去的幾個版本中，DDP的實現已經演進了好幾次。本節重點介紹PyTorch v1.5.0的當前狀態。DDP實現同時存在於 Python和C++文件，Python 部分包括公開API和非性能關鍵的組件，C++提供核心梯度歸並算法。Python API 通過Pybind11來調用C++核心。
5.1 Python前端
DDP nn.module在distributed.py中實現，它包含面向用戶的組件。組件包括構造函數、forward 函數和 no_sync 上下文管理器。除了在第4節中強調的一般思想外，Python前端中還有幾個塑造DDP行為的實現細節。
DDP構造器API中公開了Configurable Knobs，包括
-
process_group，用於指定DDP運行AllReduce的流程組實例，這有助於避免和默認流程組混淆；
-
bucket_cap_mb，用於控制AllReduce bucket大小，應用程序應調整此以優化訓練速度，
-
find_unused_parameters，來切換DDP是否應檢測未使用的參數，DDP是通過遍歷autograd圖來完成檢測的。
本地模型中的模型設備關聯性（Model Device Affinity ）也控制DDP的行為，特別是當模型跨越多個設備時，這在模型太大而無法裝入單個設備時很常見。對於大型模型，應用程序可以將模型的不同層放置在不同的設備上，並使用Tensor.to(device) API將中間輸出從一個設備移動到另一個設備。DDP也適用於多設備模型。只要將 device_ids參數設置為None或空列表，DDP就會檢查模型，執行健全性檢查並相應地應用配置。然后，將多設備模型視為一個整體。
當一個層需要跟蹤運行方差和運行平均值（例如BatchNorm）等狀態時，模型緩沖區（Model Buffers）是必要的。DDP通過讓rank 0 的進程獲得支持模型緩沖區的權限。如果模型包含緩沖區，DDP在本地模型上開始前向傳遞之前，將緩沖區值從rank 0進程廣播到所有其他進程。此行為也與no_sync模式兼容。當啟用no_sync模式時，它會在正向過程中正確設置一個標志，以指示它是否期望在下一個反向過程中執行梯度規約。如果通信發生，DDP將在隨后的前向傳遞之前廣播緩沖區。
5.2 Core Gradient Reduction
主要的開發工作花費在gradient reduction上，因為這是DDP中與性能最相關的步驟。該實現存在於reducer.cpp中，它由四個主要組件組成，即：
- 構建參數到桶的映射。
- 安裝autograd hook。
- 啟動bucket AllReduce
- 檢測全局未使用的參數。
我們接下來闡述這四個組成部分。
參數到桶映射（Parameter-to-Bucket Mapping）對DDP速度有相當大的影響。在每次向后傳播中，將所有參數梯度中的張量復制到桶中，並在AllReduce之后將平均梯度復制回桶中。為了加速復制操作，存儲桶始終與參數在同一設備上創建。如果模型跨越多個設備，DDP會考慮設備關聯性，以確保同一存儲桶中的所有參數都位於同一設備上。AllReduce的順序也會對結果產生影響，因為它決定了多少通信可以與計算重疊。DDP按model.parameters()的相反順序啟動AllReduce。
Autograd Hook是DDP在后向傳播中的切入點。在構建過程中，DDP遍歷模型中的所有參數，在每個參數上找到梯度累加器，並為每個梯度累加器安裝相同的post hook函數。梯度累加器將在相應的梯度准備就緒時，會觸發post hooks，DDP將計算出整個桶何時全部就緒，這樣可以啟動AllReduce操作。然而，由於無法保證梯度准備的順序，DDP不能選擇性地選擇安裝掛鈎的參數。在當前的實現中，每個bucket都保留一個掛起的梯度計數。每個post-hook函數都會遞減計數，當計數為零時，DDP會將一個桶標記為就緒。在下一次向前傳播中，DDP會為每個桶補齊待定的累積計數。
Bucket AllReduce是DDP中通信開銷的主要來源。一方面，在同一個桶中裝入更多的梯度將減少通信開銷的攤銷系統。另一方面，由於每個桶需要等待更多的梯度，因此使用較大的桶尺寸將導致更長的歸並等待時間。因此，桶大小是關鍵的權衡。默認情況下，每個存儲桶的大小為25MB。應用程序應該根據經驗測量其影響，並將其設置為其用例的最佳值。
全局未使用參數（Globally Unused Parameters）的梯度在向前和向后過程中應保持不變。檢測未使用的參數需要全局信息，因為在一個DDP過程中，一個參數可能在一次操作中不存在，但可能在另一個過程的同一次迭代中參與訓練。因此DDP在位圖中維護本地未使用的參數信息，並啟動額外的AllReduce以收集全局位圖。由於位圖比張量尺寸小得多，因此模型中的所有參數共享同一位圖，而不是創建每桶位圖（per-bucket bitmaps）。位圖位於CPU上，以避免為每次更新啟動專用CUDA內核。但是，某些ProcessGroup后端可能無法在CPU 張量上運行AllReduce。例如，ProcessGroupNCCL僅支持CUDA張量。此外，由於DDP應該與任何定制的ProcessGroup后端一起工作，它不能假設所有后端都支持CPU張量。為了解決這個問題，DDP在同一設備上維護另一個位圖作為第一個模型參數，並調用非阻塞拷貝操作（non-blocking copy）將CPU位圖移動到設備位圖以進行集合通信。