Pytorch 分散式訓練 DistributedDataParallel — 概念篇

李謦伊
謦伊的閱讀筆記
7 min readMar 5, 2022

分散式訓練就是指將模型放置在很多台機器且每台機器有多個 GPU 上進行訓練,之所以使用分散式訓練的原因有兩種:第一、模型在一塊 GPU 上放不下,第二、使用多塊 GPU 進行平行計算能夠加速訓練。但需要注意的是隨著使用的 GPU 數量增加,各個設備之間的通訊會越複雜,導致訓練速度下降。

分散式訓練主要分為兩種類型:資料平行化 (Data Parallel)、模型平行化 (Model Parallel)。

資料平行化 (Data Parallel)

當數據量非常大,並且模型架構能夠放置在單個 GPU 上時,就可以採用資料平行化的方式來進行分工合作。

做法是依照一些規則將數據分配到不同的 GPU 上,並且每個 GPU 都有相同的模型架構,也就是會在每個 GPU 上複製一份相同的模型,各自進行訓練後,將計算結果合併,再進行參數更新。

source

參數更新的方式又分為同步及異步:

  • 同步 (synchronous): 所有的 GPU 在訓練時會等待其他 GPU 計算完畢後,才會進行一次參數更新,因此訓練速度上會比使用異步的方式來得慢。但因為在更新參數時會合併其他計算結果,相當於增加了 batch size 的大小,對於訓練結果有一定的提升。
  • 非同步、異步 (asynchronous): 每個 GPU 各自進行訓練和參數更新,不須等待其他 GPU 計算完畢。能夠提升訓練速度,但可能會產生 Slow and Stale Gradients (梯度失效、梯度過期) 問題,收斂過程較不穩定,訓練結果會比使用同步的方式差。
source

模型平行化 (Model Parallel)

當模型架構太大以至於在一個 GPU 放不下時,可以採用模型平行化的方式來將模型拆解並分配到不同的 GPU 上。

由於模型層與層之間通常有依賴性,也就是指在進行前向傳播、反向傳播時,前面及後面的層會作為彼此的輸入和輸出,在訓練速度上會有一定的限制,因此若非使用較大的模型較不建議採用模型平行化。若想提升訓練速度,可以選擇較容易進行平行運算的 module,ex: Inception。

source

資料平行化 + 模型平行化

當然也可以採用結合資料平行化與模型平行化的方式,如下圖所示,做法是將模型拆解並分配到單個機器中不同的 GPU 上,並且在其他機器中的每個 GPU 都複製一份相同的模型架構 (拆解後的部分模型),而資料數據也會依照一些規則分配到不同的 GPU 上。

source

分散式訓練系統架構

Parameter Server (PS) vs Ring-All-Reduce

這些平行化運算的方式是如何進行加速的呢?各個設備又是如何將計算結果相互交換?分散式訓練之所以能達到這兩件事是藉由系統架構的設計來實現,目前的系統架構主要分成 Parameter Server (PS)、Ring-All-Reduce。

Parameter Server (PS)

Parameter Server 架構分為 sever group 和多個 worker group 兩大部分以及 resource manager,其中 resource manager 負責整體的資源分配。

sever group 用於初始化及維護所有的模型參數,由 sever manager 和多個 sever node (master node) 組成,其中 sever manager 負責分配資源並維護各個 sever node 的狀態;sever node 則是負責維護被分配的一部分參數。

worker group 用於維護從 training data 分配到的數據及對應的 sever node 所分配到的部分參數,每個 worker group 由 task scheduler 和多個 worker node 組成,其中 task scheduler 負責分配各個 worker node 的任務並進行監控;worker node 則是負責進行訓練和更新各自分配到的參數。此外,worker node 只與對應的 sever node 進行通訊,不論是 worker group 和 worker group、worker group 內部的每個 worker node 之間都不會互相交流。

source

整體流程如下圖所示,每個 worker group 會各自進行訓練、計算梯度、更新分配到的部分參數,再將結果傳送給 server node。接著 server node 把收到的參數彙總並進行全局更新 (ex: 平均參數),再傳回給 worker node。

由於 server node 需要跟所有對應的 worker node 進行溝通,容易導致 sever node 的計算過重以及網路阻塞的狀況,從而使得訓練速度下降。

source

Ring-All-Reduce

Ring-All-Reduce 架構是將所有 GPU (假設為 N 個) 的通訊方式呈現環狀的形式,並且每個 GPU 中的數據會分為 N 份,GPU 之間會根據環狀的方式來傳送數據。這樣的架構所產生的通訊量與設備的多寡無關,因此能夠大幅地減少通訊的開銷、更好地平衡設備間通訊的使用程度。另外,在訓練過程中,可以利用 Backpropagation 的特性 (後面層的梯度會先於前面層計算),在計算前面層梯度的同時,進行後面層梯度的傳送,以提升訓練效率。

整體流程如下圖所示,首先第 k 個 GPU 把第 k 份的數據傳給下一個 GPU,接著下一個 GPU 會將收到的第 k 份數據與自己的第 k 份數據進行加總整合後,再傳送給下一個 GPU。經過 N 次的迭代循環後,第 k 個 GPU 會收到所有第 k 份的最終整合數據,再以這份數據來更新各自的參數。

source

Pytorch 分散式訓練方法

DataParallel (DP) vs DistributedDataParallel (DDP)

Pytorch 提供了兩種分散式訓練方法 — DataParallel (DP)、 DistributedDataParallel (DDP)。

  • DP 採用 Parameter Server (PS) 架構,並且僅支援單台機器多 GPU,任務中只會有一個進程
  • DDP 則是採用 Ring-All-Reduce 架構,可支援多台機器多 GPU,能夠使用多進程。

DP 的 code 寫法較為簡單,但由於 DDP 的執行速度更快,官方更推薦使用 DDP。

下一篇會來進行 Pytorch 的 DistributedDataParallel 實作~

--

--