高效訓練百萬億參數預訓練模型的系統(tǒng)挑戰(zhàn)和對策

馬子軒/MA Zixuan，翟季冬/ZHAI Jidong，韓文弢/HAN Wentao，陳文光/CHEN Wenguang，鄭緯民/ZHENG Weimin

（清華大學，中國北京100083）

1 大規(guī)模預訓練模型的發(fā)展背景

近年來，深度學習在計算機視覺（CV）、自然語言處理（NLP）、推薦系統(tǒng)、決策模型等各個領域都發(fā)揮了重要作用。同時，大規(guī)模預訓練模型技術已經成為深度學習領域最先進的技術，并在多個領域的下游應用中表現出優(yōu)秀的性能。近年來，谷歌的推薦系統(tǒng)、搜索引擎，阿里巴巴的推薦系統(tǒng)、圖像生成等任務均采用了預訓練模型。而在預訓練模型方面，學術界已達成共識：模型規(guī)模和模型準確度有明顯的正相關關系。表1展示了近年來預訓練模型的發(fā)展趨勢：從最早的1.1億參數量的第1代預訓練GPT發(fā)展到最新的萬億參數量的GPT-3、Switch Transformer等模型。探索更大參數量的模型具有重要的科學意義。

表1 近年發(fā)布的預訓練模型

從計算的角度看，隨著模型規(guī)模的增長，訓練模型的數據會變多。此時，單機已經無法滿足大規(guī)模模型訓練的計算需求。在模型訓練的過程中，計算分為3個步驟：前向、反向、更新。在前向過程中，模型使用輸入數據和參數進行計算，得到預測結果并和標注數據進行對比，從而計算出該次預測的損失；在反向過程中，模型對損失進行傳播，計算所有參數的梯度；在更新階段，模型再使用計算出的梯度來更新參數。該過程在訓練中不斷迭代，最終達到模型的收斂狀態(tài)。當模型規(guī)模較大時，單機串行執(zhí)行已無法滿足模型的計算需求。這時，我們需要使用一些并行策略來輔助訓練，如數據并行和模型并行。

顧名思義，使用數據并行策略時，模型對輸入的數據進行拆分。如圖1所示，數據被劃分到不同的節(jié)點后再進行計算，每個節(jié)點上都保存一份完整的模型。模型在正向和反向的過程中，均不會進行通信；而在正向結束計算梯度以及更新時，則會引入全局的All-Reduce通信。

使用模型并行策略時，模型對模型參數本身進行拆分，這可以理解成對參數矩陣進行切分。在執(zhí)行過程中，模型執(zhí)行被切分的計算時，會根據切分方式的不同引入不同的通信方式。如圖1所示，在兩種不同的模型并行執(zhí)行模式中，不同的切分策略產生了兩種不同的通信模式：All-Reduce和All-Gather。

圖1 數據并行與模型并行示意圖

無論是數據并行還是模型并行，在執(zhí)行過程中，都需要通過對數據進行切分來減少每個節(jié)點的計算量和存儲量，從而達到并行訓練的目的。在模型訓練中，參數量和訓練速度是兩個重要的衡量指標。經驗證明，越大規(guī)模的模型所需的訓練樣本數就越多。因此，大模型訓練需要同時擴展參數量和訓練吞吐量。在傳統(tǒng)的并行模式中，通過對輸入進行切分，數據并行可以有效擴展訓練吞吐量。而由于數據并行需要在每個節(jié)點上都存儲一份參數，這種并行模式不能擴展參數量。模型并行可以對參數進行切分，從而有效擴展參數量。切分后多個機器使用同一組輸入數據，因此不能有效擴展訓練吞吐量。傳統(tǒng)的大模型訓練通常使用混合數據并行和模型并行的方法進行訓練。

混合專家（MoE）模型是近年來一種新的預訓練模型。這種模型將一個大模型切分為多個小模型（專家）。在訓練過程中，數據先通過一個小規(guī)模網絡（一般稱之為網關）進行路由，再選擇適合的一部分專家進行計算，最后加權求和得到輸出。

對MoE模型而言，一個整體的大模型由多個小模型構成。這樣一來，這類大模型的參數量可以和傳統(tǒng)的稠密大模型一樣大或更大。同時，由于每次運算只是從離散的模型中稀疏地選取一部分來激活，MoE模型的整體計算量與稠密小模型相當。因此，MoE模型可以在擴展模型參數、增大模型規(guī)模的同時，在相同的訓練時間內，達到優(yōu)于稠密大模型的準確度。谷歌最新的Switch Transformer就采用了MoE的架構。

在MoE大模型中，并行訓練會引入兩種劃分：數據劃分和專家劃分。在訓練過程中，每一組數據需要選擇合適的專家，在不同節(jié)點間交換數據。這個過程會引入全局的Allto-All通信。MoE的并行訓練很好地滿足了大模型訓練對吞吐量和模型規(guī)模的需求。在MoE并行訓練中，不同機器存儲不同的專家，有效擴展了模型的參數規(guī)模。同時，由于不同機器選取不同的輸入，該方法也有效擴展了模型訓練的吞吐量。因此，MoE模型非常適合對大模型進行擴展。

2 大規(guī)模預訓練模型面臨的挑戰(zhàn)

當模型擴展到超大規(guī)模時，會遇到多方面的挑戰(zhàn)。例如，如何選取高效的并行策略，如何進行高效數據存儲，如何選取合適數據精度以及如何實現動態(tài)負載均衡。

2.1 如何選取高效的并行策略

在大模型訓練中，有多種并行策略可供選擇，如數據并行、模型并行、流水線并行以及MoE并行等各種并行模式。這些并行策略將大模型訓練按照不同的方式進行劃分，從而使任務可以很好地在多機上執(zhí)行。而不同的劃分方法，對底層有不同的計算與通信需求。例如，一類并行策略的通信需求雖然小，但可能會引入更多的計算；另一類并行策略雖然不引入額外計算，但可能會引入更多的通信量。在訓練過程中，如何選擇合適的并行策略是一個重要的挑戰(zhàn)。

以16個節(jié)點為例，這些節(jié)點在訓練過程中需要被切分的部分包括輸入數據與模型參數。如果使用簡單的數據并行，那么模型參數并不被劃分，每臺機器上都有一份完整復制的模型。那么，數據并行的核心是把輸入數據簡單地切分成16份。與此類似，模型并行把模型切成16份。如果是MoE并行，則有16個專家，每個節(jié)點1個專家。在此基礎上，不同并行策略之間還可以進行混合，例如數據并行加模型并行的混合方式。

這些策略的組合給并行設計帶來了巨大的選擇空間。在此基礎上，在很多系統(tǒng)中，網絡拓撲會帶來不同節(jié)點間通信性能的不同。例如，在胖樹結構中，通常會存在網絡裁剪的問題，這導致一部分節(jié)點之間的通信延遲更高、帶寬更低。這種差異導致相同并行策略在不同情況下的性能有所不同，從而使并行策略選擇問題更加復雜。因此，如何選擇一個高效的并行策略具有很大的挑戰(zhàn)性。

2.2 如何高效存儲和劃分數據

以萬億參數量的模型為例，如果模型精度是32位，模型參數本身就會占用4 TB的空間。與此同時，模型的梯度需要占用4 TB的空間，優(yōu)化器的更新參數需要占用8 TB的空間，計算中間值需要約8 TB的空間。為了保證模型的訓練性能，這些數據需要存儲在內存中。如果使用NVIDA Tesla V100訓練該模型，僅在顯存中存儲這些數據就需要768塊顯卡。在這種情況下，不同的劃分方式對底層的計算和通信也會產生不同的影響。因此，用高效劃分數據來支持高效訓練同樣非常具有挑戰(zhàn)性。

2.3 如何選取合適的存儲精度

在計算過程中，使用更低的精度（如16位）訓練時，模型需要的內存量更小且性能更好，但準確度有所下降，訓練輪數有所增加；使用更高精度（如32位或64位）訓練時，模型需要的內存量更大，計算的要求更高，但模型準確度高，訓練輪數少。因此，如何在模型中選擇合適的精度進行計算，以平衡模型精度與模型訓練時間，給模型擴展帶來了新的挑戰(zhàn)。

2.4 如何實現動態(tài)負載均衡

MoE模型會通過網關對不同的輸入樣本選擇不同的專家進行計算。在實際運行過程中，樣本的選擇存在嚴重的負載不均衡問題。例如，在語言處理中，常用語言元素出現的頻率高，對應的專家負載更高，如圖2所示的專家0。此時，如果訓練規(guī)模非常大，部分節(jié)點的高負載就會降低整體系統(tǒng)的執(zhí)行效率。因此，如何動態(tài)改善負載均衡，以提高模型的訓練效率成為大模型訓練的重要挑戰(zhàn)。

圖2 混合專家模型中的負載均衡問題

3 新一代國產超級計算機

新一代國產神威超級計算機采用新一代神威處理器芯片——申威26010-Pro，其架構如圖3所示。其中，每個節(jié)點包括6個核組（CG），每個核組包含1個主核（計算控制核心）和64個從核（計算核心CPE）。加速計算主要使用從核陣列。如果使用中央處理器-圖形處理器（CPU-GPU）系統(tǒng)來做類比，我們可以認為主核相當于其中的CPU，從核陣列相當于GPU。神威處理器將6個核組通過1個環(huán)形網絡連接起來，并封裝到1個芯片中。同時，每個節(jié)點上包含2個自主研發(fā)的高速網絡芯片，所有的節(jié)點通過網絡再進行連接。神威網絡的架構是一個雙層胖樹，最下面的一層把256個節(jié)點組織成一個超節(jié)點。超節(jié)點的內部節(jié)點間完全互連；超節(jié)點之間通過頂層網絡互聯(lián)，跨超節(jié)點的通信需要經過頂層網絡。

圖3 新一代國產超級計算機架構

出于成本考慮，兩層的網絡拓撲會帶來網絡裁剪。當通信跨越超節(jié)點時，帶寬會降低到超節(jié)點內通信帶寬的3/16。正如前文所述，因為訓練對數據和模型的需求不同，將并行模式擴展到超大規(guī)模的異構系統(tǒng)時，需要考慮如何實現高效的映射。

大模型對內存計算和通信都有非常大的需求，新一代的超級計算機有PB級的內存空間、強大的計算能力、自主可控的高速網絡以及靈活的通信接口。因此，新一代的國產超級計算機為大規(guī)模模型訓練奠定了基礎。

4 在國產超算平臺上訓練大模型

根據國產超算平臺的特點，我們將大模型訓練擴展到了千萬核心規(guī)模。我們的工作核心在于驗證當模型做到一定規(guī)模時，并行加速還面臨哪些挑戰(zhàn)，以及如何應對這些挑戰(zhàn)。當前，我們的訓練參數規(guī)模達到了174萬億，已達到人腦神經元突觸數量的規(guī)模。針對前文所述的4個挑戰(zhàn)，我們提出了一些解決方案。盡管這些策略并不一定是最優(yōu)的，但仍可以為接下來進一步的優(yōu)化擴展提供參考。

4.1 根據通信帶寬采用合適的并行策略

新一代神威超級計算機的網絡拓撲分為兩層，其中上層是3/16的裁剪網絡。此時，采用單一的并行策略，如數據并行或模型并行，會帶來全局通信。受限于頂層帶寬和參與通信的高進程數，全局通信帶寬性能相比于小規(guī)模（如一個超節(jié)點以內）有顯著的下降。那么，如何選擇合適的并行策略來規(guī)避這類問題呢？

本質上，該問題屬于高性能計算的基本問題。我們將應用程序的通信模式與底層的網絡拓撲相結合，構建了一個性能模型，并通過性能模型判斷并行策略在大規(guī)模下的執(zhí)行效率。最終，如圖4所示，我們在全機規(guī)模下的解決方案是在超節(jié)點內（256個節(jié)點）做數據并行，在跨超節(jié)點做MoE并行。相比于對稱的方案（即超節(jié)點內做MoE并行，超節(jié)點間做數據并行），該方案可以將性能提高1.6倍。

圖4 全機規(guī)模并行策略設計

4.2 實現高效不重復的參數存儲和劃分

前文提到的數據存儲問題，在本質上為如何將參數進行高效劃分。在訓練過程中，參數主要分為3部分：模型參數、梯度參數和更新參數。其中，模型參數指訓練需要的計算參數；梯度參數指反向傳播時每個模型參數的梯度；更新參數指優(yōu)化器所需的參數，例如Adam優(yōu)化器中需要的參數。這些參數的規(guī)模都和模型參數的大小相當。此外，在混合精度訓練中，我們還需要考慮混合精度策略需要存儲的額外參數，如主參數與主梯度。因此，如何對這些參數做切分是一個復雜的問題。

我們以APEX O2的混合精度訓練為例。在訓練過程中，模型的更新參數會占用75%的空間。如果簡單地采用數據并行，模型參數在參與數據并行的節(jié)點中將被重復存儲，浪費了寶貴的存儲資源。微軟于2019年提出的工作ZerO就對數據并行的參數存儲進行了優(yōu)化，其核心思想是把參數分布在不同節(jié)點來進行更新，并在不影響通信量的前提下，將參數分布式存儲到不同機器中。如圖5所示，分布式參數更新有效降低了模型參數的內存。

圖5 分布式參數更新方法

我們對ZerO的方法進行了深度擴展。在裁剪網絡中，采用分層訓練策略，并針對全局和部分通信，擴展了通信算法，從而在不影響通信量的前提下，在分層并行策略中實現了分布式參數更新。

4.3 設計合適的混合精度策略

在新一代神威超級計算機中，神威CPU同時支持雙精度和半精度計算。單精度計算使用雙精度計算模擬，因此半精度性能為單精度的4倍。在機器學習訓練中，廣泛使用的是單精度與半精度類型。此時，如何針對國產超級計算機設計合適的混合精度策略，成為影響性能優(yōu)化的重要因素。

以Multibox SSD模型為例，對該模型參數分布的分析如圖6所示，橫坐標表示參數數值的指數分布。紅色區(qū)域表示FP16可以表示的范圍。如果該模型的參數直接采用半精度存儲，那么大部分參數會失真或無法表示。因此，我們采用了混合精度訓練中常用的動態(tài)損失縮放技術：在正向計算結束后對損失進行縮放，即乘以一個系數，然后在反向計算結束后再進行反縮放，并動態(tài)調整縮放系數，保證在不發(fā)生上溢出的情況下，數值盡量大。理想情況下，縮放后的表示范圍如圖6中綠色區(qū)域所示。此時，大部分參數都可以被半精度表示。

圖6 Multibox SSD模型的參數分布示意

NVIDIA的APEX系統(tǒng)是針對NVIDIA GPU的混合精度框架。APEX提供了4種訓練策略，分別是O0（單精度訓練）、O1（只優(yōu)化計算）、O2（優(yōu)化除更新外的部分）、O3（全部優(yōu)化）。在NVIDIA平臺中，使用較多的是O2策略。我們將APEX的策略移植到了神威平臺。遺憾的是，這4種策略并不能直接在神威平臺的大模型訓練中使用。測試結果表明：O0和O1策略可以正常收斂，但性能差；O2、O3性能好，但不能收斂。

因此，我們提出了一種分層的混合精度策略，如圖7所示。通過對不同層進行小規(guī)模驗證，我們發(fā)現不同層對精度的要求不同。例如，訓練中前饋網絡（FFN）層和注意力層對精度不敏感，可以使用半精度計算、單精度更新；而其他層對精度敏感，因此可以使用單精度訓練。這種策略在保證收斂的前提下，可以獲得較高的訓練性能。

圖7 分層的混合精度策略

此外，混合精度訓練還包含更多的探索方向，例如，訓練時間對訓練精度的敏感性等。另外，在大的科學計算程序中，不同迭代使用不同精度。在迭代的開始、中間和最后，迭代都可以選擇不同的混合精度策略。以動態(tài)的方式使計算和內存達到最高效率也是一個比較前沿的研究方向。

4.4 實現高效均衡負載

為了解決MoE模型負載不均衡的問題，我們提出了一個算法SWIPE。算法的核心思想是當一個專家非常受歡迎，且樣本數多于均值時，就將多余樣本分配給其他專家，以保證全局的絕對均衡。這個策略雖然簡單直接，但驗證結果顯示模型訓練能夠有效收斂。

4.5 國產系統(tǒng)底層優(yōu)化

此外，為了適配新一代神威超級計算機,我們針對神威平臺的軟件系統(tǒng)、算子庫、計算框架和基礎設施等進行了大量的優(yōu)化工作。例如，針對神威系統(tǒng)動態(tài)模式下內存分配的性能問題，根據機器學習應用的特點，我們設計了高效內存分配器SWAlloc；針對神威平臺，實現了一套高性能算子庫SWTensor；同時深度重構了SWPyTorch，以支持各類自定義算子表示。該部分工作包含約60 000行C/C++代碼，實現了100余個定制算子，有效地支持了模型訓練。

5 實驗結果

最終，我們在新一代神威平臺上整合了上述優(yōu)化方法，促成了百萬億參數模型的訓練。

我們的模型基于阿里巴巴的中文預訓練處理模型M6，模型輸入包含文本和圖像數據。如圖8所示，模型將輸入圖像拆分為多個塊，并使用ResNet等預訓練模型來提取特征，然后將圖像特征與詞向量連接起來形成一個序列，再由Transformer模型處理并生成高級表示。

圖8 模型結構示意圖

該模型在中文最大的多模態(tài)數據集M6-Corpus上進行訓練。數據來源包括百科全書、電子商務平臺和其他類型的網頁。最終處理的數據集的詳細統(tǒng)計情況如表2所示，其中，其中，圖像部分提取特征后大小為16 TB。

表2 預訓練數據集規(guī)模

我們共測試了3個不同的參數量模型的性能，模型超參數如表3所示。其中，d表示模型的隱藏層規(guī)模，d表示FFN內部層規(guī)模。在訓練的3個模型中，MoDa-174T規(guī)模達到了百萬億參數量（我們獲得的最大規(guī)模模型）。

表3 實驗用到的測試模型超參數

訓練過程與性能實驗均在新一代神威超級計算機上進行。其中，通過對PyTorch的前向、反向和更新時間的檢測，可以得到時間數據。浮點運算次數（FLOPs）通過神威系統(tǒng)的性能計數器得到，主核和從核分別計數，累加結果作為總浮點運算次數。

性能測試結果如表4所示。在1.93萬億參數規(guī)模時，模型訓練的計算性能達到了1.18 EFLOPS（百億億次浮點計算每秒）；當模型增大至14.5萬億時，混合精度性能達到1.00 EFLOPS；當模型擴展到174萬億時，由于模型規(guī)模太大，模型訓練無法進行混合并行，并行策略只有一維，因此通信性能會有明顯下降，只達到0.230 EFLOPS。

表4 性能測試結果

圖9和圖10分別展示了MoDa-1.93T模型的弱擴展性與強擴展性測試結果。在實驗中，我們將問題規(guī)模定義為專家個數和樣本個數。在弱擴展性測試中，兩者隨測試規(guī)模的變化發(fā)生變化，即在測試規(guī)模為全系統(tǒng)的1/n的測試中，模型規(guī)模和樣本個數均為全系統(tǒng)測試的1/n。因此，每個計算節(jié)點處理的樣本數和模型規(guī)模均保持不變。測試結果表明，我們的系統(tǒng)可以做到接近線性的可擴展性。

圖9 弱擴展性測試結果

圖10 強擴展性測試結果

而在強擴展性測試中，我們固定專家個數和樣本個數。此時，在測試規(guī)模為全系統(tǒng)的1/n時，每個節(jié)點處理的樣本為全系統(tǒng)測試中每個節(jié)點的n倍。測試結果表明，從全系統(tǒng)1/16擴展到整機時，性能提升約為12倍。我們的系統(tǒng)表現出很好的擴展性。

為了驗證系統(tǒng)的正確性，我們將MoDa-1.93T模型訓練了500步，收斂曲線如圖11所示。損失在500步后下降到3.46，根據參考文獻[21]，我們認為訓練接近收斂，并且可以驗證系統(tǒng)的正確性。

圖11 MoDa-1.93T收斂曲線

6 結束語

我們把目前的工作命名為“八卦爐”?！鞍素誀t”是一個高性能計算（HPC）和人工智能（AI）結合得較好的例子。在系統(tǒng)方面，結合以上提到的優(yōu)化策略，我們進行了預訓練模型加速。在此過程中，我們發(fā)現HPC存在許多挑戰(zhàn)，例如網絡裁剪等。本研究為將來探索新的大模型訓練系統(tǒng)提供了寶貴的經驗。

總的來說，該項研究主要針對國產系統(tǒng)的大型模型加速訓練，并結合研究中遇到的問題提出了一系列解決方案。目前，這些方案仍處于研究階段，希望有更多學者能一起參與討論。