神經(jīng)網(wǎng)絡(luò)機(jī)器翻譯(Neural Machine Translation, NMT)模型自2013年在學(xué)術(shù)界首次被提出后，就不斷快速發(fā)展，目前在某些語(yǔ)種和場(chǎng)景下，譯文質(zhì)量甚至可以達(dá)到人工翻譯的水平。

阿里翻譯團(tuán)隊(duì)自2016年10月起正式開(kāi)始自主研發(fā)NMT模型，2016年11月首次將NMT系統(tǒng)的輸出結(jié)果應(yīng)用在中英消息通訊場(chǎng)景下的外部評(píng)測(cè)中并取得了不錯(cuò)的成績(jī)，翻譯質(zhì)量有了大幅度提升。

但是，由于NMT模型的結(jié)構(gòu)復(fù)雜，且深度神經(jīng)網(wǎng)絡(luò)模型本身的訓(xùn)練過(guò)程一般又會(huì)涉及很大量的計(jì)算，因此NMT系統(tǒng)往往需要較長(zhǎng)的訓(xùn)練周期，例如，使用3000萬(wàn)的訓(xùn)練數(shù)據(jù)在單塊GPU卡上一般需要訓(xùn)練20天以上，才能得到一個(gè)初步可用的模型。

基于上述問(wèn)題，2017年2月初開(kāi)始，阿里翻譯團(tuán)隊(duì)和阿里云Large Scale Learning的穆琢團(tuán)隊(duì)合作，共同開(kāi)發(fā)支持分布式訓(xùn)練的NMT系統(tǒng)，并于2017年3月底完成了第一個(gè)版本的分布式NMT系統(tǒng)。

項(xiàng)目成果

在2017年4月份的英俄電商翻譯質(zhì)量?jī)?yōu)化項(xiàng)目中，分布式NMT系統(tǒng)大大提高了訓(xùn)練速度，使模型訓(xùn)練時(shí)間從20天縮短到了4天，為項(xiàng)目整體迭代和推進(jìn)節(jié)省了很多時(shí)間成本。

圖一：使用不同卡數(shù)時(shí)，在中英100萬(wàn)訓(xùn)練語(yǔ)料上獲得的收斂加速比

在4張GPU卡上，可以達(dá)到3倍以上的收斂加速比，在8張GPU卡上，可以達(dá)到5倍以上的收斂加速比，在16張GPU卡上，可以達(dá)到9倍以上的收斂加速比，通過(guò)不斷累加GPU卡數(shù)，收斂加速比預(yù)期還會(huì)繼續(xù)提升。

圖二：中英2000萬(wàn)訓(xùn)練集上，利用2機(jī)4卡進(jìn)行分布式訓(xùn)練的收斂過(guò)程

圖三：多種分布式策略的加速效果對(duì)比

除了基于MA的分布式實(shí)現(xiàn)，項(xiàng)目組還嘗試了其他多種分布式實(shí)現(xiàn)方法，也都獲得了不同程度的加速效果，包括Downpour SGD、AllReduce SGD以及使用了BMUF(Blockwise Model-Update Filtering, 一種針對(duì)Model Average方法的改進(jìn)方案)策略的Model Average方法。

實(shí)現(xiàn)方案

關(guān)于多機(jī)多卡分布式方案的討論

data parallel數(shù)據(jù)并行

基于data parallel的同步SGD算法，即使用多個(gè)worker均攤整個(gè)batch的計(jì)算量，且每個(gè)GPU(worker)上存儲(chǔ)完整的模型副本。Tensorflow本身提供的PS框架可以直接用于實(shí)現(xiàn)此算法。標(biāo)準(zhǔn)的同步SGD算法每次迭代都分為三個(gè)步驟，首先，從參數(shù)服務(wù)器中把模型參數(shù)pull到本地，接著用新的模型參數(shù)計(jì)算本地mini-batch的梯度，最后后將計(jì)算出的梯度push到參數(shù)服務(wù)器。參數(shù)服務(wù)器需要收集所有workers的梯度，再統(tǒng)一處理更新模型參數(shù)。

因此，在ps框架下，同步SGD算法的總通信量 transfer data = 2 * num_of_gpus * size_of_whole_model。通過(guò)對(duì)模型的大小和一個(gè)mini-batch計(jì)算時(shí)間的評(píng)估，發(fā)現(xiàn)在單機(jī)兩卡的情況下，由于PCIe的帶寬較高(~10GB/s)，可以獲得接近2倍的加速比，但兩機(jī)兩卡的情況下，計(jì)算通信比則為1:1，而隨著機(jī)器數(shù)的增多，不但沒(méi)有帶來(lái)加速，反而比單卡更慢。

可見(jiàn)，通信優(yōu)化是同步SGD算法的主要問(wèn)題，下面的具體方案介紹中，我們也展示了相關(guān)的實(shí)驗(yàn)結(jié)果。

hybrid parallel混合并行

模型并行可以通過(guò)模型參數(shù)的存儲(chǔ)策略，利用參數(shù)的稀疏性，減少參數(shù)的實(shí)際通信量，對(duì)訓(xùn)練過(guò)程進(jìn)行加速。一般情況下，適用于稀疏模型的訓(xùn)練，但對(duì)于特殊的網(wǎng)絡(luò)結(jié)構(gòu)，模型并行可以和數(shù)據(jù)并行相結(jié)合，提高全局的計(jì)算通信比，我們這里稱之為hybrid parallel。以AlexNet舉例來(lái)說(shuō)，可以把模型結(jié)構(gòu)分為全連接層和卷積層兩部分，從而對(duì)兩個(gè)部分分別分析計(jì)算通信比。

全連接層計(jì)算通信比極低，毫無(wú)擴(kuò)展性，而卷積層的計(jì)算通信比較高，適合使用數(shù)據(jù)并行，因此hybrid parallel的方案是：多個(gè)workers使用數(shù)據(jù)并行的方式計(jì)算卷積層，而只需一個(gè)worker(或極少量workers)計(jì)算全連接層，以獲得最高的計(jì)算通信比，且卷積層和全連接層中間只有一個(gè)規(guī)模很小的feature map需要通信。

通過(guò)對(duì)NMT的模型中各個(gè)層進(jìn)行profiling，發(fā)現(xiàn)網(wǎng)絡(luò)中各部分的參數(shù)大小相差不大，計(jì)算量也相當(dāng)，不但計(jì)算通信比不高，各層結(jié)構(gòu)之間傳遞的feature map也較大，并不適合使用hybrid parallel的方式。

分布式方案的進(jìn)一步探索

基于以上的分析，我們對(duì)分布式方案進(jìn)行了進(jìn)一步的探索。為了最大程度的獲得多機(jī)多卡的擴(kuò)展性，我們選擇了Model average、BMUF、downpour等分布式方法。Model average(MA)方法較于簡(jiǎn)單的數(shù)據(jù)并行，有兩個(gè)顯著的優(yōu)勢(shì)：一、通信量可以通過(guò)同步頻次來(lái)調(diào)節(jié)，兩次同步的間隔可以以step甚至epoch為單位來(lái)控制，一般情況下，通信開(kāi)銷幾乎可以忽略不計(jì)，從而獲得線性的計(jì)算加速比;二、batch size的大小可以和單機(jī)baseline保持一致，基于數(shù)據(jù)并行的同步SGD，通常為了保持計(jì)算通信比而同比增大batch size，導(dǎo)致精度損失。

但是，MA方法、BMUF方法以及基于異步更新的downpour方法，對(duì)調(diào)參的要求都更高，不同的參數(shù)設(shè)置，會(huì)對(duì)模型的收斂結(jié)果有較大的影響，計(jì)算加速比的提高并不意味著收斂加速比的提高。下文是對(duì)各個(gè)實(shí)現(xiàn)方案的消息介紹和結(jié)果展示。

基于Model Average方法的分布式訓(xùn)練方案

方案概述

Model Average方法，即每個(gè)worker負(fù)責(zé)訓(xùn)練本地的模型，多個(gè)GPU之間通過(guò)固定(或動(dòng)態(tài))的頻率求取各個(gè)模型的均值，并以此作為繼續(xù)訓(xùn)練的基礎(chǔ)。

基于Tensorflow的Model Average設(shè)計(jì)主要分為graph構(gòu)建和分步run兩部分。

1. 首先，由于每個(gè)worker需要在各自的進(jìn)程內(nèi)進(jìn)行求解，所以需要在每個(gè)worker上分別進(jìn)行構(gòu)建forward-backward子圖，這一部分的子圖在上圖中用紅色的箭頭線表示。應(yīng)當(dāng)注意，每個(gè)worker上應(yīng)該具有獨(dú)立的op和variable。

2. 其次，需要在ps上維護(hù)一份模型拷貝，并構(gòu)建Reduce Average的子圖，這一部分的子圖在上圖中用黑色的箭頭線表示。

3. 然后，需要構(gòu)建broadcast的子圖，這一部分的子圖在上圖中用藍(lán)色的線表示。(這里由于美觀，在圖中沒(méi)有畫出全部的broadcast依賴關(guān)系，ps的variable 1和variable 2都應(yīng)該和worker 1和worker 2中的variable 1和variable 2建立依賴關(guān)系)

4. 在構(gòu)建上述依賴關(guān)系后，我們需要通過(guò)client啟動(dòng)Model Average的機(jī)制。用戶只需要通過(guò)控制session run的順序來(lái)指定每個(gè)子圖的求解即可完成。在這里，紅色的子圖需要每輪都run，而到了一定輪數(shù)間隔后開(kāi)始依次run黑色的子圖和藍(lán)色的子圖即可完成Model Average的機(jī)制。

在實(shí)現(xiàn)上，可以用主worker完成所有子圖的構(gòu)建，其他worker處于join狀態(tài)。在運(yùn)行時(shí)，各個(gè)worker會(huì)根據(jù)圖中的依賴關(guān)系驅(qū)動(dòng)起來(lái)。

超參的調(diào)整

當(dāng)兩機(jī)四卡的參數(shù)保持和單機(jī)單卡baseline參數(shù)一致的情況下，從多組實(shí)驗(yàn)中可以觀察到，不同的同步頻次對(duì)模型的質(zhì)量有影響，加速比最快只達(dá)到1.5。通過(guò)推導(dǎo)，我們發(fā)現(xiàn)，MA方法中每次計(jì)算多個(gè)worker的模型均值，都相當(dāng)于把每個(gè)訓(xùn)練樣本對(duì)模型梯度的貢獻(xiàn)縮減為原來(lái)的1/num_of_gpus, 大大抵消了多卡帶來(lái)的加速。因此，MA中的本地訓(xùn)練learning rate(lr)應(yīng)該調(diào)整為原baseline lr的num_of_gpus倍。

實(shí)驗(yàn)結(jié)果

其中bs表示batch size，lr表示learning rate，橫坐標(biāo)是訓(xùn)練時(shí)間，縱坐標(biāo)代表翻譯質(zhì)量。

使用BMUF方法對(duì)Model Average方案進(jìn)行改進(jìn)

方案概述

上面提到的Model Average方法在兩卡、四卡、八卡上均獲得了較高的加速比，在不斷并行擴(kuò)展的過(guò)程中，需要根據(jù)機(jī)器的個(gè)數(shù)對(duì)lr進(jìn)行調(diào)整，然而lr的大小一定有一個(gè)適當(dāng)?shù)姆秶?，并不能無(wú)限的增大，隨著機(jī)器數(shù)的增多，這種naive MA方法也逐漸暴露出它的局限性。

BMUF(Blockwise Model-Update Filtering)是Model average(naive MA)的一種優(yōu)化算法，最主要的區(qū)別是引入了梯度的歷史量。MA算法等價(jià)于在每次同步時(shí)，把所有workers上的模型增量加和平均，作為“block gradients”，用于更新參數(shù)服務(wù)器上的全局參數(shù)，而B(niǎo)MUF的“block gradients”則在此基礎(chǔ)之上，增加了另外一項(xiàng)歷史梯度，且歷史梯度的權(quán)重為block momentum rate。

具體推導(dǎo)如下：

MA：一個(gè)mini-batch的訓(xùn)練數(shù)據(jù)對(duì)最后模型的貢獻(xiàn)公式為：

BMUF：一個(gè)mini-batch的訓(xùn)練數(shù)據(jù)對(duì)最后模型的貢獻(xiàn)公式為：

從公式中可以看出，MA中單個(gè)mini-batch的貢獻(xiàn)在average階段被縮小為原來(lái)的N分之1，如果要保持和baseline一樣的貢獻(xiàn)度，則需要增大lr。而B(niǎo)MUF相比baseline的貢獻(xiàn)比為：

超參的調(diào)整

上文的公式可以指導(dǎo)BMUF訓(xùn)練過(guò)程中參數(shù)的調(diào)整。BMUF最主要的參數(shù)有兩個(gè)，一是block learning rate(blr)，此超參一般設(shè)為1.0，和MA算法一致，二是block momentum rate(bmr)，可以通過(guò)調(diào)整bmr為1-1/N，從而保持lr不變，總的貢獻(xiàn)量和baseline一致。此外，引入了歷史梯度后，還可以結(jié)合Nesterov方法，在本地訓(xùn)練時(shí)先對(duì)本地的模型更新已知的歷史梯度，再進(jìn)行訓(xùn)練，從而進(jìn)一步加速了收斂。

實(shí)驗(yàn)結(jié)果

上圖橫坐標(biāo)代表已訓(xùn)練的總數(shù)據(jù)量，即每個(gè)計(jì)算節(jié)點(diǎn)訓(xùn)練數(shù)據(jù)量的總和(最理想的情況，收斂曲線應(yīng)該和baseline重合甚至更高)，縱坐標(biāo)代表模型的翻譯質(zhì)量?？梢钥吹組A和BMUF在兩機(jī)四卡上可以獲得相當(dāng)?shù)募铀俦龋珺MUF的優(yōu)勢(shì)在于，可以通過(guò)調(diào)節(jié)bmr從而把learning rate(lr)可以保持在與單機(jī)單卡的baseline一個(gè)量級(jí)，因?yàn)閘r有一個(gè)較為適當(dāng)?shù)姆秶?，不能隨著GPU的增多無(wú)限制增大。BMUF的效果在更大的lr以及十六卡的實(shí)驗(yàn)中已得到體現(xiàn)。

基于Downpour SGD方法的分布式訓(xùn)練方案

方案概述

和MA的設(shè)計(jì)類似，Downpour SGD的設(shè)計(jì)也比較自然，同樣分為graph構(gòu)建和分布run兩部分。在這里需要注意的是，Downpour SGD方法是一種ASGD的異步分布式訓(xùn)練方法，其在Apply gradients時(shí)是不需要多個(gè)worker之間同步的。

下面具體介紹Downpour的整體流程。

1. 將ps上的model weights拉到每個(gè)worker上。

2. 每個(gè)worker自己求解本地的model，在求解過(guò)程中，不但要對(duì)本地model進(jìn)行參數(shù)更新，還要將梯度累加到另一組variable中保存。

3. 當(dāng)本地worker求解到一定輪數(shù)后，將本地存儲(chǔ)的累積梯度異步apply到ps的model參數(shù)上。

這種方式在實(shí)現(xiàn)上，需要在每個(gè)worker上再保存一份variable用來(lái)存儲(chǔ)梯度的累積量，并且在構(gòu)圖時(shí)需要對(duì)optimizer進(jìn)行更改，主要體現(xiàn)在需要將AdaDelta計(jì)算的梯度在apply到自己的模型上之前先累加到本地的variable中，然后再apply到自身參數(shù)中。

這里還存在一個(gè)問(wèn)題關(guān)于累加梯度的問(wèn)題，因?yàn)樵贏daDelta的實(shí)現(xiàn)中，實(shí)際上apply到model上的量包含梯度和delta量?jī)蓚€(gè)部分，因此在累加上有可能需要將這些delta量也考慮進(jìn)去。在實(shí)現(xiàn)時(shí)，無(wú)需重寫新的optimizer，只需要將調(diào)用apply gradients前后的model weights相減即可得出正確的結(jié)果。

Downpour SGD會(huì)帶來(lái)一些其他的超參數(shù)，這其中就包括push gradients和pull weights的間隔輪數(shù)，我們稱之為step。另外，若push和pull的step過(guò)大，也會(huì)導(dǎo)致這期間的累積梯度過(guò)大，如果此時(shí)直接將累積梯度apply到weights上很可能會(huì)出現(xiàn)NAN的情況，因此需要對(duì)累積梯度做clipping操作，所以對(duì)累加梯度進(jìn)行clipping的norm值也是一個(gè)額外的超參數(shù)。

超參的調(diào)整

由于ASGD帶來(lái)的異步性，導(dǎo)致調(diào)參的難度相對(duì)于同步SGD更為困難，而且Downpour SGD的超參數(shù)確實(shí)比較多。為了加快訓(xùn)練的速度，減少每個(gè)epoch的迭代時(shí)間，我們采用了比較激進(jìn)的weak scaling的方式，即每個(gè)worker都使用較大的batch size。

但是應(yīng)當(dāng)注意，雖然增加batch size會(huì)使過(guò)數(shù)據(jù)的速度加快，但也會(huì)讓每個(gè)數(shù)據(jù)的學(xué)習(xí)效果降低，這是因?yàn)楹托atch size相比，我們更新模型的頻率大大降低了，因此我們需要同時(shí)加大learning rate。另外，在異步條件下，每個(gè)worker計(jì)算的累積梯度都有可能成為stale的gradients，這些gradients的表示的方向并不一定最優(yōu)，相反，有時(shí)候會(huì)很差。

倘若我們不對(duì)這些累積的梯度進(jìn)行clip，那么很可能出現(xiàn)NAN的情況，因此我們需要調(diào)整累積梯度的clipping norm值。但是clipping norm值越小，梯度值就被削弱的越狠，模型的學(xué)習(xí)效果就越不好，它與learning rate的調(diào)參是一對(duì)trade off。最后，用于push和pull的step間隔值不宜過(guò)大，因?yàn)檫^(guò)大的step會(huì)導(dǎo)致過(guò)大的累積梯度，如果此時(shí)使用較小的clipping norm對(duì)累積梯度進(jìn)行削減，那么這將意味著如此長(zhǎng)輪數(shù)計(jì)算出來(lái)的累積梯度效果將被削減。在實(shí)驗(yàn)中我們發(fā)現(xiàn)這個(gè)step間隔輪數(shù)設(shè)置在20左右比較理想。

綜上所述，調(diào)參的基本目標(biāo)是，固定push和pull的step間隔輪數(shù)后，選取較大的batch size并適當(dāng)增加learning rate，逐步調(diào)整增加clipping norm的大小，使學(xué)習(xí)效果達(dá)到最大。

實(shí)驗(yàn)結(jié)果

其中bs表示batch size，lr表示learning rate，norm將累積梯度做clipping的最大值，fetch為downpour中push和pull的間隔值，縱坐標(biāo)代表翻譯質(zhì)量值，橫坐標(biāo)是訓(xùn)練時(shí)間。

基于Ring-allReduce SGD方法的分布式訓(xùn)練方案

方案概述

上文提到，同步SGD算法每次迭代都需要pull參數(shù)和push梯度這兩次通信，這樣的一個(gè)push+pull的過(guò)程對(duì)于整個(gè)系統(tǒng)來(lái)說(shuō)，等價(jià)于一次allreduce操作，通信量相當(dāng)于2 * nGPUs * model_size，隨著GPU個(gè)數(shù)的增多，allreduce成為最主要的性能瓶頸。

Ring-allReduce op

Allreduce聚合通信已有較成熟的優(yōu)化方法，最常用的方法之一有Ring allreduce。其基本思想是將allreduce分拆成reduce_scatter和allgather兩步。

1. 首先，根據(jù)節(jié)點(diǎn)數(shù)對(duì)待通信的message進(jìn)行分片，分片的數(shù)量與通信的節(jié)點(diǎn)數(shù)相同(不需要額外的參數(shù)服務(wù)器)。

2. 開(kāi)始進(jìn)行reduce_scatter操作。對(duì)worker(i)來(lái)說(shuō)，第j次通信是把自己的第j片message發(fā)送給worker(i+1)，并把收到的消息累加到message的對(duì)應(yīng)片段。如此經(jīng)過(guò)n-1輪的環(huán)狀流水線通信后，每個(gè)worker上都有一個(gè)分片是reduce了所有workers的結(jié)果。

3. 然后開(kāi)始allgather操作。第j次通信是把自己的第j+1片message發(fā)送給worker(i+1)，如此經(jīng)過(guò)n-1輪通信，所有worker的整個(gè)message就都經(jīng)過(guò)了reduce操作，即完成了all_reduce操作。

ring-allreduce方法可以使得集群內(nèi)的通信總量保持在一個(gè)常數(shù)，略小于2 * mdoel_size，不隨并行規(guī)模的增大而增多，有很好的擴(kuò)展性。

gRPC vs MPI

Ring-allreduce方法能夠帶來(lái)性能提升的前提是：latency的開(kāi)銷相對(duì)于bandwidth來(lái)說(shuō)可以忽略不計(jì)，因?yàn)閞ing allreduce雖然降低了通信總量，卻增加了通信次數(shù)。當(dāng)前tensorflow的通信協(xié)議只支持gRPC，通信的latency較高不能被簡(jiǎn)單忽略，而使用基于支持RDMA的cuda-aware MPI來(lái)實(shí)現(xiàn)allreduce op，效果則更加顯著。

此外，allreduce操作同樣存在于Model average算法中，同樣可以帶來(lái)一定的性能收益。

超參的調(diào)整

因?yàn)閷?duì)于同步SGD分布式算法，如果保持Total Batch size不變，每塊GPU卡上的mini batch size會(huì)隨著GPU卡數(shù)的增多而減小，其計(jì)算時(shí)間并不是線性減少，因?yàn)楫?dāng)batch size足夠小后，計(jì)算時(shí)間會(huì)逐漸趨于平穩(wěn)，雖然通信已經(jīng)通過(guò)優(yōu)化而大幅減少，計(jì)算的拓展性依然有限。因此，在使用多機(jī)多卡的同步SGD時(shí)，batch size與GPU個(gè)數(shù)同比例增大，同理，為了保持單個(gè)訓(xùn)練樣本的貢獻(xiàn)，lr也同比增大。

實(shí)驗(yàn)結(jié)果

上圖中，多卡同步SGD的總batch size設(shè)置為1280，是baseline(batch size=160)的8倍，橫坐標(biāo)代表已訓(xùn)練的總數(shù)據(jù)量，縱坐標(biāo)代表模型的翻譯質(zhì)量。在InfiniBand(10GB/s)網(wǎng)絡(luò)上，獲得4倍的計(jì)算加速比，在萬(wàn)兆(1GB/s)以太網(wǎng)上,獲得2.56倍計(jì)算加速比。但batch size的增大會(huì)影響收斂精度，上圖可以看到，收斂的最高點(diǎn)比baseline略低，有明顯的精度損失。

未來(lái)工作

上一階段工作主要集中在模型訓(xùn)練階段的加速策略上，接下來(lái)的工作主要分為兩方面：一方面是繼續(xù)挖掘分布式訓(xùn)練的加速潛力，通過(guò)系統(tǒng)與算法相結(jié)合的優(yōu)化策略，最大化利用硬件資源，提升收斂加速比，并將分布式優(yōu)化策略和算法模型本身解耦，實(shí)現(xiàn)復(fù)雜DL模型分布式加速功能的組件化和通用化。

另一方面，需要在現(xiàn)有的服務(wù)化方案的基礎(chǔ)上，進(jìn)一步通過(guò)模型精度壓縮、網(wǎng)絡(luò)結(jié)構(gòu)簡(jiǎn)化等方式，在保證模型效果的同時(shí)，提高解碼速度，降低線上延時(shí)，進(jìn)而增強(qiáng)線上服務(wù)能力，節(jié)約服務(wù)化所需的硬件成本。

【本文為51CTO專欄作者“阿里巴巴官方技術(shù)”原創(chuàng)稿件，轉(zhuǎn)載請(qǐng)聯(lián)系原作者】

戳這里，看該作者更多好文