摘要：千億參數量的中文大規模預訓練語言模型時代到來。
本文分享自華爲雲社區《 MindSpore開源框架加持，如何「煉出」首個千億參數、TB級內存的中文預訓練語言模型？》，原文做者：chengxiaoli。html
千億參數量的中文大規模預訓練語言模型時代到來。算法
近段時間，中文大規模預訓練語言模型圈有些熱鬧。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 版第一，這意味着第一次基於國內軟硬件協同優化，以及超大規模分佈式訓練，結果是使人振奮的，咱們本身也具備了足夠強的基礎設施。」
固然，也誠如以上所言，盤古只是對超大規模分佈式訓練、超大規模中文預訓練模型的一次探索，將來還須要更多的研究工做者投入到通用智能與大規模分佈式計算的研究工做中。