【NLP】華為推出了全球最大的中文語(yǔ)言模型-盤古

作者：金雪鋒
鏈接：https://www.zhihu.com/question/456443707/answer/1856014437

這次HDC大會(huì)，華為聯(lián)合鵬城發(fā)布了兩個(gè)千億的NLP大模型（其中一個(gè)模型還和循環(huán)智能進(jìn)行了合作），都是盤古命名的。

一個(gè)是4月25日發(fā)布的Transformer encoder-decoder的中文理解模型；另外一個(gè)4月26日發(fā)布的是Transformer decoder的中文預(yù)訓(xùn)練模型。這兩個(gè)模型是不同的定位，也是分別訓(xùn)練出來(lái)的。其中第二個(gè)預(yù)訓(xùn)練模型現(xiàn)在已經(jīng)開(kāi)源出來(lái)了。

這種超大的模型考驗(yàn)的是全棧和全流程的能力：

1、豐富的數(shù)據(jù)集和良好的數(shù)據(jù)清洗

2、硬件系統(tǒng)：芯片、互聯(lián)、整機(jī)

3、硬件使能：高性能的算子庫(kù)和融合能力

4、AI框架：分布式并行、編譯優(yōu)化

5、AI使能平臺(tái)：資源管理和調(diào)度

6、優(yōu)秀的算法和模型設(shè)計(jì)

7、全面的系統(tǒng)工程：包括精度調(diào)優(yōu)、性能調(diào)優(yōu)、模型切分、集群可靠性等等

只有這些能力都達(dá)到一流/沒(méi)有短板，你才能訓(xùn)練出這樣的大模型，否則即便給你K級(jí)集群，你也用不起來(lái)。所以這兩個(gè)千億模型訓(xùn)練出來(lái)是鵬城、華為EI/智能計(jì)算/海思/諾亞實(shí)驗(yàn)室/中央軟件院等通力協(xié)作的結(jié)果。

MindSpore作為AI框架有幸同時(shí)支撐了這兩個(gè)千億模型的訓(xùn)練，過(guò)程其實(shí)非常艱辛，但總算熬過(guò)來(lái)。下面就簡(jiǎn)單介紹一下作為AI框架支撐千億模型的挑戰(zhàn)和解決方案：

挑戰(zhàn)： 

千億參數(shù)，

TB級(jí)顯存的模型以盤古2000億（基于Transformer decoder結(jié)構(gòu)的預(yù)訓(xùn)練模型）為例，如果我們訓(xùn)練時(shí)權(quán)重都用標(biāo)準(zhǔn)的 FP32 數(shù)據(jù)格式，那么算下來(lái)，權(quán)重占的空間就達(dá)到了 750GB，訓(xùn)練過(guò)程中內(nèi)存開(kāi)銷還會(huì)數(shù)倍上升。這 750GB 參數(shù)，不是放在硬盤上，也不是加載到內(nèi)存中，而是需要移到昇騰基礎(chǔ)硬件平臺(tái) HBM（High Bandwidth Memory高帶寬存儲(chǔ)器）內(nèi)存中，以利用昇騰基礎(chǔ)軟硬件平臺(tái)進(jìn)行模型訓(xùn)練。

模型大 ，意味著數(shù)據(jù)也大，而且都需要是高質(zhì)量數(shù)據(jù)。為了滿足數(shù)據(jù)需求，研發(fā)團(tuán)隊(duì)從互聯(lián)網(wǎng)爬取了80TB文本，并最后清洗為1TB的中文數(shù)據(jù)集。

這樣的模型與數(shù)據(jù)，已經(jīng)不是我們幾臺(tái)服務(wù)器能加載上的了，更不用說(shuō)進(jìn)行訓(xùn)練。好在研發(fā)團(tuán)隊(duì)會(huì)提供 API，一般算法工程師直接調(diào)用接口就能試試效果。
圖注：盤古千億模型架構(gòu)。圖注：盤古 2 千億模型架構(gòu)。

超大規(guī)模并行

如果給你足夠的計(jì)算力，你能想到如何訓(xùn)練這么大的模型嗎？我們最常用的分布式訓(xùn)練方式數(shù)據(jù)并行，單獨(dú)這么做肯定是不行的，因?yàn)闆](méi)有哪個(gè) 計(jì)算硬件能放下 800GB 的參數(shù)。那么再加上模型并行呢？又產(chǎn)生了新問(wèn)題，我們?cè)撊绾尾鸱秩绱司薮蟮摹副P古」？硬件產(chǎn)品（如GPU等）之間的梯度流、數(shù)據(jù)流通信又是什么樣的？

顯然訓(xùn)練如此龐大的模型，遠(yuǎn)比我們想象中的復(fù)雜，需要大量的工程化操作，并保證這些操作不會(huì)或極少影響到模型最終收斂效果。

如果手動(dòng)來(lái)寫分布式訓(xùn)練邏輯，那么需要綜合考慮計(jì)算量與類型、集群帶寬、拓?fù)浣Y(jié)構(gòu)、樣本數(shù)量等等一大堆復(fù)雜的東西，然后再設(shè)計(jì)出性能比較優(yōu)秀的并行切分策略，并編寫大量并行切分和節(jié)點(diǎn)間的通信代碼。如果系統(tǒng)環(huán)境變了，還要重新設(shè)計(jì)并修改算法，想想就覺(jué)得頭大。

MindSpore的解決方案

5 大維度的并行能力

MindSpore提供了5維的并行方式：數(shù)據(jù)并行、算子級(jí)模型并行、Pipeline模型并行、優(yōu)化器模型并行和重計(jì)算，并且在圖編譯階段，有機(jī)融合了5個(gè)維度的并行。這5維并行方式組合起來(lái)構(gòu)成了盤古的并行策略。
1、數(shù)據(jù)并行

數(shù)據(jù)并行是最基本，應(yīng)用最廣的并行方式，其將訓(xùn)練數(shù)據(jù)（mini-batch）切分，每臺(tái)設(shè)備取得其中一份；每臺(tái)設(shè)備擁有完整的模型。在訓(xùn)練時(shí)，每臺(tái)設(shè)備經(jīng)過(guò)梯度計(jì)算后，需要經(jīng)過(guò)設(shè)備間的梯度同步，然后才能進(jìn)行模型參數(shù)的更新。

2、算子級(jí)模型并行

算子級(jí)模型并行是對(duì)模型網(wǎng)絡(luò)中的每個(gè)算子涉及到的張量進(jìn)行切分。MindSpore對(duì)每個(gè)算子都獨(dú)立建模，每個(gè)算子可以擁有不同的切分策略。

以矩陣乘算子MatMul(x, w)為例，x是訓(xùn)練數(shù)據(jù)，w是模型參數(shù)，兩者都是二維矩陣。并行策略((4, 1), (1, 1))表示將x按行切4份，保持w不切，如果一共有4臺(tái)設(shè)備，那么每臺(tái)設(shè)備擁有一份x的切片，和完整的w。

3、Pipeline 模型并行

Pipeline模型并行將模型的按層分成多個(gè)stage，再把各個(gè)sage映射到多臺(tái)設(shè)備上。為了提高設(shè)備資源的利用率，又將mini-batch劃分成多個(gè)micro-batch,這樣就能夠使得不同設(shè)備在同一時(shí)刻處理不同micro-batch的數(shù)據(jù)。

一種Pipeline并行方式(Gpipe) 要求反向計(jì)算要等所有設(shè)備的正向計(jì)算完成后才開(kāi)始，而反向計(jì)算可能依賴于正向的輸出，導(dǎo)致每個(gè)卡正向計(jì)算過(guò)程中累積的activation內(nèi)存與micro-batch數(shù)量成正比，從而限制了micro-batch的數(shù)量。MindSpore的Pipeline并行中，將反向提前，每個(gè)micro-batch計(jì)算完成后，就開(kāi)始計(jì)算反向，有效降低activation存儲(chǔ)時(shí)間，從而提升整體并行效率。4、優(yōu)化器模型并行

優(yōu)化器模型并行將優(yōu)化器涉及到的參數(shù)和梯度切分到多臺(tái)設(shè)備上。以Adam優(yōu)化器為例，其內(nèi)部可能有多份與權(quán)重同等大小的“動(dòng)量”需要參與計(jì)算。在數(shù)據(jù)并行的情況下，每個(gè)卡都擁有完整的“動(dòng)量”，它們?cè)诿總€(gè)卡上都重復(fù)計(jì)算，造成了內(nèi)存及計(jì)算的浪費(fèi)。通過(guò)引入優(yōu)化器并行，每個(gè)卡只保存權(quán)重及“動(dòng)量”的切片，能降低每個(gè)卡的靜態(tài)內(nèi)存及提升計(jì)算效率。
5、重計(jì)算

重計(jì)算(Rematerialization)針對(duì)正向算子的輸出累計(jì)保存在內(nèi)存中，導(dǎo)致內(nèi)存峰值過(guò)大的問(wèn)題，舍棄了部分正向算子的輸出，而是在反向階段用到時(shí)再重新計(jì)算一遍。這樣做有效地降低了訓(xùn)練過(guò)程中的內(nèi)存使用峰值。如下圖所示，第一個(gè)內(nèi)存峰值通過(guò)重計(jì)算消除，第二個(gè)內(nèi)存峰值可以通過(guò)前面講到的優(yōu)化器并行消除。

有了這5維的并行維度后，如何將其組合起來(lái)作用于盤古，并且如何將切分后的模型分片分配到每臺(tái)設(shè)備上仍然是難題。MindSpore自動(dòng)并行，把這5個(gè)維度并行有機(jī)組合起來(lái)，可以實(shí)現(xiàn)非常高效的大模型分布式訓(xùn)練能力。

下圖(b)是一典型的樹(shù)形的硬件拓?fù)浣Y(jié)構(gòu)，其帶寬隨著樹(shù)深度的增加而降低，并且會(huì)產(chǎn)生一些流量沖突。為了利用此特征，MindSpore的目標(biāo)是最大化計(jì)算通信比，將通信量大的并行方式（算子級(jí)并行）放置在服務(wù)器內(nèi)部的多卡之間；將通信量較?。≒ipeline并行）的放置在同一機(jī)架內(nèi)的服務(wù)器間；將數(shù)據(jù)并行（優(yōu)化器并行）的部分放置在不同機(jī)架間，因?yàn)樵撏ㄐ趴梢院陀?jì)算同時(shí)執(zhí)行(overlap)，對(duì)帶寬要求較低。
在盤古2000億模型中，MindSpore將64層（layer）劃分為16個(gè)stage，每個(gè)stage包含4層。在每層中，利用算子級(jí)并行的方式對(duì)張量進(jìn)行切分。

如下圖中的Q,K,V的參數(shù)在實(shí)際中（按列）被切了8份，輸入張量（按行）被切了16份，輸出張量因此被切了128份（8*16）。重計(jì)算配置是配置在每層內(nèi)的，也就是重計(jì)算引入的多余的計(jì)算量不會(huì)超過(guò)一層的計(jì)算量。總計(jì)，MindSpore使用了2048塊昇騰處理器來(lái)訓(xùn)練盤古。
MindSpore對(duì)外屏蔽了復(fù)雜并行實(shí)現(xiàn)的細(xì)節(jié)，使得用戶像編寫單機(jī)模型腳本那樣簡(jiǎn)單。用戶在單機(jī)腳本的基礎(chǔ)上，僅通過(guò)少了配置就能實(shí)現(xiàn)多維度的混合并行。下圖是簡(jiǎn)化版的盤古腳本，其中紅色加粗字體表示的在MindSpore中的并行策略。將紅色加粗字體去掉，則是單機(jī)腳本。
圖算跨層聯(lián)合優(yōu)化，發(fā)揮硬件極致性能

除了跨節(jié)點(diǎn)間的大規(guī)模自動(dòng)外，在單卡節(jié)點(diǎn)內(nèi)，MindSpore通過(guò)圖層和算子層的跨層協(xié)同優(yōu)化，來(lái)進(jìn)一步發(fā)揮昇騰算力。

在傳統(tǒng)的NN網(wǎng)絡(luò)中，不同算子承載的計(jì)算量和計(jì)算復(fù)雜度也各不相同。如LayerNorm由11個(gè)基本算子組成，而Add則只有1個(gè)基本算子。這種基于用戶角度的算子定義，通常是無(wú)法充分發(fā)揮硬件資源計(jì)算能力的。因?yàn)橛?jì)算量過(guò)大、過(guò)復(fù)雜的算子，通常很難生成切分較好的高性能算子。從而降低設(shè)備利用率；而計(jì)算量過(guò)小的算子，由于計(jì)算無(wú)法有效隱藏?cái)?shù)據(jù)搬移開(kāi)銷，也可能會(huì)造成計(jì)算的空等時(shí)延，從而降低設(shè)備利用率。

為了提升硬件利用率，MindSpore使用了圖算融合優(yōu)化技術(shù)，通過(guò)圖層和算子層聯(lián)合優(yōu)化，并將“用戶使用角度的易用性算子”進(jìn)行重組融合，然后轉(zhuǎn)換為“硬件執(zhí)行角度的高性能算子”，從而充分提升硬件資源利用率，進(jìn)而提升整網(wǎng)執(zhí)行性能。具體優(yōu)化流程如下圖所示：
以LayerNorm算子為例，通過(guò)算子拆分和重組，11個(gè)小算子，組成了1個(gè)單算子和2個(gè)融合算子。這些重組后的算子可以生成更加高性能的算子，從而大大降低了整體網(wǎng)絡(luò)運(yùn)行時(shí)間。
在盤古模型中，圖算融合幫助整體訓(xùn)練時(shí)間減少了20%以上。除此之外，對(duì)于其它 NLP、CV等任務(wù)，圖算融合在優(yōu)化性能方面都有不錯(cuò)的表現(xiàn)。
總結(jié)一下

即使給我們足夠的算力，超大模型的訓(xùn)練還是異常復(fù)雜，遠(yuǎn)比想象中的困難。對(duì)于我們一般算法工程師來(lái)說(shuō)，針對(duì)某個(gè)任務(wù)，上億參數(shù)量已經(jīng)算大的了，但是并不會(huì)感到訓(xùn)練上會(huì)有什么困難，因?yàn)楦鱾€(gè)深度學(xué)習(xí)框架直接調(diào)用數(shù)據(jù)并行接口就能搞定。

但是如果模型繼續(xù)增大到百億級(jí)、千億級(jí)甚至萬(wàn)億級(jí)，并行與優(yōu)化策略的復(fù)雜度猛然上升，算法工程師一點(diǎn)點(diǎn)地編寫與優(yōu)化代碼可太難了。MindSpore通過(guò)編譯器并行優(yōu)化能力，把計(jì)算邏輯和并行邏輯解耦，單卡串行代碼自動(dòng)實(shí)現(xiàn)分布式并行，從而使得算法科學(xué)家將精力都解放到模型本身上。

為了從預(yù)訓(xùn)練獲取更多的知識(shí)， GPT-3 與盤古 這樣的模型會(huì)越來(lái)越大，畢竟到現(xiàn)在我們還沒(méi)看到大模型預(yù)訓(xùn)練效果的極限在哪。屆時(shí)，這類模型對(duì)基礎(chǔ)設(shè)施的需求會(huì)更大，并行與優(yōu)化策略也會(huì)更加復(fù)雜。預(yù)訓(xùn)練的效果，看的也已經(jīng)不是模型本身，而是基礎(chǔ)設(shè)施構(gòu)建得是否足夠優(yōu)秀。大規(guī)模計(jì)算集群及軟硬件協(xié)同優(yōu)化，這次在盤古 的訓(xùn)練上得到了充分的完美體現(xiàn)。

當(dāng)然，也誠(chéng)如以上所言，盤古只是對(duì)超大規(guī)模分布式訓(xùn)練、超大規(guī)模中文預(yù)訓(xùn)練模型的一次探索，未來(lái)還需要更多的研究工作者投入到通用智能與大規(guī)模分布式計(jì)算的研究工作中。