s18 大语言模型 — exercise.py 练习指南

Download exercise.py

练习目标

通过动手实现 Scaling Law 最优配比计算、DPO 损失函数和 LoRA 配置，建立对 LLM 核心概念的量化直觉。完成后你将能够：

1. 给定计算预算，找出 Chinchilla 最优的参数量和数据量配比
2. 独立写出 DPO 损失函数的 PyTorch 实现
3. 理解 LoRA 配置中的关键参数

预备知识

• Scaling Law：$L(N,D)=a/N^{\alpha }+b/D^{\beta }+c$，其中 $\alpha \approx 0.076$, $\beta \approx 0.095$
• Chinchilla 最优配比：$D\approx 20N$，训练计算量 $C\approx 6ND$
• DPO 损失：$-\log \sigma (\beta \cdot ({\mathrm{\Delta }}_{\text{policy}}-{\mathrm{\Delta }}_{\text{ref}}))$
• LoRA：$h=Wx+\frac{\alpha }{r}BAx$，其中 $B\in {\mathbb{R}}^{d\times r}$，$A\in {\mathbb{R}}^{r\times k}$

---

任务清单

练习 1：实现 Chinchilla 最优配比计算

目标：给定计算预算 $C$，找到使损失最小的参数量 $N$ 和数据量 $D$。

核心约束：训练计算量 $C\approx 6ND$（前向传播约 $2ND$ FLOPs，反向传播约 $4ND$ FLOPs）。

求解方法：在给定的 $C$ 约束下，问题变为找到最小化 $L(N,D)$ 的 $N$ 和 $D$：

$$\underset{N,D}{min}\frac{a}{N^{\alpha }}+\frac{b}{D^{\beta }}+c\text{s.t.}C=6ND$$

使用拉格朗日乘数法可以推导出解析解，但练习中使用更直观的对数空间网格搜索：

def find_optimal_ND(compute_budget, a=1.5, b=2.0, alpha=0.076, beta=0.095):
    # 在对数空间中搜索 N
    N_candidates = np.logspace(6, 12, 200)          # 1M → 1T 参数
    best_loss = float('inf')
    best_N, best_D = None, None

    for N in N_candidates:
        D = compute_budget / (6 * N)                 # 由约束确定 D
        if D <= 0:
            continue
        loss = a / (N ** alpha) + b / (D ** beta) + 1.0
        if loss < best_loss:
            best_loss = loss
            best_N = N
            best_D = D

    return best_N, best_D, best_loss

关键步骤：

1. 对 $N$ 做对数均匀采样（np.logspace(6, 12, 200)）
2. 对每个 $N$，由 $C=6ND$ 计算 $D$（从约束推导）
3. 计算对应的损失 $L(N,D)$
4. 找到使损失最小的 $(N,D)$ 对

检验标准：最优配比 $D/N$ 应该约为 20（即 D=20N）。

预期输出：

[练习1] Chinchilla 最优配比:
  给定计算预算 C=5.88e+22 FLOPs
  最优参数量 N=7.00e+09 (~7.0B)
  最优数据量 D=1.40e+12 tokens (~1.40T)
  最优配比 D/N=200.0 (期望: ≈20)

---

练习 2：实现 DPO 损失函数

目标：补全 dpo_loss() 函数，实现 DPO 损失的计算。

核心公式：

$${\mathcal{L}}_{\text{DPO}}=-\frac{1}{N}\sum _{i=1}^N\log \sigma (\beta \cdot [(\log {\pi }_{\theta }(y_w|x)-\log {\pi }_{\theta }(y_l|x))-(\log {\pi }_{\text{ref}}(y_w|x)-\log {\pi }_{\text{ref}}(y_l|x))])$$

TODO 步骤：

def dpo_loss(pi_logps_chosen, pi_logps_rejected,
             ref_logps_chosen, ref_logps_rejected, beta=0.1):
    # 步骤 1: 计算策略模型的"偏好差异"
    pi_log_ratio = pi_logps_chosen - pi_logps_rejected
    # 步骤 2: 计算参考模型的"偏好差异"
    ref_log_ratio = ref_logps_chosen - ref_logps_rejected
    # 步骤 3: 策略模型相对于参考模型的改善
    logits = beta * (pi_log_ratio - ref_log_ratio)
    # 步骤 4: -log σ(logits)
    loss = -F.logsigmoid(logits).mean()
    return loss

逐步理解：

1. pi_log_ratio：如果策略模型正确偏好（给好回答更高的 log 概率）→ 此值为正 → 好信号
2. ref_log_ratio：参考模型的基准偏好。即使参考模型的偏好也是对的，策略模型仍可以比它做得更好
3. logits = beta * (pi_log_ratio - ref_log_ratio)：策略模型相对于参考模型的改善程度。这个值越大，DPO 损失越小
4. -F.logsigmoid(logits)：logsigmoid(x) = log(1/(1+e^{-x}))，其范围为 $(-\mathrm{\infty },0)$。取负号后损失为正，当 logits 很大时损失趋近于 0

logsigmoid vs log(sigmoid(.))：F.logsigmoid 在数值上比 torch.log(torch.sigmoid(x)) 更稳定（前者直接计算 -softplus(-x)，避免了 sigmoid 可能出现的数值溢出）。

预期输出：

[练习2] DPO 损失:
  好回答log P=[-1.5, -2.0, -1.8], 差回答log P=[-4.0, -5.0, -4.5]
  DPO Loss = XXXXX (期望: 一个较小的正数)

当策略模型的 $P(y_w|x)\gg P(y_l|x)$ 时，损失应很小（模型已经很好地学会区分好/差回答）。

---

练习 3：配置 LoRA 适配器

目标：创建 LoRA 配置字典，理解每个参数的含义。

LoRA 的核心参数：

参数	典型值	含义
r	8, 16, 32, 64	低秩分解的秩。越大表达能力越强，但参数量也越大
lora_alpha	16, 32	缩放系数。越大 LoRA 更新的影响越大。实际缩放 = alpha/r
target_modules	["q_proj", "v_proj"]	对哪些模块应用 LoRA。Qwen/Llama 系列通常加在 Q 和 V 投影上
lora_dropout	0.0 - 0.1	LoRA 的 dropout 概率。小数据集上可设为 0.05-0.1
bias	"none"	LoRA 不训练 bias，保持与原始模型的 bias 一致
task_type	"CAUSAL_LM"	任务类型：因果语言模型（GPT 系列）

TODO 步骤：

def create_lora_config(r=8, lora_alpha=16.0, target_modules=None,
                       lora_dropout=0.1):
    if target_modules is None:
        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"]

    config = {
        "r": r,
        "lora_alpha": lora_alpha,
        "target_modules": target_modules,
        "lora_dropout": lora_dropout,
        "bias": "none",
        "task_type": "CAUSAL_LM",
    }
    return config

常见配置建议：

• 仅 Q+V 投影：参数量最小，效果通常够用（["q_proj", "v_proj"]）
• Q+K+V+O 投影：参数量翻倍，效果更好（["q_proj", "k_proj", "v_proj", "o_proj"]）
• 全注意力+FFN：参数量最多，但某些任务可能需要（加上 ["gate_proj", "up_proj", "down_proj"]）

为什么主要在 Attention 投影上加 LoRA？ 注意力层负责"查找"相关信息，任务适配的核心在于改变"模型关注什么"。FFN 层更多是"知识存储"，改变 FFN 可能导致原有知识被覆盖。但实践中，在 FFN 上加 LoRA 有时效果也很好。

预期输出：

[练习3] LoRA 配置: {'r': 16, 'lora_alpha': 32.0, 'target_modules': [...], ...}

---

三个练习的关系

练习	LLM 概念	在 LLM 创业/应用中的位置
练习 1: Chinchilla 配比	如何分配算力	训练新模型的决策依据
练习 2: DPO 损失	如何对齐人类偏好	微调模型使其"更好用"
练习 3: LoRA 配置	如何高效微调	消费级硬件上微调大模型

这三个概念在 LLM 产品开发中都有直接应用：决定模型规模和数据量（练习 1）、通过偏好数据优化模型行为（练习 2）、在有限算力下微调模型（练习 3）。

检查要点

运行 python exercise.py，确认：

• [ ] 练习 1 D/N 配比接近 20
• [ ] 练习 2 DPO 损失为正数，正确偏好场景损失 < 错误偏好场景
• [ ] 练习 3 LoRA 配置包含所有必需字段

完成练习后，返回 demo.py 观察这些概念的完整实现和可视化。

完整代码

# -*- coding: utf-8 -*-
"""
s18 大语言模型 — 练习题
==============================================
请补全以下 TODO 部分，完成后运行验证。
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math


# ============================================================
# 练习 1：实现 Chinchilla 最优配比计算
# ============================================================

def find_optimal_ND(
    compute_budget: float,  # 计算预算 C（FLOPs）
    a: float = 1.5,
    b: float = 2.0,
    alpha: float = 0.076,
    beta: float = 0.095,
) -> tuple:
    """
    TODO: 对于给定的计算预算 C，找到最优的参数 N 和数据 D
    使得在约束 C ≈ 6ND 下最小化损失 L(N, D) = a/N^α + b/D^β + c

    用拉格朗日乘数法可以推出:
    N_opt ∝ C^(β/(α+β))
    D_opt ∝ C^(α/(α+β))

    参数：
        compute_budget: 计算预算 C（FLOPs）
        a, b, alpha, beta: Kaplan scaling law 参数
    返回：
        (N_opt, D_opt, L_opt): 最优参数量、最优数据量、对应的损失
    """
    # TODO: 实现找到最优 N 和 D
    # 步骤：
    #   1. 设置约束: C = 6 * N * D
    #   2. 将 D = C / (6N) 代入损失函数
    #   3. 使用数值方法（如简单网格搜索）找到最小化损失的 N
    #   4. 或者使用拉格朗日乘数推导的比例关系
    #
    # 提示（简化数值解法）:
    #   对 N 在合理范围内做搜索：
    #     N 从 1e6 到 1e12 对数搜索
    #     D = compute_budget / (6 * N)
    #     计算 L = a/N^α + b/D^β + c
    #     找到最小 L 对应的 N, D
    # ===== 你的代码在这里 =====
    N_opt = 1e9  # 占位值
    D_opt = compute_budget / (6 * N_opt)
    L_opt = a / (N_opt ** alpha) + b / (D_opt ** beta) + 1.0  # c=1.0
    # ==========================
    return N_opt, D_opt, L_opt


# 测试
# 假设 C = 6 * 7e9 * 1.4e12 FLOPs（LLaMA 7B 规模的计算预算）
C_test = 6 * 7e9 * 1.4e12
try:
    N_opt, D_opt, L_opt = find_optimal_ND(C_test)
    print(f"[练习1] Chinchilla 最优配比:")
    print(f"  给定计算预算 C={C_test:.2e} FLOPs")
    print(f"  最优参数量 N={N_opt:.2e} (~{N_opt/1e9:.1f}B)")
    print(f"  最优数据量 D={D_opt:.2e} tokens (~{D_opt/1e12:.2f}T)")
    print(f"  最优配比 D/N={D_opt/N_opt:.1f} (期望: ≈20)")
    print(f"  预测损失 L={L_opt:.4f}")
except Exception as e:
    print(f"[练习1] 未完成实现: {e}")


# ============================================================
# 练习 2：实现 DPO 损失函数
# ============================================================

def dpo_loss(
    pi_logps_chosen: torch.Tensor,    # (batch,) 策略模型对偏好回答的 log 概率
    pi_logps_rejected: torch.Tensor,   # (batch,) 策略模型对较差回答的 log 概率
    ref_logps_chosen: torch.Tensor,    # (batch,) 参考模型对偏好回答的 log 概率
    ref_logps_rejected: torch.Tensor,  # (batch,) 参考模型对较差回答的 log 概率
    beta: float = 0.1,                 # KL 惩罚系数
) -> torch.Tensor:
    """
    TODO: 实现 DPO（Direct Preference Optimization）损失函数

    公式:
    L_DPO = -log σ( β*(log π_θ(y_w|x) - log π_ref(y_w|x))
                   - β*(log π_θ(y_l|x) - log π_ref(y_l|x)) )

    即: -log σ( β * (pi_diff - ref_diff) )
    其中 pi_diff = log π_θ(y_w) - log π_θ(y_l)
          ref_diff = log π_ref(y_w) - log π_ref(y_l)

    参数：
        pi_logps_chosen: log π_θ(y_w | x)
        pi_logps_rejected: log π_θ(y_l | x)
        ref_logps_chosen: log π_ref(y_w | x)
        ref_logps_rejected: log π_ref(y_l | x)
        beta: KL 惩罚系数
    返回：
        loss: 标量 DPO 损失值
    """
    # TODO: 实现步骤
    #   1. pi_log_ratio = pi_logps_chosen - pi_logps_rejected
    #   2. ref_log_ratio = ref_logps_chosen - ref_logps_rejected
    #   3. logits = beta * (pi_log_ratio - ref_log_ratio)
    #   4. loss = -F.logsigmoid(logits).mean()
    # ===== 你的代码在这里 =====
    loss = torch.tensor(0.0)
    # ==========================
    return loss


# 测试 DPO 损失
chosen = torch.tensor([-1.5, -2.0, -1.8])    # 好回答 log 概率
rejected = torch.tensor([-4.0, -5.0, -4.5])  # 差回答 log 概率
ref_chosen = torch.tensor([-2.5, -2.5, -2.5])  # 参考模型均匀
ref_rejected = torch.tensor([-2.5, -2.5, -2.5])
try:
    loss = dpo_loss(chosen, rejected, ref_chosen, ref_rejected, beta=0.1)
    print(f"\n[练习2] DPO 损失:")
    print(f"  好回答log P={chosen.tolist()}, 差回答log P={rejected.tolist()}")
    print(f"  DPO Loss = {loss.item():.4f} (期望: 一个较小的正数)")
    # 好的模型应该给好回答高概率，差回答低概率 → 损失应该小
except Exception as e:
    print(f"[练习2] 未完成实现: {e}")


# ============================================================
# 练习 3：配置 LoRA 适配器
# ============================================================

def create_lora_config(
    r: int = 8,               # LoRA 秩
    lora_alpha: float = 16.0, # LoRA 缩放系数
    target_modules: list = None,  # 需要应用 LoRA 的模块名列表
    lora_dropout: float = 0.1,    # LoRA dropout
) -> dict:
    """
    TODO: 创建一个 LoRA 配置字典（模拟 PEFT 库的 LoraConfig）

    参数：
        r: LoRA 秩
        lora_alpha: LoRA alpha 缩放参数
        target_modules: 需要应用 LoRA 的目标模块（如 ["q_proj", "v_proj"]）
        lora_dropout: LoRA dropout 概率
    返回：
        config: LoRA 配置字典
    """
    if target_modules is None:
        # Qwen/Llama 系列常见的目标模块
        target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"]

    # TODO: 构建配置字典
    # 包含以下键:
    #   "r": r
    #   "lora_alpha": lora_alpha
    #   "target_modules": target_modules
    #   "lora_dropout": lora_dropout
    #   "bias": "none" (LoRA 不训练 bias)
    #   "task_type": "CAUSAL_LM" (因果语言模型)
    # ===== 你的代码在这里 =====
    config = {}
    # ==========================
    return config


# 测试
config = create_lora_config(r=16, lora_alpha=32.0,
                            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"])
print(f"\n[练习3] LoRA 配置: {config}")

print("\n所有练习测试完成！请对比 demo.py 查看参考实现。")
print("""
提示:
  - Chinchilla 最优: 用拉格朗日乘数法或网格搜索找到最小化损失的配比
  - DPO: 核心是让好回答概率 > 差回答概率，公式: -log σ(β·ΔlogP)
  - LoRA: 秩 r 通常 8-64, alpha 通常 16-32, 常见的 target_modules 是 q_proj 和 v_proj
""")