大模型工業(yè)化的方法論，都藏在 GPU 里

本文轉(zhuǎn)載自DataFunTalk

編輯推薦

大模型已經(jīng)是人工智能的基礎(chǔ)設(shè)施了，如何進(jìn)一步普惠商用并盈利是當(dāng)務(wù)之急。近日，NVIDIA 邀請(qǐng)百度、字節(jié)跳動(dòng)火山翻譯、騰訊微信等 6 位專(zhuān)家，圍繞“AI 大模型技術(shù)路線和工業(yè)化落地實(shí)踐”展開(kāi)分享，以下為分享內(nèi)容精華匯總。

數(shù)據(jù)智能時(shí)代，計(jì)算是剛需，也是痛點(diǎn)，最主要的特點(diǎn)就是一個(gè)字——大。

將“大”拆分成三個(gè)特點(diǎn)，則包括：

• 數(shù)據(jù)大批量

• 計(jì)算串行依賴(lài)

• 計(jì)算復(fù)雜性高

這囊括了數(shù)據(jù)和算法的復(fù)雜性，而數(shù)據(jù)、算法、算力是智能時(shí)代三要素，前兩者的復(fù)雜性最后都要由算力來(lái)承載。

這使得業(yè)界對(duì)算力的需求在空間和時(shí)間上都極速膨脹，有如海嘯之勢(shì)。

GPU 卻有辦法將海嘯在空間上細(xì)分成萬(wàn)千條涓涓細(xì)流，在時(shí)間上縮短流水的路徑，在結(jié)構(gòu)上精簡(jiǎn)路徑分支，將大規(guī)模任務(wù)層層拆解成小規(guī)模任務(wù)，舉重若輕般承載了海量算力需求，成為智能時(shí)代的算力基礎(chǔ)。

針對(duì)上述三個(gè)特點(diǎn)，GPU 基于吞吐量、顯存等指標(biāo)，分別采用并行、融合、簡(jiǎn)化三種方法，在算子層面進(jìn)行加速。

GPU 加速的主要方法論，亦適用于大模型的工業(yè)化。

伴隨底層芯片、算力、數(shù)據(jù)等基礎(chǔ)設(shè)施的完善和進(jìn)步，全球 AI 產(chǎn)業(yè)正逐步從運(yùn)算智能走向感知智能、認(rèn)知智能，并相應(yīng)形成“芯片、算力設(shè)施、AI 框架&算法模型、應(yīng)用場(chǎng)景”的產(chǎn)業(yè)分工、協(xié)作體系。2019 年以來(lái)，AI 大模型帶來(lái)問(wèn)題泛化求解能力大幅提升，“大模型+小模型”逐步成為產(chǎn)業(yè)主流技術(shù)路線，驅(qū)動(dòng)全球 AI 產(chǎn)業(yè)發(fā)展全面加速。

從 NVIDIA、百度、字節(jié)跳動(dòng)火山翻譯、騰訊微信等幾位專(zhuān)家的分享，我們也不難發(fā)現(xiàn)，他們?cè)跇I(yè)界落地大模型的時(shí)候，很大程度上都采用了并行、融合與簡(jiǎn)化的方法，并且還從訓(xùn)練、推理層面延伸到了算法建模層面。

01

并行

并行方法是一種空間換時(shí)間的方法，它將海嘯細(xì)分成涓涓細(xì)流。具體而言，對(duì)于數(shù)據(jù)大批量的計(jì)算，每個(gè)計(jì)算步驟的耗時(shí)都比較長(zhǎng)。GPU 利用并行計(jì)算，也就是把沒(méi)有計(jì)算依賴(lài)的數(shù)據(jù)盡可能并行，把大批量拆分成小批量，以減少每個(gè)計(jì)算步驟的 GPU 空閑等待時(shí)間，提高計(jì)算吞吐量。

為了實(shí)際完成大模型的訓(xùn)練，就需要高效率的軟件框架，使得訓(xùn)練在單個(gè) GPU、單個(gè)節(jié)點(diǎn)甚至在大規(guī)模集群上都有很好的計(jì)算效率。

因此，NVIDIA 開(kāi)發(fā)出了 Megatron 訓(xùn)練框架。

Megatron 采用了模型并行、Sequence 并行等優(yōu)化手段以高效地訓(xùn)練 Transformer 大模型，可訓(xùn)練萬(wàn)億級(jí)別參數(shù)量的模型。

在算法建模層面，火山翻譯和百度主要對(duì) MoE 模型等建模方法進(jìn)行了探索。

1. 模型并行

模型并行可分為 Pipeline 并行和 Tensor 并行。

Pipeline 并行也就是層間并行（圖中上半部分），將不同的層劃分到不同的 GPU 進(jìn)行計(jì)算。這種模式的通信只發(fā)生在層的邊界，通信次數(shù)和通信數(shù)據(jù)量較少，但會(huì)引入額外的 GPU 空間等待時(shí)間。

Tensor 并行也就是層內(nèi)并行（圖中下半部分），將一個(gè)層的計(jì)算劃分到不同的 GPU 上計(jì)算。這種模式的實(shí)現(xiàn)更加容易，對(duì)于大矩陣的效果更好，更好實(shí)現(xiàn) GPU 之間的負(fù)載均衡，但通信次數(shù)和數(shù)據(jù)量都比較大。

為了充分利用 GPU 資源，Megatron 將每個(gè)訓(xùn)練 batch 劃分成了更小的 micro batch。

由于不同的 micro batch 之間沒(méi)有數(shù)據(jù)依賴(lài)，因此可以互相覆蓋等待時(shí)間，從而能夠提高 GPU 的利用率，進(jìn)而提高整體的訓(xùn)練性能。

Tensor 并行把每個(gè)算子的計(jì)算劃分到不同的 GPU 上，對(duì)于一個(gè)矩陣層，存在橫切和縱切兩種方式。

如圖所示，Megatron 在 Transformer block 的 attention 和 MLP 部分都引入了這兩種切分方式。

在 Tensor 并行模式下，每個(gè) Transformer 層的前向反向加起來(lái)總共需要四次 All-reduce 通信，由于 All-reduce 的通信量較大，因此 Tensor 并行更適合單卡內(nèi)部使用。

結(jié)合 Pipeline 并行和 Tensor 并行，Megatron 可以將在 32 個(gè) GPU 上訓(xùn)練 1700 億參數(shù)模型，擴(kuò)展到在 3072 個(gè) GPU 上訓(xùn)練 1 萬(wàn)億參數(shù)規(guī)模的模型。

2. Sequence 并行

Tensor 并行其實(shí)并沒(méi)有對(duì) Layer-norm 以及 Dropout 做拆分，因此這兩個(gè)算子在每個(gè) GPU 之間是復(fù)制的。

然而，這些操作本身不需要大量計(jì)算，卻非常占用激活顯存。

為此，Megatron 又提出了 Sequence 并行的優(yōu)化方法。Sequence 并行的好處在于不會(huì)增加通信量，并且可以大大減少顯存占用

由于 Layer-norm 和 Dropout 沿著序列的維度是獨(dú)立的，因此可以按照 Sequence 維度進(jìn)行拆分。

使用了 Sequence 并行之后，對(duì)于超大規(guī)模的模型而言，其實(shí)顯存占用量還是很大的。因此，Megatron 又引入了激活重計(jì)算技術(shù)。

Megatron 的做法是，找到一些計(jì)算量很少但顯存占用很大的算子，比如 Attention 里的 Softmax、Dropout 等算子，對(duì)這些算子進(jìn)行激活重計(jì)算就可以顯著減少顯存，并且計(jì)算開(kāi)銷(xiāo)增加不大。

Sequence 并行和選擇性激活重計(jì)算的結(jié)合可以將顯存占用降低為原來(lái)的 1/5 左右。相對(duì)于原本直接將所有激活進(jìn)行重計(jì)算的方案，其顯存也只有其兩倍，同時(shí)計(jì)算開(kāi)銷(xiāo)顯著降低，并且隨著模型規(guī)模增大，計(jì)算開(kāi)銷(xiāo)的占比也會(huì)逐漸降低。到萬(wàn)億規(guī)模模型的時(shí)候，重計(jì)算的開(kāi)銷(xiāo)只占整體的 2% 左右。

3. 算法并行

MoE 模型因其設(shè)計(jì)思想簡(jiǎn)潔、可擴(kuò)展性強(qiáng)等特點(diǎn)，使其在業(yè)界得到了越來(lái)越多的關(guān)注。

MoE 模型提出了這樣的設(shè)計(jì)思想，也就是將大模型拆分成多個(gè)小模型。每個(gè)樣本只需要激活部分專(zhuān)家模型進(jìn)行計(jì)算，從而大大節(jié)省計(jì)算資源。

目前最常用的稠密大模型是 BERT、T5、GPT-3，最常用的稀疏 MoE 模型是 T5+MoE，MoE 正成為大模型構(gòu)建的趨勢(shì)。

可以說(shuō)，MoE 在算法建模層面，結(jié)合了并行計(jì)算的思想。

大模型的通用性，體現(xiàn)在多個(gè)方面，除了我們已經(jīng)熟知的幾點(diǎn)，比如注意力機(jī)制歸納偏置更弱，模型容量大，模型數(shù)據(jù)大等等，還可以在任務(wù)建模方式上做優(yōu)化，MoE 就是典型的代表。

對(duì)于火山翻譯而言，MoE 的基本思路是通過(guò)寬度換取深度，因?yàn)槟Ｐ蜕疃仍缴睿?jì)算層數(shù)越多，進(jìn)而推理時(shí)間越長(zhǎng)。

比如，對(duì)于擁有 4 層 Encoder、4 層 Decoder 的 Transformer 模型，每次計(jì)算必須經(jīng)過(guò)所有 8 個(gè) FFN 的計(jì)算。如果是混合專(zhuān)家模型，則可以把 FFN 平行放置，最終把計(jì)算路徑減半，因而推理時(shí)間也減半。

而在相同推理時(shí)間下，也就是模型深度相近的時(shí)候，由于 MoE 可以增加模型寬度，在機(jī)器翻譯的最終效果上也有所提升。

針對(duì) 24 種非洲語(yǔ)言和英、法語(yǔ)言的多語(yǔ)言翻譯任務(wù)，火山翻譯開(kāi)發(fā)出了擁有 128 層 Transformer、24 個(gè)專(zhuān)家層的 MoE 模型，相比傳統(tǒng)架構(gòu)實(shí)現(xiàn)了更好的翻譯效果。

但 Sparse MoE 中的“專(zhuān)家模型”可能有些名不副實(shí)，因?yàn)閷?duì)于一句話，比如其每個(gè) Token 經(jīng)過(guò)的專(zhuān)家都有可能是不同的。

火山翻譯因此開(kāi)發(fā)了 Hard Gate MoE，使得句子經(jīng)過(guò)的專(zhuān)家由語(yǔ)種確定，這使得模型結(jié)構(gòu)更加簡(jiǎn)單，實(shí)驗(yàn)結(jié)果也表明其翻譯效果更好。

在算法建模的并行化探索中，百度也在知識(shí)增強(qiáng)跨模態(tài)生成大模型 ERNIE-ViLG 2.0 中采用了混合專(zhuān)家擴(kuò)散模型框架。

為何要對(duì)擴(kuò)散模型采用專(zhuān)家模型？

其實(shí)是因?yàn)樵诓煌纳呻A段，模型建模要求的不同。比如在初始階段，模型著重學(xué)習(xí)從高斯噪聲中生成有語(yǔ)義的圖像，在最后階段，模型著重從含噪圖像中恢復(fù)圖像細(xì)節(jié)。

實(shí)際上，在 ERNIE 3.0 的早期版本中就融合了自編碼和自回歸，其可在通用的語(yǔ)義表示上，針對(duì)具體的生成任務(wù)和理解任務(wù)，結(jié)合兩種建模方式。

融合自編碼和自回歸的基本思想其實(shí)與專(zhuān)家模型的建模方法論類(lèi)似。

具體來(lái)說(shuō)，是在通用表示的基礎(chǔ)上，根據(jù)理解任務(wù)適合自編碼網(wǎng)絡(luò)結(jié)構(gòu)，生成任務(wù)適合自回歸網(wǎng)絡(luò)結(jié)構(gòu)，來(lái)進(jìn)行建模。此外，這種建模方式通常還能學(xué)習(xí)到更好的通用表示。

此外，在 ERNIE-UniX2 模型中，百度通過(guò)將對(duì)比學(xué)習(xí)、語(yǔ)言模型等預(yù)訓(xùn)練范式進(jìn)行融合，將多語(yǔ)言、多模態(tài)的理解和生成任務(wù)進(jìn)行了統(tǒng)一。

訓(xùn)練完 MoE 模型后，推理部署也是非常重視效率的環(huán)節(jié)。

在進(jìn)行超大規(guī)模模型推理部署方案選擇的時(shí)候，首先會(huì)根據(jù)模型的參數(shù)規(guī)模、模型結(jié)構(gòu)、GPU 顯存和推理框架，以及對(duì)模型精度和推理性能的權(quán)衡來(lái)決定是使用單卡推理還是多卡推理。如果顯存不足，則會(huì)考慮模型壓縮或多卡推理的方案。

多卡推理包括 Tensor 并行、Pipeline 并行、Expert 并行等模式。

對(duì) MoE 超大模型采用不同模式會(huì)遇到不同的挑戰(zhàn)。其中，MoE 模型的 Tensor 并行和稠密模型類(lèi)似。

如果選擇 Expert 并行模式，每個(gè) MoE Layer 的 Expert 就會(huì)被劃分到不同的 GPU 上，這可能帶來(lái)負(fù)載均衡問(wèn)題，從而導(dǎo)致大量的 GPU 是空閑的，最終使得整體吞吐量不高。這是 MoE 多卡推理中需要關(guān)注的重點(diǎn)。

對(duì)于 Tensor 并行和 Pipeline 并行，除了通過(guò)微調(diào)減少卡間通訊以外，更直接的方法是提升卡間帶寬。而當(dāng)對(duì) MoE 模型使用 Expert 并行導(dǎo)致負(fù)載均衡問(wèn)題的時(shí)候，可以通過(guò) Profiling 分析優(yōu)化。

多卡推理方案增加了通信開(kāi)銷(xiāo)，對(duì)模型推理延時(shí)有一定影響。

02

融合

融合是解決并行計(jì)算中遇到的天然矛盾的方法，并行計(jì)算和串行計(jì)算是兩種基本的計(jì)算模式。而在應(yīng)用并行計(jì)算的時(shí)候，最典型的困難就是大量的串行依賴(lài)，以及因此產(chǎn)生的中間值顯存占用問(wèn)題，而 GPU 顯存通常會(huì)成為大模型訓(xùn)練和推理的硬件性能瓶頸之一。

對(duì)于海量計(jì)算串行依賴(lài)問(wèn)題，最主要的方法是將細(xì)流的路徑縮短，也就是減少中間停留過(guò)程。具體來(lái)說(shuō)，就是利用算子融合，把次序存在先后依賴(lài)關(guān)系的算子進(jìn)行合并，以減少顯存占用。

算子融合不僅在計(jì)算層面，也可以在算子設(shè)計(jì)層面實(shí)現(xiàn)。

1. 1F1B

Pipeline 并行中如果將前向和反向過(guò)程分開(kāi)，就會(huì)出現(xiàn)顯存占用過(guò)多問(wèn)題。

因此，Megatron 又提出了 Pipeline 并行的新模式 1F1B，每個(gè) GPU 以交替的方式執(zhí)行每個(gè) micro batch 的正向和反向過(guò)程，以盡早釋放其占用的顯存，進(jìn)而減少顯存占用。

1F1B 并不能減少 bubble time，為了進(jìn)一步減少 bubble time，Megatron 又提出了 interleaved 1F1B 模式。也就是原本每個(gè) GPU 負(fù)責(zé)連續(xù) 4 個(gè)層的計(jì)算，現(xiàn)在變成負(fù)責(zé)連續(xù)兩個(gè)層的計(jì)算，只有原來(lái)的一半，從而 bubble time 也變成了原來(lái)的一半。

2. Kernel 融合

當(dāng)在做 GPU 計(jì)算的時(shí)候，每個(gè)計(jì)算流程都可以封裝成一個(gè) GPU 的 Kernel，放到 GPU 上執(zhí)行，并且是順序性的。傳統(tǒng)的算子庫(kù)為了通用性，會(huì)把算子設(shè)計(jì)的非?；?，因此數(shù)量也非常多，帶來(lái)的弊端是顯存占用多，因?yàn)樾枰鎯?chǔ)大量的中間隱藏表示，另外這對(duì)帶寬的要求也比較高，最終可能造成延遲或者性能損失。

火山翻譯基于 CuBLAS 乘法接口將其他非矩陣乘法算子進(jìn)行了融合，包括了 Softmax、LayerNorm 等。

除了比較通用算子的融合，火山翻譯還針對(duì)一些特定算子比如 Beam Search 無(wú)法很好利用 GPU 并行性的特點(diǎn)，優(yōu)化其計(jì)算依賴(lài)問(wèn)題，從而實(shí)現(xiàn)加速。

在四種主流 Transformer 模型上，LightSeq 算子融合在 PyTorch 的基礎(chǔ)上取得了最高 8 倍的加速。

03

簡(jiǎn)化

簡(jiǎn)化是一種比較簡(jiǎn)單直觀的加速方式，在細(xì)微處將流水分支精簡(jiǎn)。具體而言，就是對(duì)于計(jì)算高復(fù)雜性，在保證性能的前提下將算子復(fù)雜度簡(jiǎn)化，最終減少計(jì)算量。

超大規(guī)模模型的單卡推理一般會(huì)涉及模型壓縮。

常見(jiàn)的模型壓縮方案是量化、蒸餾和剪枝。量化是業(yè)內(nèi)最常用的模型壓縮方案之一。雖然量化的計(jì)算采用了更低的精度，但可以保持模型的參數(shù)量級(jí)，在某些情況下或許能更好地保證模型整體的精度。

1. 量化

目前有兩種量化方法，一種是訓(xùn)練后量化，一種是量化感知訓(xùn)練。后者通常比前者對(duì)模型的精度保持更好。

完成量化后，可以通過(guò) TensorRT 或 FasterTransformer 等推理加速框架，進(jìn)一步加速超大模型的推理。

LightSeq 在訓(xùn)練過(guò)程的量化中采用了真 int8 量化，也就是在矩陣乘法之前，會(huì)執(zhí)行量化操作，并且在矩陣乘法之后才執(zhí)行反量化操作。而不像過(guò)去的偽量化那樣，在矩陣乘法之前就執(zhí)行了量化和反量化操作，以讓模型適應(yīng)量化所帶來(lái)的損失和波動(dòng)。后者在實(shí)際計(jì)算中并不能帶來(lái)加速，反而可能增大延時(shí)，或者使得顯存占用上升。而真 int8 量化在實(shí)際應(yīng)用中也帶來(lái)了很好的加速效果。

2. 蒸餾

第二種模型壓縮方式是蒸餾。蒸餾可以針對(duì)不同應(yīng)用場(chǎng)景采用不同的策略對(duì)超大模型進(jìn)行壓縮，在某些情況下，蒸餾可以讓超大模型擁有更好的泛化能力。

3. 剪枝

最后一種模型壓縮方案是剪枝。剪枝可分為全模型剪枝和部分層剪枝，對(duì)于超大模型，了解模型關(guān)鍵層非常重要，需要避開(kāi)這些對(duì)精度影響最大的部分的剪枝，這對(duì)于稀疏 MoE 模型也是適用的。

4. 大模型工業(yè)化

大模型的研究和落地已成趨勢(shì)，預(yù)計(jì)在 2022 年，關(guān)于大規(guī)模語(yǔ)言模型和 Transformers 的論文超過(guò) 1 萬(wàn)篇，比五年前 Transformers 剛提出的時(shí)候增長(zhǎng)了 7 倍。另外，大模型也有非常廣泛的應(yīng)用，比如圖片生成、推薦系統(tǒng)、機(jī)器翻譯，甚至是生命科學(xué)、代碼生成等。

OpenAI 也曾經(jīng)在 2020 年發(fā)表過(guò)兩篇論文，就展示過(guò)一個(gè)模型的表現(xiàn)基本上和三個(gè)主要的因素掛鉤，即算力、數(shù)據(jù)集大小、模型參數(shù)量，以這三個(gè)指標(biāo)就能很好地預(yù)測(cè)模型的效果。

Richard Sutton 曾經(jīng)說(shuō)過(guò)，在過(guò)去 70 年的 AI 發(fā)展中，一個(gè)反復(fù)出現(xiàn)的趨勢(shì)是一個(gè)通用的能夠高效利用計(jì)算資源的方法，總會(huì)是最后的贏家。

根據(jù) Richard Sutton 的“贏家定律”，深度學(xué)習(xí)在過(guò)去近十年是贏在了通用性。

但如今，大模型訓(xùn)練的挑戰(zhàn)已不言而喻。以 GPT-3 為例，其在訓(xùn)練的時(shí)候如果使用原始的混合精度，需要保存訓(xùn)練時(shí)的參數(shù)和梯度以及 FP 32 的主參數(shù)，如果使用 Adam 優(yōu)化器，還要保存兩個(gè)優(yōu)化器的動(dòng)量信息，則最終總共需要 2.8 個(gè) TB 的顯存，這遠(yuǎn)遠(yuǎn)超出了單卡的顯存容量，需要超過(guò) 35 張 A100 才能承載。

NVIDIA 2021 年的論文“Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM”中得出一個(gè)經(jīng)驗(yàn)公式，表明單次迭代參數(shù)量為 1750 億的 GPT-3 模型就需要 4.5 億 FLOPs 的算力。如果整個(gè)訓(xùn)練周期包含 95000 次迭代的話，就需要 430 ZettaFLOPs。換句話說(shuō)，需要一塊 A100 訓(xùn)練 16000 天，這還是不考慮計(jì)算效率的結(jié)論。

也就是說(shuō)，光靠堆積這三個(gè)指標(biāo)在大模型工業(yè)化時(shí)代將極大浪費(fèi)資源。

DeepMind 在 2022 年發(fā)表的 ChinChilla 的論文中曾表示，實(shí)際上 GPT-3、OPT、PaLM 等大模型，基本都屬于欠擬合模型。如果基于同樣的計(jì)算資源，調(diào)小模型參數(shù)量，并訓(xùn)練更多步驟，最終的模型效果才能更好一些。這也是微信在 WeLM 大規(guī)模語(yǔ)言模型中遵循的設(shè)計(jì)思想。

業(yè)界各企業(yè)基本都在開(kāi)始將注意力從規(guī)模上放松，轉(zhuǎn)而關(guān)注大模型落地時(shí)的效率問(wèn)題。

比如，從整體執(zhí)行效率來(lái)看，經(jīng)過(guò) Megatron 優(yōu)化的幾乎所有模型都有 30% 的吞吐量提升，并且隨著模型大小的增加，可以實(shí)現(xiàn)更高的 GPU 利用率。在 1750 億參數(shù)的 GPT-3 模型上，GPU 利用率可以達(dá)到 52.8%。而在 5300 億參數(shù)規(guī)模以上的模型上，利用率可以達(dá)到 57%。

也就是說(shuō)，根據(jù) Richard Sutton 的“贏家定律”，效率，將成為大模型工業(yè)化的主基調(diào)。