平衡流水線平行化與詞彙平行化以提升大型語言模型訓練效率

核心概念

本文提出了一種稱為「詞彙平行化」的新方法，用於解決大型語言模型訓練中，因詞彙層造成流水線平行化效能瓶頸的問題，透過將詞彙層計算分散到所有流水線設備，並優化計算順序和通訊策略，有效提升了模型訓練的吞吐量和記憶體使用效率。

摘要

大型語言模型訓練的效能瓶頸

近年來，隨著 Transformer 模型規模不斷擴大，模型平行化技術成為深度學習領域的重要研究方向。其中，流水線平行化（Pipeline Parallelism）由於其低通訊成本和高算術強度，成為訓練大型模型的常用方法。然而，流水線平行化面臨著兩個主要挑戰：流水線氣泡（Pipeline Bubbles）和高記憶體消耗。

詞彙層造成的不平衡問題

本文關注的是一個經常被忽視的問題：詞彙層（Vocabulary Layers）造成的不平衡問題。傳統上，Transformer 層會被平均分配到不同的流水線階段，而第一階段包含額外的輸入層，最後階段包含額外的輸出層。這種不平衡的設置會嚴重影響效能：

• 由於工作負載較少，其他流水線階段會出現流水線氣泡。
• 第一階段的額外輸入層加劇了記憶體瓶頸。
• 隨著詞彙量增加，這種不平衡會更加明顯。

詞彙平行化：解決方案

為了徹底解決這個問題，本文提出了詞彙平行化（Vocabulary Parallelism）方法，其設計原則如下：

• 將詞彙層按詞彙維度進行分割，並平均分配到所有流水線設備。
• 將詞彙層的計算表示為類似於 Transformer 層的前向/後向傳遞的傳遞過程，使其原生於流水線平行化。
• 將詞彙傳遞整合到流水線排程中，不應顯著影響原始排程的記憶體和效率。

詞彙傳遞的建構與優化

本文詳細介紹了如何將詞彙層分割成多個計算傳遞，並將其分配到不同的流水線設備上。同時，為了減少通訊障礙，本文提出了兩種新穎的演算法：

• 前向階段優化：將計算 softmax 所需的全局最大值和總和的通訊延遲到計算 softmax 之後，有效減少了通訊次數。
• 後向階段優化：將所有三個 all-reduce 通訊延遲到計算輸入梯度的矩陣乘法之後，進一步減少了通訊次數。

流水線排程

本文採用 Qi 等人提出的框架來建構流水線排程，將詞彙傳遞整合到現有的流水線排程中，例如 1F1B 和 V-Half。通過在最後一個 Transformer 層的前向和後向傳遞之間插入適當的時間間隔，為詞彙層計算創造空間，同時確保排程滿足所有依賴關係。

實驗結果

實驗結果表明，詞彙平行化方法可以顯著提高模型訓練的吞吐量，最高可達 51%，同時降低了峰值記憶體消耗。此外，當與記憶體平衡的排程（如 V-Half）結合使用時，詞彙平行化可以實現完全的記憶體平衡。

總結

詞彙平行化方法有效解決了大型語言模型訓練中詞彙層造成的不平衡問題，顯著提高了模型訓練的效率。

統計資料

在 Gemma2 9B 模型中，詞彙量為 256k 時，輸出層的計算量和參數記憶體大約是 Transformer 層的 5 倍。
與 na¨ıve 方法相比，詞彙平行化方法的吞吐量提高了 5% 到 51%。
詞彙平行化方法僅增加了少量且恆定的激活記憶體開銷。
與交錯流水線相比，詞彙平行化方法在 32 個 GPU 上的端到端訓練速度提高了約 11%。
在 24 個設備上，將詞彙量填充為 256032（48 的倍數）比使用原始值 256008 的效能提高了約 8%。
對於 21B 模型，詞彙平行化在 MFU 中的效能比交錯流水線高出 6.7% 到 8.2%。
在 V-Half 排程上，詞彙平行化方法的吞吐量比基準方法高出 7.2% 到 143%。
在 V-Half 排程上，基準方法在不同流水線設備之間的記憶體使用不平衡高達 45GB，而詞彙平行化方法實現了平衡的記憶體使用。

引述

從以下內容提煉的關鍵洞見

Balancing Pipeline Parallelism with Vocabulary Parallelism

by Man Tsung Ye... 於 arxiv.org 11-11-2024

https://arxiv.org/pdf/2411.05288.pdf

深入探究

除了詞彙層之外，還有哪些因素會導致流水線平行化訓練中出現不平衡？如何解決這些問題？

除了詞彙層之外，還有其他因素會導致流水線平行化訓練中出現不平衡，主要可以分為計算不平衡和記憶體不平衡兩個方面：
計算不平衡:

不同 Transformer 層的計算複雜度不同: 某些 Transformer 層，例如包含 Multi-Head Attention 的層，計算複雜度可能高於其他層，例如 Feed-Forward Network 層。

解決方案: 可以根據每一層的計算複雜度，動態調整分配到每個流水線階段的層數，或者將計算複雜度高的層進一步拆分到多個流水線階段。

資料集的樣本長度不一致: 在處理變長序列時，不同批次的樣本長度差異可能會導致計算時間不同，進而影響流水線的平衡。

解決方案: 可以使用動態批次大小調整策略，根據當前批次的樣本長度動態調整批次大小，或者使用 padding 將所有樣本填充到相同長度。

硬體異構性: 如果流水線中的不同設備具有不同的計算能力，例如使用不同型号的 GPU，則可能會導致計算不平衡。

解決方案: 可以根據設備的計算能力，調整分配到每個設備的計算量，或者使用更先進的流水線调度算法，例如 V-Half，來平衡不同設備的負載。
記憶體不平衡:

不同 Transformer 層的激活值大小不同: 某些 Transformer 層的激活值大小可能大於其他層，例如 Multi-Head Attention 層。

解決方案: 可以使用激活值檢查點技術，只儲存部分層的激活值，或者使用激活值壓縮技術，減少激活值的大小。

梯度累積步數: 更大的梯度累積步數會導致更大的激活值記憶體佔用。

解決方案: 可以根據模型大小和設備記憶體容量，調整梯度累積步數，或者使用混合精度訓練，減少激活值的大小。

詞彙平行化方法是否適用於其他模型架構，例如基於 RNN 的語言模型？

詞彙平行化方法主要針對 Transformer 模型中詞彙量過大導致的計算和記憶體不平衡問題。雖然 RNN 模型也包含詞彙層，但其架構與 Transformer 不同，因此直接套用詞彙平行化方法可能並不適用。

RNN 模型的詞彙層計算與隱藏狀態相關: RNN 模型的詞彙層計算需要用到每個時間步的隱藏狀態，而詞彙平行化方法需要對詞彙層進行切分，這會導致不同設備上的隱藏狀態不完整，影響模型訓練。
RNN 模型的並行化策略不同: RNN 模型通常採用時間步上的並行化策略，而不是像 Transformer 模型那樣在模型結構上進行並行化。
因此，要將詞彙平行化方法應用於 RNN 模型，需要針對 RNN 模型的特性進行調整和優化。例如，可以考慮將詞彙層的計算與隱藏狀態的更新解耦，或者探索新的 RNN 並行化策略。

如何在不影響模型精度的情況下，進一步優化詞彙平行化方法的效能？

在不影響模型精度的情況下，可以從以下幾個方面進一步優化詞彙平行化方法的效能：

減少通訊開銷: 詞彙平行化方法需要在不同設備之間進行通訊，例如 All-Reduce 操作。可以通過使用更高效的通訊原語，例如 NCCL，或者優化通訊拓撲結構來減少通訊開銷。
與其他模型並行化技術結合: 可以將詞彙平行化方法與其他模型並行化技術，例如張量平行化和流水線平行化，結合起來使用，進一步提高模型訓練效率。
硬體優化: 可以針對特定硬體平台，例如 NVIDIA GPU，對詞彙平行化方法進行優化，例如使用混合精度訓練和 GPU 核函數融合等技術。
探索更優的詞彙層切分策略: 可以根據模型架構和硬體資源，探索更優的詞彙層切分策略，例如使用非均匀切分或者動態切分等方法。
通過以上優化方法，可以進一步提高詞彙平行化方法的效能，同時保持模型精度。