---
# System prepended metadata

title: '利用多張 GPU 訓練大型語言模型：從 0 開始介紹 DeepSpeed, Flash Attention, Liger Kernel 和 Quantization'
tags: [分散式訓練]

---

## 訓練 LLM 所面臨的根本挑戰：記憶體與運算瓶頸
訓練大型語言模型最主要的瓶頸在於記憶體使用量與訓練穩定性，以一個 80 億參數（8B）模型為例，若使用傳統 fp32 精度格式，光是儲存模型參數本身就需要 32GB VRAM，搭配 Adam 優化器的 momentum 與 variance 需要各 32GB。

訓練過程還會產生一份梯度資料（gradient tensor），用來反映 loss 函數對各個模型參數的偏導數，這份梯度資料的精度通常與前向傳播相同，並且可以使用較低精度的 fp16，因此大小為 16GB。

![image](https://hackmd.io/_uploads/Ski66SA6yg.png)
資料來源：[【生成式AI時代下的機器學習(2025)】助教課](https://www.youtube.com/watch?v=mpuRca2UZtI&t=1718s)

除了參數與梯度外，activations（即前向傳播中每層的中間值）更是一大記憶體殺手，其大小會根據 batch size、context length（token 長度）與 hidden size（例如 embedding 維度）而有所不同。圖中顯示，若一個 32 層模型，每層的 attention 占用約 40GB，Feedforward 模組約 2.2GB，還有 LayerNorm 等累積下來，activations 需求將達 1.35TB，為了解決這個問題，實務中常使用 Gradient Checkpointing 技術，選擇性儲存關鍵激活值，其餘反向傳播時再重算，以節省大量記憶體。

![image](https://hackmd.io/_uploads/rJFWbICTJx.png)
資料來源：[【生成式AI時代下的機器學習(2025)】助教課](https://www.youtube.com/watch?v=mpuRca2UZtI&t=1718s)

## DeepSpeed、Flash Attention、Liger Kernel 與 Quantization 背後的運作原理
![image](https://hackmd.io/_uploads/r1MOML0TJg.png)

### DeepSpeed
DeepSpeed 是由微軟開發的分散式訓練框架，它的核心價值在於讓超大規模模型能夠在有限 GPU 資源下完成訓練，這項技術的核心機制是 ZeRO（Zero Redundancy Optimizer），透過將模型參數、梯度與優化器狀態切分並分散儲存到不同的 GPU 上，避免了傳統資料並行方式中每張 GPU 上需要完整複製整個模型的高記憶體消耗。

#### ZeRO 三階段優化：
Stage 1：將優化器狀態（Optimizer States 如Adam的動量和方差）分片並分配到各個GPU，而梯度和參數仍由所有GPU完整持有。
Stage 2：進一步將梯度（Gradients）也進行分片，每個GPU僅儲存與其計算相關的部分，並在反向傳播後動態通信以更新參數。
Stage 3：將模型參數本身也分片，每個GPU僅持有部分參數，並在需要時通過通信獲取其他部分資訊，這使得單張GPU能訓練數百億甚至千億參數的模型，記憶體效率提升至極致。
![image](https://hackmd.io/_uploads/rykp7IRa1l.png)
:::info
#### 支援混合精度訓練
結合FP16（半精度浮點數）和FP32（單精度浮點數）來加速計算並減少記憶體使用，為了解決精度下降可能導致的梯度消失問題，採用動態損失縮放技術，確保訓練穩定性。
#### Offload 
專為降低 GPU 顯示記憶體使用壓力而設計的功能，從昂貴且資源有限的 GPU VRAM 中搬移到 CPU 的主記憶體（RAM），甚至是更大容量但更慢速的 NVMe 磁碟中。然而，從 GPU 傳輸資料到 CPU RAM 已需耗費數百微秒，若是搬到 NVMe（即磁碟），I/O 時間更可能達到毫秒等級，這種傳輸延遲會嚴重拖慢訓練過程中的運算。
:::

### Flash Attention
Flash Attention 針對 Transformer 模型中的注意力機制進行優化，解決長序列場景下的記憶體瓶頸，傳統注意力計算需將查詢（Query）、鍵（Key）和值（Value）矩陣完整載入，並生成龐大的注意力分數矩陣，導致記憶體需求隨序列長度平方增長。
![image](https://hackmd.io/_uploads/H1x6y27AJe.png)
#### Flash Attention 主要特點包括：
* 區塊處理（Tiling）：把輸入序列分成多個小區塊，一次只處理一塊，並透過逐步累積 softmax 結果與加權值，避免一次性構建完整矩陣。
* Fused kernel：所有運算（如矩陣乘法、遮罩、softmax、dropout 等）串接在單一 GPU 核心（kernel）內執行，這種做法允許將中間結果（注意力計算）保存在 GPU 的 register 和 L1/L2 cache 中，減少了對高頻寬記憶體（HBM）的訪問次數，大幅減少資料搬移與計算中斷。
* 重新計算（Recomputation）：在反向傳播過程中，透過重新計算部分中間結果，減少了對中間激活值的存儲需求，進一步降低了記憶體使用量。

![image](https://hackmd.io/_uploads/Bk0eQ2QAke.png)

### Liger Kernel
Liger Kernel 是由 LinkedIn 開發的 Triton 內核集合，專為LLM訓練設計，通過高效內核替代 PyTorch 中的慢速層（如RMSNorm、RoPE和SwiGLU）來提升性能，它使用 Triton 語言編寫，直接生成高效的 GPU 代碼。

* 提升訓練吞吐量：透過高效的核心實現，Liger Kernel 能夠將多 GPU 訓練的吞吐量提高超過 20%，加速模型訓練過程。
* 減少記憶體使用：透過核心融合（Kernel Fusion）和就地替換（In-place Replacement）等技術，Liger Kernel 將記憶體使用量降低了 60%，允許在有限的硬體資源下訓練更大的模型。
* 易於整合：​Liger Kernel 可與 Hugging Face 的模型、PyTorch 的 FSDP 和 Microsoft 的 DeepSpeed 無縫整合，方便用戶在現有框架中使用。

### Quantization
模型量化是一種透過將模型的權重和激活值從高精度（如 32 位元浮點數）轉換為低精度（如 8 位元整數）來減少計算和記憶體需求的技術。主要方法包括：
* 後訓練量化（Post-Training Quantization, PTQ）：​在模型訓練完成後，對其進行量化，這種方法簡單易行，但可能會導致輕微的準確性下降。 ​
* 量化感知訓練（Quantization-Aware Training, QAT）：​在訓練過程中考慮量化的影響，透過模擬低精度運算來調整模型參數，以減少量化帶來的準確性損失。

![image](https://hackmd.io/_uploads/rJ8md2Q0kx.png)

## 總結
DeepSpeed、Flash Attention、Liger Kernel和 Quantization 各有側重，但它們在實際應用中常相互配合。例如，DeepSpeed 的 ZeRO 與 Flash Attention 結合，可在分布式環境中高效處理長序列注意力計算；Liger Kernel 則能進一步加速 DeepSpeed 中的特定層；Quantization 則為整體流程提供記憶體壓縮基礎，這些技術的共同目標是讓大模型訓練更高效、可擴展，並適應有限硬體資源的場景。