多GPU訓練大型模型：資源分配與優化技巧｜英偉達將推出面向中國的改良芯片HGX H20、L20 PCIe、L2 PCIe

在人工智能領域，大型模型因其強大的預測能力和泛化性能而備受矚目。然而，隨著模型規模的不斷擴大，計算資源和訓練時間成為制約其發展的重大挑戰。特別是在英偉達禁令之后，中國AI計算行業面臨前所未有的困境。為了解決這個問題，英偉達將針對中國市場推出新的AI芯片，以應對美國出口限制。本文將探討如何在多個GPU上訓練大型模型，并分析英偉達禁令對中國AI計算行業的影響。

如何在多個 GPU 上訓練大型模型？

神經網絡的訓練是一個反復迭代的過程。在每次迭代中，數據首先向前傳播，通過模型的各層，為每個訓練樣本計算輸出。然后，梯度向后傳播，計算每個參數對最終輸出的影響程度。這些參數的平均梯度和優化狀態被傳遞給優化算法，如Adam，用于計算下一次迭代的參數和新的優化狀態。隨著訓練的進行，模型逐漸發展以產生更準確的輸出。

然而，隨著大模型的到來，單機難以完成訓練。并行技術應運而生，基于數據并行性、管道并行性、張量并行性和混合專家等策略，將訓練過程劃分為不同的維度。此外，由于機器和內存資源的限制，還出現了混合精度訓練、梯度累積、模型卸載CPU、重算、模型壓縮和內存優化版優化器等策略。

為進一步加速訓練過程，可以從數據和模型兩個角度同時進行并行處理。一種常見的方式是將數據切分，并將相同的模型復制到多個設備上，處理不同數據分片，這種方法也被稱為數據并行。另外一種方法是模型并行即將模型中的算子劃分到多個設備上分別完成（包括流水線并行和張量并行）。當訓練超大規模語言模型時，需要對數據和模型同時進行切分，以實現更高級別的并行，這種方法通常被稱為混合并行。通過這些并行策略，可以顯著提高神經網絡的訓練速度和效率。

一、數據并行

在數據并行系統中，每個計算設備都有完整的神經網絡模型副本，在進行迭代時，每個設備僅負責處理一批數據子集并基于該子集進行前向計算。假設一批次的訓練樣本數為N，使用M個設備并行計算，每個設備將處理N/M個樣本。完成前向計算后，每個設備將根據本地樣本計算誤差梯度Gi（i為加速卡編號）并進行廣播。所有設備需要聚合其他加速卡提供的梯度值，然后使用平均梯度（ΣN i=1Gi）/N來更新模型，完成該批次訓練。

數據并行訓練系統通過增加計算設備，可以顯著提高整體訓練吞吐量和每秒全局批次數。與單計算設備訓練相比，最主要的區別在于反向計算中梯度需要在所有計算設備中進行同步，以確保每個計算設備上最終得到所有進程上梯度平均值。

二、模型并行

模型并行可以從計算圖的角度出發，采用流水線并行和張量并行兩種方式進行切分。

1、流水線并行

流水線并行（Pipeline Parallelism，PP）是一種計算策略，將模型的各層劃分為多個階段，并在不同計算設備上進行處理，實現前后階段的連續工作。PP廣泛應用于大規模模型的并行系統，以解決單個設備內存不足問題。下圖展示了由四個計算設備組成的PP系統，包括前向計算和后向計算。其中F1、F2、F3、F4代表四個前向路徑，位于不同設備上；B4、B3、B2、B1代表逆序后向路徑，位于四個不同設備上。然而，下游設備需要等待上游設備計算完成才能開始計算任務，導致設備平均使用率降低，形成模型并行氣泡或流水線氣泡。

樸素流水線策略會導致并行氣泡，使系統無法充分利用計算資源，降低整體計算效率。為減少并行氣泡，可以將小批次進一步劃分為更小的微批次，并利用流水線并行方案處理每個微批次數據。在完成當前階段計算并得到結果后，將該微批次的結果發送給下游設備，同時開始處理下一微批次的數據，在一定程度上減少并行氣泡。如下圖所示，前向F1計算被拆解為F11、F12、F13、F14，在計算設備1中完成F11計算后，會在計算設備2中開始進行F21計算，同時計算設備1中并行開始F12的計算。與原始流水線并行方法相比，有效降低并行氣泡。

2、張量并行

張量并行需要針對模型結構和算子類型處理參數如何在不同設備上進行切分，并確保切分后的數學一致性。大語言模型以Transformer結構為基礎，包含三種算子：嵌入表示、矩陣乘和交叉熵損失計算。這三種算子具有較大差異，因此需要設計相應的張量并行策略，以便將參數分配到不同設備上。對于嵌入表示層參數，可按照詞維度進行劃分，每個計算設備只存儲部分詞向量，然后通過匯總各個設備上的部分詞向量來獲得完整的詞向量。

矩陣乘的張量并行可以利用矩陣分塊乘法原理來優化計算。以矩陣乘法Y = X × A為例，其中X是M × N維的輸入矩陣，A是N × K維的參數矩陣，Y是M × K維的結果矩陣。當參數矩陣A過大超出單張卡的顯存容量時，可以將A切分到多張卡上，并通過集合通信匯集結果，確保最終結果的數學計算等價于單計算設備的計算結果。參數矩陣A有兩種切分方式：

1）按列切分

將矩陣A按列切成A1和A2，分別放置在兩個計算設備上。兩個計算設備分別計算Y1 = X × A1和Y2 = X × A2。計算完成后，多計算設備間進行通信，拼接得到最終結果矩陣Y，其數學計算與單計算設備結果等價。

2）按行切分

將矩陣A按行切成B1,B2,...,Bn，每個Bi為N*（K/n）即（K/n）N維。將這n個切分后的矩陣分別放到n個GPU上，則可并行執行矩陣乘法Y=XB1，Y=X*（B1+B2），...，Y=X*（B1+B2+...+Bn）。每步并行計算完成后，各GPU間進行通信，拼接得到最終結果矩陣Y。

在Transformer中FFN結構包含兩層全連接（FC）層，每層都涉及兩個矩陣乘法。這兩個矩陣乘法分別采用上述兩種切分方式。對于第一個FC層的參數矩陣，采用按列切塊方式，而對于第二個FC層參數矩陣，則采用按行切塊方式。這樣的切分方式使得第一個FC層輸出能夠直接滿足第二個FC層輸入要求（按列切分），從而省去了第一個FC層后匯總通信操作。

多頭自注意力機制張量并行與FFN類似，由于具有多個獨立的頭，因此相較于FFN更容易實現并行。其矩陣切分方式如圖所示。

在分類網絡最后一層，通常會使用Softmax和Cross_entropy算子來計算交叉熵損失。然而，當類別數量非常大時，單計算設備內存可能無法存儲和計算logit矩陣。針對這種情況，可以對這類算子進行類別維度切分，并通過中間結果通信來獲得最終的全局交叉熵損失。首先計算的是softmax值，其公式如下：

在計算交叉熵損失時，可以采用張量并行的方式，按照類別維度對softmax值和目標標簽進行切分，每個設備計算部分損失。最后再進行一次通信，得到所有類別的損失。整個過程中，只需要進行三次小量的通信，就可以完成交叉熵損失的計算。

3、管道并行

管道并行性將模型按層“垂直”分割。同時，還可以“水平”分割層內的某些操作，稱為張量并行訓練。對于現代模型（如Transformer）的計算瓶頸，即將激活批矩陣與大權重矩陣相乘，可以在不同GPU上計算獨立的點積或每個點積的一部分并對結果求和。無論采用哪種策略，都可以將權重矩陣分割成均勻大小的分片，托管在不同的GPU上，并使用分片計算整個矩陣乘積的相關部分，再通過通信組合結果。Megatron-LM是一個例子，在Transformer自注意力層和MLP層中實現矩陣乘法的并行化。PTD-P結合張量、數據和管道并行性，通過為每個設備分配多個非連續層以減少氣泡開銷，但增加了網絡通信成本。有時，輸入可以跨維度并行化，并通過更細粒度的示例進行計算，以減少峰值內存消耗。序列并行是一種思想，將輸入序列在時間上分割成多個子示例，從而按比例減少內存消耗。

四、混合專家 (MoE)

隨著研究人員試圖突破模型大小限制，混合專家(MoE) 方法引起廣泛關注。其核心思想是集成學習，即多個弱學習器組合可生強大的學習器。使用 MoE 方法時，僅需使用網絡一小部分即可計算任何輸入的輸出。一種示例方法是擁有多組權重，網絡可以在推理時通過門控機制選擇使用哪一組權重。這可以在不增加計算成本的情況下啟用更多參數。每組權重都被稱為“專家”，希望網絡能夠學會為每個專家分配專門的計算和技能。不同專家可以托管在不同 GPU 上，從而提供一種清晰方法來擴展模型所使用的 GPU 數量。恰好一層 MoE 包含作為專家前饋網絡 {E_i}^n_{i=1} 和可訓練門控網絡 G 學習概率分布 n “專家”，以便將流量路由到少數選定的 “專家”。當 “專家” 數量過多時，可以考慮使用兩級分層 MoE。

GShard（Google Brain團隊開發的一款分布式訓練框架

）通過分片將MoE變壓器模型擴展至6000億個參數。MoE變壓器用MoE層替換所有其他前饋層。分片MoE變壓器僅具有跨多臺機器分片的MoE層，其他層只是簡單地復制。Switch Transformer（Transformer類的萬億級別模型

）通過稀疏開關FFN層替換密集前饋層（其中每個輸入僅路由到一個專家網絡），將模型大小擴展到數萬億個參數，并具有更高的稀疏性。

五、其他節省內存的設計

1、混合精度計算（Mixed Precision Training）

混合精度訓練（Mixed Precision Training）是指在訓練模型時同時使用16位和32位浮點類型，以加快運算速度和減少內存使用。在NVIDIA GPU上，使用float16進行運算比使用float32快一倍多，大大提高了算力的上限。然而，將模型的運算轉換為FP16并不能完全解決問題，因為FP16的數值范圍遠小于FP32和TF32，限制模型的運算能力。為確保模型能夠收斂到與FP32相同結果，需要采用額外的技巧。

1）權重備份（Weight Backup）

其中一種避免以半精度丟失關鍵信息的技術是權重備份。在訓練時，權重、激活值和梯度都使用FP16進行計算，但會額外保存TF32的權重值。在進行梯度更新時，對TF32的權重進行更新。在下一步訓練時，將TF32的權重值轉換為FP16，然后進行前向和反向計算。

2）損失縮放（Loss Scaling）

在訓練模型時，由于梯度量級往往非常小，使用FP16格式可能會導致一些微小梯度直接被歸零。大部分非零梯度實際上并不在FP16表示范圍內。由于FP16格式右側部分并未被充分利用，我們可以通過將梯度乘以一個較大系數，使整個梯度分布向右移動并完全落在FP16表示范圍內。一種簡單方法是在計算梯度之前先將損失乘以一個較大值，以此放大所有梯度。在進行梯度更新時，再將其縮小回原來的并使用TF32進行更新。

3）精度累加（Precision Accumulation）

在FP16模型中，一些算術運算如矩陣乘法需要用TF32來累加乘積結果，然后再轉換為FP16。例如，Nvidia GPU設備中的Tensor Core支持利用FP16混合精度加速，同時保持精度。Tensor Core主要用于實現FP16的矩陣相乘，并在累加階段使用TF32大幅減少混合精度訓練的精度損失。

2、梯度累積（Gradient Accumulation）

梯度累積是一種神經網絡訓練技術，通過將數據樣本按批次拆分為幾個小批次，并按順序計算。在每個小批次中，計算梯度并累積，在最后一個批次后求平均來更新模型參數。神經網絡由許多相互連接的神經網絡單元組成，樣本數據通過所有層并計算預測值，然后通過損失函數計算每個樣本的損失值（誤差）。神經網絡通過反向傳播算法計算損失值相對于模型參數的梯度，并利用這些梯度信息來更新網絡參數。梯度累積每次獲取一個批次的數據，計算一次梯度（前向），不斷累積梯度，累積一定次數后根據累積的梯度更新網絡參數，然后清空所有梯度信息進行下一次循環。

3、卸載CPU（CPU Offloading）

CPU Offloading是指將未使用的數據暫時卸載到CPU或不同的設備之間，并在需要時重新讀取回來。由于CPU存儲相比GPU存儲具有更大的空間和更低的價格，因此實現雙層存儲可以大大擴展訓練時的存儲空間。然而，簡單的實現可能會導致訓練速度降低，而復雜的實現需要實現預取數據以確保設備無需等待。ZeRO是一種實現這一想法的方式，它將參數、梯度和優化器狀態分配到所有可用的硬件上，并根據需要進行具體化。

4、激活重新計算（Activation Recomputation）

Recompute是一種在前向計算中釋放tensor，在反向傳播時需要重新計算的方法，適用于占用內存大但重新計算量小的tensor。重新計算的方式有三種：

Speed Centric會保留計算出的tensor以備后續使用；

Memory Centric會在計算完成后釋放tensor，需要時再重新計算；

Cost Aware會在計算完成后判斷是否保留tensor，若可能導致內存峰值則釋放。

可以將swap和recompute結合使用，針對特定op采用不同方式。還可以預先迭代幾次，收集內存和運行時間信息，判斷哪些tensor該swap，哪些該recompute。

5、模型壓縮（Compression）

模型壓縮是通過裁剪、權重共享等方式處理大模型，以減少參數量。然而，這種方式容易降低模型精度，因此使用較少。常見的模型壓縮方法包括修剪、權重共享、低秩分解、二值化權重和知識蒸餾。

修剪可以采用對連接、kernel、channel進行裁剪的方式；權重共享是通過共享模型參數來減少參數量；低秩分解將矩陣分解為低秩形式，從而減少參數量；二值化權重是將權重從32位降至8位或16位，實現混合精度訓練；知識蒸餾是使用訓練好的教師模型指導學生模型訓練。

6、高效內存優化器（Memory Efficient Optimizer）

優化器在模型訓練中的內存消耗是一個重要問題。以Adam優化器為例，它需要存儲動量和方差，與梯度和模型參數規模相同，內存需求增加。為減少內存占用，已經提出了幾種優化器，如Adafactor和SM3，采用不同的方法估計二階矩或大幅減少內存使用。

ZeRO優化器是一種針對大型模型訓練的內存優化方法。通過觀察模型狀態和激活臨時緩沖區及不可用碎片內存的消耗，采用兩種方法：ZeRO-DP和ZeRO-R。ZeRO-DP通過動態通信調度來減少模型狀態上的冗余，而ZeRO-R則使用分區激活重新計算、恒定緩沖區大小和動態內存碎片整理來優化殘留狀態的內存消耗。

英偉達禁令之后，中國AI計算何去何從？

在10月17日，美國強化對中國市場的AI芯片禁令，將性能和密度作為出口管制標準，禁止出口單芯片超過300teraflops算力、性能密度超過每平方毫米370gigaflops的芯片。由于限制AMD、英特爾等公司的高端AI芯片，尤其是英偉達的主流AI訓練用GPU A100和H100，該禁令又被稱為“英偉達禁令”。

針對新的芯片禁令，AI產業議論紛紛，焦點主要集中在實施時間、緩沖地帶、涉及的GPU型號和禁令期限等方面。盡管存在爭議，但針對中國的高端AI芯片禁令仍在堅定推行。

現在，AI行業必須形成共識應對挑戰。與其過分關注被禁的GPU，我們應更深入思考在芯片鐵幕時代下中國AI計算未來發展路徑。下面將探討當前產業形勢并共同探討AI計算前行之路。

一、目前現狀

與之前情況相比，英偉達禁令出臺后大眾輿論與AI行業反應似乎更為冷靜。僅在消費級顯卡RTX 4090是否被禁問題上引發游戲玩家和商家爭論。盡管行業不希望看到高端AI芯片被禁售，但對此局面已有預期。美國對華芯片封鎖已持續多年，英偉達部分高端GPU已被禁止出售，產業界的反應也從驚訝轉變為冷靜應對。加上ChatGPT的火爆導致全球高端GPU市場行情上漲，美國方面多次表示要推動對華整體性的高端AI芯片禁售。

為應對禁令并受到大模型發展的推動，去年年底到今年上半年，眾多中國科技、金融、汽車等企業集中購買英偉達高端GPU，導致市場上GPU供不應求。對于許多中國中小型科技企業和AI創業公司來說，原本就很難買到高端GPU，禁售并未帶來太大變化。實際上，國內AI芯片產業在貿易摩擦初期便開始加速發展，雖然英偉達的高端GPU在AI訓練需求方面難以替代，但并非不可替代。

此外，AI芯片與手機芯片不同，并不關乎大眾消費者。華為已在手機芯片領域取得突破。因此，無論是大眾還是行業，對禁令都持坦然態度，甚至有些習以為常。然而，必須承認的是，禁令對中國AI行業仍造成了一定程度的傷害：短期內更換英偉達GPU面臨芯片產能和生態兼容性等難題；禁令還將直接損害使用英偉達產品的AI服務器等領域的廠商。

長期禁令可能使中國AI計算與全球高端芯片脫鉤，可能帶來復雜的負面影響，包括：中國AI算力發展可能落后于英偉達高端GPU的更新迭代；在底層算力發展分歧下，中國AI產業可能在軟件技術方面掉隊；科技封鎖可能從AI芯片擴展到通用算力、存儲、基礎軟件等數字化基礎能力。因此，需要制定三項同時發力的“突圍方案”：加快國產AI芯片的自主研發和生態建設；加大力度投資大模型等軟件技術，降低對英偉達等公司的依賴；加強與國際科技合作，推動中國AI計算的全球化發展。

二、解決方案一：用好買家身份

作為全球芯片市場最大買家，中國企業應該利用好這個身份，擺脫中美科技貿易中的思維誤區。我們往往認為游戲規則是由美國政府和企業制定的，只能被動接受，但實際上作為買家應該擁有更多話語權。針對中國市場的AI芯片禁令，最直接傷害的是以英偉達為代表的美國科技巨頭，因為中國市場對他們的AI芯片需求最大。英偉達CEO黃仁勛曾表示，如果被剝奪了中國市場，他們將沒有應急措施，世界上沒有另一個中國。因此，我們應該認識到作為買家的力量，并利用好這個身份來維護自己的利益。

我們可以看到美國科技公司和政府之間的矛盾?？萍脊咀非笊虡I利益，而政府則追求政治利益。美國科技公司一直在嘗試反對和繞過禁令，例如英偉達推出針對中國市場的特供版GPU。

三、解決方案2：以云代卡，算力集中

在可見的較長時間里，美國對中國AI芯片封禁只會加強，這給AI大模型發展帶來挑戰。許多業內人士認為，大模型發展雖快，但沒有呈現此前科技風口的迅猛局面，投資缺錢、計算缺卡是主要原因。

為解決中國AI產業在禁令之下的算力缺口問題，企業需要加大云端AI算力配置和投入，推動以云代卡。事實上，在高端AI芯片可能被禁的大趨勢下，中國幾大公有云廠商都開始加強囤積英偉達高端GPU。這不僅因為自身需要加大大模型投入，打開MaaS市場，也對AI算力有直接需求。此外，GPU轉化為云資源池后可以長期復用，對云廠商來說具有進可攻、退可守的優勢。因此，今年上半年出現高端AI芯片流向云廠商、中小企業難以獲得芯片的局面。

客觀來看，這種高端AI芯片集中向云的舉動有利于中國市場統籌應對AI芯片禁令，也符合東數西算戰略思路。另一趨勢是，隨著大模型參數和使用數據量不斷加大，本地化卡池訓練已經越來越緊張，在云端進行千卡、萬卡訓練成為未來主要發展方向，因此企業用戶會更加積極地走向云端。

同時，云端AI算力不僅限于囤積英偉達GPU。隨著政策推動和自主AI芯片采購力度的加大，云端化和自主化結合的AI算力將成為發展趨勢。根據IDC數據，2023上半年中國AI服務器已經使用50萬塊自主開發的AI加速器芯片。華為已經推出昇騰AI云服務，提供自主AI算力服務。在東數西算背景下，各地建立一批采用自主AI算力的AI計算中心，保障云端AI算力穩定可靠供給。

然而，很多企業仍然傾向于采購本地AI算力。一方面是因為英偉達GPU市場緊缺，保值性高，甚至可以作為企業的核心資產。另一方面是因為云端AI算力存在排隊、宕機、軟件服務缺失等問題，影響開發者體驗。為進一步提高開發者的云端AI算力使用體驗，公有云廠商需要進一步努力。

四、方案三：讓國產AI算力爆發式成長

面對新一輪AI芯片禁令，中國AI產業并非依賴英偉達的高端GPU，而是經過多年發展，AI芯片產業已經得到巨大發展。雖然英偉達市場份額仍占主導，國產AI算力已經具備一定市場占比，但在核心性能、軟件生態和出貨能力方面仍需不斷提升?？陀^上，禁令的倒逼將加速國產AI算力的成長與成熟周期。

為了實現這個目標，有幾件事非常重要：

1、形成產業共識，避免概念混淆

雖然AI芯片市場呈現出眾多品牌和類型參與者，但其中存在的問題也不容忽視。對于類腦芯片等前沿技術，目前仍處于暢想階段，而一些AI芯片廠商僅能自用，無法面向市場出貨，同時還有大量廠商處于早期建設階段，短期內對AI計算自主化貢獻有限。

為應對英偉達高端GPU禁售問題，需要將關注點集中在可行、有效的GPU替代方案上，避免過多的聯想和發散。只有形成產業共識，才能更好地解決問題。

2、走向規?；逃?，避免PPT造芯

目前國內能夠出貨的AI芯片廠商主要集中在華為、百度、燧原科技和海光信息等少數幾家。大量半導體廠商與AI企業還停留在打造芯片的計劃與愿景上，導致政策支持與投資市場期待的國產AI芯片發展停滯，甚至有些企業可能只是在這一階段享受金融市場紅利而缺乏實質性進展。

為推動產業發展，未來的產業導向應該重將AI芯片從計劃轉向出貨，幫助廠商獲得直接商業回饋，讓產品與產能接受市場檢驗，逐步塑造正向現金流。

3、加強軟件生態，強化遷移能力

英偉達GPU重要性不僅在于硬件性能，更在于其CUDA和PyTorch等軟件生態的強大能力。因此，發展國產AI芯片不能忽視軟件能力的提升。在加強自主軟件生態建設的同時，還需要關注基于英偉達生態的AI模型遷移能力和遷移成本。

許多廠商已經在這方面進行探索，例如海光信息的DCU與CUDA在生態和編程環境上高度相似，使得CUDA用戶能夠以較低代價快速遷移到海光的ROCm平臺。此前，PyTorch2.1版本宣布支持華為昇騰，顯示出國產AI芯片已經具備一定的規?；绊懥?，可以更多地融入全球軟件生態。未來要實現國產AI計算的爆發，離不開國產AI基礎軟件生態的蓬勃發展。

4、加大對“主品牌”支持，形成規?；?/strong>