如何訓練2457億引數量的中文巨量模型“源1.0”

浪潮人工智慧研究院

從2018年的BERT到2020年的GPT-3，NLP語言模型經歷了爆發式的發展過程，其中BERT模型的引數量為3.4億，而GPT-3的模型引數量達到了1750億。2021年9月，浪潮人工智慧研究院釋出了“源1.0”，它是目前規模最大的中文AI單體模型，引數規模高達2457億，訓練採用的中文資料集達5TB。“源1.0”在語言智慧方面表現優異，獲得中文語言理解評測基準CLUE榜單的零樣本學習和小樣本學習兩類總榜冠軍。測試結果顯示，人群能夠準確分辨人與“源1.0”作品差別的成功率低於50%。

海量的引數帶來了模型訓練和部署上的巨大挑戰。本文將聚焦“源1.0”背後的計算挑戰以及我們採取的訓練方法。

1.“源1.0”的模型結構

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

圖1 模型結構示意圖(左圖為LM，右圖為PLM)

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

2. 如何訓練“源1.0”

2.1.源1.0訓練面對的挑戰

“源1.0”的訓練需要面對的第一個挑戰就是資料和計算量的挑戰。

資料方面，如果把訓練一個巨量模型的訓練過程比作上異常戰役的話，那麼資料就是我們的彈藥。資料量的多少，決定了我們可以訓練模型的規模，以及最後的效果。針對這一方面，我們構建了一個全新的中文語料庫，清洗後的高質量資料規模達到了5TB，是目前規模最大的中文語料庫。

圖2 資料預處理流程圖

算力方面，根據OpenAI提出的PetaFlop/s-day衡量標準，我們可以估算“源1.0”訓練的計算需求情況。根據Wikipedia提供的資料()，GPT-3的計算需求約為3640 PetaFlop/s-day, 約等於64個A100 GPU訓練1年時間。而“源1.0”的計算需求達到了4095 PetaFlop/s-day。

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

在模型訓練時一般最常用的是採用資料並行分散式計算策略，但這隻能滿足小模型的訓練需求。對於巨量模型來說，由於其模型引數量過大，遠遠超過常用計算裝置比如GPU卡的視訊記憶體容量，因此需要專門的演算法設計來解決巨量模型訓練的視訊記憶體佔用問題，同時還需要兼顧訓練過程中的GPU計算效能的利用率。

2.2“源1.0”的訓練策略

為了解決視訊記憶體不足的問題，我們採用了張量並行、流水並行、資料並行相結合的並行策略，實現了在2128個GPU上部署“源1.0”，並完成了1800億tokens的訓練。

一、張量並行

針對單個GPU裝置不能完整的承載模型訓練，一個解決方案就是張量並行+資料並行的2D並行策略。具體來說，使用多個GPU裝置為1組，比如單個伺服器內的8個GPU為1組，組內使用張量並行策略對模型進行拆分，組間(伺服器間)採用資料並行。

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

圖3 張量並行示意圖

在訓練過程中，tensor經過每一層的時候，計算量與通訊資料量之比如下：

其中，S為輸入序列的長度，h為隱藏層的大小(hidden size)。

二、流水並行

圖4 流水線並行示意圖

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

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

與此同時，在流水並行過程中，節點間的計算量與通訊資料量之比為：

根據上面的公式，流水線中節點的計算效率與h和S呈線性關係，這與張量並行類似。

三、資料並行

圖6 資料並行示意圖

採用資料並行時，全域性批次大小(global batch size)按照流水線分組進行分割。每個流水線組都包含模型的一個副本，資料在組內按照區域性批次規模送入模型副本。資料並行時的計算量與通訊資料量的比值可用如下公式近似：

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

根據這一公式，我們可以看出資料並行的計算效率與全域性批次大小B和序列長度S呈正比關係。由於模型對記憶體的需求與S的平方成正比，與B成線性關係，因此增加全域性批次大小可以更有效的提升資料並行的效率。

當全域性批次大小過大的時候，模型很容易出現不收斂的問題，為了保證模型訓練過程的穩定性，我們將全域性批次大小限制在了個token內。

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

· 儘可能增加序列長度，因為它有利於張量並行、流水線並行和資料並行。由於記憶體佔用與序列長度的平方成正比，因此有必要在反向傳播時重新計算啟用函式，以節省記憶體開銷。

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

· 增加隱藏層大小可以提高張量並行和流水線並行的效能。

· 增加節點中的微批次大小可以提高流水線並行效率，增加全域性批次大小可以提升資料並行的效率。

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

Model

Layers

Hidden size

Global BS

Micro BS

Sequence

Length

t

p

d

GPUs

Yuan 1.0

76

16384

3360

1

2048

8

38

7

2128

結合模型結構的特性以及我們使用叢集的硬體特性，我們如下的節點配置和分散式策略選擇：

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

· 模型分成了7組，每組38臺AI伺服器，裡面放置一個完整的“源1.0”模型，7組之間採用資料並行

· 每組的38個伺服器，採用流水並行每個伺服器放置1/38的模型(2個Transformer Layer)，一共76層

· 在每臺伺服器內採用張量並行，按照Transformer結構的每一層進行均勻切分

模型收斂曲線如下圖：

關於“源1.0”的更多資訊，大家可以參照浪潮人工智慧研究院釋出在arxiv上的論文：。

另外，我們也整理了更詳細的《NLP模型訓練解決方案白皮書》。感興趣的讀者可以關注“浪潮AIHPC”公眾號，回覆“NLP白皮書”即可下載。