生成式AI時代下的機器學習 - 生成式 AI 時代下的大型語言模型訓練工具

生成式AI時代下的機器學習 - 生成式 AI 時代下的大型語言模型訓練工具¶

在當前的生成式 AI 浪潮中，大型語言模型 (LLMs) 扮演著核心的角色。為了訓練這些龐大的模型，往往需要多張好的 GPU，並且經常會遇到模型因為「塞不進 GPU 裡面」而無法訓練的挑戰。

本文將介紹 DeepSpeed、Flash Attention、Liger Kernel 與 Quantization 等關鍵工具，說明如何善用多張 GPU 的能力，克服訓練 LLM 過程中的巨大資源限制。

LLM 訓練的資源困境：為何模型會「塞不進去」？

訓練 LLM 之所以困難，主因在於模型本身實在是太大了。

1. 核心訓練元件的記憶體需求

如果訓練方法僅依賴 PyTorch 等框架，在大型模型情境下通常無法運行。我們必須先了解訓練一個 LLM 所需的記憶體總量。

以一個 80 億（8B）參數的模型為例，光是儲存其模型權重（使用 32 個 bit 的浮點數），就需要大約 32GB 的空間。然而，訓練所需的總記憶體遠超於此，主要來自以下幾個核心訓練元件：

這些核心訓練元件加總起來，一個 8B 模型總共需要大約 128GB 的記憶體。這個數字已經超過現今幾乎所有單張 GPU 的儲存大小。

1. 啟用值 (Activation) 與長輸入的挑戰

除了上述靜態參數外，訓練過程中最大的記憶體佔用者其實是 Activation（啟用值）。啟用值是指模型每一層對於輸入所產生的輸出或反應。在反向傳播 (Back Propagation) 算梯度時，需要儲存每一層的輸出和損失 (Loss)。

最大的問題出在 Attention 機制：

• Attention 是計算兩個序列之間的矩陣。

• 它的空間複雜度會從原本的線性 O(N) 變成 平方 O(N2)，其中 N 是輸入序列的長度。

• 當模型處理長輸入（例如 16k 或 32k context）時，啟用值佔用的空間會變得非常龐大。對於一個 8B 模型而言，在長輸入情境下，啟用值可能佔用高達 1.35TB 的記憶體。

1. Batch Size

為了製造足夠的梯度(Gradient)，batch size也需要對應的大小，也會導致空間增大。

不同元件佔的空間大小示意圖如下，因此有三種解決方案

解決方案一：多 GPU 平行化與 DeepSpeed

由於 128GB 的模型狀態無法被單張 GPU 容納，我們必須將訓練所需的組件切片，並分配給多張 GPU。

DeepSpeed 是由微軟開發的一套軟體套件，專門用於優化多 GPU 的平行運算，以最大化多張 GPU 的使用效益。

DeepSpeed 的核心演算法是 Zero Redundancy Optimizer (Zero)。Zero 演算法有三個級別 (Level)，級別越高，節省的記憶體越多：

• Zero-1： 優先將佔用空間最大的 優化器狀態 (Optimizer States) 切片，分配到各個 GPU 上。

• Zero-2： 除了優化器狀態外，也將 梯度 (Gradients) 和 FP16 模型權重 切片。

• Zero-3： 將所有模型狀態（包括 FP32 模型權重）都切開。這是將 LLM 訓練記憶體需求壓到最小的狀態。

雖然切片會增加 GPU 間的傳輸開銷 (transfer overhead)，但 NVIDIA GPU 之間的 NVLink 技術傳輸速度非常快，每秒可達 900 GB。DeepSpeed 也會進行傳輸排程 (scheduling) 來減少延遲。透過 Zero-3 演算法，即使是 8B 模型在 80GB 的 H100 GPU 上，也能訓練長度達 16k 的序列。差異如下

CPU Offload

如果 GPU 記憶體實在太小，DeepSpeed 還提供了 Offload 功能，允許將部分訓練組件（如優化器或模型參數）搬到容量更大的 CPU RAM 上儲存。CPU RAM 的容量通常比 GPU RAM 大上一個數量級（約 10 倍）。

然而，應盡量避免使用 Offload，因為 CPU 與 GPU 之間的傳輸速度相對慢得多。實驗顯示，使用 Offload 會導致訓練速度顯著變慢，可能慢上 10 倍以上（例如，一個步驟從 7.3 秒增加到 74 秒）。

解決方案二：優化長輸入與啟用值計算

當模型的輸入序列非常長時，啟用值佔用記憶體太多會成為主要瓶頸。解決方法是重寫 Attention 實作程式碼（即 GPU 上的 Kernel Function）來提高計算效率。從效率與高低階語言對照如下

1. Flash Attention

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

• 加速機制： 傳統 Attention 包含多個複雜步驟，如矩陣乘法、縮放、SoftMax 等。Flash Attention 將這四五個步驟整合為單一、優化過的核心函數 (Fused Kernel)。在 GPU 上，SoftMax 或 Masking 等運算比矩陣乘法更耗時，融合核心能有效提升整體速度。

• 記憶體節省機制： 它將大部分資料（如 QKV 矩陣）存放在 CPU 上，只在計算時快速載入（Flash）到 GPU 記憶體中，從而將記憶體空間需求降至接近線性 O(N)。

1. Liger Kernel

Liger Kernel 是一套利用 Triton（一個允許在 Python 中編寫 GPU 核心功能的介面）實作的優化核心函數。

Liger Kernel 將 LLM 運算中常用的一些函數打包成優化的核心，能夠讓 LLM 算得更快，並且使用更少的記憶體。其優點是使用方式非常簡單：只需將載入模型的方式從 AutoModelForCausalLM 替換為 AutoLigerKernelForCausalLM，即可為許多開源模型進行加速。

解決方案三：量化 (Quantization)

量化是一種有損壓縮 (lossy compression) 技術，主要目的是減少模型儲存所需的記憶體空間。

• 原理： 將高精度浮點數（如 32bit 或 16bit）壓縮成更低精度（如 8bit 或 4bit）。儘管這是種壓縮，但模型還原後與原始張量仍非常近似。

• 主要應用： 量化技術主要應用於 Inference（推論）階段。

• 記憶體效益： 透過量化演算法（例如 GGUF family 演算法中的 Q8, 8-bit 量化），一個 8B 模型只需要約 8GB 的記憶體即可載入。這使得使用者可以在記憶體較小的 GPU 上（例如 Colab 上的 15GB T4 GPU）載入並運行大型模型。

總結來說，這些開源工具讓在有限的 GPU 資源下訓練和部署大型模型成為可能。

概念比喻：

我們可以將訓練 LLM 想像成建造一棟摩天大樓：123

• DeepSpeed 就像是將建造大樓所需的龐大材料（128GB 的參數和狀態）分割成小塊，分配給一組工人（多個 GPU）同時處理，確保所有材料都能進入並被利用，而不是因為單一倉庫太小而被卡住。

• Flash Attention 與 Liger Kernel 則像是替大樓中特別複雜且耗能的生產線（Attention 機制）安裝了客製化的、高效能的模組，讓該生產線的運作更快，且在運作時佔用的空間（啟用值）更少。

• Quantization 則負責在建築完成後（推論階段），將大樓的設計圖（模型權重）進行高密度壓縮（例如從 32bit 壓縮到 4bit），使其可以用更小的隨身碟（資源有限的硬體）攜帶和運行。

參考資源:

【生成式AI時代下的機器學習(2025)】助教課：利用多張GPU訓練大型語言模型

學習影片 https://youtu.be/rhypibF_WI4