Distributed Training/zh
| Article | |
|---|---|
| Topic area | Deep Learning |
| Prerequisites | Stochastic Gradient Descent, Backpropagation |
概述
分布式訓練是將單個神經網絡訓練作業拆分到多個處理器、加速器或機器上,使工作負載並行運行的實踐。它解決了單設備訓練的兩個限制:模型可能太大,無法裝入單個設備的內存;以及在大數據集上收斂所需的實際時間可能令人無法接受。具有數千億參數的現代基礎模型如果沒有它根本無法訓練,即使是較小的模型也通常使用它來將訓練時間從數周縮短到數小時。該領域結合了數值優化、高性能計算和系統工程的思想:如何在工作節點之間協調隨機優化,如何在設備之間高效地移動張量,以及如何調度計算以使加速器保持忙碌。
兩個基本的軸是數據並行,其中每個工作節點持有模型的副本並處理批次的不同分片;以及模型並行,其中模型本身被拆分到工作節點之間。現實世界的系統將兩者與流水線並行和優化器狀態分片相結合,從而擴展到數千個加速器,同時保持較高的算術強度和可承受的通信開銷。
動機
三種壓力將訓練推離單一設備。首先是參數數量:具有 $ N $ 個參數的 16 位 精度模型至少需要 $ 2N $ 字節用於權重,梯度和優化器狀態也需要類似的量,因此一個 700 億參數的模型在混合精度下,一旦加入 Adam 矩,就輕易超過 1 TB 內存。其次是數據量:在單設備吞吐量下用數萬億個 token 訓練一個大型語言模型需要數年時間。第三是樣本效率:更大的有效批次大小通過數據並行變得 made 可行,可以穩定梯度並改善泛化,直到臨界批次大小,超過該值後回報會遞減。
然而,擴展並非免費。工作節點之間的通信與計算相互競爭,同步屏障可能使設備閒置。分布式訓練的核心工程問題是如何在保留原始算法優化動態的同時,重疊、壓縮或消除該通信。
數據並行
在同步數據並行訓練中,$ K $ 個工作節點中的每一個都持有參數 $ \theta $ 的相同副本,並處理大小為 $ B $ 的不同微批次。每個工作節點計算局部梯度 $ g_k = \nabla_\theta L_k(\theta) $,然後工作節點通過 all-reduce 操作交換梯度,使每個工作節點都持有平均值 $ \bar{g} = \frac{1}{K}\sum_{k=1}^{K} g_k $,並且每個工作節點在本地應用相同的優化器更新:
$ {\displaystyle \theta \leftarrow \theta - \eta \, \bar{g}.} $
由於所有副本以相同狀態開始,並對相同的平均梯度應用相同的更新,它們保持逐位一致(除歸約中的非確定性外)。有效批次大小為 $ K B $,這通常需要學習率預熱和線性縮放規則來保持穩定性。
all-reduce 原語可用環形或樹形算法實現;ring all-reduce每個工作節點傳輸梯度大小的 $ 2(K-1)/K $ 倍,在均勻互聯上是帶寬最優的。NCCL、Gloo 和 MPI 都提供了經過調優的實現。
一種較舊的替代方案是參數服務器,即由一個中心進程持有規範的參數並聚合工作節點的更新。異步變體允許工作節點無需等待即可推送陳舊的梯度,以優化質量換取吞吐量;這種方法在密集深度學習中大多已成為歷史,但在某些推薦器和嵌入工作負載中仍然存在,因為稀疏性和負載不平衡有利於這種方法。
模型並行
當單層的參數或激活超過設備內存時,該層本身必須被分區。在張量並行中,矩陣乘法 $ Y = X W $ 沿隱藏維度被拆分:每個工作節點持有列條 $ W_k $,計算 $ Y_k = X W_k $,部分輸出根據層類型進行拼接或求和。Megatron-LM 推廣了這種方案用於transformer 塊,通過對注意力和 MLP 投影進行分片,使得每個塊在前向和反向傳遞中只需要兩次all-reduce。
張量並行在一個節點內擴展良好,在該節點內 NVLink 等高帶寬鏈路連接着一小組 GPU(通常為 4 到 8 個)。在跨節點的情況下,每步通信成本快速增長,通常與其他形式的並行結合使用,而不是單獨使用。
流水線並行
流水線並行按深度將網絡在工作節點之間分區:工作節點 $ k $ 持有連續的層塊,並將激活轉發到工作節點 $ k+1 $。簡單地一次運行一個微批次會使大多數工作節點空閒(稱為"流水線氣泡")。GPipe 通過將小批次拆分為 $ M $ 個微批次並重疊它們的前向傳遞來解決這個問題;對於 $ P $ 個階段,氣泡比例大致為 $ (P-1)/(M+P-1) $。PipeDream-Flush 和 1F1B 調度更激進地交錯前向和反向傳遞,以減少激活內存。
流水線並行對負載均衡敏感:單個慢速階段會停滯整條流水線。從業者會分析每個階段的耗時並調整層的分配,有時會在各階段之間轉移工作或複製開銷小的層。
混合併行與 3D 並行
在集群規模下,任何單一軸都不夠用。常見的安排——有時稱為 3D 並行——組合了數據並行、張量並行和流水線並行。張量並行處理節點內最大的層,流水線並行將模型在節點之間分區,而數據並行在副本組之間複製整條流水線以獲得吞吐量。序列並行沿序列維度對激活進行分片,是用於超長上下文的第四個軸。
選擇分區是一個離散優化問題,由內存預算、互聯拓撲以及模型的注意力和前饋形狀決定。Alpa、GSPMD 以及 DeepSpeed 和 Megatron 中的分區例程會自動搜索這一空間。
ZeRO 與分片優化器狀態
Zero Redundancy Optimizer(ZeRO)通過對每個副本存儲的內容進行分片來降低數據並行的內存開銷。ZeRO-1 將優化器狀態在數據並行組之間分區;ZeRO-2 也對梯度進行分區;ZeRO-3(也稱為 Fully Sharded Data Parallel,FSDP)則對參數本身進行分區,在每一層按需聚合。由於 gather 和 reduce-scatter 原語可以與反向計算組合,通信可以在很大程度上與算術運算重疊,從而以數據並行更簡單的編程模型實現接近模型並行的內存占用。
結合混合精度存儲和 CPU 或 NVMe 卸載,ZeRO 已使得在商品級規模的 GPU 集群上訓練萬億參數的模型成為可能。
通信與同步
分布式訓練中占主導地位的開銷是移動張量。有幾種技術可以減少或隱藏這一開銷:
- 梯度壓縮,包括 1-bit SGD、signSGD 和 PowerSGD,以引入必須加以控制的偏差或方差為代價,減少了每次 all-reduce 的字節數。
- Local SGD 在每個工作節點上執行 $ H $ 步後再求平均,以略微的優化漂移為代價降低了通信頻率。
- 計算與通信的重疊利用了這樣一個事實:早期層的梯度在後續層仍在計算時就已經就緒;PyTorch DDP 和 Horovod 等框架以桶為單位調度異步 all-reduce,使得網絡在反向傳遞期間始終在傳輸數據。
- 拓撲感知的集合通信優先使用節點內高帶寬鏈路進行早期歸約,只在最後才跨越慢速鏈路。
實踐考量與局限
分布式訓練會放大每一個小故障。在數千個加速器規模下,落後節點、鏈路錯誤和靜默數據損壞變得司空見慣,因此檢查點頻率、確定性重啟和彈性重分片都需要大量工程關注。諸如損失尖峰不穩定之類的數值問題在大批次大小下更為常見,可能需要梯度裁剪、更低的學習率,或從最近的檢查點恢復。
並行策略的選擇很少具有可移植性:為某種集群拓撲調優的分區在另一種拓撲上可能遠非最優,改變模型形狀可能需要重新分片優化器狀態。通信庫、內核啟動開銷和每設備的內存碎片都成為頭等關注的問題。儘管存在這些工程負擔,分布式訓練仍是通向當今前沿規模模型的唯一已知途徑,集合通信算法、分片策略和加速器互聯方面的改進將繼續拓寬其實際可行的邊界。