大規模分布式訓練的挑戰

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

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

數據并行和模型并行并不相互排斥，它們可以一起使用來加速大規模模型訓練過程。通過對模型進行分區應用模型并行后，模型的一個分區可以復制并分發到多個 GPU 并應用數據并行。

到目前為止一切看起來都很好？現在讓我們來談談大規模分布式訓練的挑戰。在模型和計算分布在數百或數千個 GPU 上之后，如何充分利用這些大規模計算資源是一個巨大的挑戰（還有其他挑戰，例如如果一臺機器在訓練期間崩潰了怎么辦，因為單個集群中有數千個節點，節點崩潰的概率高于小集群，如何從檢查點恢復這部分，例如從空閑的 HA 池中添加新的 GPU / 節點回集群或重新平衡模型和訓練任務以適應拓撲等。在這篇博客中，我們重點介紹如何充分利用資源。）。正如我在上一節中所描述的，前向計算和后向計算之間存在依賴關系，層的計算之間存在依賴關系，模型已跨不同節點進行分區。需要跨節點進行數據傳輸和通信。由于依賴于其他節點的計算結果，某些節點可能處于空閑狀態，而計算結果仍在進行中；梯度、某一層的輸出和其他數據可能需要從一個節點傳輸到另一個節點，這可能會達到網絡帶寬瓶頸并導致某些節點空閑。可能有些步驟（例如，優化器步驟）需要同步，也會導致某些節點空閑。

零氣泡流水線并行性

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

1F1B管道時間表

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

零泡沫管道計劃

DeekSeek 中的 DualPipe

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

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