# [2025李宏毅ML] 助教課：利用多張GPU訓練大型語言模型—從零開始介紹DeepSpeed、Liger Kernel、Flash Attention及Quantization
:::success
:+1: 筆記全集Book 請至: https://hackmd.io/@4j/r1U_UJ_pye/
:::
:::info
2025生成式AI時代下的機器學習_李宏毅
課程網站 https://speech.ee.ntu.edu.tw/~hylee/ml/2025-spring.php

video: https://youtu.be/mpuRca2UZtI?si=GQu-6l4WENr-5zqQ
[[Excalidraw]](https://drive.google.com/file/d/1pKgY4EuyXh7SuTbI_r7it1CK5DCFUKYZ/view)

▪ Instructor
Hsiu-Hsuan Wang(王秀軒)
Find more at https://anthony-wss.github.io/
(20250329)
:::

:::spoiler 目錄
[TOC]
:::

## Overview
在 LLM 時代，一般 pytorch 訓練 NN 的方法會無法訓練
可利用以下 package (場景)
1. **DeepSpeed**: 多張 GPU 平行運算
2. **Liger Kernel, Flash Attention**: 輸入 context 太長 ex: 32k 128k
3. **Quantization**: 如想用 colab 跑作業時

## Introduction
### 1. Forward & Backward pass

▪ **LLM訓練為什麼大**

cross entropy (output, label)算loss
再 backpropagation 推到輸入
得到 gradient
經由 optimizer 更新 weights

訓練上困難 主要是 **"LLM weights"** 很大!

### 2. Memory needed, assuming LLM is a 8b model
▪ **要訓練8B模型**
fp32: 一個float=32bit, 就需要 32GB
在訓練時, 常轉成fp16 (nv對fp16也有特別加速)
▪ **adam**
他的設計是每個權重都有個 lr, 故 momentum 跟 variance 大小同 w
gradients 作為輸入 更新w

⇒ 儲存空間(不含input)全部加起來為 128GB!!!
但 **input** 才是佔最多的

### 3. Activations
▪ **activation**: 模型每層對 input 產生的反應
"attention" 在其中佔最大, 空間複雜度為 n平方

▪ **3.1 Activation recomputation (gradient checkpointion)**
訓練時不把 activation 都存下來, **存重要的就好**,
在 backward 需要時再算(訓練時會慢一點)

### 4. Batch size
Batch size 要夠大才不會有 bias
基本上要 4-60M token
(DeepSeek V3: bz=1920)

解法: **Gradient accumulation**
把1920切小批 都過完後再一起合併更新

### 5. Challenges
(紫紅綠) para,grad,opti 在不同輸入長度時的 memory 會是一樣大小
只有 **activations** 會依輸入長度指數型增長 memory

## Part I. parameters, gradients and optimizer states
### 1. How to leverage multiple GPUs?
(16/32 bit 對應到前2的圖)
一個GPU放不下, 就分到多個去運算

### 2. DeepSpeed - Zero Redundancy Optimizer (ZeRO)
DeepSpeed的套件
他使用 **Zero Redundancy Optimizer (ZeRO)** 演算法

▪ **3 個 Level**
-zero1: opt 佔最大先切分他, 會多一些GPU間的傳輸
-zero2: 再切第二大的部分 grad, LLM
-zero3: 再切 LLM
(會從32bit 開始切是因為他們 佔體大 且 用到的時間較少)

其實GPU間傳輸是很快的, 如NVLink: 900GB/s

### 3. Memory Reduction
DP=8 切八片, 虛線: H100 為 80GB
橘色不會減少, 紫紅綠依序縮小八倍
若使用 zero3, 就可訓練到16k長度的8B模型 :D

▪ **3.1 Traning Speed**
速度還取決於 network bandwidth, model size,...等
但這些都不至於慢太多~

### 4. Zero - Offload
用了 Zero 還是放不下的話, 借用 CPU 運算

▪ **4.1 optimizer offload**
把跟 opt 相關的部分都搬到 CPU
▪ **4.2 optimizer + model parameter offload**
全部都搬到 CPU, 要用時再搬過來 GPU 做

但是 CPU GPU間的傳輸是很慢的!! 訓練會變非常慢

### 5. Simple Experiment on TWCC(國網中心)
實驗有多慢

V100(32G)*4  → 128G
input開很小: bz=1, max_len=512

實測 1GPU 需要 15GB, 1 step 要 74 秒 (非 常 慢)
如果有 8張GPU 就裝得下 不用跑offload, 變 7.3 秒 快了 10 倍

▪ **結論**
使用 8張V100 是可以做到 fully finetuning 8B的模型, 速度約 7.3 s/step
盡量不要使用 offload 慢 不好用

## Part II. activations
想改善attention計算的效率
⇒ 重寫實作的 code

### 1. kernel
Kernel: 跑在 GPU 上的 function

### 2. Flash Attention Algorithm
Flash Attention Algorithm (做一個新kernel function)
算得更快 又更少記憶體 (會佔到CPU RAM)

▪ **2.1.1 Fused Kernel**
右圖可發現在attention裡矩陣乘法並不是佔最久的
反而是dropout, softmax,mask

▪ **2.2 Memory**
主要都放在 CPU, 要計算時才放到 GPU
▪ **2.3 Performance**
H100上 不同方法實作的速度

### 3. Liger Kernel
`Liger: [音]賴ger`
是用 Triton 實作的加速套件
用法簡單直接改載入的方式即可
速度記憶體皆有優化

## Part III. Quantization
Quantization 僅用在 inference
### 1. Lossy compression
壓縮後再還原, 用壓縮過後的儲存, 要計算時再還原
常見的演算法: GGML family,...
Q8: 指把每個tensor壓到8 bit儲存

## Take away
You may need packages like DeepSpeed, Flash Attention, or Liger Kernel to train
LLM with multiple gpus and limited memory.

### Recommended Reading
- https://huggingface.co/spaces/nanotron/ultrascale-playbook
- https://huggingface.co/docs/transformers/deepspeed

-- END --