[原始碼解析] PyTorch 分散式(8) -------- DistributedDataParallel之論文篇

羅西的思考發表於2021-11-21
原始碼PyTorch分散式Parallel

[原始碼解析] 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 有部分出入。

本系列其他文章如下：

深度學習利器之自動微分(1)

深度學習利器之自動微分(2)

[原始碼解析]深度學習利器之自動微分(3) --- 示例解讀

[原始碼解析]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通訊庫之間的比較。結果表明：

  1. 通訊是影響訓練延遲的主要因素，其影響隨模型尺寸的增大而增大；
  2. 儲存桶大小對通訊效率有很大影響，如果配置正確，可能會導致2倍以上的加速；
  3. 適當跳過同步將顯著減少分攤的通訊開銷，而不會顯著降低收斂速度。

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點陣圖移動到裝置點陣圖以進行集合通訊。

0xFF 參考

http://www.vldb.org/pvldb/vol13/p3005-li.pdf

相關文章