随着人工智能與深度學習的發展，大規模和超大規模的模型越來越受到業界的推崇。以NLP行業為例，從最開始的Bert-base隻有1億左右的參數量，到千億級别的GPT-3，再到今年6月釋出的目前全球最大預訓練模型“悟道2.0”，參數規模達到驚人的1.75萬億，整個業界都由一種向更大模型發展的趨勢。面對如此龐大的模型，必然也需要龐大的資料量才能進行訓練，如果沒有分布式訓練的大算力加持，一個Epoch可能就要訓練到天荒地老。抛開業界淬煉超大模型的場景，對于一個AI行業的普通算法工程師，面對日常的工作，分布式訓練也可以大大加速模型的訓練、調參的節奏、以及版本的疊代更新，在時間如此珍貴的當下，相信沒有工程師會抗拒分布式訓練帶來的收益。是以，我們今天就聊聊深度學習中關于分布式訓練的那些事兒。
分布式訓練政策按照并行方式不同，可以簡單的分為資料并行和模型并行兩種方式。
資料并行是指在不同的GPU上都copy儲存一份模型的副本，然後将不同的資料配置設定到不同的GPU上進行計算，最後将所有GPU計算的結果進行合并，進而達到加速模型訓練的目的。由于資料并行會涉及到把不同GPU的計算結果進行合并然後再更新模型，根據跟新方式不同，又可以分為同步更新和異步更新。在資料并行中，每個GPU隻計算一個batch中的一部分資料，同步更新指的就是在等待所有的GPU都計算完成之後，然後再統一合并和更新網絡的權重，并廣播到所有的GPU中，随後進行下一輪的計算。而異步跟新不同，異步更新中每個GPU在獨立計算完成之後，都無需等待其他GPU，可以立即更新整體權重，然後廣播到其他GPU中，随後馬上進入下一輪的計算。由此可見，同步更新需要等待所有的GPU都計算完成才能更新，如果叢集中某一個GPU訓練慢了，或者叢集中的通信出現抖動，都會影響到整個網絡的訓練速度，類似木桶效應，最短闆決定了最大的容量。而異步更新由于不用等待其他GPU節點，是以總體訓練速度會快一些，但是會有一個嚴重的梯度失效的問題。即在異步的情況下，每一個節點完成訓練之後，都會馬上去更新，這會造成其他節點現在的模型參數和這一輪訓練前采用的模型參數可能不一緻，進而導緻此時的梯度過期。是以，異步更新雖然快，但是由于梯度失效問題，模型往往會陷入到次優解中。
與資料并行不同，分布式訓練中的模型并行是指将整個神經網絡模型拆解分布到不同的GPU中，不同的GPU負責計算網絡模型中的不同部分。這通常是在網絡模型很大很大、單個GPU的顯存已經完全裝不下整體網絡的情況下才會采用。由于深度學習的模型通常包含很多層，層與層之間的運作有先後訓練，前向傳播和反向梯度計算的時候，前面的層和後面的層都會彼此依賴作為輸入輸出，是以這種串行的邏輯對加速造成了一定的限制。但是相比起來，我們也算可以通過模型并行的方式把一個超大模型訓練起來，不然對于一個單GPU的話，超大模型是完全沒辦法work的。 是以，對比起來，模型并行由于各個GPU隻加載了模型的一部分網絡結構，存在一定的依賴關系，造成了規模的伸縮性比較差，不能随意的增減GPU的數量，是以在實際中運用的并不多。而資料并行的方式，由于各個GPU互相獨立，友善GPU的擴縮容，同時加速效果好，是以在實際中運用較多，但是在某些時候，我們也可以同時結合資料并行和模型并行兩種方式。
在Pytorch中為我們提供了兩種多GPU的分布式訓練方案：torch.nn.DataParallel（DP）和 torch.nn.parallel.Distributed Data Parallel（DDP）。
DP模式使用起來非常容易，隻需要對單GPU的代碼修改其中一行就可以運作了，由于DP模式采用的是PS架構，存在負載不均衡問題，主卡往往會成為訓練的瓶頸，是以訓練速度會比DDP模式慢一些。而且DP隻支援單機多卡的方式，一般一台機器隻能安裝最多8張卡，當我們要訓練特别大型的任務時，8卡就會顯得特别吃緊，是以會有一定的限制。
與DP模式不同，DDP模式本身是為多機多卡設計的，當然在單機多卡的情況下也可以使用。DDP采用的是all-reduce架構，基本解決了PS架構中通信成本與GPU的數量線性相關的問題。雖然在單機多卡情況下，可以使用DP模式，但是使用DDP通常會比DP模式快一些，是以DDP模式也是官方推薦大家使用的方式。改造現有的代碼使用DDP也非常友善，通過下面幾個步驟就可以輕松搞定。
除了Pytorch原生提供的DP和DDP方式以外，也有很多優秀的由第三方提供的分布式訓練工具，其中Horovod就是比較常用的一款。Horovod是Uber開源的跨平台分布式訓練架構（horovod名字來源于俄羅斯一種民間舞蹈，舞者手拉手站成一個圓圈跳舞，類比了GPU裝置之間的通信模式，如果該架構是中國人或者華人開發的話，我估計可能就叫“鍋莊”了吧^-^），從名字可以看出來，Horovod采用all-reduce架構來提高分布式裝置的通信效率。同時，Horovod不僅支援Pytorch，也支援TensorFlow等其他深度學習架構。訓練中如果想使用Horovod的話，其實對代碼的改動也比較少，如下所示。
另外，位元組跳動也開源了一款高性能的分布式深度學習訓練架構BytePS，提升了通信性能，據說效果可以優于Horovod。而在幾天之前，快手聯合蘇黎世理工也宣布開源了一款分布式訓練架構Bagua（八卦），Bagua專門針對分布式場景設計了特定的優化算法，實作了算法和系統層面的聯合優化，性能較同類提升60%。感興趣的同學也可以關注一下。
這裡我們對比了Pytorch原生的DP和DDP模式，同時也選擇了第三方插件Horovod進行對比。實驗選擇了基于bert-base的預訓練語言模型進行文本分類的任務。具體實驗參數如下：GPU型号: V100, learning_rate: 2e-5, batch_size: 128, max_len: 128, epochs: 1, train_set_size: 48w
由于DDP和Horovod都是采用all-reduce架構，是以性能相當，可見Pytorch原生的DDP模式也已經做得非常不錯了。而DP相比其他模式性能就會差一些。是以在實際工作中，還是比較推薦使用DDP或者Horovod進行分布式訓練。
本文探讨了深度學習中模型并行和資料并行的分布式政策，并基于Pytorch架構介紹了原生的DP和DDP模式，以及第三方Horovod分布式訓練架構。從後面的實驗對比可以看出，平時工作中比較推薦使用DDP或者Horovod的方式。分布式訓練是深度學習中非常重要的一環，除了Horovod，其他各大廠商也相繼開源了自己的分布式訓練架構，比如BytePS、DeepSpeed、Bagua等等，這些架構的開源也将進一步推動這個領域的發展，為深度學習提供更優秀的工具。