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

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

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

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

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

零氣泡流水線并行性

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

1F1B管道時(shí)間表

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

零泡沫管道計(jì)劃

DeekSeek 中的 DualPipe

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

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

雙管計(jì)劃

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

總結(jié)

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