Allen Wang
EN
Enhance LLM Training Efficiency
← 文章

Enhance LLM Training Efficiency

發佈時間

前言#

訓練一般的神經網路時，我們很少需要煩惱模型放不放得進 GPU。但訓練 LLM 完全是另一回事：模型本身就已經大到單張 GPU 裝不下，更不用說訓練過程中還會產生的一大堆中間結果。

解決方式是同時使用多張 GPU，把負擔拆開來。但在談怎麼拆之前，得先搞清楚記憶體到底被誰吃掉了。

這篇會先算一次訓練 LLM 的記憶體帳單，再依序介紹 DeepSpeed、Flash Attention、Triton kernel 與量化這幾種優化手段。

記憶體到底花在哪#

LLM 的訓練流程和一般的神經網路是一樣的，如下圖所示：

神經網路訓練流程示意圖：前向傳播計算損失，反向傳播計算梯度，再由優化器更新權重

訓練 LLM 就只是把中間的神經網路換成 LLM 而已，流程沒有改變。

真正的問題是 LLM 所需要的記憶體空間實在太大了。以一個 8B 的模型為例：

8B 模型訓練時的記憶體用量拆解：模型權重、梯度與優化器狀態各自佔用的空間

先不論 input，光是模型本身就已經塞不下了。而這還只是靜態的部分，訓練時真正的開銷來自三個地方：

  1. 模型權重：8B 參數以 FP16 儲存就要 16 GB。
  2. 梯度：每個參數都要對應一份梯度，等於再一份同樣大小的空間。
  3. 優化器狀態：Adam 需要為每個參數保存一階與二階動量，加上 FP32 的主權重副本，這一塊往往比模型本身還大。

如果再考慮每一層都會產出的 activation，一個 32 層的模型累積下來就會是 TB 等級的容量。

計算層面的優化#

在動用多張 GPU 之前，有兩個單卡就能做的優化：

  • 針對 activation：使用 activation recomputation（也叫 gradient checkpointing）。做法是前向傳播時只保留少數幾個檢查點，其餘的 activation 直接丟掉，等反向傳播需要用到時再重新算一次。這是典型的「用計算換記憶體」，通常多花 20% 到 30% 的計算時間，換來大幅下降的記憶體用量。
  • 針對 gradient：使用 gradient accumulation，把一個大 batch 拆成好幾個小 batch 依序前向與反向，梯度累加起來之後才更新一次權重。這樣可以在有限的記憶體下，得到近似大 batch size 的訓練效果。

前面講的是針對計算上的優化，接下來就會來說怎麼用多張 GPU 來優化訓練過程。

DeepSpeed#

DeepSpeed 是微軟開發的開源深度學習優化庫。它的核心方法是把前面提到的 optimizer state、gradient、model weight 這三塊，切分到多張 GPU 之間，而不是每張卡都各存一份完整副本。

因為 GPU 與 GPU 之間可以通信，需要用到別人手上那一份時再去要就好。某種程度上，它是花費了 communication 的成本，換取單一 GPU 所需空間的下降。這個技術也就是 ZeRO（Zero Redundancy Optimizer）。

它一共分成三個層次，切得越多，省得越多，通信量也越大：

  1. ZeRO-1：只切分 optimizer state，這是只有在 global step 更新時才會用到的部分，所以通信成本最低。
  2. ZeRO-2：在前一階段的基礎上，再切分 gradient。
  3. ZeRO-3：連訓練時 16 bits 的模型權重本身也一起切分，前向傳播時需要哪一層才把那一層收集回來。

底下是它的結果：

DeepSpeed ZeRO 三個階段的記憶體用量比較，每往上一階單卡記憶體需求都明顯下降

除此之外，DeepSpeed 還有 offload 的技術，也就是再加上 CPU RAM 來存放暫時用不到的部分。不過因為 CPU RAM 跟 GPU RAM 之間的通信速度太慢，會直接變成訓練的瓶頸，所以除非是真的塞不下、寧可慢也要跑起來的情況，否則建議不要用。

Flash Attention#

Flash Attention 是一種注意力機制的實作演算法，目的是提高 Transformer 模型的效率，讓它能夠處理更長的序列長度，並更快地進行訓練和推理。

它的關鍵洞察是：標準的 attention 實作瓶頸不在「算得慢」，而在「搬得慢」。計算 attention 時需要把 N × N 的注意力矩陣寫進 GPU 的 HBM（高頻寬記憶體）再讀回來，而這個矩陣的大小隨序列長度平方成長，記憶體頻寬因此成為真正的限制。

Flash Attention 的做法是把輸入切成小塊（tiling），在速度快上一個量級的 SRAM 裡分塊完成計算，全程不把完整的 N × N 矩陣寫出去。結果是記憶體用量從平方級降到線性級，速度也同時變快。

至於數學上如何在分塊的情況下仍然算出正確的 softmax，等我修完大型語言模型後會再來補充這裡的細節。

Triton Kernel#

把 PyTorch 的 kernel function 用 low-level 的 Triton 重寫，有點像是用 C++ 把 Python 的熱點改寫成 Cython 的概念。

好處除了單一運算本身更快之外，更重要的是可以做 kernel fusion：把好幾個連續的小運算合併成一個 kernel，中間結果就不需要來回寫入記憶體。以剛才提到的記憶體頻寬瓶頸來看，減少一次來回往往比讓計算本身變快更有效。事實上 Flash Attention 就是這個想法的一個具體案例。

Quantization#

Quantization 是 AI 模型的壓縮技術，把原本 FP16 或 FP32 的權重壓縮成 8 bits、4 bits 等較小的表示，需要計算時再還原回來。

以 8B 模型為例，從 FP16 換成 INT4 可以把權重從 16 GB 壓到 4 GB 左右。代價是精度損失，所以量化更常被用在推論階段；訓練時則多半採用 QLoRA 這類做法，把凍結的基礎模型量化後存放，只讓少量可訓練的參數維持在高精度。

結語#

這些方法解決的其實是同一個問題的不同面向：記憶體不夠，以及資料搬運太慢。

Activation recomputation 和量化是用計算換空間，ZeRO 是用通信換空間，而 Flash Attention 和 Triton kernel 則是直接減少搬運本身。實務上這幾種手段通常是疊加使用的，讓我們可以在有限的資源下，訓練出最好的模型。