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

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

“源1.0”的模型結構

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

圖1 模型結構示意圖（左圖為LM，右圖為PLM）

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

?

如何訓練“源1.0”

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

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

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

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

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

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

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

| “源1.0”的訓練策略

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

a. 張量并行

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

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

圖3 張量并行示意圖

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

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

b. 流水并行

圖4 流水線并行示意圖

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

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

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

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

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

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

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

當d? 1時，上面公式可以進一步簡化成：

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

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

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

盡可能增加序列長度，因為它有利于張量并行、流水線并行和數(shù)據(jù)并行。由于內存占用與序列長度的平方成正比，因此有必要在反向傳播時重新計算激活函數(shù)，以節(jié)省內存開銷；

語言模型中層數(shù)太多會對性能產生負面影響，因為這會增加在流水線氣泡上的時間消耗；

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

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

在這一設計原則的基礎上，我們設計的“源1.0”的模型結構以及分布式策略的設置如下表所示：

結合模型結構的特性以及我們使用集群的硬件特性，我們如下的節(jié)點配置和分布式策略選擇：

“源1.0”模型在訓練過程中共使用了2128個GPU；

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

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

在每臺服務器內采用張量并行，按照Transformer結構的每一層進行均勻切分；

模型收斂曲線如下圖：

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

審核編輯 ：李倩