20251108筆記 內容可能有錯誤，請參考原始影片

[李宏毅【生成式AI時代下的機器學習(2025)】](https://www.youtube.com/playlist?list=PLJV_el3uVTsNZEFAdQsDeOdzAaHTca2Gi)

[【生成式AI時代下的機器學習(2025)】助教課：利用多張GPU訓練大型語言模型—從零開始介紹DeepSpeed、Liger Kernel、Flash Attention及Quantization](https://www.youtube.com/watch?v=mpuRca2UZtI&list=PLJV_el3uVTsNZEFAdQsDeOdzAaHTca2Gi&index=6)

### 【生成式AI時代下的機器學習(2025)】助教課：利用多張GPU訓練大型語言模型—從零開始介紹DeepSpeed、Liger Kernel、Flash Attention及Quantization 大綱

今天助教課主要探討在生成式 AI 時代，訓練大型語言模型 (LLMs) 所面臨的硬體和軟體挑戰，並介紹用於解決這些挑戰的四個核心工具：DeepSpeed、Flash Attention、Liger Kernel 與 Quantization。

**I. 訓練 LLMs 的核心挑戰：記憶體限制**
* **問題概述：** LLM 模型權重、梯度、優化器狀態與 activations 佔用過多 GPU 記憶體。
* **解決方案面向：** 針對參數/梯度/優化器狀態 (Part 1) 和 Activations (Part 2) 進行優化，並討論推論時的量化 (Part 3)。

**II. Part 1：優化參數、梯度與優化器狀態 (DeepSpeed/Zero)**
* **核心工具：** DeepSpeed (Zero Redundancy Optimizer)。
* **Zero 策略：** 通過將模型狀態切分到多張 GPU 上實現平行運算。
* Zero 1：切分 Optimizer States (佔用空間最大)。
* Zero 2：切分 Optimizer States + Gradients。
* Zero 3：切分所有模型狀態 (Optimizer, Gradients, Parameters)。
* **額外技術：** 使用 CPU RAM 進行 Offload (但傳輸速度慢)。

**III. Part 2：優化 Activations 記憶體 (Liger Kernel & Flash Attention)**
* **核心挑戰：** 長輸入序列導致 Attention activations 佔用 $N^2$ 級別的記憶體。
* **通用技術：** Activation Recomputation (Gradient Checkpointing)。
* **加速技術 (Kernel Optimization)：**
* Flash Attention：透過重寫 Attention 運算的核心功能 (Fused Kernel) 來加速計算並減少記憶體使用。
* Liger Kernel：使用 Triton 實作的 Kernel 函式，用於加速 LLM 常用運算並節省記憶體。

**IV. Part 3：量化技術 (Quantization)**
* **目的：** 通過有損壓縮 (Lossy Compression) 減少模型儲存時所需的位元數 (如從 32-bit 降至 8-bit 或 4-bit)，主要用於推論 (Inference)。

---

### I. 訓練 LLMs 的核心挑戰：記憶體與資源 (The Challenges)

#### A. 訓練 LLM 的困難點：巨大的記憶體需求
訓練 LLM 困難的主因是模型本身太大。

1. **記憶體佔用計算 (以 8B 模型為例)：**
* 假設訓練一個 8 億（8B）參數的模型。
* **權重 (Parameters):** 權重通常以 32-bit 浮點數儲存。8B 參數 $\times$ 32 bit (4 bytes) = **32 GB**。
* **FP16 權重 (訓練用):** 由於 NVIDIA GPU 對 16-bit 格式有加速，訓練時通常會轉為 FP16 格式，儲存空間減半。
* **梯度 (Gradients):** 每個權重都有一個梯度，與 LM 權重大小一樣，約 **32 GB**。
* **Optimizer States (Adam):** Adam 優化器要求每個權重儲存一個 momentum (m) 和一個 variance (v)。這兩個狀態都以 FP32 儲存，因此需要 2 倍的 32 GB，約 **64 GB**。
* **總計（非 Activations 部分）：** 權重 (32G, FP32) + 梯度 (32G) + Optimizer States (64G) $\approx$ **128 GB**。
* 此 128 GB 已超過目前幾乎所有單張 GPU 的儲存容量。

2. **Activation 帶來的挑戰：長度與 $N^2$ 複雜度**
* Activation 是模型每一層對輸入產生的反應（輸出），在 Back Propagation 時需要儲存。
* Attention 機制的 Activation 佔用空間最大，因為它計算序列間的方形矩陣，空間複雜度從 $O(N)$ 變成 $O(N^2)$（N 是輸入序列長度）。
* 對於 8B 模型，當輸入長度達到 32k tokens 時，Attention 的 Activation 佔用空間可達到 **1.35 TB** 這麼誇張的數字。
* 在訓練長輸入 (例如 16k 或更長) 的模型時，Activation 是限制記憶體的主要因素。

3. **Batch Size 的要求：**
* 訓練 LLM 需要非常穩定的梯度，因此要求 Batch Size 必須很大，通常需要 4 到 60 個 Million tokens。
* 例如，訓練 32k context 時，Batch Size 可能需要 1920，總共約 61M tokens。

#### B. 解決 Batch Size 挑戰的技術
* **梯度累積 (Gradient Accumulation):** 當無法使用足夠大的 Batch Size 時，可將 Global Batch Size 切分成若干 Mini Batch。
* GPU 每次只處理 Mini Batch，計算梯度但不更新模型，直到累積到足夠的 Mini Batch 數量後，才統一進行一次 Optimizer 更新。

### II. Part 1：優化參數、梯度與優化器狀態 (DeepSpeed/Zero)

Part 1 旨在優化模型參數、梯度和 Optimizer States 等（圖表中橘色 Activation 以外的部分）所需的記憶體。

#### A. DeepSpeed 與 Zero 演算法
DeepSpeed 是由 Microsoft 開發的軟體套件，主要用於發揮多張 GPU 的平行運算能力。其核心技術是 **Zero Redundancy Optimizer (Zero)**。Zero 的精神是將訓練所需的組件（128GB 的模型狀態）切成小片，分裝到多張 GPU 中。

Zero 演算法依據節省的記憶體量分為三個等級：

| Level | 切分內容 (Partitioned Components) | 節省記憶體 | 傳輸 overhead | 選擇依據 |
| :---: | :---: | :---: | :---: | :---: |
| **Zero 1** | Optimizer States (佔用 64G) | 最大 | 較少 | Optimizer 雖然大但使用頻率低（只在 Global Step 更新時才用到）。 |
| **Zero 2** | Optimizer States + Gradients (32G) | 更多 | 增加 (中) | |
| **Zero 3** | 所有模型狀態 (Parameters + Gradients + Optimizer States) | 最多 | 增加 (高) | 能夠將訓練大小切到最小。|

* **平行化考量：** 一旦將模型狀態切分，GPU 之間就會產生額外的傳輸 (transfer) overhead。
* **傳輸速度：** NVIDIA GPU 使用 NVLink 技術，傳輸速度可達每秒 900 GB，速度非常快。DeepSpeed 還會針對傳輸進行排程 (scheduling)，盡量避免延遲訓練。
* **實例效果：** 使用 Zero 3，即使是 80GB VRAM 的 H100 也能夠訓練 16k 長度的 8B 模型。

#### B. CPU Offload 技術
如果 GPU 記憶體真的很小，連 Zero 切分後都裝不下，可以將一部分訓練組件搬到 CPU 的 RAM 中儲存。

1. **Optimizer Offload：** 將與 Optimizer 相關的組件搬到 CPU。GPU 上只保留 LLM 的權重。
2. **Parameter Offload：** 甚至將模型權重也放到 CPU。
3. **缺點：** CPU 與 GPU 之間的傳輸速度比 GPU 之間的速度慢一個數量級。
* 實測顯示，使用 Offload (V100, B=1, 長度 512) 進行一個 Step 需要 74 秒，速度極慢。
* **結論：** 除非目標是做極大模型的實驗，否則不建議使用 Offload 技術，因為它會讓訓練變慢十倍以上。

### III. Part 2：優化 Activations 記憶體 (Liger Kernel & Flash Attention)

Part 2 專注於解決長 Context 訓練時 Activation 佔用記憶體過大的問題。

#### A. Activation Recomputation / Gradient Checkpointing
這是一種解決 Activation 記憶體問題的直觀方法。

* **原理：** 在 Forward Pass 時，只儲存少量必要的 Activation。
* 在 Backward Pass 需要時，再逐層重新計算所需的 Activation。
* **代價：** 雖然節省了極大的儲存空間，但會增加計算 (computation) overhead，導致訓練速度稍微變慢。

#### B. Kernel Optimization 基礎
由於 Attention 運算在 GPU 上執行，優化方式是重寫 Attention 實作的程式碼 (Kernel) 讓它跑得更快。

* **Kernel 的定義：** 跑在 GPU 上的 Function。
* **Kernel 開發工具 (從高階到低階)：**
1. **PyTorch：** 方便但相對較慢。
2. **Torch Compile：** PyTorch 的 decorator，能無痛加速程式碼，它會編譯計算圖 (graph) 並做記憶體調度和預取 (prefetch)。
3. **Triton：** OpenAI 開發的 Python 介面，可用 Python 寫 CUDA 函式，比 PyTorch 靈活。
4. **CUDA：** NVIDIA 開發的 C/C++ 套件，最低階，速度最快，但學習曲線較陡峭。

#### C. Flash Attention
Flash Attention 的目標是加速 Attention 的計算，同時使用更少的記憶體。

1. **加速原理 (Fused Kernel)：**
* 標準 Attention 涉及矩陣乘法、Scale、Softmax 等四五個複雜函式。
* Flash Attention 將這些複雜操作壓成一個單一的 Kernel 函式（稱為 **Fused Kernel**）。
* 實務上，GPU 對矩陣乘法已做特殊加速，反而是 Softmax 或 Masking 等運算耗時較長。Fused Kernel 針對這些環節進行優化。
* **結果：** 計算速度大幅提升，尤其在長 Context 下比 CUDA 實作的 Attention 更快。

2. **記憶體優化原理：**
* Flash Attention 透過將大部分資料（如 QKV 矩陣）先放於 CPU RAM (或速度較慢的記憶體)，只在需要計算時將其 Flash 到 GPU 上運算，藉此大幅降低記憶體需求，達到接近線性的記憶體空間。

#### D. Liger Kernel
Liger Kernel 是一套使用 Triton 實作的 Kernel 集合，用於封裝 LLM 運算中常用的功能。

* **特點：** 它可以讓 LLM 算得更快，且使用更少的記憶體。
* **使用方式：** 載入模型時，只需將 `AutoModelForCausalLM` 替換為 `AutoLigerKernelForCausalLM` 即可。
* **效益：** 不論在速度還是記憶體使用上，Liger Kernel 都有顯著的優勢，允許訓練更長的 Context 模型。

### IV. Part 3：量化技術 (Quantization)

量化 (Quantization) 是一種有損壓縮 (lossy compression) 技術。

* **目的：** 將模型權重儲存的浮點數精度降低。
* **方法：** 將 32-bit 的資料壓縮成 8-bit、4-bit，甚至 1-bit。
* **機制：** 壓縮後的 Tensor 儲存空間較小，需要計算時再還原 (Dequantization)。
* **應用場景：** Quantization 主要用於**推論** (Inference)，以減少模型載入記憶體時的空間需求。
* **實例：** 在 Colab 上，免費 GPU (T4, 15GB VRAM) 可載入 8B 參數的模型。這通常是通過 GGML 家族的 Quantization 演算法實現的，例如 Q8 (8-bit 量化)，只需約 8GB 記憶體即可載入 8B 模型。

### V. 總結與建議

* 訓練 LLM 需要大量的 GPU 記憶體，但透過 DeepSpeed (Zero)、Flash Attention 和 Liger Kernel 等開源工具，可以有效將模型塞入並進行訓練。
* **DeepSpeed 建議：** 官方 API 較難使用，建議透過 Hugging Face `transformer` 的 Trainer 介面，使用 JSON 檔案來開啟 DeepSpeed 的功能。
* **CPU Offload 警告：** 雖然 Offload 可以讓極大的模型在有限的 GPU 資源下訓練，但由於 GPU 與 CPU 之間的傳輸速度太慢，會極大地拖慢訓練速度，實際實驗中並不實用。