電子發(fā)燒友App
在上一篇的介紹中，我們介紹了以Google GPipe為代表的流水線并行范式。當(dāng)模型太大，一塊GPU放不下時(shí)，流水線并行將模型的不同層放到不同的GPU上，通過(guò)切割mini-batch實(shí)現(xiàn)對(duì)訓(xùn)練數(shù)據(jù)的流水線處理，提升GPU計(jì)算通訊比。同時(shí)通過(guò)re-materialization機(jī)制降低顯存消耗。
但在實(shí)際應(yīng)用中，流水線并行并不特別流行，主要原因是模型能否均勻切割，影響了整體計(jì)算效率，這就需要算法工程師做手調(diào)。因此，今天我們來(lái)介紹一種應(yīng)用最廣泛，最易于理解的并行范式：數(shù)據(jù)并行。
數(shù)據(jù)并行的核心思想是：在各個(gè)GPU上都拷貝一份完整模型，各自吃一份數(shù)據(jù)，算一份梯度，最后對(duì)梯度進(jìn)行累加來(lái)更新整體模型。理念不復(fù)雜，但到了大模型場(chǎng)景，巨大的存儲(chǔ)和GPU間的通訊量，就是系統(tǒng)設(shè)計(jì)要考慮的重點(diǎn)了。在本文中，我們將遞進(jìn)介紹三種主流數(shù)據(jù)并行的實(shí)現(xiàn)方式：
- DP（Data Parallelism） ：最早的數(shù)據(jù)并行模式，一般采用參數(shù)服務(wù)器(Parameters Server)這一編程框架。實(shí)際中多用于單機(jī)多卡
- DDP（Distributed Data Parallelism） ：分布式數(shù)據(jù)并行，采用Ring AllReduce的通訊方式，實(shí)際中多用于多機(jī)場(chǎng)景
- ZeRO ：零冗余優(yōu)化器。由微軟推出并應(yīng)用于其DeepSpeed框架中。嚴(yán)格來(lái)講ZeRO采用數(shù)據(jù)并行+張量并行的方式，旨在降低存儲(chǔ)。
本文將首先介紹DP和DDP，在下一篇文章里，介紹ZeRO。全文內(nèi)容如下：
1、數(shù)據(jù)并行（DP）
- 1.1 整體架構(gòu)
- 1.2 通訊瓶頸與梯度異步更
2、分布式數(shù)據(jù)并行（DDP）
- 2.1 圖解Ring-AllReduce
- 2.2 DP與DDP通訊分析
推薦閱讀：
一、數(shù)據(jù)并行（DP）
1.1 整體架構(gòu)
一個(gè)經(jīng)典數(shù)據(jù)并行的過(guò)程如下：
- 若干塊 計(jì)算GPU ，如圖中GPU0~GPU2；1塊 梯度收集GPU ，如圖中AllReduce操作所在GPU。
- 在每塊計(jì)算GPU上都拷貝一份完整的模型參數(shù)。
- 把一份數(shù)據(jù)X（例如一個(gè)batch）均勻分給不同的計(jì)算GPU。
- 每塊計(jì)算GPU做一輪FWD和BWD后，算得一份梯度G。
- 每塊計(jì)算GPU將自己的梯度push給梯度收集GPU，做聚合操作。這里的聚合操作一般指 梯度累加 。當(dāng)然也支持用戶(hù)自定義。
- 梯度收集GPU聚合完畢后，計(jì)算GPU從它那pull下完整的梯度結(jié)果，用于更新模型參數(shù)W。更新完畢后，計(jì)算GPU上的模型參數(shù)依然保持一致。
- 聚合再下發(fā)梯度的操作，稱(chēng)為AllReduce。
前文說(shuō)過(guò)，實(shí)現(xiàn)DP的一種經(jīng)典編程框架叫“參數(shù)服務(wù)器”，在這個(gè)框架里，計(jì)算GPU稱(chēng)為Worker，梯度聚合GPU稱(chēng)為Server。在實(shí)際應(yīng)用中，為了盡量減少通訊量，一般可選擇一個(gè)Worker同時(shí)作為Server。比如可把梯度全發(fā)到GPU0上做聚合。需要再額外說(shuō)明幾點(diǎn)：
- 1個(gè)Worker或者Server下可以不止1塊GPU。
- Server可以只做梯度聚合，也可以梯度聚合+全量參數(shù)更新一起做 在參數(shù)服務(wù)器的語(yǔ)言體系下，DP的過(guò)程又可以被描述下圖：
1.2 通訊瓶頸與梯度異步更新
DP的框架理解起來(lái)不難，但實(shí)戰(zhàn)中確有兩個(gè)主要問(wèn)題：
- 存儲(chǔ)開(kāi)銷(xiāo)大 。每塊GPU上都存了一份完整的模型，造成冗余。關(guān)于這一點(diǎn)的優(yōu)化，我們將在后文ZeRO部分做講解。
- 通訊開(kāi)銷(xiāo)大 。Server需要和每一個(gè)Worker進(jìn)行梯度傳輸。當(dāng)Server和Worker不在一臺(tái)機(jī)器上時(shí)，Server的帶寬將會(huì)成為整個(gè)系統(tǒng)的計(jì)算效率瓶頸。
我們對(duì)通訊開(kāi)銷(xiāo)再做詳細(xì)說(shuō)明。如果將傳輸比作一條馬路，帶寬就是馬路的寬度，它決定每次并排行駛的數(shù)據(jù)量。例如帶寬是100G/s，但每秒?yún)s推給Server 1000G的數(shù)據(jù)，消化肯定需要時(shí)間。
人類(lèi)老板不愿意了：“打工系統(tǒng)里不允許有串行存在的任務(wù)！”，于是梯度異步更新這一管理層略誕生了。
上圖刻畫(huà)了在梯度異步更新的場(chǎng)景下，某個(gè)Worker的計(jì)算順序?yàn)椋?/p>
- 在第10輪計(jì)算中，該Worker正常計(jì)算梯度，并向Server發(fā)送push&pull梯度請(qǐng)求。
- 但是，該Worker并不會(huì)實(shí)際等到把聚合梯度拿回來(lái)，更新完參數(shù)W后再做計(jì)算。而是直接拿舊的W，吃新的數(shù)據(jù)，繼續(xù)第11輪的計(jì)算。 這樣就保證在通訊的時(shí)間里，Worker也在馬不停蹄做計(jì)算，提升計(jì)算通訊比 。
- 當(dāng)然，異步也不能太過(guò)份。只計(jì)算梯度，不更新權(quán)重，那模型就無(wú)法收斂。圖中刻畫(huà)的是延遲為1的異步更新，也就是在開(kāi)始第12輪對(duì)的計(jì)算時(shí)，必須保證W已經(jīng)用第10、11輪的梯度做完2次更新了。
參數(shù)服務(wù)器的框架下，延遲的步數(shù)也可以由用戶(hù)自己決定，下圖分別刻劃了幾種延遲情況：
- (a) 無(wú)延遲
- (b) 延遲但不指定延遲步數(shù) 。也即在迭代2時(shí)，用的可能是老權(quán)重，也可能是新權(quán)重，聽(tīng)天由命。
- (c) 延遲且指定延遲步數(shù)為1 。例如做迭代3時(shí)，可以不拿回迭代2的梯度，但必須保證迭代0、1的梯度都已拿回且用于參數(shù)更新。
總結(jié)一下， 異步很香，但對(duì)一個(gè)Worker來(lái)說(shuō)，只是等于W不變，batch的數(shù)量增加了而已，在SGD下，會(huì)減慢模型的整體收斂速度。 異步的整體思想是，比起讓W(xué)orker閑著，倒不如讓它多吃點(diǎn)數(shù)據(jù)，雖然反饋延遲了，但只要它在干活在學(xué)習(xí)就行。
batch就像活，異步就像畫(huà)出去的餅，且往往不指定延遲步數(shù)，每個(gè)Worker干越來(lái)越多的活，但模型卻沒(méi)收斂取效。讀懂分布式訓(xùn)練系統(tǒng)其實(shí)也不難。
二、分布式數(shù)據(jù)并行(DDP)
受通訊負(fù)載不均的影響， DP一般用于單機(jī)多卡場(chǎng)景 。因此，DDP作為一種更通用的解決方案出現(xiàn)了，既能多機(jī)，也能單機(jī)。DDP首先要解決的就是通訊問(wèn)題：將Server上的通訊壓力均衡轉(zhuǎn)到各個(gè)Worker上。實(shí)現(xiàn)這一點(diǎn)后，可以進(jìn)一步去Server，留Worker。
前文我們說(shuō)過(guò)，聚合梯度 + 下發(fā)梯度這一輪操作，稱(chēng)為AllReduce。 接下來(lái)我們介紹目前最通用的AllReduce方法：Ring-AllReduce 。它由百度最先提出，非常有效地解決了數(shù)據(jù)并行中通訊負(fù)載不均的問(wèn)題，使得DDP得以實(shí)現(xiàn)。
2.1 Ring-AllReduce
如下圖，假設(shè)有4塊GPU，每塊GPU上的數(shù)據(jù)也對(duì)應(yīng)被切成4份。AllReduce的最終目標(biāo)，就是讓每塊GPU上的數(shù)據(jù)都變成箭頭右邊匯總的樣子。
Ring-ALLReduce則分兩大步驟實(shí)現(xiàn)該目標(biāo)： Reduce-Scatter和All-Gather 。
Reduce-Scatter
定義網(wǎng)絡(luò)拓?fù)潢P(guān)系，使得每個(gè)GPU只和其相鄰的兩塊GPU通訊 。每次發(fā)送對(duì)應(yīng)位置的數(shù)據(jù)進(jìn)行 累加 。每一次累加更新都形成一個(gè)拓?fù)洵h(huán)，因此被稱(chēng)為Ring?？吹竭@覺(jué)得困惑不要緊，我們用圖例把詳細(xì)步驟畫(huà)出來(lái)。
一次累加完畢后，藍(lán)色位置的數(shù)據(jù)塊被更新，被更新的數(shù)據(jù)塊將成為下一次更新的起點(diǎn)，繼續(xù)做累加操作。
3次更新之后，每塊GPU上都有一塊數(shù)據(jù)擁有了對(duì)應(yīng)位置完整的聚合（圖中紅色）。此時(shí)，Reduce-Scatter階段結(jié)束。進(jìn)入All-Gather階段。目標(biāo)是把紅色塊的數(shù)據(jù)廣播到其余GPU對(duì)應(yīng)的位置上。
All-Gather
如名字里Gather所述的一樣，這操作里依然按照“相鄰GPU對(duì)應(yīng)位置進(jìn)行通訊”的原則，但對(duì)應(yīng)位置數(shù)據(jù)不再做相加，而是直接替換。All-Gather以紅色塊作為起點(diǎn)。
以此類(lèi)推，同樣經(jīng)過(guò) 3輪迭代后 ，使得每塊GPU上都匯總到了完整的數(shù)據(jù)，變成如下形式：
建議讀者們手動(dòng)推一次，加深理解(注：最后一圖箭頭畫(huà)錯(cuò)，公眾號(hào)不許修改
2.2 Ring-AllReduce通訊量分析
假設(shè)模型參數(shù)W的大小為，GPU個(gè)數(shù)為。則梯度大小也為，每個(gè)梯度塊的大小為。
對(duì)單卡GPU來(lái)說(shuō)：
- Reduce-Scatter階段，通訊量為
- All-Gather階段，通訊量為
總通訊量為，隨著N的增大，可以近似
而對(duì)前文的DP來(lái)說(shuō)，它的Server承載的總通訊量也是。 雖然通訊量相同，但搬運(yùn)相同數(shù)據(jù)量的時(shí)間卻不一定相同 。DDP把通訊量均衡負(fù)載到了每一時(shí)刻的每個(gè)Worker上，而DP僅讓Server做勤勞的搬運(yùn)工。當(dāng)越來(lái)越多的GPU分布在距離較遠(yuǎn)的機(jī)器上時(shí)，DP的通訊時(shí)間是會(huì)增加的。
但這并不說(shuō)明參數(shù)服務(wù)器不能打（有很多文章將參數(shù)服務(wù)器當(dāng)作old dinosaur來(lái)看）。事實(shí)上，參數(shù)服務(wù)器也提供了多Server方法，如下圖：
在多Server的模式下，進(jìn)一步，每個(gè)Server可以只負(fù)責(zé)維護(hù)和更新某一塊梯度（也可以某塊梯度+參數(shù)一起維護(hù)），此時(shí)雖然每個(gè)Server仍然需要和所有Worker通訊，但它的帶寬壓力會(huì)小非常多。經(jīng)過(guò)調(diào)整設(shè)計(jì)后，依然可以用來(lái)做DDP。雖然這篇文章是用遞進(jìn)式的方式來(lái)介紹兩者，但不代表兩者間一定要決出優(yōu)劣。 我想表達(dá)的觀點(diǎn)是，方法是多樣性的 。對(duì)參數(shù)服務(wù)器有興趣的朋友，可以閱讀參考的第1個(gè)鏈接。
最后，請(qǐng)大家記住Ring-AllReduce的方法，因?yàn)樵谥蟮腪eRO，Megatron-LM中，它將頻繁地出現(xiàn)，是分布式訓(xùn)練系統(tǒng)中重要的算子。
三、總結(jié)
1、在DP中，每個(gè)GPU上都拷貝一份完整的模型，每個(gè)GPU上處理batch的一部分?jǐn)?shù)據(jù)，所有GPU算出來(lái)的梯度進(jìn)行累加后，再傳回各GPU用于更新參數(shù)。
2、DP多采用參數(shù)服務(wù)器這一編程框架，一般由若個(gè)計(jì)算Worker和1個(gè)梯度聚合Server組成。Server與每個(gè)Worker通訊，Worker間并不通訊。因此Server承擔(dān)了系統(tǒng)所有的通訊壓力?；诖薉P常用于單機(jī)多卡場(chǎng)景。
3、異步梯度更新是提升計(jì)算通訊比的一種方法，延遲更新的步數(shù)大小決定了模型的收斂速度。
4、Ring-AllReduce通過(guò)定義網(wǎng)絡(luò)環(huán)拓?fù)涞姆绞?，將通訊壓力均衡地分到每個(gè)GPU上，使得跨機(jī)器的數(shù)據(jù)并行（DDP）得以高效實(shí)現(xiàn)。
5、DP和DDP的總通訊量相同，但因負(fù)載不均的原因，DP需要耗費(fèi)更多的時(shí)間搬運(yùn)數(shù)據(jù)。
審核編輯：郭婷
工商網(wǎng)監(jiān)
評(píng)論