摘要：千億參數量的中文大規模預訓練語言模型時代到來。
本文分享自華為雲社區《 MindSpore開源框架加持，如何「煉出」首個千億參數、TB級內存的中文預訓練語言模型？》，原文作者：chengxiaoli。
千億參數量的中文大規模預訓練語言模型時代到來。
近段時間，中文大規模預訓練語言模型圈有些熱鬧。26 億參數量的「悟道 · 文源」， 270 億參數量的 PLUG，以及昨天華為雲發布的千億級別「盤古」NLP 大模型，預訓練語言模型已經成長到僅加載就需要 TB 級的內存或顯存。
我們可以直觀地想到，「盤古」效果理應更好，但計算量需求也更大，訓練起來更困難。
然而「盤古」實際上是這樣一次探索：開源框架 MindSpore，昇騰基礎軟硬件平台，加上超大規模中文預訓練模型，意味着基礎設施已然完善了。
這項工作由華為以及北京大學相關技術團隊聯手完成，在昇騰基礎軟硬件平台，以及 MindSpore 框架自動並行等黑科技的幫助下，訓練出當前最大的中文預訓練模型。
那么量級不斷拔高的盤古大模型是如何訓練出來的？接下來，讓我們細致解讀下「盤古」背后的關鍵技術。
千億參數，TB 級內存的模型
以盤古 2000 億為例，如果我們訓練時權重都用標准的 FP32 數據格式，那么算下來，權重占的空間就達到了 750GB，訓練過程中內存開銷還會數倍上升。這 750GB 參數，不是放在硬盤上，也不是加載到內存中，而是需要移到昇騰Atlas訓練服務器 HBM（High Bandwidth Memory 高帶寬存儲器）內存中，以利用昇騰Atlas訓練服務器訓練模型。
模型大 ，意味着數據也大，而且都需要是高質量數據。為了滿足數據需求，研發團隊從互聯網爬取了 80 TB 文本，並最后清洗為 1TB 的中文數據集。
這樣的模型與數據，已經不是我們幾台服務器能加載上的了，更不用說進行訓練。好在研發團隊會提供 API，一般算法工程師直接調用接口就能試試效果。
可以說，目前盤古是業界首創的千億規模中文預訓練模型，其中最高參數量達 2000 億。
超大規模自動並行，算法工程師的福音
先考慮一個問題，你想到如何訓練這樣的大模型了嗎？
如果給你足夠的計算力，你能想到如何訓練這么大的模型嗎？我們最常用的分布式訓練方式數據並行，單獨這么做肯定是不行的，因為沒有哪個計算硬件能放下 800GB 的參數。那么再加上模型並行呢？又產生了新問題，我們該如何拆分如此巨大的「盤古」？硬件產品（如 NPU、GPU 等）之間的梯度流、數據流通信又是什么樣的？
顯然訓練如此龐大的模型，遠比我們想象中的復雜，需要大量的工程化操作，並保證這些操作不會或極少影響到模型最終收斂效果。
難道盤古真得靠手動並行優化？
如果手動來寫分布式訓練邏輯，那么需要綜合考慮計算量與類型、集群帶寬、拓撲結構、樣本數量等等一大堆復雜的東西，然后再設計出性能比較優秀的並行切分策略，並編寫大量並行切分和節點間的通信代碼。如果系統環境變了，還要重新設計並修改算法，想想就覺得頭大。
倘若我們用 TensorFlow 或其他類似框架，MirroredStrategy 這一系列自帶的分布式策略完全用不上，看起來自行寫並行策略是必不可少的。然而，盤古 真正的訓練是一種軟硬件協同的方式，MindSpore 計算框架、CANN 異構計算架構、昇騰基礎軟硬件平台整套基礎設施。其中，MindSpore 提供的，就包含了至關重要的自動並行能力。
融合 5 大維度，強大的自動並行
MindSpore 自動並行提供了 5 維的並行方式：數據並行、算子級模型並行、Pipeline 模型並行、優化器模型並行和重計算，並且在圖編譯階段，有機融合了 5 個維度的並行。這 5 維並行方式組合起來構成了盤古的並行策略。
a. 數據並行
數據並行是最基本，應用最廣的並行方式，其將訓練數據（mini-batch）切分，每台設備取得其中一份；每台設備擁有完整的模型。在訓練時，每台設備經過梯度計算后，需要經過設備間的梯度同步，然后才能進行模型參數的更新。
b. 算子級模型並行
算子級模型並行是對模型網絡中的每個算子涉及到的張量進行切分。MindSpore 對每個算子都獨立建模，每個算子可以擁有不同的切分策略。
以矩陣乘算子 MatMul(x, w)為例，x 是訓練數據，w 是模型參數，兩者都是二維矩陣。並行策略 ((4, 1), (1, 1)) 表示將 x 按行切 4 份，保持 w 不切，如果一共有 4 台設備，那么每台設備擁有一份 x 的切片，和完整的 w。
c.Pipeline 模型並行
Pipeline 模型並行將模型的按層分成多個 stage，再把各個 sage 映射到多台設備上。為了提高設備資源的利用率，又將 mini-batch 划分成多個 micro-batch, 這樣就能夠使得不同設備在同一時刻處理不同 micro-batch 的數據。
一種 Pipeline 並行方式(Gpipe) 要求反向計算要等所有設備的正向計算完成后才開始，而反向計算可能依賴於正向的輸出，導致每個卡正向計算過程中累積的 activation 內存與 micro-batch 數量成正比，從而限制了 micro-batch 的數量。MindSpore 的 Pipeline 並行中，將反向提前，每個 micro-batch 計算完成后，就開始計算反向，有效降低 activation 存儲時間，從而提升整體並行效率。
d. 優化器模型並行
優化器模型並行將優化器涉及到的參數和梯度切分到多台設備上。以 Adam 優化器為例，其內部可能有多份與權重同等大小的「動量」需要參與計算。在數據並行的情況下，每個卡都擁有完整的「動量」，它們在每個卡上都重復計算，造成了內存及計算的浪費。通過引入優化器並行，每個卡只保存權重及「動量」的切片，能降低每個卡的靜態內存及提升計算效率。
e. 重計算
重計算 (Rematerialization) 針對正向算子的輸出累計保存在內存中，導致內存峰值過大的問題，舍棄了部分正向算子的輸出，而是在反向階段用到時再重新計算一遍。這樣做有效地降低了訓練過程中的內存使用峰值。如下圖所示，第一個內存峰值通過重計算消除，第二個內存峰值可以通過前面講到的優化器並行消除。
有了這 5 維的並行維度后，如何將其組合起來作用於盤古，並且如何將切分后的模型分片分配到每台設備上仍然是難題。MindSpore 自動並行，把這 5 個維度並行有機組合起來，可以實現非常高效的大模型分布式訓練能力。
下圖 (b) 是一典型的樹形的硬件拓撲結構，其帶寬隨着樹深度的增加而降低，並且會產生一些流量沖突。為了利用此特征，MindSpore 的目標是最大化計算通信比，將通信量大的並行方式（算子級並行）放置在服務器內部的多卡之間；將通信量較小（Pipeline 並行）的放置在同一機架內的服務器間；將數據並行（優化器並行）的部分放置在不同機架間，因為該通信可以和計算同時執行(overlap)，對帶寬要求較低。
在盤古 2000 億模型中，MindSpore 將 64 層（layer）划分為 16 個 stage，每個 stage 包含 4 層。在每層中，利用算子級並行的方式對張量進行切分。
如下圖中的 Q,K,V 的參數在實際中（按列）被切了 8 份，輸入張量（按行）被切了 16 份，輸出張量因此被切了 128 份（8*16）。重計算配置是配置在每層內的，也就是重計算引入的多余的計算量不會超過一層的計算量。總計，MindSpore 使用了 2048 塊昇騰處理器來訓練盤古。
MindSpore 對外屏蔽了復雜並行實現的細節，使得用戶像編寫單機模型腳本那樣簡單。用戶在單機腳本的基礎上，僅通過少了配置就能實現多維度的混合並行。下圖是簡化版的盤古腳本，其中紅色加粗字體表示的在 MindSpore 中的並行策略。將紅色加粗字體去掉，則是單機腳本。
圖算跨層聯合優化，發揮硬件極致性能
除了跨節點間的大規模自動外，在單卡節點內，MindSpore 通過圖層和算子層的跨層協同優化，來進一步發揮昇騰算力。
在傳統的 NN 網絡中，不同算子承載的計算量和計算復雜度也各不相同。如 LayerNorm 由 11 個基本算子組成，而 Add 則只有 1 個基本算子。這種基於用戶角度的算子定義，通常是無法充分發揮硬件資源計算能力的。因為計算量過大、過復雜的算子，通常很難生成切分較好的高性能算子。從而降低設備利用率；而計算量過小的算子，由於計算無法有效隱藏數據搬移開銷，也可能會造成計算的空等時延，從而降低設備利用率。
為了提升硬件利用率，MindSpore 使用了圖算融合優化技術，通過圖層和算子層聯合優化，並將「用戶使用角度的易用性算子」進行重組融合，然后轉換為「硬件執行角度的高性能算子」，從而充分提升硬件資源利用率，進而提升整網執行性能。具體優化流程如下圖所示：
以 LayerNorm 算子為例，通過算子拆分和重組，11 個小算子，組成了 1 個單算子和 2 個融合算子。這些重組后的算子可以生成更加高性能的算子，從而大大降低了整體網絡運行時間。
在盤古模型中，圖算融合幫助整體訓練時間減少了 20% 以上。除此之外，對於其它 NLP、CV 等任務，圖算融合在優化性能方面都有不錯的表現。
總結：訓練超大模型的完美體現
即使給我們足夠的算力，超大模型的訓練還是異常復雜，遠比想象中的困難。對於我們一般算法工程師來說，針對某個任務，上億參數量已經算大的了，但是並不會感到訓練上會有什么困難，因為各個深度學習框架直接調用數據並行接口就能搞定。
但是如果模型繼續增大到百億級、千億級甚至萬億級，並行與優化策略的復雜度猛然上升，算法工程師一點點地編寫與優化代碼可太難了。MindSpore 通過自動並行，把計算邏輯和並行邏輯解耦，單卡串行代碼自動實現分布式並行，從而使得算法科學家將精力都解放到模型本身上。
為了從預訓練獲取更多的知識， GPT-3 與盤古這樣的模型會越來越大，畢竟到現在我們還沒看到大模型預訓練效果的極限在哪。屆時，這類模型對基礎設施的需求會更大，並行與優化策略也會更加復雜。只有擁有足夠優秀的基礎設施，大規模預訓練的效果才會更好，從而在知識問答、知識檢索、知識推理、閱讀理解等場景發揮更大作用，實現智慧客服、營銷、文案生成等商業價值。
大規模計算集群及軟硬件協同優化，這次在盤古的訓練上得到了充分的完美體現。正如開發團隊所言，「基於 Mindspore 和昇騰基礎軟硬件平台在千億參數模型上的實踐也是一次探索，大模型的分布式訓練、超參調優、數據集組成、模型結構適應性等都存在太多的未知。現在，盤古模型效果很好，刷新了 clue 版第一，這意味着第一次基於國內軟硬件協同優化，以及超大規模分布式訓練，結果是令人振奮的，我們自己也具有了足夠強的基礎設施。」
當然，也誠如以上所言，盤古只是對超大規模分布式訓練、超大規模中文預訓練模型的一次探索，未來還需要更多的研究工作者投入到通用智能與大規模分布式計算的研究工作中。