04 DPO 与训练优化

学习理念：DPO 是 RLHF 的简化替代方案——不需要训练奖励模型，直接用偏好数据优化。核心是让模型"比参考模型更倾向于 chosen 而非 rejected"。后半部分的 GPU 优化技术（梯度累积/ZeRO/NF4）是实际训练中解决"显存不够"问题的必杀技。

海外对标：Anthropic RLHF、DPO（Stanford 2023）、ZeRO（Microsoft DeepSpeed）

本节 AI 替代率：~65% | 人工干预率：~35%

角色	能力范围
🤖 AI 擅长	生成 DPO 损失函数代码、优化方案对比
👤 人类需理解	DPO 损失中的"当前模型 - 参考模型"差值含义、ZeRO 各 stage 的内存分布差异

---

📌 来源说明：以下内容提取自原始笔记第7章（DPO实现）+ 第8章（GPU优化），代码来自 03_dpo_train.py 和 04_single_gpu_optimizations_pseudocode.py。

📖 阅读优先级

等级	章节	说明
🔴 必须深入	三、ZeRO + NF4	Agent 部署时要知道显存瓶颈在哪、QLoRA 省多少
🟡 理解即可	二、单 GPU 优化	了解梯度累积、检查点这些概念就行
🟢 直接跳过	一、DPO 全部	做模型对齐的工程师才需要手写 DPO

---

一、DPO：直接偏好优化

1.1 与 SFT 的对比

维度	SFT	DPO
数据	单条"输入-标准回答"	成对"chosen vs rejected"
目标	拟合标准答案	让 chosen 概率 > rejected 概率
参考模型	不需要	需要——提供训练基线
Loss	交叉熵	基于偏好比对的 sigmoid loss

1.2 DPO 数据格式

每条数据包含三部分：prompt + chosen（优回答）+ rejected（差回答）

prompt:   "如何培养一个习惯？"
chosen:   "21天培养法，每天坚持xxx..."
rejected: "不知道，你自己查一下。"

数据集：HuggingFaceH4/ultrafeedback_binarized（GPT-4 打分选出的二元偏好数据）

1.3 DPO 损失函数

🔥 【P0 必须理解】 DPO loss 的公式可以用一句话理解：

"让当前模型比参考模型更喜欢 chosen 胜过 rejected"

核心差值：
  [π_θ(chosen) - π_θ(rejected)] - [π_ref(chosen) - π_ref(rejected)]
   └── 当前模型偏好 ──┘     └── 参考模型偏好 ──┘
   └────────── 当前模型相比参考模型的"额外偏好提升" ──────────┘

最终 loss：
  loss = -log(sigmoid(β × 差值))

def compute_loss(preferred_logits, rejected_logits,
                 ref_preferred_logits, ref_rejected_logits,
                 preferred_labels, rejected_labels,
                 preferred_answer_mask, rejected_answer_mask, beta=0.5):

    # 当前模型：chosen 和 rejected 的平均 log 概率
    preferred_log_prob = _compute_average_log_probability(
        preferred_logits, preferred_labels, preferred_answer_mask)
    rejected_log_prob = _compute_average_log_probability(
        rejected_logits, rejected_labels, rejected_answer_mask)

    # 参考模型：chosen 和 rejected 的平均 log 概率
    ref_preferred_log_prob = _compute_average_log_probability(
        ref_preferred_logits, preferred_labels, preferred_answer_mask)
    ref_rejected_log_prob = _compute_average_log_probability(
        ref_rejected_logits, rejected_labels, rejected_answer_mask)

    # DPO 核心差值
    final = (preferred_log_prob - rejected_log_prob) - \
            (ref_preferred_log_prob - ref_rejected_log_prob)

    loss = -torch.nn.functional.logsigmoid(final * beta)
    return loss.mean()

1.4 训练要点

def train(model, ref_model, tokenizer, config):
    model.train()
    ref_model.eval()               # 参考模型不更新
    optimizer = torch.optim.AdamW(model.parameters(), lr=config.max_lr)

    for step in range(total_steps):
        # 当前模型：分别计算 chosen 和 rejected 的 logits
        preferred_logits = model(preferred_input_ids).logits
        rejected_logits = model(rejected_input_ids).logits

        # 参考模型：同样计算，但不反传梯度
        with torch.no_grad():
            ref_preferred_logits = ref_model(preferred_input_ids).logits
            ref_rejected_logits = ref_model(rejected_input_ids).logits

        # 计算 loss（只更新 model）
        loss = compute_loss(preferred_logits, rejected_logits,
                          ref_preferred_logits, ref_rejected_logits, ...)
        loss.backward()
        optimizer.step()

---

二、单 GPU 训练优化

2.1 显存消耗构成

训练大模型时，显存主要被四部分占用：

参数 (Weights)  +  梯度 (Gradients)  +  优化器状态 (Optimizer States)  +  激活值 (Activations)
   FP16: 2B        FP16: 2B            Adam: 8B (fp32+momentum+var)     与 batch_size × seq_len 成正比

例如 7B 模型：参数量 14GB（FP16）+ 梯度 14GB + 优化器 56GB（Adam 32bit）= 84GB，远超单卡显存。

2.2 梯度累积（Gradient Accumulation）

🟡 【P1 看注释就行】 显存放不下 batch_size=12？拆成 3 个 micro batch，每步前向+反向但不更新，累积 3 步后再 optimizer.step()。

per_device_batch_size = 4
gradient_accumulation_steps = 3
effective_batch_size = 4 * 3 = 12  # 等效 batch size

for micro_step in range(gradient_accumulation_steps):
    loss = compute_loss(micro_batch)
    loss.backward()         # 梯度累加到 .grad
optimizer.step()            # 3 步后统一更新
optimizer.zero_grad()       # 清空梯度

2.3 梯度检查点（Gradient Checkpointing）

核心：前向时不保存中间激活值，反向时重新计算。以时间换空间——节省 ~50% 显存，增加 ~20% 时间。

model.gradient_checkpointing_enable()  # 一行代码启用

2.4 混合精度训练（AMP/Mixed Precision）

# 参数用 FP16/BF16 存储，但关键计算（梯度更新）用 FP32
with torch.amp.autocast("cuda", dtype=torch.bfloat16):
    loss = compute_loss(batch)

scaler = torch.cuda.amp.GradScaler()  # FP16 需要梯度缩放
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

BF16 vs FP16：BF16 动态范围和 FP32 相同，训练更稳定，是当前主流选择。

2.5 CPU 卸载（Offload）

优化器状态 → 移到 CPU 内存（不占显存），仅当需要更新对应参数时才加载到 GPU：

ZeRO-Offload: 优化器状态放 CPU，参数和梯度放 GPU
              → 显存占用减半，但 CPU↔GPU 传输增加延迟

---

三、分布式训练（多 GPU）

3.1 四种并行方式

数据并行：每张卡一份完整模型，分一批数据                 → 减速最小，显存不变
张量并行：单层切分到多卡                                     → 适合单层太大的大模型
流水线并行：不同层放不同卡                                   → 适合层数极深的模型
ZeRO：参数/梯度/优化器状态 分片到各卡                       → 主流方案

3.2 ZeRO（Zero Redundancy Optimizer）🔥 主流

Stage	分片内容	显存节省	通信量	典型场景
ZeRO-1	优化器状态	4x	低	单卡优化器存不下时
ZeRO-2	优化器状态 + 梯度	8x	中	常见训练配置
ZeRO-3	优化器状态 + 梯度 + 参数	16x	高	超大模型训练

3.3 4-bit 量化（NF4）

🔥 【P0 必须理解】 QLoRA 的核心：用 4-bit 量化加载 base model（不动），仅对 LoRA 层用 16-bit 训练。

NormalFloat4 (NF4)：专门为神经网络权重设计的 4-bit 数据类型
  特点：对权重分布（零均值正态分布）更优的量化映射

from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",           # NF4 量化
    bnb_4bit_use_double_quant=False       # 双重量化（更省但精度更低）
)
quantized_model = AutoModelForCausalLM.from_pretrained(
    "model/Qwen3-8B", quantization_config=bnb_config
)

3.4 优化方案对比

方案	显存需求（7B 模型）	速度	实现复杂度
全参 FP16	~84GB	快	低
+ 梯度检查点	~42GB	中	低
+ ZeRO-2	~21GB	中	低
+ 4-bit QLoRA	~8GB	中	中
+ CPU 卸载	~4GB	慢	中

---

四、本阶段文件索引

优先级	文件	路径
🔥 P0	03_dpo_train.py	3.代码/fine_tune_proj/03_dpo_train.py
🟡 P1	04_single_gpu_optimizations_pseudocode.py	3.代码/fine_tune_proj/04_single_gpu_optimizations_pseudocode.py