來一個快問快答，如何使用性能有限的GPU卡從頭訓練出一個ChatGPT？
在目前，随着模型參數量不斷暴漲，大家對算力的要求也水漲船高。對比GPT-3當年使用了10000塊GPU，花了30天訓練了1750億參數，現實條件下我們一沒法組織起無限算力，二來手裡的計算卡硬體裝置還可能差異很大，顯存、帶寬水準參差不齊，在卡不多條件又比較艱苦的背景下，我們該如何進行高效的大模型訓練呢？
我們在這篇論文中找到了答案。
論文标題:
YUAN 2.0: A Large Language Model with Localized Filtering-based Attention
論文連結:
https://arxiv.org/ftp/arxiv/papers/2311/2311.15786.pdf
項目位址:
https://github.com/IEIT-Yuan/Yuan-2.0
模型下載下傳:
https://huggingface.co/IEITYuan
在上個月，浪潮資訊釋出了千億級别的開源大模型源2.0，作為一個開源模型，源2.0在多個評測名額上超越了 GPT-3.5 直逼 GPT-4.0。作為源2.0的技術報告，這篇論文核心的關注點還是放在了源2.0對模型結構的創新，也就是标題中提到的新型注意力機制結構LFA（局部注意力過濾增強機制，Localized Filtering-based Attention），但是我們注意到了源2.0在分布式訓練政策上的創新同樣有可以深挖之處。
從Scaling Law出發，當下“有限的算力”與幾乎看不到上限的“模型參數”之間存在着天然的沖突，而在有限的資源下進行高效的大模型訓練與開發是一個世界範圍内的關鍵技術難點與研究問題，而源2.0則在“分布式訓練政策”上下了功夫做了文章，最終很好的解決了在有限且多元的算力背景下高效進行大模型訓練的關鍵工程難題。
那麼源2.0是如何做到在有限的GPU環境下，從頭訓練出來了一個優質高效的大模型的呢？OK，那我們就細細展開從“大模型是如何誕生的”出發慢慢說起。
一、大模型誕生記：3D并行
大模型要訓練，自然離不開“并行”。而在大模型訓練中為什麼要考慮并行？顯然是因為“量大而力不足”。
大模型訓練過程之中，“量大”可以分為兩個方面，一個是“模型大”，模型參數軍備競賽直沖萬億，而另一個則是“資料多”，訓練資料量級成了天文數字。而“力不足”也可以從兩個方面進行了解，一個“算力不足”算太慢，另一個則是“顯存不足”存不下。舉一個具體一點的例子，GPT-3有1750億個參數在570GB的語料資料中進行訓練，而要訓練這樣一個“龐然大物”，使用 8 張 V100 顯示卡，訓練時長預計需要36 年。
顯然，這不是一個合理的時間，讓我們回顧模型訓練的基本流程，一個神經網絡，輸入資料我們首先要跑一遍前向Forward的過程，算算wx+b，算算激活函數，之後計算Loss function利用損失函數進行Backward對參數求導得到梯度Grad，拿到Grad後扔給優化器Optimizer更新模型權重，如此反複。
事實上，制約訓練複雜度的無非是模型的大小以及輸入資料的多少，由此衍生出兩種經典的分布式訓練政策，模型并行與資料并行。
先來說模型并行，大模型訓練的一個痛點也在于模型太大參數量太多，一種直接的思路就是模型參數量大那麼我們就拆模型。因而模型并行（Model Parallel, MP）思想應運而生，而模型并行又有兩種思路，分别是流水線并行（Pipeline Parallel, PP）與張量并行（Tensor Parallel, TP）。
先來看流水線并行，一般而言深度學習的模型總是分層的，假設有一個8層Layer的模型,，那麼最樸素的一種流水線并行方法就是按照Layer切分進行并行。譬如使用4個GPU 進行流水線并行，其中GPU 0負責計算
，
而GPU 1則負責
,
依次類推GPU 3負責
。
這其中GPU 0首先計算中間值1以張量形式傳遞給GPU 1，GPU 1再将計算值傳遞至GPU 2，直到GPU 3 接收到模型的輸出後進行反向傳播，完成一次訓練過程。
而除了按Layer層拆，還有一種模型并行的思路就是“直接拆張量”，這其實源于矩陣乘法的性質，假設有矩陣乘法，現在我們可以把按列給拆成，這樣做矩陣乘法，如果有N個GPU，這樣就可以把一個大的矩陣乘法拆成N個小乘法。
除卻模型并行，輸入資料的大小也直接決定了記憶體的使用多少與計算量是以模型大我們拆模型，資料大那我們就拆資料，因而一種分布式政策的思想也就呼之欲出：資料并行（Data Parallel, DP），簡而言之一張卡放不下的資料算不過來的資料我們按Batch進行切分，将資料配置設定到多張卡中進行模型訓練。
此時，綜合前面提到的流水線并行（PP）、張量并行（TP）以及資料并行（DP），就構成了當下大模型訓練的基本方法：3D并行。通過在訓練資料、模型張量、模型層三個方面應用分布式訓練與并行技術，3D并行技術成功訓練了數萬億參數規模的模型。3D并行的應用從計算量與顯存兩個方面規避了大模型訓練的資源瓶頸，讓模型規模達到了一個在過去無法設想的高度。
其實早在浪潮資訊2021年對标GPT-3推出的源1.0 中，就已經使用了傳統的資料+流水線+張量的3D并行政策，從結果導向來看，源1.0 交出了一份不錯的答卷。據公開資料表明，2457億參數的源1.0的訓練算力效率達到了44.8%，而1750億參數的GPT-3訓練算力效率為21.3%。
二、打破3D并行：創造大模型訓練的另一種可能
3D并行解決了所有問題嗎？顯然沒有，舉個例子，當硬體環境差異較大時，不同GPU卡顯存與帶寬本身存在着顯著差異，有時即便是同一種GPU卡，卡的負載也并不完全一樣。在這種多元算力背景下，3D并行往往很容易出現中途顯存超标OOM的問題，而現實場景下在多元算力晶片上開發大模型又往往是不得已而為之的必要手段，那麼如何在現有3D并行的技術基礎之上“求解”這種複雜現實問題就成為迫在眉睫的難點痛點。
時間來到兩年後，源2.0的參數規模是1026億，不到源1.0的一半，但是明顯的，源2.0不僅在各方面能力上“遠超”源1.0，對計算力的使用效率也由極大提升。那麼在分布式訓練政策之上，源2.0對比源1.0最大的差別何在呢？其實答案就是一套創新性的分布式訓練政策——非均勻流水線并行+優化器參數并行（ZeRO-1）+ 資料并行 + Loss計算分塊，這套分布式訓練政策也正是源2.0成功的關鍵！
具體來看，回想流水線并行的工作流程，其中有一個問題就在于傳統流水線并行給每一個工作GPU配置設定的模型層數都是相同的。這個思路乍一看沒什麼問題，但是在實際模型訓練過程中，往往前幾個階段需要緩存更多的激活值以進行反向傳播計算；而在後幾個流水并行階段，隻需要緩存少量的激活值。也就是說很可能在前幾個stage顯存幾乎用到上限的時候，後面的顯存還剩餘很多。假設一個24層的Transformer用8個GPU做流水線并行，如果是均勻配置設定則一個GPU中存在3層網絡，此時根據反向傳播從後前的方向将使得最後方的GPU要儲存更多的激活值，進而非常有可能使得GPU超過顯存上限使得訓練失敗，而在GPU卡性能不高的場景下這個問題更甚。
而源2.0在這一顯存瓶頸中做出了突破，如上圖(b)所示，不是說均勻的流水線并行這種“絕對的平均主義”會導緻“多勞不多得”嗎，那麼不如就“按勞配置設定”，将流水線并行中的每個GPU中承擔的Layers不均勻風格，以降低峰值顯存開銷，實作存儲和計算的總體最佳性能。比如一個12層的模型，分到4個GPU上，GPU 1的任務中隻分2層，GPU 2 分 3 層，GPU 3 分 4 層，GPU 4 分 3 層。
與此同時，為了更進一步降低損耗，源2.0又提出了一種“逐塊交叉熵計算”的方法，舉個例子，輸入tokens數為2048。逐塊交叉熵計算先對tokens進行切分，切分為16個塊，每個塊長度為128個tokens，每計算一次即得128個Loss，通過将所有的Loss連接配接為一個張量，釋放中間的臨時變量，使得模型訓練可以在不需要額外計算與通信的情況下解決前文所述最後一個流水線階段的顯存瓶頸。
此外，如果要繼續将資料并行的思路發揮到極緻，仔細分析可以看到，在資料并行中每個工作GPU内還儲存有一份完整的模型參數，這樣直接使得GPU的顯存成為了模型規模的瓶頸。那麼，真的有必要每個GPU内都儲存有一份完整的模型參數嗎？答案顯然是不必要的，這就引出了大名鼎鼎的ZeRO。
拆解一下這幾個字母，ZERO代表 Zero Redundancy Optimizer，即零備援優化，ZeRO首先對“模型參數”的構成進行了分析，由上文模型的訓練的過程可以得到，所謂模型參數主要包含三大部分，分别是模型本身參數、梯度參數以及優化器參數三塊，對于模型訓練而言，事實上并不需要工作GPU“人手一份”資料，相反整個訓練過程其實僅需要“一份參數”就夠了。
從這個角度出發，那麼十分自然，假設有N塊GPU，無論是模型參數的哪一部分，其實目前工作GPU僅需要1/N份參數就足夠可以完成訓練，由此，源2.0利用了ZeRO的思想，通過在切分三種參數的“顯存下降”與“通信成本升高”之間權衡，源2.0選擇了ZeRO-1即優先對優化器參數進行并行。從上圖來看，使用優化器參數并行對比Baseline将執行個體中的顯存由120GB降低至31.4GB，大大優化了GPU顯存使用。
除了前面提到的ZeRO-1與非均勻流水線并行，源2.0還取消了張量并行政策。拆張量的張量并行由于直接從“源頭”進行拆分，導緻其對模型參數的切分粒度較細，而細粒度的計算就意味着GPU間就必須得進行頻繁的通信，使得晶片間的P2P通信帶寬成為訓練的主要瓶頸。舉個例子，晶片間的P2P通信帶寬為900GB/s時，通信時間占比為10%，而如果P2P通信帶寬為400GB/s時，通信時間占比會上升到23%。進一步換算，可以得到P2P通信帶寬的降低，導緻了約17%的性能損失。
注：浪潮資訊“源1.0” 3D并行訓練政策通信分析
依據這一發現，源2.0取消了張量并行，應用非均勻流水線并行、Loss計算分塊、ZeRO-1、資料并行等技術重組了分布式訓練政策，使得晶片間通信帶寬也不再成為限制模型訓練的阻礙，如下圖所示，當晶片間的P2P帶寬從400GB/s降低到100GB/s，這套分布式訓練算法的性能也幾乎不會産生太大的變動（0.23%），可以更好地應用在性能水準差異大的各類算力叢集中。
三、一點總結
2023又到年底，回望這大模型狂飙突進的一年，從 ChatGPT 問世技驚四座到 AI Agent 火熱浪潮洶湧，從 GPT-4 問鼎天下睥睨群雄到當下大模型百花齊放百家争鳴……短短一年，大模型就賦予了我們最接近未來“智能時代”的狂熱想象，各種技術噴薄而出各類行業應用泉湧而下，哪怕是最近這短短一個月，我們都可以看到諸如谷歌的 Gemini，微軟的Phi-2，Mistral AI的Mixtral等等模型，
而這一批又一批的大模型來來去去。如果僅僅關注這些模型在某個基準中上升幾名或在某個資料集上名額如何，或許在一段時間内大家茶餘飯後還能熱火朝天的讨論兩句，但是在大模型超高速發展的當下，很快又會被大家抛之腦後成為過眼雲煙。
抛開刷榜抛開SOTA，事實上面向實際場景對訓練技術進行創新相當重要，我們常常期望國内各種技術“趕英超美”，但是所謂“工程優勢”從來都不是突然冒出來的一篇論文或一項技術所賦予的，真正的大國工程優勢的出現，往往就在于類似Yuan 2.0分布式訓練這種創新技術的一點一滴的累積。
在論文中提出一個Fancy的概念或許容易，在一個Toy example中刷出兩個漂亮結果或許簡單，但是能紮紮實實一步一步的把各種理論組裝搭建，面向實際問題兵來将擋水來土掩，把一個又一個小的技術創新串聯，直至撐起來一個巨型工程甚至開源出去，或許是一條頗為困難但我們不得不前進的道路。
而在大模型遍地開花的當下，也許就是為什麼像源2.0這樣的工作顯得尤為難得的答案吧！