來源：時序人

國產之光 Deepseek 大模型成功出圈，在效果比肩非開源模型ChatGPT的同時，其運行成本很低。那么 Deepseek 是如何實現這么高效的訓練呢？文本將簡要介紹 Deepseek 母公司幻方開源的 HAI-platform 大模型訓練工具，來一窺極限 AI 工程的秘密。

為了更好地發(fā)揮 GPU 集群的計算能力，訓練具有驚人功能的、強大的萬億參數模型，一個高效簡潔的大模型訓練工具十分必要?；梅?深度求索研發(fā)了一個能對 GPU 資源進行統一高效利用的 AI 平臺，同時實現了四種并行訓練方式：ZeRO 支持的數據并行、流水線并行、張量切片模型并行和序列并行。這種方式可以極大優(yōu)化集群的使用，提高模型訓練的顯存效率和計算效率。

【項目地址】

https://github.com/HFAiLab/hai-platform

顯存與計算效率

首先，我們先了解訓練巨大模型的顯存和計算效率的挑戰(zhàn)。

01

顯存效率

訓練具有萬億參數的模型對顯存的需求極為龐大，遠超過單個 GPU 的承載能力。舉例來說，當采用 Adam 優(yōu)化器進行混合精度訓練時，僅保存模型的狀態(tài)信息（包括參數、梯度及優(yōu)化器狀態(tài)）就需要大約12TB的顯存空間。以英偉達A100 GPU為例，其顯存容量為40GB，這意味著僅為了存儲模型狀態(tài)，就需要多達400張 A100 GPU 的顯存資源。

在模型的前向傳播過程中，中間激活值會被保留，直至反向傳播完成并計算出損失函數后才會被釋放。這部分額外的顯存消耗與批量大小成正比。即便將批量大小設為1，訓練一個萬億參數的模型也會產生超過 400GiB 的激活顯存占用。盡管可以通過 Checkpoint 技術（以計算時間換取顯存空間）來提前處理這些激活顯存，但整體上，訓練過程中的顯存需求依然極為巨大。

因此，為了確保這類大規(guī)模模型能夠在顯存不溢出的情況下進行訓練，必須將模型狀態(tài)量和激活顯存有效地分散到多個GPU設備上。這種分布式策略是實現大規(guī)模模型訓練的關鍵。

02

計算效率

據估算，端到端訓練一個擁有萬億參數的模型大約需要消耗 5000 Zflops 的計算量（即5乘以10的24次方；這一估算基于OpenAI的研究《Scaling Laws for Neural Language Models》）。這一龐大的計算需求意味著，若使用4000張英偉達 A100 GPU 并以50%的計算效率運行，大約需要100天的時間才能完成訓練。

雖然大型超級計算 GPU 集群可能配備超過4000個 GPU，但由于批量大?。╞atch size）的限制，要在如此規(guī)模上實現高計算效率仍然面臨巨大挑戰(zhàn)。計算效率與計算時間對通信時間的比率成正比，而這一比率又與批量大小直接相關。然而，訓練模型時批量大小的設置存在一個上限：超過該上限，模型的收斂性能將顯著下降。

以目前最大的模型之一 GPT-3 為例，其訓練批量大小約為 1500。如果使用約 4000 張 GPU 進行訓練，即使將批量大小設置為 4000，每張 GPU 上的批量大小也僅為1，這將嚴重限制模型訓練的擴展性。因此，如何在保證模型收斂性能的同時優(yōu)化批量大小和計算效率，是實現大規(guī)模模型訓練的關鍵難題。

AIGC 集群：

整合 GPU 資源

構建一個能管理大規(guī)模 GPU 集群資源的 AI 平臺，主要會遇到如下幾個痛點問題：

• 資源調度：算力規(guī)模不斷擴大，而訓練任務的計算需求又多種多樣，如何處理任務和算力的關系以最大化集群資源使用？

• 使用效率：集群使用有峰谷差異，同時又要應對突發(fā)任務的需求，如何兼顧時效性和集群整體效率？

• 迭代適配：集群會有升級迭代，在調整算力規(guī)模、類型、使用規(guī)則等場景下，如何讓平臺快速適配，盡可能降低切換成本？

幻方提出以任務級分時調度共享 AI 算力的理念，將集群零散資源進行整合再分配，成功支持在上千計算節(jié)點上穩(wěn)定運行深度學習訓練和其他多類型任務，日常算力占用率 95% 以上，日常 GPU 使用率 75% 以上。

01

任務調度

任務是集群使用的基本單位，而非用戶。所有用戶提交的任務都由集群統一管理和調度。

用戶獨占 vs 分時調度

幻方提出分時調度理念對集群資源進行管理。用戶提交任務，如運行 python / bash 代碼，啟動開發(fā)容器等，由平臺根據當前資源需求、集群忙閑程度等進行任務的中斷和加載。任務代碼需要遵循平臺編碼規(guī)則以確?？梢詳帱c續(xù)跑，具體流程如下：

• 接受集群的打斷信號；

• 保存 checkpoint（模型參數，優(yōu)化器參數等）；

• 通知集群打斷；

• 從 checkpoint 恢復，繼續(xù)運行。

這里，集群不會將 GPU 資源池化，而是以計算節(jié)點為基本單位，根據資源類型、網絡區(qū)域等條件進行分類標記。AIGC 集群鼓勵用戶一次性用滿多張 GPU，進行并行訓練。用戶提交任務時需選定節(jié)點數量 n，則該任務可獲得 n 的整數倍個 GPU，比如在8卡計算節(jié)點上提交使用4個節(jié)點的任務，則該任務會獲得32張 GPU 進行并行訓練。

02

用戶管理

集群提供以配額的方式記錄集群的各類資源，通過優(yōu)先級管理用戶的使用權限。例如：某用戶在 NORMAL 優(yōu)先級上擁有10個計算節(jié)點的配額，其在集群上提交若干個任務，則最多同時以 NORMAL 優(yōu)先級調度10個計算節(jié)點運行。算力資源優(yōu)先滿足高優(yōu)先級用戶的訓練需求，同優(yōu)先級內則交替使用集群算力資源。

同時，每個用戶都可以歸屬于某一個用戶組。計算資源和優(yōu)先級配額是按用戶為單位進行分配的，而同組的用戶可以共享自定義環(huán)境、私有數據存儲。

03

數據管理

AIGC 集群會將訓練數據存儲在文件系統中。訓練中的任務會實時從文件系統中讀取數據，流轉到計算節(jié)點。這里，優(yōu)良的文件系統將是提高 GPU 利用率，進而提升集群整體效率的決定性因素之一。

為了最大化 GPU 利用率，降低數據讀取的成本，幻方自研了高速文件系統 3FS，其專門針對深度學習訓練中樣本讀取的場景進行優(yōu)化，能支持計算和存儲節(jié)點間 7TB/s 以上的數據吞吐。3FS 需要用戶將原始數據中較多的小文件聚合成大文件進行存儲，從而實現優(yōu)秀的性能。

并行訓練

01

3D并行

設計高效的并行方式對訓練大模型至關重要。deepseek 結合集群的特性，可以實現三種并行訓練方式的靈活組合：數據并行、流水并行、張量并行，并使用序列并行方式對 Transformer 進一步并行優(yōu)化，極大地提升顯存利用和計算效果。

彩色塊顯示有 32 張 GPU （八個計算節(jié)點，每個節(jié)點有四個 GPU）的系統上進行 3D 并行訓練的 GPU 映射。同一顏色的 GPU 在同一節(jié)點上

02

數據并行

為了減少顯存的占用，Deepseek 采用 ZeRO 數據并行的方式，把訓練優(yōu)化器的狀態(tài)均分到 1/N 個 GPU 上 （N 是數據并行的數量），如下圖所示：

在做前向和反向傳播時，ZeRO 數據并行會先做 allgather 獲得完整的參數，然后在前向和反向傳播結束后釋放掉，只保留 1/N 的參數和梯度。

03

流水并行

流水并行把模型切分成 M 個階段，分別放在 M 個 GPU 上，能夠大幅節(jié)省顯存占用和提升擴展性能。訓練的時候每個階段接收上一階段的輸出作為輸入，然后傳給下一個階段。如下圖所示：

可以看到，這種并行方式會造成大量的 GPU 空閑（Idle）。為了減少浪費，Gpipe 和 PipeDream 的概念被提出：如果同時進行多個迭代，每個節(jié)點在同一時刻負責不同的迭代的計算，就可以避免數據依賴，不用在原地干等了：

在 DeepSeek 中，模型的切分是在 Builder 里由用戶切分的，用戶可以繼承 并實現自己的 build_model 方法。 Builder 創(chuàng)建出來的模型包含一個 forward_backward 方法，用戶可以傳進輸入、損失函數、標簽，該方法會執(zhí)行前向傳播和反向傳播并返回 Loss。

04

張量并行

張量并行把全連接層的參數和計算切分到 K 個 GPU 上，能夠節(jié)省顯存占用和提升擴展性能。假設 K=2，張量并行把矩陣乘法 X @ A @ B 變成 X @ A1 @ B1 + X @A2 @ B2 ，其中 A 沿著縱軸切成 A1 和 A2，B 沿著橫軸切成 B1 和 B2；GPU 0 會存放參數 A1 和 B1 并負責 X @ A1 @ B1 的計算，GPU 1 則存放參數 A2 和 B2 并負責 X @ A2 @ B2 的計算。

05

序列并行

Megatron-LM 提出的序列并行方案旨在通過一些方式分攤張量并行中無法分攤的顯存。

如上圖所示，在張量并行的基礎上，將 Transformer 核的 LayerNorm 以及 Dropout 層的輸入按序列長度維度進行了切分，使得每個 GPU 上面只需要做一部分的 Dropout 和 LayerNorm 。這樣做進一步降低了計算資源和顯存的開銷。