大規(guī)模分布式訓(xùn)練的挑戰(zhàn)

當(dāng)模型較小并且可以裝入單個 GPU 內(nèi)存時，它非常簡單高效：輸入樣本、前向計算、后向計算等，GPU 計算資源可以得到充分利用。為了加速使用更多數(shù)據(jù)訓(xùn)練這種小型到中型模型，可以使用數(shù)據(jù)并行（例如PyTorch Distributed — 2020 Meta AI）等技術(shù)將整個模型復(fù)制到多個計算資源（GPU），獨立生成梯度，然后在每次迭代時傳達這些梯度以保持模型副本的一致性。

對于大型模型，比如具有數(shù)千億個參數(shù)的 LLM，該模型太大而無法放入單個 GPU，我們需要對模型進行分區(qū)并將其分布到具有數(shù)百或數(shù)千個 GPU 的大型集群中（例如，據(jù)報道，具有 175B 個參數(shù)的 GPT-3 模型在 10,000 個 A100 GPU 上進行訓(xùn)練），這種技術(shù)稱為模型并行性。我們可以按層對模型進行分區(qū)（例如，GPT-3 175B 有大約 96 層），并將不同的層分布到不同的 GPU。如果單個層仍然太大而無法放入單個 GPU 內(nèi)存（此外，參數(shù)，優(yōu)化器狀態(tài)和梯度在訓(xùn)練期間也需要內(nèi)存），仍然可以將該單層分成幾部分，并將每部分分配給專用的 GPU。還記得我在第一篇博客中介紹的專家混合嗎？ Mixture-of-Experts 已經(jīng)將前饋層分成了獨立的分區(qū)，我們可以將這些專家分布在不同的 GPU 上，每個 Attention 層都有多個 head（GPT-3 175B 有 96 個 head），我們可以將這些 head（及其 MatMul 計算）分布在不同的 GPU 上。而且我們可以從張量級別進一步劃分模型。

數(shù)據(jù)并行和模型并行并不相互排斥，它們可以一起使用來加速大規(guī)模模型訓(xùn)練過程。通過對模型進行分區(qū)應(yīng)用模型并行后，模型的一個分區(qū)可以復(fù)制并分發(fā)到多個 GPU 并應(yīng)用數(shù)據(jù)并行。

到目前為止一切看起來都很好？現(xiàn)在讓我們來談?wù)劥笠?guī)模分布式訓(xùn)練的挑戰(zhàn)。在模型和計算分布在數(shù)百或數(shù)千個 GPU 上之后，如何充分利用這些大規(guī)模計算資源是一個巨大的挑戰(zhàn)（還有其他挑戰(zhàn)，例如如果一臺機器在訓(xùn)練期間崩潰了怎么辦，因為單個集群中有數(shù)千個節(jié)點，節(jié)點崩潰的概率高于小集群，如何從檢查點恢復(fù)這部分，例如從空閑的 HA 池中添加新的 GPU / 節(jié)點回集群或重新平衡模型和訓(xùn)練任務(wù)以適應(yīng)拓?fù)涞取Ｔ谶@篇博客中，我們重點介紹如何充分利用資源。）。正如我在上一節(jié)中所描述的，前向計算和后向計算之間存在依賴關(guān)系，層的計算之間存在依賴關(guān)系，模型已跨不同節(jié)點進行分區(qū)。需要跨節(jié)點進行數(shù)據(jù)傳輸和通信。由于依賴于其他節(jié)點的計算結(jié)果，某些節(jié)點可能處于空閑狀態(tài)，而計算結(jié)果仍在進行中；梯度、某一層的輸出和其他數(shù)據(jù)可能需要從一個節(jié)點傳輸?shù)搅硪粋€節(jié)點，這可能會達到網(wǎng)絡(luò)帶寬瓶頸并導(dǎo)致某些節(jié)點空閑。可能有些步驟（例如，優(yōu)化器步驟）需要同步，也會導(dǎo)致某些節(jié)點空閑。

零氣泡流水線并行性

為了在深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練期間充分利用集群的計算資源，引入了幾種管道并行技術(shù)。PipeDream ：快速高效的管道并行 DNN 訓(xùn)練 — 2018（微軟、卡內(nèi)基梅隆大學(xué)、斯坦福大學(xué)）這篇論文引入了“一前一后”（1F1B）調(diào)度策略，通過管道重疊通信和計算來提高集群的 GPU 利用率。下圖中的 1、2、3、4 表示不同的訓(xùn)練數(shù)據(jù)小批量。

1F1B管道時間表

零氣泡流水線并行性——2023 Sea AI指出，后向計算實際上包含兩個部分：計算關(guān)于輸入 x (B) 和層參數(shù) W (W) 的梯度。1F1B 策略將 B 和 W 合并為 B，但不必要地增加了順序依賴的計算。因此，零氣泡流水線將 B 和 W 分成不同的階段，以減少流水線中的氣泡。并用更新后驗證替換事前同步，以進一步減少優(yōu)化器步驟中的氣泡（下圖底部）。

零泡沫管道計劃

DeekSeek 中的 DualPipe

DeekSeek（從 V3 開始）引入了 DualPipe Schedule，其思想與 Zero Bubble Pipeline Schedule 類似，但進行了一些額外的更改，以進一步提高計算與通信的比率和效率：

細(xì)粒度階段：將每個塊分為 4 個部分：注意力、全對全調(diào)度（處理設(shè)備之間的通信）、MLP（多層感知器）、全對全組合（跨設(shè)備合并輸出）。對于后向塊，注意力和 MLP 進一步分為兩個部分：后向輸入（B）和后向權(quán)重（W），如 Zero Bubble。
雙向流水線調(diào)度可同時從流水線兩端提供微批次，并且大部分通信可以完全重疊（參見下圖中的 2 個黑色箭頭）。為了支持雙向流水線調(diào)度，DualPipe 需要保留兩份模型參數(shù)。如果我們有 8 個設(shè)備和 8 層模型，則在零氣泡調(diào)度中，每個設(shè)備都有一個相應(yīng)的層。但在 DualPipe 調(diào)度中，為了處理雙向流水線，設(shè)備 0 應(yīng)該具有模型的 layer0 和 layer7，而設(shè)備 7 應(yīng)該具有模型的 layer7 和 layer0。

雙管計劃

為了保證 DualPipe 有足夠的計算性能，DeepSeek 還定制了高效的跨節(jié)點全對全通信內(nèi)核（包括調(diào)度和合并），以節(jié)省專用于通信的 SM 數(shù)量。更多詳細(xì)信息請參閱DeepSeek-V3 技術(shù)報告。

總結(jié)

富有創(chuàng)意的 DualPipe 調(diào)度，加上他們出色的基礎(chǔ)工程優(yōu)化，使得 DeepSeek 能夠充分利用集群的計算資源（GPU）。從這部分，我可以看到這個團隊的聰明才智和出色的工程精神。這可能部分是因為他們的資源與其他 LLM 巨頭（如 OpenAI、Meta、Google 等）相比有限。