多 GPU 訓練大型語言模型（LLM）

tags: DeepSpeed ZeRO FlashAttention LLM GPU 模型訓練

前言

大型語言模型（LLM）的訓練需要龐大的計算資源，特別是 GPU 記憶體。本文介紹如何利用多 GPU 技術有效訓練大型語言模型，從原理到實作，包含常見工具和技術。

參考資源

• 課程影片：利用多張GPU訓練大型語言模型—從零開始介紹DeepSpeed、Liger Kernel、Flash Attention與Quantization (李宏毅2025課程)
• 課程投影片及更多資料：利用多張GPU訓練大型語言模型 (Lily's AI學習筆記)

1. 正向與反向傳播流程（Forward & Backward Pass）

大型語言模型的訓練流程包含以下步驟：

• 輸入：token序列，例如 "The dog chased its ..."
• 模型輸出：token分佈（output distribution），例如可能是 "tail"、"shadow"、"ball" 等
• 正確標籤：例如 "tail"
• 損失計算：使用交叉熵（cross-entropy）損失函數計算誤差
• 反向傳播：計算 LLM 的梯度（gradients）
• 參數更新：使用 Adam 優化器更新模型參數：
  • 動量（Momentum）
  • 變異數（Variance）

這個過程與一般神經網路訓練類似，差別在於大語言模型的參數量極大，導致訓練困難。

大語言模型記憶體架構流程圖

[mermaid 圖表 — 原始 HackMD 版本可正常渲染]

flowchart TD subgraph CPU["CPU區域(32-bit 精度)"] W32[LLM weights: 32GB] M32[Momentum: 32GB] V32[Variance: 32GB] Opt32[Adam optimizer] W32 --> M32 W32 --> V32 M32 --> Opt32 V32 --> Opt32 end

subgraph GPU["GPU區域(16-bit 精度)"] Input[Input: ?? GB] W16[LLM weights: 16GB] G16[LLM gradients: 16GB] Input --> W16 W16 --> G16 G16 -->|"loss function\nbackpropagation"| W16 end

W32 -->|"Utilize optimized fp16 computation on GPU"| W16

style CPU fill:#ffffff,stroke:#2f5597,stroke-width:2px style GPU fill:#fef3c7,stroke:#b45309,stroke-width:2px style W32 fill:#ffffff,stroke:#2f5597,stroke-width:2px style M32 fill:#ffffff,stroke:#2f5597,stroke-width:2px style V32 fill:#ffffff,stroke:#2f5597,stroke-width:2px style Opt32 fill:#ffffff,stroke:#2f5597,stroke-width:2px style Input fill:#fef3c7,stroke:#b45309,stroke-width:2px style W16 fill:#fef3c7,stroke:#b45309,stroke-width:2px style G16 fill:#fef3c7,stroke:#b45309,stroke-width:2px

上圖展示了大語言模型訓練中的記憶體使用結構：

• CPU 區域：通常存儲 32 位元精度的模型權重和優化器狀態
• GPU 區域：使用 16 位元精度進行高效運算，包含模型權重的副本和梯度

2. 記憶體需求分析

以一個 80 億（8B）參數的模型為例：

主要記憶體消耗：

項目	32位元 (FP32)	16位元 (FP16)	說明
LLM 參數	32GB	16GB	模型權重
LLM 梯度	32GB	16GB	每個參數的梯度
Adam 優化器狀態	64GB	N/A	每個參數的動量和變異數，固定使用 FP32
啟用記憶體 (activations)	可達TB級	可達TB級	隨序列長度增加，尤其是自注意力機制

合計基礎需求：約 128GB（不含 activations）

啟用記憶體（Activations）分析：

• 8B 模型通常有 32 層
• 每層都保存 activations，特別是自注意力層
• 自注意力機制的空間複雜度是 O(N²)，N 是序列長度
• 當輸入從 256 tokens 增加到 16K 或更多時，記憶體需求呈平方級增長
• 某些模型（如 DeepSeek V3、Google Gemini）的上下文長度可達 128K-200K

記憶體需求示例：

• 單層 8B 模型的自注意力機制在長序列時可佔用 40GB
• 32層合計可達 1.35TB 記憶體（若全部保存）

### 批次大小（Batch Size）考量：

為了獲得穩定的梯度，需要足夠大的批次：

• 通常需要 4-60M tokens/batch
• 例如 DeepSeek V3 訓練時使用 1920 x 32K = 61M tokens/batch
• 解決方案：梯度累積（Gradient Accumulation）
  • 將大批次切成多個小批次（mini-batch）
  • 每個 mini-batch 計算一次前向和反向傳播
  • 累積多個 mini-batch 的梯度後再更新模型

3. 多 GPU 的需求與挑戰

為什麼需要多 GPU 訓練？

• 模型大小：單一 GPU 無法容納完整模型及訓練所需狀態
• 記憶體爆炸：自注意力機制導致記憶體隨序列長度平方增長
• 訓練效率：訓練需要處理巨量資料和進行數十億次迭代

4. DeepSpeed 與 ZeRO 優化

DeepSpeed 是 Microsoft 開發的框架，基於 ZeRO（Zero Redundancy Optimizer）演算法，能有效解決多 GPU 訓練問題。

ZeRO 分層解法：

• ZeRO-1：僅優化器狀態分散存儲（32GB → 8GB/GPU，使用4GPU）

  • 最小化通信開銷
  • 在需要時從其他 GPU 獲取優化器狀態
  • 優化器狀態（optimizers）在訓練中用到的頻率較低，因此首先被選擇分散
  • 數學表示：單卡顯存消耗降至 $2\Phi+2\Phi + \frac{K*\Phi}{N_d}$（約原始需求的 26%）

• ZeRO-2：優化器狀態 + 梯度分散

  • 進一步減少單 GPU 記憶體需求
  • 稍微增加通信開銷
  • 數學表示：單卡顯存消耗降至 $2\Phi + \frac{(2+K)*\Phi}{N_d}$（約原始需求的 13.8%）

• ZeRO-3：優化器狀態 + 梯度 + 模型參數分散

  • 最大程度減少記憶體使用
  • 最高通信開銷，但仍可接受
  • 完整模型狀態（權重、梯度、優化器狀態）都被分散式存儲
  • 數學表示：單卡顯存消耗降至 $\frac{(2+2+K)*\Phi}{N_d}$（僅原始需求的 1.58%）

其中，$\Phi$表示模型參數量，$N_d$表示 GPU 數量，$K$為常數。

NVIDIA GPU 間通信：

• NVLink 技術每秒可傳輸 900GB
• DeepSpeed 提供智能調度，減少通信延遲

ZeRO 效能比較：

以 8B 模型，8 張 GPU 訓練為例：

• 不使用 ZeRO：每 GPU 需 >80GB
• ZeRO-1：每 GPU 需約 40GB
• ZeRO-2：每 GPU 需約 25GB
• ZeRO-3：每 GPU 需<20GB

ZeRO Offload：

CPU RAM 通常比 GPU 記憶體大 10 倍，可作為額外存儲：

• Optimizer Offload：

  • 將優化器狀態轉移到 CPU RAM
  • GPU 僅保留模型參數和梯度

• 完全 Offload：

  • 優化器狀態和部分模型參數都放在 CPU RAM
  • 需要時傳回 GPU

**實際測試**（使用 V100 32GB）： - 使用 Offload：每 GPU 使用 _{15GB，但一步需 74 秒 - 不使用 Offload，8 GPU：每 GPU 使用} 24GB，一步僅需 7.3 秒

結論：

• 能用 GPU 就不要用 CPU Offload（慢 10 倍）
• 8 張 V100 足以全面微調（fully fine-tune）8B 模型
• CPU RAM 設計優先考慮容量而非速度，而 GPU RAM 設計優先考慮速度

DeepSpeed 配置：

DeepSpeed 配置相對複雜，但可通過 Hugging Face Transformers 簡化：

• 只需編寫一個簡單的 JSON 配置文件
• 使用 Transformers 的 Trainer API 進行訓練
• 參考: Hugging Face 的 DeepSpeed 整合文檔

5. 啟用記憶體（Activations）優化

啟用記憶體再計算（Activation Recomputation）：

又稱梯度檢查點（Gradient Checkpointing）：

• 前向傳播時僅保存關鍵 checkpoints
• 反向傳播時重新計算需要的啟用記憶體
• 以計算成本換取記憶體節省
• 訓練略微變慢，但可顯著減少記憶體需求

FlashAttention 和 LargerKernel 技術：

FlashAttention：

• 重新實現注意力機制的 GPU kernel
• 核心技術：
  • 融合核心（Fused Kernel）：將多個操作合併為單一 GPU 核心
  • 記憶體管理：將部分中間結果存放在 CPU，需要時再載入 GPU
  • 效能提升：相比標準實現，速度提升 2-4 倍
  • 記憶體優化：將 O(N²) 空間複雜度近似降至線性
  • 計算重點：在注意力機制中，矩陣乘法已被 GPU 高度優化，真正耗時的反而是 dropout、softmax 和 mask 操作

LargerKernel：

• 由臺灣工程師開發的工具（在 LinkedIn 工作）
• 使用 Triton 語言實現高效 GPU kernel
• 集成到 Transformers 庫，使用極為簡單：

  # 標準寫法
  model = AutoModelForCausalLM.from_pretrained("gpt2")
  
  # LargerKernel 寫法
  model = AutoLargerKernel.from_pretrained("gpt2")
  

• 支持多種開源模型
• 提供顯著的速度和記憶體改進

6. Kernel 開發層級

為實現高效 GPU 計算，可選擇不同層級的工具（由高階到低階）：

1. PyTorch：高階，易用但較慢

   • 容易使用
   • 自動處理基本 GPU 運算
   • 靈活性較低

2. Torch Compile：透過 @torch.compile 裝飾器簡單加速

   • 在 PyTorch 基礎上無痛加速
   • 自動優化計算圖和記憶體調度

3. Triton：OpenAI 開發的 Python GPU 編程框架

   • 比 PyTorch 更靈活
   • 允許自定義 kernel 函數
   • 在 Python 中直接編寫 GPU 代碼

4. CUDA：NVIDIA 的低階 C/C++ GPU 編程工具

   • 最完整的控制
   • 學習曲線最陡峭
   • 最高效能但開發成本高

7. 量化技術（Quantization）

量化是一種有損壓縮技術，適用於推理階段：

• 原理：將高精度浮點數（如 FP32）轉換為低精度表示（8-bit、4-bit 等）
• 常見技術：
  • GGML/GGUF 系列（Llama.cpp 使用）
  • GPTQ
  • BitsAndBytes
  • AWQ
  • QLoRA

量化效益：

• 8B 模型使用 8-bit 量化僅需 ~8GB 記憶體
• 可在消費級 GPU（如 T4 15GB）上運行
• 使得 Google Colab 免費版（T4 GPU）能夠運行較大語言模型

## 8. 訓練挑戰的三大來源

根據 HuggingFace UltraScale Playbook：

1. 模型規模：參數、梯度和優化器狀態

   • 解決方案：DeepSpeed ZeRO

2. 啟用記憶體：尤其是長序列自注意力

   • 解決方案：FlashAttention、LargerKernel、梯度檢查點

3. 精度與批次：訓練穩定性與記憶體平衡

   • 解決方案：混合精度訓練、梯度累積

9. GPU 型號參考

GPU 型號	記憶體容量	適用模型規模
T4	15GB	量化後 8B-13B
RTX 4090	24GB	量化後 70B，訓練 7B
A100	40GB	訓練 13B，量化後 70B+
H100	80GB	訓練 70B

10. PyTorch Distributed 與 DeepSpeed 的差異

• PyTorch Distributed：

  • 專注於多 GPU 間的通信機制
  • 將相同代碼部署到多個 GPU 上執行
  • 處理主節點（Master Node）和工作節點（Worker Node）間的數據傳輸
  • 適合數據並行處理

• DeepSpeed：

  • 在 Distributed 基礎上增加了模型並行能力
  • 能將單個模型切分到多個 GPU 上
  • 處理跨 GPU 參數訪問與計算
  • 兩者通常結合使用：Distributed 處理通信，DeepSpeed 處理模型分割

11. Megatron-DeepSpeed：大規模訓練解決方案

Megatron-DeepSpeed 是由微軟將其 DeepSpeed 庫整合到 NVIDIA 的 Megatron-LM 框架中開發的強大工具。這個整合框架特別適合使用多 GPU 集群進行超大型語言模型的預訓練和微調。

• Megatron-LM：NVIDIA 開發的專為大型 Transformer 模型設計的框架
• DeepSpeed：微軟的優化庫，簡化並增強分布式訓練和推理
• 整合優勢：
  • 支持 3D 並行化（數據並行、模型並行和流水線並行）
  • 更有效的記憶體管理
  • 更好的計算負載均衡
  • 適應不同硬體架構（包括 AMD GPU）

使用 Megatron-DeepSpeed 訓練超大模型時，通常需要配置以下幾個關鍵部分：

1. 模型並行度（模型如何被拆分到不同 GPU）
2. 數據並行度（數據如何在 GPU 之間分配）
3. 流水線並行度（模型層如何分配到不同 GPU 組）
4. ZeRO 階段設置（根據硬體情況選擇最合適的優化級別）

12. 多 GPU 訓練的最佳實踐

基於最新研究和工業實踐，以下是使用多 GPU 訓練大型語言模型的一些最佳實踐：

1. 硬體選擇與配置：

   • 使用高帶寬互連的 GPU 集群（如 NVLink、InfiniBand）
   • 保證充足的 CPU 記憶體用於數據加載和可能的 offload
   • 考慮使用高速 SSD 進行 ZeRO-Infinity（NVMe offload）

2. 訓練策略優化：

   • 使用漸進式訓練：從短上下文開始，逐步增加到目標長度
   • 梯度累積時，選擇合適的微批次大小（micro-batch size）
   • 使用混合精度訓練（AMP）但注意數值穩定性

3. 記憶體使用優化：

   • 總是使用梯度檢查點（gradient checkpointing）
   • 警惕和監控記憶體尖峰（memory spikes）
   • 考慮使用優化過的 Attention 實現：FlashAttention 或 xFormers
   • 使用監控工具如 nvidia-smi、PyTorch Profiler 或 Weights & Biases 追蹤資源使用

4. 故障恢復機制：

   • 實施頻繁的檢查點保存策略
   • 使用分布式檢查點（distributed checkpointing）
   • 針對大型檢查點實施增量保存策略

進階閱讀與深入探討
如果你對多GPU訓練的數學原理與更深入的理論感興趣，可以參考：

• ZeRO 論文中的公式推導
• FlashAttention的IO感知算法設計
• GPU內存模型與NVLink頻寬計算

這些內容需要較強的數學和系統架構基礎，但能幫助你理解為何這些技術如此有效。

## 推薦資源

• HuggingFace UltraScale Playbook - GPU 訓練細節指南（推薦閱讀時間：2-3天）
• DeepSpeed with Transformers - 簡化 DeepSpeed 配置
• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models - 原始 ZeRO 論文
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness - FlashAttention 論文
• DeepSpeed ZeRO 理論與實踐 - 詳細解釋 ZeRO 各階段的數學原理
• Fine-Tuning Large Language Models with DeepSpeed - 使用 DeepSpeed 微調 LLM 的實用指南

---

本文最初發布於 HackMD @BASHCAT。