---
title: 多 GPU 訓練大型語言模型（LLM）
description: 從原理到實踐，了解如何利用多 GPU 技術有效訓練大型語言模型
tags: DeepSpeed, ZeRO, FlashAttention, GPU, LLM, 模型訓練
image: https://i.imgur.com/mErPwqL.png
lang: zh-tw
---
# 多 GPU 訓練大型語言模型（LLM）
###### tags: `DeepSpeed` `ZeRO` `FlashAttention` `LLM` `GPU` `模型訓練`
[TOC]
## 前言
大型語言模型（LLM）的訓練需要龐大的計算資源，特別是 GPU 記憶體。本文介紹如何利用多 GPU 技術有效訓練大型語言模型，從原理到實作，包含常見工具和技術。
### 參考資源
- 課程影片：[利用多張GPU訓練大型語言模型—從零開始介紹DeepSpeed、Liger Kernel、Flash Attention與Quantization (李宏毅2025課程)](https://www.youtube.com/watch?v=mpuRca2UZtI)
- 課程投影片及更多資料：[利用多張GPU訓練大型語言模型 (Lily's AI學習筆記)](https://lilys.ai/notes/878530)
## 1. 正向與反向傳播流程（Forward & Backward Pass）
大型語言模型的訓練流程包含以下步驟：
- **輸入**：token序列，例如 "The dog chased its ..."
- **模型輸出**：token分佈（output distribution），例如可能是 "tail"、"shadow"、"ball" 等
- **正確標籤**：例如 "tail"
- **損失計算**：使用交叉熵（cross-entropy）損失函數計算誤差
- **反向傳播**：計算 LLM 的梯度（gradients）
- **參數更新**：使用 Adam 優化器更新模型參數：
- 動量（Momentum）
- 變異數（Variance）
這個過程與一般神經網路訓練類似，差別在於大語言模型的參數量極大，導致訓練困難。
### 大語言模型記憶體架構流程圖
```mermaid
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
:::info
**記憶體需求示例**：
- 單層 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 倍，可作為額外存儲：
:::warning
- **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 庫，使用極為簡單：
```python
# 標準寫法
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
:::success
**量化效益**：
- 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）
- 針對大型檢查點實施增量保存策略
:::spoiler 進階閱讀與深入探討
如果你對多GPU訓練的數學原理與更深入的理論感興趣，可以參考：
- [ZeRO 論文](https://arxiv.org/abs/1910.02054)中的公式推導
- [FlashAttention](https://arxiv.org/abs/2205.14135)的IO感知算法設計
- GPU內存模型與NVLink頻寬計算
這些內容需要較強的數學和系統架構基礎，但能幫助你理解為何這些技術如此有效。
:::
## 推薦資源
- [HuggingFace UltraScale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook) - GPU 訓練細節指南（推薦閱讀時間：2-3天）
- [DeepSpeed with Transformers](https://huggingface.co/docs/transformers/main/deepspeed) - 簡化 DeepSpeed 配置
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://arxiv.org/abs/1910.02054) - 原始 ZeRO 論文
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://arxiv.org/abs/2205.14135) - FlashAttention 論文
- [DeepSpeed ZeRO 理論與實踐](https://zhuanlan.zhihu.com/p/675360966) - 詳細解釋 ZeRO 各階段的數學原理
- [Fine-Tuning Large Language Models with DeepSpeed](https://medium.com/@yxinli92/fine-tuning-large-language-models-with-deepspeed-a-step-by-step-guide-2fa6ce27f68a) - 使用 DeepSpeed 微調 LLM 的實用指南
Sign in with Wallet
Connect another wallet