PyTorch如何加速數據并行訓練?分佈式秘籍大揭秘

PyTorch 在學術圈裡已經成為最為流行的深度學習框架,如何在使用 PyTorch 時實現高效的并行化?

在芯片性能提升有限的今天,分佈式訓練成為了應對超大規模數據集和模型的主要方法。本文將向你介紹流行深度學習框架 PyTorch 最新版本( v1.5)的分佈式數據并行包的設計、實現和評估。

論文地址:https://arxiv.org/pdf/2006.15704.pdf

PyTorch 是深度學習研究和應用中廣泛使用的科學計算包。深度學習的最新進展證明了大型數據集和大型模型的價值,這需要擴展模型訓練更多計算資源的能力。

同時,由於簡明的原理和廣泛的適用性,數據并行已經成為了分佈式訓練的一種熱門方案。通常,分佈式數據并行技術會在每個計算資源上複製模型以獨立生成梯度,然後在每次迭代時傳遞這些梯度以保持模型副本的一致性。儘管該技術在概念上很簡單,但計算和通信之間的細微依賴關係使優化分佈式訓練的效率變得不簡單。

因此,在這篇論文中,來自 Facebook AI 和華沙大學的研究者介紹了 PyTorch 分佈式數據并行模型的設計、實現以及評估

從 v1.5 開始,PyTorch 自身提供了幾種加速分佈數據并行的技術,包括分桶梯度(bucketing gradients)、通信重疊計算(overlapping computation with communication)以及跳過梯度同步(skipping gradient synchronization)。相關評估結果显示,在配置正確的情況下,PyTorch 分佈式數據并行模型可以用 256 個 GPU 達到接近線性的可擴展性。

接下來,我們來看 PyTorch 分佈式數據并行訓練的模型設計、具體實現和效果評估。

系統設計

PyTorch 提供了一個數據分佈式并行(DistributedDataParalle, DDP)模型來幫助實現在多個進程和機器的并行訓練。在分佈訓練期間,每個模型都有自己的本地模型副本和本地優化器。就糾錯而言,分佈式數據并行訓練和本地訓練在數學上必須是等價的。

下圖 1 描述了 DDP 構造塊的組成,其中包含一個 Python API 前端和 C++ 梯度下降核心算法,並採用了 c10d 聚合通信庫。

Python API 前端

在設計 API 時,研究者制定了以下兩個設計目標來達到必要的功能:

非侵入式:對應用提供的 API 必須是非侵入式的;

攔截式:API 需要允許攔截各種信號並立即觸發適當的算法。

分佈式數據并行化旨在使用更多的計算資源來加速訓練。

根據以上需求,研究者用 nn.Module 實現了分佈式數據并行。nn.Module 採用本地模型作為構造函數的參數,並在反向傳播中透明地同步梯度。下面的代碼是使用 DDP 模型的示例:

梯度下降

研究者闡述了在 PyTorch 上進行分佈式數據并行訓練的幾種梯度降低技術。DDP 中的梯度下降算法已經有了新的改進。為了介紹當前實現的結構,研究者從一個簡單的初始方案(naive solution)開始,逐步介紹更多複雜的版本,最終在 PyTorch v1.5.0 上使用當前版本。

初始方案

DDP 首先校正了所有的訓練進程,以保證各個進程:

從相同的模型狀態開始;

每次迭代花費同樣多的梯度。

為了完成第二點,初始方案在進行本地反向傳播之後、更新本地參數之前插入了一個梯度同步環節。幸運的是,PyTorch 的 autograd 引擎能夠接受定製的 backward 鈎子(hook)。DDP 可以註冊 autograd 鈎子來觸發每次反向傳播之後的計算。然後,它會使用 AllReduce 聚合通信來號召計算所有進程中每個參數的平均梯度,並且把結果寫回梯度 tensor。

初始方案足以完成想要的目標,但存在兩項性能缺陷。聚合通信在小型 tensor 上性能表現很差,這種缺陷在帶有大量小參數的大型模型上尤為突出。由於兩者之間存在界限,分別進行梯度計算和同步化會造成通信重疊計算機會的缺失。

梯度分桶(bucketing )

梯度分桶的觀點是受聚合通信在大型 tensor 上更加高效的啟發而提出的。

下圖 2(a)和 (b) 給出的定量視圖展示了在每個 AllReduce 中參數數目不同的情況下,AllReduce 60M torch 的 float32 參數的完整執行時間:

這些實驗表明,不用等到每個梯度 tensor 都可用時再啟動 AllReduce,DDP 在等待較短的時間並將多個梯度存儲到一個 AllReduce 操作中時,就可以實現更高的吞吐量和更短的延遲。

通信重疊計算

在使用分桶的情況下,DDP 只需在啟動通信之前在同一個 bucket 中等待所有的內容。在這樣的設置下,在反向傳播的最後觸發 AllReduce 就顯得不足了。因此需要對更加頻繁的信號做出相應,並且更加迅速地啟動 AllReduce。因此,DDP 為每個梯度累加器都註冊了 autograd 鈎子。

下圖 3(a)的示例中,兩個豎直軸表示時間,虛線代表梯度準備就緒的時間。進程 1 中,4 個梯度按順序計算。進程 2 中,g_2 在 g_3 和 g_4 之後計算;圖 3(b)的示例中,梯度 g_3 對應的參數在一次迭代中被跳過了,導致 g_3 的就緒信號缺失。

為了解決這個問題,DDP 遍歷了前向傳播的輸出 tensor 中的 autograd 圖以找到涉及到的所有參數。涉及到 tensor 的就緒狀態足以充當反向傳播完成的信號。

以下算法 1 給出了 DDP 的偽代碼:

下圖 4 展示了 DDP 在前向傳播和反向傳播過程中如何與本地模型交互:

梯度累加

此外,DDP 無法分辨應用程序是計劃在反向傳播之後立即調用 optimizer.step()還是通過多次迭代累加梯度。因此,研究者需要為這個用例再引入一個接口(即 no sync)。以下是樣例代碼片段:

聚合通信

DDP 是在集合通信庫基礎上建立的,包括 3 個選項 NCCL、Gloo 和 MPI。DDP 採用了來自這三個庫的 API,並將它們封裝進同一個 ProcessGroup API 中。

由於所有的通信都是聚合操作,因此所有的 ProcessGroup 實例上的後續操作必須和其類型匹配並遵循相同的順序。對所有的庫使用同一個 ProcessGroup API 允許研究者在相同的 DDP 實現上試驗不同的通信算法。

如果單一 NCCL、Gloo 或 MPI 的 ProcessGroup 無法使鏈路容量達到飽和,通過使用循環的 ProcessGroups,DDP 可以獲得更高的帶寬利用率。

具體實現

DDP 的實現在之前的幾個版本中已經改進了多次。研究者介紹了當前 PyTorch v1.5.0 的狀態。DDP 同時在 Python 和 C++ 上都可以實現,Python 開放了 API 並組成了非性能關鍵因素組件,而 C++ 提供了核心梯度下降算法。Python API 通過 Pybind11 的 API 調用了 C++ 內核。

Python 前端

Python 前端中的實現細節決定了 DDP 的行為。可配置的 Knobs 在 DDP 構造函數 API 中開放。具體包括:

分組處理以找出 DDP 中運行 AllReduce 的進程組實例,它能夠幫助避免與默認進程組混淆;

bucket_cap_mb 控制 AllReduce 的 bucket 大小,其中的應用應調整 knob 來優化訓練速度;

找出沒有用到的參數以驗證 DDP 是否應該通過遍歷 autograd 圖來檢測未用到的參數。

本地模型中的 Model Device Affinity 也能控制 DDP 的行為,尤其是當模型因為太大而需要跨越多個設備運行時,更是如此。對於大型模型,模型的每一層可以放在不同的設備上,使用 Tensor.to(device) API 可以將中間輸出從一個設備轉移到另一個上。DDP 也可以在多個模型上運行。

當層(例如 BatchNorm)需要跟蹤狀態,例如運行方差和均值時,模型緩衝器(buffer)是非常必要的。DDP 通過讓 rank 為 0 的進程獲得授權來支持模型緩衝器。

核心梯度下降

開發過程中的主要工作就是梯度降低,它也是 DDP 中決定性能的關鍵步驟。這個在 reducer.cpp 中的實現有 4 個主要的組成部分:構建 parameter-to-bucket map、安裝 autograd 鈎子,啟動 bucket AllReduce 以及檢測全局未用過的參數。

Parameter-to-Bucket Mapping 已經對 DDP 的速度有了相當大的影響。在每次反向傳播中,tensor 從全部的參數梯度到 bucket 被複制,平均梯度在 AllReduce 之後又被複制回 tensor。

Autograd Hook 是 DDP 反向傳播的進入點。在構造期間,DDP 遍歷模型中的所有參數,找出每個參數的梯度累加器,並且為每個梯度累加器安裝相同的 post-hook 函數。當相應的梯度準備就緒時,梯度累加器會啟用 post hook,並且當整個 bucket 準備好啟動 AllReduce 操作時,DDP 會確定啟用。

Bucket Allreduce 是 DDP 中通信開銷的主要來源。默認情況下,bucket 的大小是 25MB。

實驗評估

研究者展示了使用專屬 32-GPU 集群和共享權限時 PyTorch DDP 的評估結果,其中 GPU 部署在 4 台服務器,並通過邁絡思 MT27700 ConnectX-4 100GB/s 的網卡連接。每台服務器配有 8 個英偉達 Tesla V100 GPU。

下圖 5 展示了一台服務器上 8 個 GPU 的互連方式:

接下來,研究者利用 ResNet50 和 BERT 這兩個流行的模型度量了 PyTorch DDP 在每次迭代時的延遲和可擴展性,並且大多數實驗使用隨機生成的合成輸入和標籤,這對於比較每次迭代時的延遲來說足夠了。

延遲故障

為了驗證通信重疊計算的有效性,下圖 6 展示了 ResNet50 和 BERT 模型分別使用 NCCL 和 Gloo 反向傳遞時的延遲故障。所有實現都用到了 4 台服務器上的 32 個 GPU。

結果显示,在 PyTorch DDP 訓練時,反向傳遞是耗時最長的步驟,這是因為 AllReduce 通信(即是梯度同步)在這一過程中完成。

Bucket 大小

bucket 大小是一個重要的配置選項。根據經驗,出於最大努力估計,bucket_cap_mb 的默認值是 25MB。研究者使用兩台機器上的 16 個 GPU 比較不同 bucket 大小下每次迭代的延遲。另一個極端是在短時間內傳遞全部的梯度,結果如下圖 7 所示。

下圖 8 給出了相同設置下、32 個 GPU 上的實驗結果。在這種情況下,離群值(outlier)的跨度更大,這並不意外。因為在有更多參与者的情況下,同步必然要花費更長的時間,並且 strangler 的影響更明顯。

可擴展性

為了理解 DDP 的可擴展性,研究者用多達 256 個 GPU 上的 NCCL 和 Gloo 後端來度量 ResNet50 和 BERT 每次迭代的訓練延遲。結果如下圖 9 所示。

下圖 10 給出了每 1、2、4 和 8 次迭代進行梯度下降時每次迭代的平均延遲。

除了每次迭代延遲,測量收斂速度以驗證加速度是否會因收斂放緩而被消除也非常關鍵。實驗採用 MNIST 數據集來訓練 ResNet。學習率設置為 0.02,批處理大小是 8。結果如下圖 11(a)所示;圖 11(b)是將批處理大小設為 256,學習率設為 0.06 的測量結果。

循環分配(Round-Robin)進程組

PyTorch 分佈式包支持將 Round-Robin 進程組和多個 NCCL 或 Gloo 進程組組合在一起,從而按照 Robin-Robin 順序向各個進程組實例分配聚合通信。

下圖 12 展示了使用 1、3 和 5 個 NCCL 或 Gloo 進程組的 Round-Robin 進程組每次迭代的延遲。最顯著的加速是使用 NCCL 後端的 BERT 模型。