训练流程

通用文本大模型的主流预训练范式

大规模原始语料
    ↓
清洗、去重、质量筛选
    ↓
Tokenizer 编码
    ↓
Decoder-only Transformer
    ↓
Next Token Prediction
    ↓
Pretrained Base Model
    ↓
Fine-Tuning
   ↓
Chat / Instruct Model

Pretraining

GPT3：[2005.14165] Language Models are Few-Shot Learners

预训练阶段主要使用大规模无标注文本，通过 Next Token Prediction 学习语言、知识和通用能力

对于序列：

$$ x=(x_1,x_2,\ldots,x_T) $$

Decoder-only LLM 优化：

$$ \mathcal{L}_{\mathrm{PT}} = -\sum_{t=1}^{T} \log P_\theta(x_t\mid x_{<t}) $$

Continued Pretraining，CPT：在已有 Base Model 上继续使用领域原始文本进行下一 Token 预测，主要改变模型的领域知识和语言分布适应能力，广义上也属于微调

Supervised Fine-tuning

使用高质量的 Instruction → Response 数据，让 Base Model 学习如何遵循指令并按照期望格式回答

Training Data通常是

1
2
3

Instruction / User
        ↓
Response / Assistant

例如：

User:
解释一下 Transformer

Assistant:
Transformer 是一种基于 Self-Attention 的神经网络架构……

多轮对话则类似：

System:
You are a helpful assistant.

User:
什么是 LoRA？

Assistant:
LoRA 是一种参数高效微调方法……

User:
为什么能减少显存？

Assistant:
因为训练时只更新低秩矩阵……

经过 Chat Template 和 Tokenizer 后，最终仍然变成一个 Token Sequence

本质上仍然是 Next Token Prediction

$$ \mathcal{L}_{\mathrm{SFT}} = -\sum_{t=1}^{T} \log \pi_\theta(y_t\mid x,y_{<t}) $$

但只对 Assistant Response 计算 Loss

区别可以简单记为：

阶段	数据	主要目的
Pre-training	大规模普通文本	学习语言、知识和通用能力
CPT	特定领域普通文本	补充领域知识
SFT	Instruction-Response	学习如何完成用户指令

Preference Alignment

SFT 让模型学会“怎么回答”，Preference Alignment 进一步让模型学习“什么回答更好”

典型 Preference Dataset 包含

Prompt
+
Chosen Response
+
Rejected Response

RLHF

[2203.02155] Training language models to follow instructions with human feedback

基于人类反馈的强化学习，是 RL 在大模型对齐中的一种应用范式

经典流程：

SFT Model
    ↓
生成多个 Response
    ↓
Human Ranking
    ↓
Reward Model
    ↓
Reinforcement Learning
    ↓
Aligned Model

需要根据偏好训练一个 Reward Model

然后 Reward Model 对回答输出一个标量奖励 $r_\phi(x,y)$

成对偏好训练常写为：

$$ \mathcal{L}_{\mathrm{RM}} = -\log \sigma \left( r_\phi(x,y_w)-r_\phi(x,y_l) \right) $$

目标是使：

$$ r_\phi(x,y_w)>r_\phi(x,y_l) $$

PPO

[1707.06347] Proximal Policy Optimization Algorithms

核心思想是：提高高优势动作的概率，同时限制新策略相对于旧策略一次更新过大

对于状态 $s_t $ 下采取的动作 $a_t$，定义新旧策略的概率比：

$$ \rho_t(\theta)= \frac{ \pi_\theta(a_t\mid s_t) }{ \pi_{\theta_{\mathrm{old}}}(a_t\mid s_t) } $$

$\theta$：当前策略模型的参数
$a_t$：当前生成的 token
$s_t$：Prompt + 前面已经生成的 token

但是 PPO 不希望策略一次更新得太激进，引入裁剪目标：

$$ J_{\mathrm{PPO}}(\theta) = \mathbb{E}_t \left[ \min \left( \rho_t(\theta)\hat A_t, \operatorname{clip} \left( \rho_t(\theta),1-\epsilon,1+\epsilon \right)\hat A_t \right) \right] $$

$\hat A$ 表示在状态 $s_t$ 下采取 $a_t$ 相比平均表现好多少

在 RLHF 中，Reward Model 对完整 Response 给出 Reward，再结合 Value Model 等方法估计每个 timestep 的 $\hat A_t$

核心是逐 token / 时间步优化

Reward Model + PPO 是经典的 RLHF 对齐方案

示例

Prompt:
解释 Transformer 是什么
Answer:
Transformer / uses / attention / EOS

π_old:
Transformer   0.40
uses          0.50
attention     0.60
EOS           0.80

π_theta:
Transformer   0.48
uses          0.45
attention     0.72
EOS           0.72

Transformer: 0.48 / 0.40 = 1.20
uses:        0.45 / 0.50 = 0.90
attention:   0.72 / 0.60 = 1.20
EOS:         0.72 / 0.80 = 0.90

假设每个时间步 Advantage

1	A = [1.0, 0.5, 1.2, -0.2]

t1: ratio=1.20 × advantage=1.0
t2: ratio=0.90 × advantage=0.5
t3: ratio=1.20 × advantage=1.2
t4: ratio=0.90 × advantage=-0.2

DPO

[2305.18290] Direct Preference Optimization: Your Language Model is Secretly a Reward Model

DPO 没有 PPO 那种 token-level 优化信号，直接利用

1
2
3

Prompt
Chosen
Rejected

优化模型对 Chosen 和 Rejected 的相对偏好

常见目标：

$$ \mathcal{L}_{DPO} = - \log \sigma \left[ \beta \left( \log \frac{ \pi_\theta(y_w\mid x) }{ \pi_{\mathrm{ref}}(y_w\mid x) } - \log \frac{ \pi_\theta(y_l\mid x) }{ \pi_{\mathrm{ref}}(y_l\mid x) } \right) \right] $$

直接看整条回答的概率提升，关心的是相对于 Reference 变化了多少

示例

Prompt:
解释 Transformer 是什么

chosen:
Transformer / uses / self-attention / to / model / relationships / EOS
长度 = 7

rejected:
Transformer / is / a / CNN / EOS
长度 = 5

假设当前模型 $\pi_\theta$ 对 chosen 的 token 概率是：

Transformer      0.8
uses             0.7
self-attention   0.6
to               0.9
model            0.7
relationships    0.6
EOS              0.8

$$ \log \pi_\theta(y_w\mid x) = \log0.8+\log0.7+\log0.6+\log0.9+\log0.7+\log0.6+\log0.8 $$

假设 rejected

Transformer   0.7
is            0.6
a             0.8
CNN           0.2
EOS           0.6

$$ \log \pi_\theta(y_l\mid x) = \log0.7+\log0.6+\log0.8+\log0.2+\log0.6 $$

然后 reference model 也分别计算

最后比较

$$ \left[ \log\pi_\theta(y_w|x) - \log\pi_{\mathrm{ref}}(y_w|x) \right] - \left[ \log\pi_\theta(y_l|x) - \log\pi_{\mathrm{ref}}(y_l|x) \right] $$

对比 chosen 答案和 rejected 答案在新模型中概率的变化量

GRPO

[2402.03300] DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models

对于同一个问题，从旧策略采样一组回答：

$$ y_1,y_2,\ldots,y_G $$

分别获得 Reward：

$$ r_1,r_2,\ldots,r_G $$

GRPO 本身不限定 Reward 来源，可以使用 Reward Model、Rule-based Reward 或 Verifier

Rule-based Reward ：

答案正确：Reward = 1
答案错误：Reward = 0

根据同一 Prompt 下多个 Response 的组内 Reward 估计相对 Advantage

$$ \hat A_i = \frac{ r_i-\operatorname{mean}(r) }{ \operatorname{std}(r)+\epsilon } $$

与 PPO 相比，GRPO 不需要独立的 Value Model，减少了额外模型带来的显存开销

总结

对比项	DPO	GRPO
主要用途	通用偏好对齐	推理能力强化
典型场景	风格、语气、安全偏好	数学、代码、逻辑推理
训练数据	`Prompt + Chosen + Rejected`	`Prompt + Group Response`
Reward 来源	Chosen / Rejected 偏好	Rule-based Reward

训练显存与数值精度

训练大模型时，显存主要由以下部分组成：

$$ M_{\mathrm{train}} \approx M_{\mathrm{param}} + M_{\mathrm{grad}} + M_{\mathrm{opt}} + M_{\mathrm{act}} + M_{\mathrm{temp}} $$

其中：

Parameters：模型权重
Gradients：反向传播产生的梯度
Optimizer States：优化器保存的额外状态，例如 Adam 的一阶矩、二阶矩
Activations：Forward 中保存、供 Backward 使用的中间结果
Temporary Buffers：Attention、矩阵乘法等计算过程中的临时空间

数值精度

数据类型	位数	Sign	Exponent	Mantissa	每个元素显存	主要用途
FP32	32 bit	1	8	23	4 Byte	高精度计算，显存占用大
FP16	16 bit	1	5	10	2 Byte	混合精度训练
BF16	16 bit	1	8	7	2 Byte	数值范围接近 FP32，当前大模型训练常用
FP8	8 bit	1	4/5	3/2	1 Byte	低精度训练
INT8	8 bit	--	--	--	1 Byte	权重量化、推理
INT4	4 bit	--	--	--	0.5 Byte	低比特量化

如果考虑模型参数

$$ M_{\text{param}} = N_{\text{param}} \times \frac{\text{bit}}{8} $$

例如一个 (7B) 参数模型

FP32：$7\times10^9\times4\approx28\text{ GB}$
FP16 / BF16：$7\times10^9\times2\approx14\text{ GB}$
INT8：$7\times10^9\times1\approx7\text{ GB}$

这些只是模型权重本身的显存，并不是完整训练显存

Mixed Precision

[1710.03740] Mixed Precision Training

基本思想：使用 FP16 进行神经网络训练，在降低显存和计算开销的同时，尽量保持与 FP32 训练相同的精度

FP16
├── Weights（Forward / Backward 使用）
├── Activations
└── Gradients

FP32
├── Master Weights
└── 部分累加操作

FP32 Master Weights

Forward 和 Backward 使用 FP16 权重，但额外维护一份 FP32 Master Copy

        FP32 Master Weights
                │
        cast FP32 → FP16
                ▼
           FP16 Weights
                │
                ▼
           Forward Pass
                │
                ▼
               Loss
                │
          × Loss Scale S
                ▼
           Scaled Loss
                │
                ▼
          Backward Pass
             /      \
            /        \
           ▼          ▼
Activation Grad.   Weight Grad.
传给前一层         用于参数更新
                       │
                     ÷ S
                       ▼
              Unscaled Weight Grad.
                       │
                       ▼
                FP32 Weight Update
                       │
                       ▼
              Updated Master Weights
                       │
                       └──→ 下一轮

原因是 FP16 精度有限，较小的 Weight Update 可能无法表示，甚至在更新时直接变成 0

使用 FP32 Master Weight 可以保留这些小更新

也就是存储使用FP32，但是计算用FP16

Loss Scaling

FP16 的动态范围比 FP32 小，较小的 Gradient 可能发生 Underflow

因此在 Backward 前将 Loss 乘一个 Scale

$$ L'=SL\qquad g'_t = \nabla_\theta L' = S\nabla_\theta L = Sg_t $$

这样可以把小梯度移动到 FP16 可表示范围内

Unscale 将放大的 Gradient 除以 $S$ 恢复原始尺度

$$ g_t = \frac{g'_t}{S} $$

回到FP32，再用这个梯度更新参数

FP16 Multiply + FP32 Accumulation

对于矩阵乘法：

$$ C_{ij}=\sum_k A_{ik}B_{kj} $$

采用：

FP16 × FP16
     ↓
FP32 Accumulation
     ↓
FP16 Storage

因为大量 FP16 数值直接累加容易产生较大的数值误差

分布式训练与并行

当单张 GPU 无法提供足够的计算能力或显存时，需要把训练拆到多张 GPU

核心可以分为两类

数据并行：模型复制，数据切分
模型并行：数据可以相同或分成 micro-batch，模型本身被切分

实际大模型训练通常会组合多种并行方式

并行方式	切分目标
DP	Data
TP	Tensor
PP	Model Layers
EP	Experts
ZeRO / FSDP	Training States

通信方式

AllGather

将各 GPU 上的 Shard 收集起来，使每张 GPU 都得到完整 Tensor

GPU0: [A]  
GPU1: [B]

AllGather →

GPU0: [A, B]  
GPU1: [A, B]

Shard → Full Tensor

AllToAll

常用于 Expert Parallel

每张 GPU 向其他 GPU 发送不同的数据块，用于数据重分布；各 GPU 通常同时持有本地 Token 和部分 Experts

Router 后通过 AllToAll 将不同 GPU 上的 Token 重新分发到对应 Expert 所在设备，计算完成后再通过一次 AllToAll 返回

GPU0: [A0 → GPU0, A1 → GPU1]  
GPU1: [B0 → GPU0, B1 → GPU1]

AllToAll →

GPU0: [A0, B0]  
GPU1: [A1, B1]

每个参与通信的 GPU 都既是发送方，也是接收方

ReduceScatter

AllReduce = 聚合后每张卡都有完整结果

ReduceScatter = 聚合后每张卡只保留一个 shard

数据维度

Data Parallel

每张 GPU 保存一份完整模型，但读取不同的数据

GPU 0: batch 0 → gradient G0 ─┐
GPU 1: batch 1 → gradient G1 ─┤
GPU 2: batch 2 → gradient G2 ─┼→ AllReduce → 同步梯度
GPU 3: batch 3 → gradient G3 ─┘

如果共有 $P$ 张 GPU：

$$ g = \frac{1}{P} \sum_{p=1}^{P}G_p $$

PyTorch 中常用 DistributedDataParallel，DDP（这里是AllReduce ）

DDP 的限制是：每张 GPU 都需要保存完整模型以及对应训练状态

因此模型本身必须能够放入单卡

ZeRO

[1910.02054] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

Data Parallel 中，每张 GPU 都保存相同的：

1
2
3

P：Parameters
G：Gradients
O：Optimizer States

核心思想：将 Data Parallel 中冗余复制的训练状态分片存储

以下举例

GPU 0：
参数      p0 p1 p2 p3 p4 p5 p6 p7
梯度      g0 g1 g2 g3 g4 g5 g6 g7
优化器    O0 O1 O2 O3 O4 O5 O6 O7

GPU 1：
参数      p0 p1 p2 p3 p4 p5 p6 p7
梯度      g0 g1 g2 g3 g4 g5 g6 g7
优化器    O0 O1 O2 O3 O4 O5 O6 O7

ZeRO-1

只分片 Optimizer States

GPU 0：[完整参数] [完整梯度] [O0 O1]
GPU 1：[完整参数] [完整梯度] [O2 O3]
GPU 2：[完整参数] [完整梯度] [O4 O5]
GPU 3：[完整参数] [完整梯度] [O6 O7]

ZeRO-2

额外分片 Gradients

GPU 0：[完整参数] [g0 g1] [O0 O1]
GPU 1：[完整参数] [g2 g3] [O2 O3]
GPU 2：[完整参数] [g4 g5] [O4 O5]
GPU 3：[完整参数] [g6 g7] [O6 O7]

使用 ReduceScatter聚合完成后直接拆开

ZeRO-3

把参数也分片

GPU 0：[p0 p1] [g0 g1] [O0 O1]
GPU 1：[p2 p3] [g2 g3] [O2 O3]
GPU 2：[p4 p5] [g4 g5] [O4 O5]
GPU 3：[p6 p7] [g6 g7] [O6 O7]

计算某一层之前，需要通过 AllGather 临时获取当前层所需参数（因为参数长期分片）

方法	主要分片内容	典型通信
DDP	不分片	AllReduce 梯度
ZeRO-1	Optimizer State	ReduceScatter + AllGather 参数
ZeRO-2	Optimizer State + Gradient	ReduceScatter 梯度 + AllGather 参数
ZeRO-3	Optimizer + Gradient + Parameter	ReduceScatter 梯度 + 按需 AllGather 参数

FSDP

核心思想与 ZeRO-3 类似，将 Parameters、Gradients、Optimizer States 分片到不同 GPU，并在计算时按需 AllGather 参数

FSDP 是 PyTorch 原生提供的 Full Sharding 训练方案

核心行为基本就是：

1
2
3

Parameters   → Shard
Gradients    → Shard
Optimizer    → Shard

计算某一层前：

Parameter Shards
        ↓
AllGather
        ↓
Forward / Backward

Backward 后：

Gradients
    ↓
ReduceScatter
    ↓
每张 GPU 只保留 Gradient Shard

模型维度

Tensor Parallel

[1909.08053] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism

TP：将一个大的矩阵运算拆到多张 GPU

例如线性层，将权重矩阵按列切分：

$$ W= \left[ W_1\quad W_2 \right] $$

不同 GPU 分别计算：

$$ Y_1=XW_1,\qquad Y_2=XW_2 $$

在 Transformer 中通常会对：

Attention Heads
Q/K/V Projection
MLP Linear Layers

进行 Tensor Parallel

TP 能解决单个 Layer 太大的问题，但会带来较频繁的 GPU 间通信

Transformer Layer

       ┌── GPU 0 ──┐
       ├── GPU 1 ──┤
Input ─┼── GPU 2 ──┼→ Communication → Output
       └── GPU 3 ──┘

Pipeline Parallel

[1811.06965] GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

PP：按照 Transformer Layer 切分模型

1
2
3

GPU 0            GPU 1            GPU 2            GPU 3

Layer 0~7   →   Layer 8~15   →   Layer 16~23   →   Layer 24~31

简单流水线会产生 Pipeline Bubble

时间      0    1    2    3    4    5    6    7
         ---------------------------------------
GPU 0    F    --   --   --   --   --   --   B
GPU 1    --   F    --   --   --   --   B    --
GPU 2    --   --   F    --   --   B    --   --
GPU 3    --   --   --   F    B    --   --   --

F  = 整个 batch 的前向传播
B  = 整个 batch 的反向传播
-- = GPU 空闲

大部分时间只有一张 GPU 在工作，设备利用率很低

因此通常把一个 Batch 切分为多个 Micro-batch：

时间      0    1    2    3    4    5    6    7    8    9   10   11   12   13
         ---------------------------------------------------------------------
GPU 0    F0   F1   F2   F3   --   --   --   --   --   --   B3   B2   B1   B0
GPU 1    --   F0   F1   F2   F3   --   --   --   --   B3   B2   B1   B0   --
GPU 2    --   --   F0   F1   F2   F3   --   --   B3   B2   B1   B0   --   --
GPU 3    --   --   --   F0   F1   F2   F3   B3   B2   B1   B0   --   --   --

若有 $K$ 个 pipeline stages，$M$ 个batch，简单情况下设备利用率大致为：

$$ \text{Utilization} \approx \frac{M}{M+K-1} $$

增加 Micro-batch 数量可以摊薄 Pipeline 的填充和排空气泡，提高设备利用率

但 Micro-batch 过多会增加调度、通信开销

Expert Parallel

MoE 中通常存在大量 Experts，但每个 Token 只激活少量专家

EP：将不同专家放到不同 GPU

GPU 0: Expert 0, 1
GPU 1: Expert 2, 3
GPU 2: Expert 4, 5
GPU 3: Expert 6, 7

Router 根据 Token 决定目标 Expert 后，需要把 Token 发送到对应 GPU

Token
 ↓
Router
 ↓
AllToAll
 ↓
Expert GPU
 ↓
AllToAll
 ↓
返回原设备

因此，MoE 可以降低单 Token 的激活参数量，但性能通常受 通信开销 和 Expert 负载均衡 影响

Hybrid Parallelism

                Hybrid Parallelism
                        │
       ┌────────────────┼────────────────┐
       │                │                │
      PP               TP               EP
Pipeline Parallel  Tensor Parallel  Expert Parallel
  按 Layer 切分      按 Tensor 切分      按 Expert 切分
       │                │                │
       └────────────────┼────────────────┘
                        │
                       DP
                 Data Parallel
              不同 GPU 处理不同数据
                        │
            ┌───────────┴───────────┐
            │                       │
       Replicated DP           Sharded DP
        普通 DDP              ZeRO / FSDP
     完整模型副本          分片模型训练状态

DP + TP + PP：Dense LLM 很经典的一种组合

DP + TP + PP + EP：MoE模型

总 GPU 数通常可表示为：

$$ \text{Total GPUs} = DP × TP × PP (\times EP) $$

参数高效微调

LoRA

[2106.09685] LoRA: Low-Rank Adaptation of Large Language Models

LoRA：Low-Rank Adaptation

传统 Full Fine-tuning 需要更新模型全部参数，随着模型规模增大，会带来较高的训练、显存和模型存储成本

LoRA 的核心假设：大模型适配下游任务时，权重更新具有较低的 intrinsic rank，因此不需要学习完整的参数更新矩阵。

对于预训练权重：

$$ W_0 \in \mathbb{R}^{d \times k} $$

Full Fine-tuning 直接学习完整的 $\Delta W$，而 LoRA 将其限制为低秩分解

$$ \boxed{\Delta W = BA} $$

其中：

$$ B\in\mathbb{R}^{d\times r},\qquad A\in\mathbb{R}^{r\times k},\qquad r\ll\min(d,k) $$

训练时冻结预训练参数 $W_0$，只更新低秩矩阵 $A$ 和 $B$，前向传播变为：

$$ \boxed{h = W_0 x + BAx} $$

即在原始线性层旁增加一条**低秩更新分支**，用少量可训练参数完成模型适配

论文实验表明，在部分模型和任务中，即使使用很小的 rank，例如 $r=1$、$r=2$ 或 $r=4$，也可以获得较好的效果

Parameter Reduction

对一个 $d\times k$ 的权重矩阵

Full Fine-tuning 训练参数量为 $dk$

LoRA 只需要训练 $A$ 和 $B$，参数量为：

$$ rk+dr=r(k+d) $$

例如

$$ d = k = 4096,\qquad r = 8 $$

Full Fine-tuning：$4096^2 = 16,777,216$

LoRA：$4096 \times 8 + 8 \times 4096 = 65,536$

仅约为原参数量的：$\frac{65,536}{16,777,216}\times 100% \approx 0.39%$

Initialization

LoRA 使用：

$$ A\sim\mathcal{N}(0,\sigma^2),\qquad B=0 $$

因此初始化时

$$ \Delta W=BA=0 $$

模型开始训练时仍等价于原始预训练模型，不会因为加入 LoRA 而立即改变模型输出

Scaling

实际前向传播中，LoRA 会对低秩更新加入缩放系数

$$ \boxed{ h=W_0x+\frac{\alpha}{r}BAx } $$

这种 scaling 可以减少改变 $r$ 时重新调整其他超参数的需要

Applying to Transformer

LoRA 原理上可以应用于任意 Dense / Linear Layer

Transformer Self-Attention 中主要包含：

$$ W_q,\quad W_k,\quad W_v,\quad W_o $$

原论文主要研究将 LoRA 应用于 **Attention 权重**，并冻结 MLP。

在大多数实验中将 LoRA 应用于 $W_q$ 和 $W_v$

$$ W_q'=W_q+B_qA_q \qquad W_v'=W_v+B_vA_v $$

Merge at Inference

训练完成后，可以将 LoRA 更新直接合并进原始权重

$$ \boxed{ W'=W_0+BA } $$

如果考虑 scaling，则实际为：

$$ \boxed{W'=W_0+\frac{\alpha}{r}BA} $$

之后推理直接使用：

$$ h=W'x $$

因此 LoRA 权重 Merge 后不需要额外的低秩分支计算，不会引入额外的推理延迟

Practical Issues

LoRA 能显著减少可训练参数量，但实际使用中仍存在几个需要重点关注的问题

Rank $r$ 的选择

$r$ 决定了 LoRA 分支的表达能力

$r$ 太小：参数更新空间受限，可能无法充分适配下游任务，出现欠拟合；
$r$ 太大：会增加可训练参数和显存占用，也可能增加过拟合风险。

本质上是在调参效率和适配能力之间进行权衡

Target Modules 的选择

LoRA 可以应用到 Transformer 中不同的线性层，也可以应用到 MLP 中的 Linear Layer

不同任务需要调整的参数方向不同，因此 LoRA 插入的位置会直接影响效果

Target Modules 本身也是 LoRA 的重要超参数

训练显存

LoRA 主要减少可训练参数和优化器状态，前向传播仍然需要经过完整模型，因此模型权重和中间 Activation 仍然会占用大量显存

对于大模型训练，通常还需要结合

Mixed Precision
Gradient Checkpointing
Quantization

进一步降低显存占用

QLoRA

[2305.14314] QLoRA: Efficient Finetuning of Quantized LLMs

QLoRA：Quantized Low-Rank Adaptation

QLoRA 在 LoRA 基础上进一步将预训练模型量化为 4-bit，冻结量化后的 Base Model，只训练 LoRA Adapter。

Quantization

量化就是把高精度数据表示成更低 bit 的数据

以 FP32 → Int8 为例

$$ X^{\mathrm{Int8}} = \operatorname{round} \left( \frac{127} {\operatorname{absmax}(X^{\mathrm{FP32}})} X^{\mathrm{FP32}} \right) $$

其中量化 scale 为：

$$ c= \frac{127} {\operatorname{absmax}(X)} $$

因为

$$ \mathrm{Int8}\in[-128,127] $$

FP32 张量中找到最大值后可以认为有效范围为 $X\in[-M,M]$

而后做一个线性映射到 Int8 的范围，于是就得到上式

1
2
3

FP32:  -2     -1      0      1      2
        ↓      ↓      ↓      ↓      ↓
Int8: -127    -64     0     64     127

FP32 中，两个很接近的数可以被区分，但Int8 只有有限的整数

所以原来连续范围里的很多 FP32 数值，都会被映射到同一个 Int8 整数

那么 Int8 相邻两个量化等级对应的 FP32 间隔，也就是 quantization step

$$ \Delta=\frac{M}{127} $$

FP32 数值实际上被划分成一个个宽度约为 $\Delta $ 的区间

一个数最多偏离最近量化点半个 step

$$ |x-\hat x| \le \frac{\Delta}{2} = \frac{M}{254} $$

NF4

4-bit NormalFloat，NF4

预训练神经网络权重通常近似 zero-centered normal distribution(零中心正态分布)

因此不是简单地把数值范围均匀划成 16 个区间，而是根据正态分布的 quantile 设计 4-bit 表示，使每个量化区间预期包含相近数量的权重值

论文实验中，NF4 的量化精度优于 FP4 和 Int4

Double Quantization

量化之后并不是只有 4-bit Weight

还需要保存每个 block 对应的 Quantization Scale

1
2
3

Block 1 → scale c₁
Block 2 → scale c₂
Block 3 → scale c₃

假设 block = 64，那么每个 block 需要 32 bit

平均到每个模型参数上

$$ \frac{32}{64}=0.5\text{ bits/parameter} $$

权重本身已经压到 4 bit了，加上保存的 scale 大约为 4.5 bits/parameter

QLoRA 提出继续量化第一次量化产生的量化常数

Weights
   │
   ▼
第一次量化
   │
   ├── 4-bit Weight
   │
   └── FP32 Quantization Constants
                  │
                  ▼
             第二次量化
                  │
                  ▼
            8-bit Constants

把这些 scale 量化成 8 bit，假设第二级 block size 为 256

$$ \frac{8}{64} + \frac{32}{64\times256} = 0.127 \text{ bits/parameter} $$

平均节省：

$$ \boxed{0.373\text{ bits/parameter}} $$

4-bit weights
+
8-bit scales
+
极少量 FP32 scales

但主要贡献还是来自于 Base Model 和训练过程中的 activations

为了尽量恢复 16-bit 微调效果，在 Transformer Block 的所有 Linear Layers 上加入 LoRA 很重要

Paged Optimizer

训练过程中某些长序列或 Gradient Checkpointing 场景可能突然产生显存峰值

Paged Optimizer 使用 NVIDIA Unified Memory

当 GPU 显存不足时：

GPU
 │
 │ memory pressure
 ▼
CPU RAM

需要时
 │
 ▼
GPU

Optimizer State 可以暂时 page 到 CPU RAM，之后再移回 GPU，从而避免偶发的 OOM

LLM 推理

LLM 推理的核心特点是 Autoregressive Generation

$$ P(x) = P(x_1) P(x_2\mid x_1) \cdots P(x_n\mid x_1,\ldots,x_{n-1}) $$

一个请求通常分为两个阶段：

Prompt
  ↓
Prefill
  ↓
生成第一个 Token
  ↓
Decode
  ↓
Token-by-Token Generation

KV Cache

如果不使用 KV Cache，每生成一个新 Token，都需要重新计算历史 Token 的 $K,V$

因此缓存：

$$ K_1,\ldots,K_T,\qquad V_1,\ldots,V_T $$

生成第 $T+1$ 个 Token 时，只计算：

$$ K_{T+1},\qquad V_{T+1} $$

对于标准 MHA，如果：

Transformer 层数为 $L$
Context Length 为 $T$
KV Head 数为 $H_{KV}$
Head Dimension 为 $d_h$
每个元素占 $b$ Bytes

单条 Sequence 的 KV Cache 约为：

$$ M_{\mathrm{KV}}=2LTH_{KV}d_hb $$

前面的 $2$ 表示 Key 和 Value

因此长上下文、大 Batch 下，KV Cache 会成为推理显存的重要瓶颈

FlashAttention

[2205.14135] FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness

标准 Self-Attention 的时间和显存开销会随序列长度 $N$ 呈平方增长，很多已有方法只关注减少 FLOPs，却忽略了 GPU 中真正影响 wall-clock time 的 memory access

FlashAttention 因此提出 IO-aware Attention

Standard Attention

$$ Q,K,V\in\mathbb{R}^{N\times d} $$

标准 Attention：

$$ S=QK^T\qquad P=\operatorname{softmax}(S)\qquad O=PV $$

从 HBM 读取 Q、K
    ↓
计算 S = QKᵀ
    ↓
把完整 S 写入 HBM
    ↓
从 HBM 重新读取 S
    ↓
计算 Softmax，得到 P
    ↓
把完整 P 写入 HBM
    ↓
从 HBM 读取 P 和 V
    ↓
计算 O = PV

标准实现会把 $S$ 和 $P$ 以及注意力矩阵反复在 HBM 和片上存储之间搬运

GPU Memory Hierarchy

存储	位置	容量	速度 / 带宽	延迟	存储对象
CPU DRAM	电脑内存	几十到几百GB	最慢	最高	未放入 GPU 的模型/数据
GPU HBM	GPU 显存	几十GB	很快	中等	Model Weights、Activations KV Cache、Gradient
GPU SRAM	GPU 芯片内部	KB～MB 级	非常快	很低	正在计算的小块数据、中间结果

FlashAttention 的核心问题：怎么尽可能把计算留在 SRAM 里完成，减少访问 HBM

Tiling

把 $Q,K,V$ 划分成多个小 block，每次只把一小块从 HBM 加载进 SRAM

就不会在 HBM 中构造完整的 Attention Matrix

Softmax分块

普通 Softmax

$$ \operatorname{softmax}(x_i)=\frac{e^{x_i}} {\sum_j e^{x_j}} $$

为了数值稳定，通常减去整行最大值：

$$ m(x)=\max_i x_i\qquad \ell(x) = \sum_i e^{x_i-m(x)} $$

那么：

$$ \operatorname{softmax}(x) = \frac{e^{x-m(x)}}{\ell(x)} $$

普通 Softmax 的问题在于：依赖整行所有元素，如果只拿到一块不能直接得到最终 Softmax

FlashAttention 使用可增量更新的 Softmax 统计量

$m$：当前最大值
$l$：当前指数和
$u$：输出累积向量

就可以一个 block 一个 block 地计算 Softmax，而不需要一次把整行元素全部放进 SRAM

因为 Softmax 的基准最大值发生改变，旧统计量需要重新缩放：

$$ \boxed{\ell_{\text{new}} = e^{m_1-m_{\text{new}}}\ell_1 + e^{m_2-m_{\text{new}}}\ell_2} $$

同时在线更新 Attention Output

$$ u_{\text{new}} = e^{m-m_{\text{new}}}u + \sum_{j\in\text{block}} e^{s_j-m_{\text{new}}}v_j $$

所有 Key / Value Block 处理完成后

$$ o=\frac{u}{\ell} $$

FlashAttention 不是 Approximate Attention，和普通 Attention 的数学结果相同，但是额外内存只需要随 $N $ 线性增长

Recomputation

标准 Attention 为了 backward，需要保存 $S, P$ 和注意力矩阵

FlashAttention 不保存它们，所以需要通过保存的 $O, m, \ell$ 重新计算来获得梯度

用额外计算换取更少 HBM 访问和更少显存

虽然 FLOPs 增加了，但由于减少了 HBM IO，实际反而更快

vLLM

[2309.06180] Efficient Memory Management for Large Language Model Serving with PagedAttention

vLLM 是一个面向大语言模型的 LLM Serving Engine

LLM 推理过程中，KV Cache 占用大量 GPU 显存，并且会随着生成长度动态增长。

传统连续显存分配方式通常需要提前预留较大的 KV Cache 空间，同时要求显存连续，容易产生显存浪费和碎片，从而限制 Batch Size 和 Serving Throughput。

KV Cache

LLM 使用 Autoregressive Generation

$$ P(x) = P(x_1) P(x_2|x_1) \cdots P(x_n|x_1,\ldots,x_{n-1}) $$

每生成一个新 token，都需要使用之前 token 的 Key 和 Value

因此将之前已经计算过的 $K,V$ 缓存下来：

$$ K_{\text{cache}},\quad V_{\text{cache}} $$

生成新 token 时只需要计算当前 token 的 $K, V$，之前 token 的直接从 KV Cache 中读取。

PagedAttention

传统 LLM Serving 通常为每个 Request 分配连续的 KV Cache 空间，并可能按照较大的最大 Sequence Length 提前预留显存

PagedAttention 借鉴操作系统中的 Virtual Memory / Paging

将一个 Request 的 KV Cache 划分为多个固定大小的 KV Block，而不再要求 KV Cache 在物理显存中连续存储

逻辑上：

Request A

Block 0
Block 1
Block 2

在 GPU 显存中可以实际存储为：

1
2
3

Logical Block 0 → Physical Block 7
Logical Block 1 → Physical Block 1
Logical Block 2 → Physical Block 3

通过 Block Table 保存 Logical Block 到 Physical Block 的映射关系

token 的逻辑位置
       ↓
算出属于哪个 Logical Block
       ↓
查 Block Table
       ↓
得到 Physical Block Number
       ↓
去 GPU DRAM 读取真正的 K/V

因此一个 Sequence 的 KV Cache 可以分散存储在不连续的 Physical Blocks 中，从而利用零散的空闲显存

Dynamic Allocation

vLLM 不需要提前按照最大 Sequence Length 分配完整 KV Cache，而是随着 token 生成按需分配新的 Physical KV Block

例如一个 Block 存 4 个 token

token 1 2 3 4
↓
Block 0

token 5 6 7 8
↓
Block 1

当前 Block 被填满后，如果继续生成 token 9，才申请新的 Block 2

1
2
3

需要多少 KV Cache
        ↓
按需申请多少 Physical Block

因此通常只有最后一个未填满的 Block 会产生少量碎片。

PagedAttention 还支持多个 Sequence 共享相同的 Physical KV Block

例如多个输出拥有相同 Prompt 时，Prompt 对应的 KV Cache 可以只保存一份

当某个 Sequence 需要修改共享 Block 时，vLLM 使用 Copy-on-Write

这样可以进一步减少重复 KV Cache 的显存占用

大语言模型训练与推理

训练流程

Pretraining

Supervised Fine-tuning

Preference Alignment

RLHF

PPO

示例

DPO

示例

GRPO

总结

训练显存与数值精度

数值精度

Mixed Precision

FP32 Master Weights

Loss Scaling

FP16 Multiply + FP32 Accumulation

分布式训练与并行

通信方式

AllGather

AllToAll

ReduceScatter

数据维度

Data Parallel

ZeRO

ZeRO-1

ZeRO-2

ZeRO-3

FSDP

模型维度

Tensor Parallel

Pipeline Parallel

Expert Parallel

Hybrid Parallelism

参数高效微调

LoRA

Parameter Reduction

Initialization

Scaling

Applying to Transformer

Merge at Inference

Practical Issues

Rank $r$ 的选择

Target Modules 的选择

训练显存

QLoRA

Quantization

NF4

Double Quantization

Paged Optimizer

LLM 推理

KV Cache

FlashAttention

Standard Attention

GPU Memory Hierarchy

Tiling

Recomputation

vLLM

KV Cache

PagedAttention

Dynamic Allocation

KV Cache Sharing