PyTorch如何加速數據並行訓練？分佈式祕籍大揭祕

本文轉載自機器之心。

選自arXiv

作者：Shen Li等

機器之心編譯

參與：小舟、杜偉

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 的互連方式：