OpenAI 大型神經網路訓練技術

訓練大型神經網路需要協調 GPU 叢集以執行單一同步計算。由於模型與叢集規模的增長，實務者使用各種平行化技術將計算負載與記憶體需求分散到多個硬體加速器上。

資料平行

資料平行透過將相同的模型參數複製到每個工作節點，並指派不同的訓練批次子集同時處理，使多個 GPU 能夠被使用。雖然這樣可以利用許多 GPU 的計算能力，但模型仍必須能夠容納於單一 GPU 的記憶體中。

為了在工作節點之間保持一致性，系統必須協調參數更新。標準的同步方法包含三個步驟：

  1. 在每個工作節點上獨立計算梯度。
  2. 在所有工作節點之間平均梯度（阻塞式通訊步驟）。
  3. 在每個工作節點上獨立計算新參數。

雖然存在非同步同步方案以減少阻塞平均的開銷，但它們往往降低學習效率，導致大多數實務者偏好同步方法。

管線平行

管線平行將模型的連續區塊（相鄰層）分割至不同的 GPU 上，降低每個裝置的記憶體占用。

解決「氣泡」問題

簡單的管線實作會產生「氣泡」——即工作節點等待前一個機器的輸出而閒置的時間。為了最小化這些氣泡，批次會被切分成更小的微批次。每個工作節點在微批次可用時立即開始處理下一個微批次，將計算與等待時間重疊。梯度在微批次之間進行平均，且參數更新僅在所有微批次完成後才執行。

排程策略

管理前向與反向傳遞時，主要使用兩種排程方案：

  • GPipe：工作節點連續處理前向與反向傳遞，於最後同步聚合多個微批次的梯度。
  • PipeDream：工作節點交替處理前向與反向傳遞，雖可提升效率，但可能導致部分計算使用過時的參數。

張量平行

張量平行將層內的個別運算「水平」切分至多個 GPU。以 Transformer 等現代架構為例，主要瓶頸在於將激活批次矩陣與大型權重矩陣相乘。張量平行將權重矩陣切分為等大小的碎片，分別放置於不同的 GPU 上計算整體矩陣乘積的部分，然後再通訊合併結果。

具體實作包括：

  • Megatron-LM：在 Transformer 的 MLP 與自注意力層內平行化矩陣乘法。
  • PTD-P：結合張量、資料與管線平行，將非相鄰層分配給每個裝置，以降低氣泡開銷。
  • Sequence Parallelism：將輸入序列在時間維度上切分為多個子樣本，以減少峰值記憶體消耗。

專家混合（MoE）

專家混合（Mixture-of-Experts，MoE）允許模型在參數數量上擴展，而不會成比例增加計算成本。其透過門控機制，只選取網路中一小部分權重（即「專家」）來計算給定輸入的輸出。由於不同的專家可以部署在不同的 GPU 上，MoE 提供了一種可擴展的方式來增加模型使用的 GPU 數量。

記憶體節省設計

除了平行化之外，還有多種策略用於減少訓練所需的裝置記憶體：

  • Checkpointing（激活重算）：不將所有原始激活保存以供梯度計算，而是只保存部分，並在反向傳遞時即時重新計算中間激活。選擇性激活重算進一步優化此方式，只對存儲成本高但計算成本低的激活進行 checkpoint。
  • Mixed Precision Training：使用較低精度的數字（通常為 FP16）進行訓練，可提升 FLOP 效率並節省裝置記憶體，且對精度的損失極小。
  • Offloading：暫時將未使用的資料移至 CPU 或其他裝置。ZeRO 實作將參數、梯度與優化器狀態分散於硬體上，並在需要時再實體化。
  • Memory Efficient Optimizers：使用如 Adafactor 等優化器，以減少執行狀態的記憶體占用。
  • Compression：壓縮中間結果，例如 Gist 壓縮反向傳遞的激活，或 DALL·E 在同步前壓縮梯度。

Sources