從2018年的BERT到2020年的GPT-3，NLP語言模型經(jīng)歷了爆發(fā)式的發(fā)展過程，其中BERT模型的參數(shù)量為3.4億，而GPT-3的模型參數(shù)量達(dá)到了1750億。2021年9月，浪潮發(fā)布了“源1.0”，它是目前規(guī)模最大的中文AI單體模型，參數(shù)規(guī)模高達(dá)2457億，訓(xùn)練采用的中文數(shù)據(jù)集達(dá)5TB。“源1.0”在語言智能方面表現(xiàn)優(yōu)異，獲得中文語言理解評測基準(zhǔn)CLUE榜單的零樣本學(xué)習(xí)和小樣本學(xué)習(xí)兩類總榜冠軍。測試結(jié)果顯示，人群能夠準(zhǔn)確分辨人與“源1.0”作品差別的成功率低于50%。

海量的參數(shù)帶來了模型訓(xùn)練和部署上的巨大挑戰(zhàn)。本文將聚焦“源1.0”背后的計(jì)算挑戰(zhàn)以及我們采取的訓(xùn)練方法。

“源1.0”的模型結(jié)構(gòu)

“源1.0”是一個(gè)典型的語言模型。語言模型通俗來講就是能夠完成自然語言理解或者生成文本的神經(jīng)網(wǎng)絡(luò)模型。對于“源1.0”，我們考慮語言模型（Language Model，LM）和前綴語言模型（Prefix Language Model，PLM）兩種模型結(jié)構(gòu)。如下圖所示：

圖1 模型結(jié)構(gòu)示意圖（左圖為LM，右圖為PLM）

我們比較了130億參數(shù)的LM和PLM在不同下游任務(wù)上的結(jié)果，注意到LM在Zero-Shot和Few-Shot上表現(xiàn)更好，而PLM在微調(diào)方面表現(xiàn)出色。微調(diào)通常會在大多數(shù)任務(wù)中帶來更好的準(zhǔn)確性，然而微調(diào)會消耗大量的計(jì)算資源，這是不經(jīng)濟(jì)的。所以我們選擇LM作為“源 1.0”模型的基礎(chǔ)模型結(jié)構(gòu)。

?

如何訓(xùn)練“源1.0”

| 源1.0訓(xùn)練面對的挑戰(zhàn)

“源1.0”的訓(xùn)練需要面對的第一個(gè)挑戰(zhàn)就是數(shù)據(jù)和計(jì)算量的挑戰(zhàn)。

數(shù)據(jù)方面，如果把訓(xùn)練一個(gè)巨量模型的訓(xùn)練過程比作上異常戰(zhàn)役的話，那么數(shù)據(jù)就是我們的彈藥。數(shù)據(jù)量的多少，決定了我們可以訓(xùn)練模型的規(guī)模，以及最后的效果。針對這一方面，我們構(gòu)建了一個(gè)全新的中文語料庫，清洗后的高質(zhì)量數(shù)據(jù)規(guī)模達(dá)到了5TB，是目前規(guī)模最大的中文語料庫。

圖2 數(shù)據(jù)預(yù)處理流程圖

算力方面，根據(jù)OpenAI提出的PetaFlop/s-day衡量標(biāo)準(zhǔn)，我們可以估算“源1.0”訓(xùn)練的計(jì)算需求情況。根據(jù)Wikipedia提供的數(shù)據(jù)（https://en.wikipedia.org/wiki/OpenAI），GPT-3的計(jì)算需求約為3640 PetaFlop/s-day，而“源1.0”的計(jì)算需求達(dá)到了4095 PetaFlop/s-day。

計(jì)算資源的巨大開銷是限制研究人員研發(fā)具有數(shù)以千萬計(jì)參數(shù)的NLP大模型的瓶頸。例如GPT-3是在由10000個(gè)GPU所組成的集群上訓(xùn)練得到的。我們在設(shè)計(jì)“源1.0”的模型結(jié)構(gòu)時(shí)，考慮到了影響大規(guī)模分布式訓(xùn)練的關(guān)鍵因素，采用了專門的分布式訓(xùn)練策略，從而加速了模型的訓(xùn)練過程。

在模型訓(xùn)練時(shí)一般最常用的是采用數(shù)據(jù)并行分布式計(jì)算策略，但這只能滿足小模型的訓(xùn)練需求。對于巨量模型來說，由于其模型參數(shù)量過大，遠(yuǎn)遠(yuǎn)超過常用計(jì)算設(shè)備比如GPU卡的顯存容量，因此需要專門的算法設(shè)計(jì)來解決巨量模型訓(xùn)練的顯存占用問題，同時(shí)還需要兼顧訓(xùn)練過程中的GPU計(jì)算性能的利用率。

| “源1.0”的訓(xùn)練策略

為了解決顯存不足的問題，我們采用了張量并行、流水并行、數(shù)據(jù)并行相結(jié)合的并行策略，實(shí)現(xiàn)了在2128個(gè)GPU上部署“源1.0”，并完成了1800億tokens的訓(xùn)練。

a. 張量并行

針對單個(gè)GPU設(shè)備不能完整的承載模型訓(xùn)練，一個(gè)解決方案就是張量并行+數(shù)據(jù)并行的2D并行策略。具體來說，使用多個(gè)GPU設(shè)備為1組，比如單個(gè)服務(wù)器內(nèi)的8個(gè)GPU為1組，組內(nèi)使用張量并行策略對模型進(jìn)行拆分，組間（服務(wù)器間）采用數(shù)據(jù)并行。

對于張量并行部分，NVIDIA在Megatron-LM中提出了針對Transformer結(jié)構(gòu)的張量并行解決方案。其思路是把每一個(gè)block的參數(shù)和計(jì)算都均勻的拆分到N個(gè)GPU設(shè)備上，從而實(shí)現(xiàn)每個(gè)GPU設(shè)備都承擔(dān)這一block的參數(shù)量和計(jì)算量的1/N效果。圖3展示了對Transformer結(jié)構(gòu)中的MLP層和self-attention層進(jìn)行張量并行拆分計(jì)算的過程示意圖。

圖3 張量并行示意圖

在訓(xùn)練過程中，tensor經(jīng)過每一層的時(shí)候，計(jì)算量與通信數(shù)據(jù)量之比如下：

其中，S為輸入序列的長度，h為隱藏層的大?。╤idden size）。

b. 流水并行

圖4 流水線并行示意圖

對于具有數(shù)千億參數(shù)的語言模型，這些參數(shù)很難被存放在單個(gè)節(jié)點(diǎn)中。流水線并行將LM的層序列在多個(gè)節(jié)點(diǎn)之間進(jìn)行分割，以解決存儲空間不足的問題，如圖5所示。每個(gè)節(jié)點(diǎn)都是流水線中的一個(gè)階段，它接受前一階段的輸出并將結(jié)果過發(fā)送到下一階段。如果前一個(gè)相鄰節(jié)點(diǎn)的輸出尚未就緒，則當(dāng)前節(jié)點(diǎn)將處于空閑狀態(tài)。節(jié)點(diǎn)的空閑時(shí)間被稱為流水線氣泡（pipline bubble）。為了提高流水行并行的性能，我們必須盡可能減少在氣泡上花費(fèi)的時(shí)間。定義流水線中氣泡的理想時(shí)間占比為如下形式：

根據(jù)這一公式，流水線氣泡的耗時(shí)隨著層數(shù)L的增加而增加，隨著微批次大?。╩icro-batch-size）的增加而減小。當(dāng)m?L/l的時(shí)候，流水并行過程中的流水線氣泡對訓(xùn)練性能的影響幾乎可以忽略。

與此同時(shí)，在流水并行過程中，節(jié)點(diǎn)間的計(jì)算量與通信數(shù)據(jù)量之比為：

根據(jù)上面的公式，流水線中節(jié)點(diǎn)的計(jì)算效率與h和S呈線性關(guān)系，這與張量并行類似。

c. 數(shù)據(jù)并行

圖6 數(shù)據(jù)并行示意圖

采用數(shù)據(jù)并行時(shí)，全局批次大小（global batch size）按照流水線分組進(jìn)行分割。每個(gè)流水線組都包含模型的一個(gè)副本，數(shù)據(jù)在組內(nèi)按照局部批次規(guī)模送入模型副本。數(shù)據(jù)并行時(shí)的計(jì)算量與通信數(shù)據(jù)量的比值可用如下公式近似：

當(dāng)d? 1時(shí)，上面公式可以進(jìn)一步簡化成：

根據(jù)這一公式，我們可以看出數(shù)據(jù)并行的計(jì)算效率與全局批次大小B和序列長度S呈正比關(guān)系。由于模型對內(nèi)存的需求與S的平方成正比，與B成線性關(guān)系，因此增加全局批次大小可以更有效的提升數(shù)據(jù)并行的效率。

當(dāng)全局批次大小過大的時(shí)候，模型很容易出現(xiàn)不收斂的問題，為了保證模型訓(xùn)練過程的穩(wěn)定性，我們將全局批次大小限制在了10^7個(gè)token內(nèi)。

根據(jù)以上的理論分析，我們確定了設(shè)計(jì)“源1.0”巨量模型結(jié)構(gòu)的基本原則：

盡可能增加序列長度，因?yàn)樗欣趶埩坎⑿?、流水線并行和數(shù)據(jù)并行。由于內(nèi)存占用與序列長度的平方成正比，因此有必要在反向傳播時(shí)重新計(jì)算激活函數(shù)，以節(jié)省內(nèi)存開銷；

語言模型中層數(shù)太多會對性能產(chǎn)生負(fù)面影響，因?yàn)檫@會增加在流水線氣泡上的時(shí)間消耗；

增加隱藏層大小可以提高張量并行和流水線并行的性能；

增加節(jié)點(diǎn)中的微批次大小可以提高流水線并行效率，增加全局批次大小可以提升數(shù)據(jù)并行的效率；

在這一設(shè)計(jì)原則的基礎(chǔ)上，我們設(shè)計(jì)的“源1.0”的模型結(jié)構(gòu)以及分布式策略的設(shè)置如下表所示：

結(jié)合模型結(jié)構(gòu)的特性以及我們使用集群的硬件特性，我們?nèi)缦碌墓?jié)點(diǎn)配置和分布式策略選擇：

“源1.0”模型在訓(xùn)練過程中共使用了2128個(gè)GPU；

模型分成了7組，每組38臺AI服務(wù)器，里面放置一個(gè)完整的“源1.0”模型，7組之間采用數(shù)據(jù)并行；

每組的38個(gè)服務(wù)器，采用流水并行每個(gè)服務(wù)器放置1/38的模型（2個(gè)Transformer Layer），一共76層；

在每臺服務(wù)器內(nèi)采用張量并行，按照Transformer結(jié)構(gòu)的每一層進(jìn)行均勻切分；

模型收斂曲線如下圖：

關(guān)于“源1.0”的更多信息，大家可以參照浪潮發(fā)布在arxiv上的論文：https://arxiv.org/abs/2110.04725

審核編輯 ：李倩