神經網絡機器翻譯(neural machine translation, nmt)模型自2013年在學術界首次被提出後,就不斷快速發展,目前在某些語種和場景下,譯文品質甚至可以達到人工翻譯的水準。
阿裡翻譯團隊自2016年10月起正式開始自主研發nmt模型,2016年11月首次将nmt系統的輸出結果應用在中英消息通訊場景下的外部評測中并取得了不錯的成績,翻譯品質有了大幅度提升。

但是,由于nmt模型的結構複雜,且深度神經網絡模型本身的訓練過程一般又會涉及很大量的計算,是以nmt系統往往需要較長的訓練周期,例如,使用3000萬的訓練資料在單塊gpu卡上一般需要訓練20天以上,才能得到一個初步可用的模型。
基于上述問題,2017年2月初開始,阿裡翻譯團隊和阿裡雲large scale learning的穆琢團隊合作,共同開發支援分布式訓練的nmt系統,并于2017年3月底完成了第一個版本的分布式nmt系統。
項目成果
在2017年4月份的英俄電商翻譯品質優化項目中,分布式nmt系統大大提高了訓練速度,使模型訓練時間從20天縮短到了4天,為項目整體疊代和推進節省了很多時間成本。
圖一:使用不同卡數時,在中英100萬訓練語料上獲得的收斂加速比
在4張gpu卡上,可以達到3倍以上的收斂加速比,在8張gpu卡上,可以達到5倍以上的收斂加速比,在16張gpu卡上,可以達到9倍以上的收斂加速比,通過不斷累加gpu卡數,收斂加速比預期還會繼續提升。
圖二:中英2000萬訓練集上,利用2機4卡進行分布式訓練的收斂過程
圖三:多種分布式政策的加速效果對比
除了基于ma的分布式實作,項目組還嘗試了其他多種分布式實作方法,也都獲得了不同程度的加速效果,包括downpour sgd、allreduce sgd以及使用了bmuf(blockwise model-update filtering, 一種針對model average方法的改進方案)政策的model average方法。
實作方案
關于多機多卡分布式方案的讨論
data parallel資料并行
基于data parallel的同步sgd算法,即使用多個worker均攤整個batch的計算量,且每個gpu(worker)上存儲完整的模型副本。tensorflow本身提供的ps架構可以直接用于實作此算法。标準的同步sgd算法每次疊代都分為三個步驟,首先,從參數伺服器中把模型參數pull到本地,接着用新的模型參數計算本地mini-batch的梯度,最後後将計算出的梯度push到參數伺服器。參數伺服器需要收集所有workers的梯度,再統一處理更新模型參數。
是以,在ps架構下,同步sgd算法的總通信量 transfer data = 2 num_of_gpus size_of_whole_model。通過對模型的大小和一個mini-batch計算時間的評估,發現在單機兩卡的情況下,由于pcie的帶寬較高(~10gb/s),可以獲得接近2倍的加速比,但兩機兩卡的情況下,計算通信比則為1:1,而随着機器數的增多,不但沒有帶來加速,反而比單卡更慢。
可見,通信優化是同步sgd算法的主要問題,下面的具體方案介紹中,我們也展示了相關的實驗結果。
hybrid parallel混合并行
模型并行可以通過模型參數的存儲政策,利用參數的稀疏性,減少參數的實際通信量,對訓練過程進行加速。一般情況下,适用于稀疏模型的訓練,但對于特殊的網絡結構,模型并行可以和資料并行相結合,提高全局的計算通信比,我們這裡稱之為hybrid parallel。以alexnet舉例來說,可以把模型結構分為全連接配接層和卷積層兩部分,進而對兩個部分分别分析計算通信比。
全連接配接層計算通信比極低,毫無擴充性,而卷積層的計算通信比較高,适合使用資料并行,是以hybrid parallel的方案是:多個workers使用資料并行的方式計算卷積層,而隻需一個worker(或極少量workers)計算全連接配接層,以獲得最高的計算通信比,且卷積層和全連接配接層中間隻有一個規模很小的feature map需要通信。
通過對nmt的模型中各個層進行profiling,發現網絡中各部分的參數大小相差不大,計算量也相當,不但計算通信比不高,各層結構之間傳遞的feature map也較大,并不适合使用hybrid parallel的方式。
分布式方案的進一步探索
基于以上的分析,我們對分布式方案進行了進一步的探索。為了最大程度的獲得多機多卡的擴充性,我們選擇了model average、bmuf、downpour等分布式方法。model average(ma)方法較于簡單的資料并行,有兩個顯著的優勢:一、通信量可以通過同步頻次來調節,兩次同步的間隔可以以step甚至epoch為機關來控制,一般情況下,通信開銷幾乎可以忽略不計,進而獲得線性的計算加速比;二、batch size的大小可以和單機baseline保持一緻,基于資料并行的同步sgd,通常為了保持計算通信比而同比增大batch size,導緻精度損失。
但是,ma方法、bmuf方法以及基于異步更新的downpour方法,對調參的要求都更高,不同的參數設定,會對模型的收斂結果有較大的影響,計算加速比的提高并不意味着收斂加速比的提高。下文是對各個實作方案的消息介紹和結果展示。
基于model average方法的分布式訓練方案
方案概述
model average方法,即每個worker負責訓練本地的模型,多個gpu之間通過固定(或動态)的頻率求取各個模型的均值,并以此作為繼續訓練的基礎。
基于tensorflow的model average設計主要分為graph建構和分步run兩部分。
1.首先,由于每個worker需要在各自的程序内進行求解,是以需要在每個worker上分别進行建構forward-backward子圖,這一部分的子圖在上圖中用紅色的箭頭線表示。應當注意,每個worker上應該具有獨立的op和variable。
2.其次,需要在ps上維護一份模型拷貝,并建構reduce average的子圖,這一部分的子圖在上圖中用黑色的箭頭線表示。
3.然後,需要建構broadcast的子圖,這一部分的子圖在上圖中用藍色的線表示。(這裡由于美觀,在圖中沒有畫出全部的broadcast依賴關系,ps的variable 1和variable 2都應該和worker 1和worker 2中的variable 1和variable 2建立依賴關系)
4.在建構上述依賴關系後,我們需要通過client啟動model average的機制。使用者隻需要通過控制session run的順序來指定每個子圖的求解即可完成。在這裡,紅色的子圖需要每輪都run,而到了一定輪數間隔後開始依次run黑色的子圖和藍色的子圖即可完成model average的機制。
在實作上,可以用主worker完成所有子圖的建構,其他worker處于join狀态。在運作時,各個worker會根據圖中的依賴關系驅動起來。
超參的調整
當兩機四卡的參數保持和單機單卡baseline參數一緻的情況下,從多組實驗中可以觀察到,不同的同步頻次對模型的品質有影響,加速比最快隻達到1.5.通過推導,我們發現,ma方法中每次計算多個worker的模型均值,都相當于把每個訓練樣本對模型梯度的貢獻縮減為原來的1/num_of_gpus, 大大抵消了多卡帶來的加速。是以,ma中的本地訓練learning rate(lr)應該調整為原baseline lr的num_of_gpus倍。
實驗結果
其中bs表示batch size,lr表示learning rate,橫坐标是訓練時間,縱坐标代表翻譯品質。
使用bmuf方法對model average方案進行改進
上面提到的model average方法在兩卡、四卡、八卡上均獲得了較高的加速比,在不斷并行擴充的過程中,需要根據機器的個數對lr進行調整,然而lr的大小一定有一個适當的範圍,并不能無限的增大,随着機器數的增多,這種naive ma方法也逐漸暴露出它的局限性。
bmuf(blockwise model-update filtering)是model average(naive ma)的一種優化算法,最主要的差別是引入了梯度的曆史量。ma算法等價于在每次同步時,把所有workers上的模型增量加和平均,作為“block gradients”,用于更新參數伺服器上的全局參數,而bmuf的“block gradients”則在此基礎之上,增加了另外一項曆史梯度,且曆史梯度的權重為block momentum rate。
具體推導如下:
ma:一個mini-batch的訓練資料對最後模型的貢獻公式為:
bmuf:一個mini-batch的訓練資料對最後模型的貢獻公式為:
從公式中可以看出,ma中單個mini-batch的貢獻在average階段被縮小為原來的n分之1,如果要保持和baseline一樣的貢獻度,則需要增大lr。而bmuf相比baseline的貢獻比為:
上文的公式可以指導bmuf訓練過程中參數的調整。bmuf最主要的參數有兩個,一是block learning rate(blr),此超參一般設為1.0,和ma算法一緻,二是block momentum rate(bmr),可以通過調整bmr為1-1/n,進而保持lr不變,總的貢獻量和baseline一緻。此外,引入了曆史梯度後,還可以結合nesterov方法,在本地訓練時先對本地的模型更新已知的曆史梯度,再進行訓練,進而進一步加速了收斂。
上圖橫坐标代表已訓練的總資料量,即每個計算節點訓練資料量的總和(最理想的情況,收斂曲線應該和baseline重合甚至更高),縱坐标代表模型的翻譯品質。可以看到ma和bmuf在兩機四卡上可以獲得相當的加速比,bmuf的優勢在于,可以通過調節bmr進而把learning rate(lr)可以保持在與單機單卡的baseline一個量級,因為lr有一個較為适當的範圍,不能随着gpu的增多無限制增大。bmuf的效果在更大的lr以及十六卡的實驗中已得到展現。
基于downpour sgd方法的分布式訓練方案
和ma的設計類似,downpour sgd的設計也比較自然,同樣分為graph建構和分布run兩部分。在這裡需要注意的是,downpour sgd方法是一種asgd的異步分布式訓練方法,其在apply gradients時是不需要多個worker之間同步的。
下面具體介紹downpour的整體流程。
将ps上的model weights拉到每個worker上。
每個worker自己求解本地的model,在求解過程中,不但要對本地model進行參數更新,還要将梯度累加到另一組variable中儲存。
當本地worker求解到一定輪數後,将本地存儲的累積梯度異步apply到ps的model參數上。
這種方式在實作上,需要在每個worker上再儲存一份variable用來存儲梯度的累積量,并且在構圖時需要對optimizer進行更改,主要展現在需要将adadelta計算的梯度在apply到自己的模型上之前先累加到本地的variable中,然後再apply到自身參數中。
這裡還存在一個問題關于累加梯度的問題,因為在adadelta的實作中,實際上apply到model上的量包含梯度和delta量兩個部分,是以在累加上有可能需要将這些delta量也考慮進去。在實作時,無需重寫新的optimizer,隻需要将調用apply gradients前後的model weights相減即可得出正确的結果。
downpour sgd會帶來一些其他的超參數,這其中就包括push gradients和pull weights的間隔輪數,我們稱之為step。另外,若push和pull的step過大,也會導緻這期間的累積梯度過大,如果此時直接将累積梯度apply到weights上很可能會出現nan的情況,是以需要對累積梯度做clipping操作,是以對累加梯度進行clipping的norm值也是一個額外的超參數。
由于asgd帶來的異步性,導緻調參的難度相對于同步sgd更為困難,而且downpour sgd的超參數确實比較多。為了加快訓練的速度,減少每個epoch的疊代時間,我們采用了比較激進的weak scaling的方式,即每個worker都使用較大的batch size。
但是應當注意,雖然增加batch size會使過資料的速度加快,但也會讓每個資料的學習效果降低,這是因為和小batch size相比,我們更新模型的頻率大大降低了,是以我們需要同時加大learning rate。另外,在異步條件下,每個worker計算的累積梯度都有可能成為stale的gradients,這些gradients的表示的方向并不一定最優,相反,有時候會很差。
倘若我們不對這些累積的梯度進行clip,那麼很可能出現nan的情況,是以我們需要調整累積梯度的clipping norm值。但是clipping norm值越小,梯度值就被削弱的越狠,模型的學習效果就越不好,它與learning rate的調參是一對trade off。最後,用于push和pull的step間隔值不宜過大,因為過大的step會導緻過大的累積梯度,如果此時使用較小的clipping norm對累積梯度進行削減,那麼這将意味着如此長輪數計算出來的累積梯度效果将被削減。在實驗中我們發現這個step間隔輪數設定在20左右比較理想。
綜上所述,調參的基本目标是,固定push和pull的step間隔輪數後,選取較大的batch size并适當增加learning rate,逐漸調整增加clipping norm的大小,使學習效果達到最大。
其中bs表示batch size,lr表示learning rate,norm将累積梯度做clipping的最大值,fetch為downpour中push和pull的間隔值,縱坐标代表翻譯品質值,橫坐标是訓練時間。
基于ring-allreduce sgd方法的分布式訓練方案
上文提到,同步sgd算法每次疊代都需要pull參數和push梯度這兩次通信,這樣的一個push+pull的過程對于整個系統來說,等價于一次allreduce操作,通信量相當于2 ngpus model_size,随着gpu個數的增多,allreduce成為最主要的性能瓶頸。
ring-allreduce op
allreduce聚合通信已有較成熟的優化方法,最常用的方法之一有ring allreduce。其基本思想是将allreduce分拆成reduce_scatter和allgather兩步。
首先,根據節點數對待通信的message進行分片,分片的數量與通信的節點數相同(不需要額外的參數伺服器)。
開始進行reduce_scatter操作。對worker(i)來說,第j次通信是把自己的第j片message發送給worker(i+1),并把收到的消息累加到message的對應片段。如此經過n-1輪的環狀流水線通信後,每個worker上都有一個分片是reduce了所有workers的結果。
然後開始allgather操作。第j次通信是把自己的第j+1片message發送給worker(i+1),如此經過n-1輪通信,所有worker的整個message就都經過了reduce操作,即完成了all_reduce操作。
ring-allreduce方法可以使得叢集内的通信總量保持在一個常數,略小于2 * mdoel_size,不随并行規模的增大而增多,有很好的擴充性。
grpc vs mpi
ring-allreduce方法能夠帶來性能提升的前提是:latency的開銷相對于bandwidth來說可以忽略不計,因為ring allreduce雖然降低了通信總量,卻增加了通信次數。目前tensorflow的通信協定隻支援grpc,通信的latency較高不能被簡單忽略,而使用基于支援rdma的cuda-aware mpi來實作allreduce op,效果則更加顯著。
此外,allreduce操作同樣存在于model average算法中,同樣可以帶來一定的性能收益。
因為對于同步sgd分布式算法,如果保持total batch size不變,每塊gpu卡上的mini batch size會随着gpu卡數的增多而減小,其計算時間并不是線性減少,因為當batch size足夠小後,計算時間會逐漸趨于平穩,雖然通信已經通過優化而大幅減少,計算的拓展性依然有限。是以,在使用多機多卡的同步sgd時,batch size與gpu個數同比例增大,同理,為了保持單個訓練樣本的貢獻,lr也同比增大。
上圖中,多卡同步sgd的總batch size設定為1280,是baseline(batch size=160)的8倍,橫坐标代表已訓練的總資料量,縱坐标代表模型的翻譯品質。在infiniband(10gb/s)網絡上,獲得4倍的計算加速比,在萬兆(1gb/s)以太網上,獲得2.56倍計算加速比。但batch size的增大會影響收斂精度,上圖可以看到,收斂的最高點比baseline略低,有明顯的精度損失。
未來工作
上一階段工作主要集中在模型訓練階段的加速政策上,接下來的工作主要分為兩方面:一方面是繼續挖掘分布式訓練的加速潛力,通過系統與算法相結合的優化政策,最大化利用硬體資源,提升收斂加速比,并将分布式優化政策和算法模型本身解耦,實作複雜dl模型分布式加速功能的元件化和通用化。
另一方面,需要在現有的服務化方案的基礎上,進一步通過模型精度壓縮、網絡結構簡化等方式,在保證模型效果的同時,提高解碼速度,降低線上延時,進而增強線上服務能力,節約服務化所需的硬體成本。
<a href="https://mp.weixin.qq.com/s/dldctlybo-nuljub9ryqbg">原文連結</a>