1. 爲何須要分佈式訓練
隨着人工智能與深度學習的發展，大規模和超大規模的模型愈來愈受到業界的推崇。以NLP行業爲例，從最開始的Bert-base只有1億左右的參數量，到千億級別的GPT-3，再到今年6月發佈的目前全球最大預訓練模型「悟道2.0」，參數規模達到驚人的1.75萬億，整個業界都由一種向更大模型發展的趨勢。面對如此龐大的模型，必然也須要龐大的數據量才能進行訓練，若是沒有分佈式訓練的大算力加持，一個Epoch可能就要訓練到天荒地老。拋開業界淬鍊超大模型的場景，對於一個AI行業的普通算法工程師，面對平常的工做，分佈式訓練也能夠大大加速模型的訓練、調參的節奏、以及版本的迭代更新，在時間如此珍貴的當下，相信沒有工程師會抗拒分佈式訓練帶來的收益。所以，咱們今天就聊聊深度學習中關於分佈式訓練的那些事兒。 
git
2. 分佈式訓練策略
分佈式訓練策略按照並行方式不一樣，能夠簡單的分爲數據並行和模型並行兩種方式。 github
2.1 數據並行
數據並行是指在不一樣的GPU上都copy保存一份模型的副本，而後將不一樣的數據分配到不一樣的GPU上進行計算，最後將全部GPU計算的結果進行合併，從而達到加速模型訓練的目的。因爲數據並行會涉及到把不一樣GPU的計算結果進行合併而後再更新模型，根據跟新方式不一樣，又能夠分爲同步更新和異步更新。在數據並行中，每一個GPU只計算一個batch中的一部分數據，同步更新指的就是在等待全部的GPU都計算完成以後，而後再統一合併和更新網絡的權重，並廣播到全部的GPU中，隨後進行下一輪的計算。而異步跟新不一樣，異步更新中每一個GPU在獨立計算完成以後，都無需等待其餘GPU，能夠當即更新總體權重，而後廣播到其餘GPU中，隨後立刻進入下一輪的計算。因而可知，同步更新須要等待全部的GPU都計算完成才能更新，若是集羣中某一個GPU訓練慢了，或者集羣中的通訊出現抖動，都會影響到整個網絡的訓練速度，相似木桶效應，最短板決定了最大的容量。而異步更新因爲不用等待其餘GPU節點，所以整體訓練速度會快一些，可是會有一個嚴重的梯度失效的問題。即在異步的狀況下，每個節點完成訓練以後，都會立刻去更新，這會形成其餘節點如今的模型參數和這一輪訓練前採用的模型參數可能不一致，從而致使此時的梯度過時。所以，異步更新雖然快，可是因爲梯度失效問題，模型每每會陷入到次優解中。 算法
2.2 模型並行
與數據並行不一樣，分佈式訓練中的模型並行是指將整個神經網絡模型拆解分佈到不一樣的GPU中，不一樣的GPU負責計算網絡模型中的不一樣部分。這一般是在網絡模型很大很大、單個GPU的顯存已經徹底裝不下總體網絡的狀況下才會採用。因爲深度學習的模型一般包含不少層，層與層之間的運行有前後訓練，前向傳播和反向梯度計算的時候，前面的層和後面的層都會彼此依賴做爲輸入輸出，所以這種串行的邏輯對加速形成了必定的限制。可是相比起來，咱們也算能夠經過模型並行的方式把一個超大模型訓練起來，否則對於一個單GPU的話，超大模型是徹底沒辦法work的。 所以，對比起來，模型並行因爲各個GPU只加載了模型的一部分網絡結構，存在必定的依賴關係，形成了規模的伸縮性比較差，不能隨意的增減GPU的數量，所以在實際中運用的並很少。而數據並行的方式，因爲各個GPU相互獨立，方便GPU的擴縮容，同時加速效果好，所以在實際中運用較多，可是在某些時候，咱們也能夠同時結合數據並行和模型並行兩種方式。網絡
3. 基於Pytorch的分佈式訓練方法
在Pytorch中爲咱們提供了兩種多GPU的分佈式訓練方案：torch.nn.DataParallel（DP）和 torch.nn.parallel.Distributed Data Parallel（DDP）。架構
3.1 Data Parallel
DP模式使用起來很是容易，只須要對單GPU的代碼修改其中一行就能夠運行了，因爲DP模式採用的是PS架構，存在負載不均衡問題，主卡每每會成爲訓練的瓶頸，所以訓練速度會比DDP模式慢一些。並且DP只支持單機多卡的方式，通常一臺機器只能安裝最多8張卡，當咱們要訓練特別大型的任務時，8卡就會顯得特別吃緊，所以會有必定的限制。框架
# use DataParallel
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = torch.nn.DataParallel(model)
model.to(device)
3.2 DistributedDataParallel
與DP模式不一樣，DDP模式自己是爲多機多卡設計的，固然在單機多卡的狀況下也可使用。DDP採用的是all-reduce架構，基本解決了PS架構中通訊成本與GPU的數量線性相關的問題。雖然在單機多卡狀況下，可使用DP模式，可是使用DDP一般會比DP模式快一些，所以DDP模式也是官方推薦你們使用的方式。改造現有的代碼使用DDP也很是方便，經過下面幾個步驟就能夠輕鬆搞定。異步
# 1. init backend nccl
torch.distributed.init_process_group(backend='nccl')
# 2. config gpu
local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)
# 3. use DistributedSampler
training_loader = DataLoader(training_set, batch_size=TRAIN_BATCH_SIZE, sampler=DistributedSampler(training_set))
# 4. move model to gpu
model.to(device)
# 5. use DistributedDataParallel
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
model = DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank)
3.3 Horovod
除了Pytorch原生提供的DP和DDP方式之外，也有不少優秀的由第三方提供的分佈式訓練工具，其中Horovod就是比較經常使用的一款。Horovod是Uber開源的跨平臺分佈式訓練框架（horovod名字來源於俄羅斯一種民間舞蹈，舞者手拉手站成一個圓圈跳舞，類比了GPU設備之間的通訊模式，若是該框架是中國人或者華人開發的話，我估計可能就叫「鍋莊」了吧^-^），從名字能夠看出來，Horovod採用all-reduce架構來提升分佈式設備的通訊效率。同時，Horovod不只支持Pytorch，也支持TensorFlow等其餘深度學習框架。訓練中若是想使用Horovod的話，其實對代碼的改動也比較少，以下所示。分佈式
import horovod.torch as hvd # 1. init horovod hvd.init() # 2. Pin GPU to be used to process local rank (one GPU per process) torch.cuda.set_device(hvd.local_rank()) # 3. Partition dataset among workers using DistributedSampler train_sampler = DistributedSampler(training_set, num_replicas=hvd.size(), rank=hvd.rank()) training_loader = DataLoader(training_set, batch_size=TRAIN_BATCH_SIZE, sampler=train_sampler) # 4. Add Horovod Distributed Optimizer optimizer = hvd.DistributedOptimizer(optimizer, named_parameters=model.named_parameters()) # 5. Horovod: broadcast parameters from rank 0 to all other processes. hvd.broadcast_parameters(model.state_dict(), root_rank=0)
另外，字節跳動也開源了一款高性能的分佈式深度學習訓練框架BytePS（項目github地址：https://github.com/bytedance/byteps），該框架沒有采用熱門的all-reduce，反而採用了PS架構，經過利用額外的CPU資源做爲Parameter Server等措施，提高了通訊性能，聽說效果能夠優於Horovod。而在幾天以前，快手聯合蘇黎世理工也宣佈開源了一款分佈式訓練框架Bagua（八卦），Bagua專門針對分佈式場景設計了特定的優化算法，實現了算法和系統層面的聯合優化，性能較同類提高60%。感興趣的同窗也能夠關注一下。項目github地址：https://github.com/BaguaSys/bagua工具
4. 實驗對比
這裏咱們對比了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等等，這些框架的開源也將進一步推進這個領域的發展，爲深度學習提供更優秀的工具。
做者簡介
Hongyu OPPO高級NLP算法工程師
主要從事NLP、知識圖譜及相關領域的工做
獲取更多精彩內容，掃碼關注[OPPO數智技術]公衆號