直接偏好优化：模型训练 - 从零开始大模型

上一章，我们构建了 DPO 的训练数据：每个样本包含一个提示词、一个接受答案和一个拒绝答案。现在，我们用这些数据对 SFT 模型进行直接偏好优化。

直接偏好优化¶

直接偏好优化的核心思想：让模型给接受答案打更高的分，给拒绝答案打更低的分。

序列对数概率¶

如何衡量模型对一个答案的打分？用序列对数概率（Sequence Log Probability）：

\log P(y|x) = \sum_{t=1}^{T} \log P(y_t | x, y_{<t})

(1)

这里：

$x$ ：提示词；
$y$ ：答案， $y_t$ 是答案的第 $t$ 个词元， $y_{<t}$ 是它前面的所有词元；
$T$ ：答案的词元数量。

也就是在提示词 $x$ 的条件下，模型生成答案中每个词元的对数概率之和。这个值越大，说明模型越“喜欢”这个答案。

DPO 损失函数¶

DPO 的损失函数基于训练模型和参考模型的对数概率比：

\mathcal{L} = -\log \sigma\left(\beta \left(\log \frac{\pi_\theta(y_c|x)}{\pi_{ref}(y_c|x)} - \log \frac{\pi_\theta(y_r|x)}{\pi_{ref}(y_r|x)} \right) \right)

(2)

这里：

$\pi_\theta$ ：训练模型（Policy）；
$\pi_{ref}$ ：参考模型（Reference），通常是 SFT 后的模型，训练时参数冻结；
$y_c$ ：接受答案（Chosen）；
$y_r$ ：拒绝答案（Rejected）；
$\beta$ ：系数，控制模型偏离参考模型的程度；
$\sigma$ ：Sigmoid 函数。

损失函数的目标：让 $\pi_\theta$ 给 $y_c$ 的概率比 $\pi_{ref}$ 给的更高，给 $y_r$ 的概率比 $\pi_{ref}$ 给的更低。

参考模型¶

如果没有参考模型，训练模型可能会走极端：给接受答案无限高的概率，给拒绝答案无限低的概率，导致生成能力下降。参考模型相当于一个锚点，防止模型偏离太远。

import os
import pickle

import numpy as np

from src.core import *
from src.gpt import *

np.random.seed(42)

数据集¶

SFT 数据集¶

和第 16 章相同的 SFT 数据集 SFTDataset，DPO 数据集在它的基础上扩展。

class SFTDataset(Dataset):

    def __init__(self, filename, context_size=64, split=0.9):
        self.filename = filename
        self.context_size = context_size
        self.split = split
        super().__init__(1)

    def load(self):
        with open(self.filename, "rb") as f:
            examples = pickle.load(f)

        split = int(len(examples) * self.split)
        self.train_data = self._pack(examples[:split])
        self.test_data = self._pack(examples[split:])

    def _pack(self, examples):
        xs, ys, masks = [], [], []
        for prompt, response in examples:
            x, y, mask = self._build_example(prompt, response)
            xs.append(x)
            ys.append(y)
            masks.append(mask)
        return xs, ys, masks

    def _build_example(self, prompt, response):
        tokens = list(prompt) + list(response)
        if len(tokens) > self.context_size + 1:
            overflow = len(tokens) - (self.context_size + 1)
            prompt = prompt[overflow:] if overflow < len(prompt) else []
            tokens = list(prompt) + list(response)
            tokens = tokens[-(self.context_size + 1):]

        response_start = len(prompt)

        x = np.array(tokens[:-1], dtype=np.int64)
        y = np.array(tokens[1:], dtype=np.int64)
        mask = np.arange(len(y)) + 1 >= response_start
        return x, y, mask

    def __getitem__(self, index):
        s = slice(index * self.batch_size, (index + 1) * self.batch_size)
        x, y, mask = self.data
        return Tensor(x[s]), Tensor(y[s]), Tensor(mask[s])

DPO 数据集¶

DPO 数据集 DPODataset 继承自 SFT 数据集，处理（提示词，接受答案，拒绝答案）三元组：

将每个三元组拆成两个：（提示词，接受答案）和（提示词，拒绝答案）；
分别将两组数据编码成特征和标签，并添加掩码；
返回六元组：（接受特征，接受标签，接受掩码，拒绝特征，拒绝标签，拒绝掩码）。

class DPODataset(SFTDataset):

    # 打包函数：把（提示词，接受，拒绝）转换成（提示词，接受）和（提示词，拒绝）两组
    def _pack(self, examples):
        # 接受特征，接受标签，接受掩码，拒绝特征，拒绝标签，拒绝掩码
        cxs, cys, cmasks, rxs, rys, rmasks = [], [], [], [], [], []
        for prompt, chosen, rejected in examples:
            x, y, m = self._build_example(prompt, chosen)
            cxs.append(x)
            cys.append(y)
            cmasks.append(m)
            x, y, m = self._build_example(prompt, rejected)
            rxs.append(x)
            rys.append(y)
            rmasks.append(m)
        return cxs, cys, cmasks, rxs, rys, rmasks

    def __getitem__(self, index):
        s = slice(index * self.batch_size, (index + 1) * self.batch_size)
        cx, cy, cmask, rx, ry, rmask = self.data
        return (Tensor(cx[s]), Tensor(cy[s]), Tensor(cmask[s]),
                Tensor(rx[s]), Tensor(ry[s]), Tensor(rmask[s]))

损失函数¶

DPO 二元交叉熵损失函数¶

DPO 损失函数 DPOLoss 实现上面的损失，也就是对边际取 Sigmoid 之后的二元交叉熵：

计算差值：

p_{diff} = policy_{chosen} - policy_{rejected}

(3)

r_{diff} = reference_{chosen} - reference_{rejected}

(4)

计算边际：

margin = \beta \cdot (p_{diff} - r_{diff})

(5)

Sigmoid：

sigmoid = \frac{1}{1 + e^{- margin}}

(6)

损失：返回 $-\log(sigmoid)$ 的均值。

这里：

$policy_{chosen}$ 、 $policy_{rejected}$ ：训练模型对接受答案、拒绝答案的序列对数概率；
$reference_{chosen}$ 、 $reference_{rejected}$ ：参考模型对接受答案、拒绝答案的序列对数概率；
$\beta$ ：系数，对应代码中的 beta，默认值为 0.1。

class DPOLoss(Loss):

    def __init__(self, beta=0.1):
        self.beta = beta

    def loss(self, pc: Tensor, pr: Tensor, rc, rr):
        # 训练模型接受和拒绝的对数概率差
        p_diff = pc.data - pr.data
        # 参考模型接受和拒绝的对数概率差
        r_diff = np.asarray(rc) - np.asarray(rr)
        # 边际：相对于参考模型，训练模型对接受答案的偏好提升了多少
        margin = self.beta * (p_diff - r_diff)
        # 创建损失值张量
        sig = 1.0 / (1.0 + np.exp(-margin))
        e = Tensor(-np.sum(np.log(np.clip(sig, 1e-10, 1.0))) / pc.shape[0])

        # 损失值张量的梯度（闭包）函数
        def gradient_fn():
            common = (sig - 1.0) / pc.shape[0] * self.beta * e.grad
            pc.grad += common
            pr.grad += -common

        # 设置损失值张量梯度（闭包）函数和父节点列表
        return e.attach(gradient_fn, parents=[pc, pr])

训练器¶

DPO 训练器¶

DPO 训练器 DPOTrainer 同样继承自 GPT 训练器 GPTTrainer，有以下主要区别：

双模型：DPO 训练器不仅有当前训练的模型 model，还有一个参考模型 reference。参考模型的参数冻结，不参与训练；
_sequence_log_prob（序列对数概率函数）：计算模型对每组（特征，标签，掩码）的对数概率（只计算掩码为 1 的部分）；
train（训练函数）：同时计算训练模型和参考模型对接受答案、拒绝答案的打分，传给 DPO 损失函数。

💡 提示：参考模型和训练模型在训练开始时参数完全相同。训练过程中，训练模型的参数不断更新，参考模型的参数保持不变，当作对照组。注意训练模型处于训练模式（启用随机失活），参考模型处于推理模式（关闭随机失活），所以即使参数相同，两者给出的打分也不会完全一样。

class DPOTrainer(GPTTrainer):

    def __init__(self, model, loss_fn, optimizer, reference):
        super().__init__(model, loss_fn, optimizer)
        self.reference = reference
        # 参考模型设置为推理模式
        self.reference.eval()

    def train(self, dataset, epochs, scheduler=None, filename=None):
        self.model.train()

        steps = 0
        for epoch in range(epochs):
            order = list(range(len(dataset)))
            np.random.shuffle(order)

            total_loss = 0.0
            for step, i in enumerate(order):
                if scheduler is not None:
                    self.optimizer.lr = scheduler.step(steps)

                cx, cy, cm, rx, ry, rm = dataset[i]
                # 训练模型的序列对数概率
                policy_chosen = self._sequence_log_prob(self.model(cx), cy, cm)
                policy_rejected = self._sequence_log_prob(self.model(rx), ry, rm)
                # 参考模型的序列对数概率
                ref_chosen = self._sequence_log_prob(self.reference(cx), cy, cm).data
                ref_rejected = self._sequence_log_prob(self.reference(rx), ry, rm).data
                # 计算损失
                loss = self.loss_fn(policy_chosen, policy_rejected, ref_chosen, ref_rejected)

                self.optimizer.zero_grad()
                loss.backward()
                total_loss += float(loss.data)
                self.optimizer.clip_grad_norm()
                self.optimizer.step()
                steps += 1

                if (step + 1) % 100 == 0:
                    lr = f" lr {self.optimizer.lr:.6f}" if scheduler is not None else ""
                    print(f"epoch {epoch + 1} step {step + 1}/{len(dataset)} loss {(total_loss / 100):.4f}{lr}")
                    total_loss = 0.0

            if filename is not None:
                self.save(filename)
                print(f"epoch {epoch + 1} saved DPO model to {filename}")

    # 测试函数：微调阶段不做测试
    def test(self, dataset):
        return None

    # 序列对数概率函数
    @staticmethod
    def _sequence_log_prob(x, y, mask):
        # 内置 Softmax
        exp = np.exp(x.data - np.max(x.data, axis=-1, keepdims=True))
        softmax = exp / np.sum(exp, axis=-1, keepdims=True)

        y = np.asarray(y.data, dtype=np.int64)
        mask = np.asarray(mask.data, dtype=softmax.dtype)
        # 创建概率值张量
        picked = np.clip(np.take_along_axis(softmax, y[..., None], axis=-1).squeeze(-1), 1e-10, 1)
        p = Tensor(np.sum(np.log(picked) * mask, axis=-1))

        # 概率值张量的梯度（闭包）函数
        def gradient_fn():
            grad = -softmax.copy()
            target_grad = np.take_along_axis(grad, y[..., None], axis=-1) + 1.0
            np.put_along_axis(grad, y[..., None], target_grad, axis=-1)
            grad *= mask[..., None]
            x.grad += p.grad[:, None, None] * grad

        # 设置概率值张量梯度（闭包）函数和父节点列表
        return p.attach(gradient_fn, {x})

    # 加载参考模型参数，并清零梯度
    def load_reference(self, filename):
        if os.path.isfile(filename):
            data = np.load(filename, allow_pickle=False)
            for i, p in enumerate(self.reference.parameters):
                p.data = data[f"param_{i}"]
                p.grad = np.zeros_like(p.data)

超参数¶

学习率¶

偏好优化阶段通常使用一个比 SFT 阶段更小的学习率。本章取 0.000001，是 SFT 阶段的百分之一。

LEARNING_RATE = 0.000001

模型参数¶

其他模型参数需要保持和预训练阶段一致。

BATCH_SIZE = 16
CONTEXT_SIZE = 256
EMBEDDING_SIZE = 384
HEADS = 6
BLOCKS = 6

建模¶

加载模型¶

创建两个模型：

model：训练模型；
reference：参考模型。

两个模型都加载 SFT 模型参数。

DATA_FILE = "../../../../data/tinyshakespeare.txt"
MODEL_FILE = "../../../../data/tinyshakespeare-sft.npz"

dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
# 训练模型
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
# 参考模型
reference = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = DPOLoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = DPOTrainer(model, loss_fn, optimizer, reference)
# 加载训练模型参数
trainer.load(MODEL_FILE)
# 加载参考模型参数
trainer.load_reference(MODEL_FILE)

加载 DPO 数据¶

DPO_SAMPLES = "../../../../data/dpo-samples.pkl"

dpo_dataset = DPODataset(DPO_SAMPLES, CONTEXT_SIZE)

训练¶

DPO_MODEL = "../../../../data/tinyshakespeare-dpo.npz"

scheduler = WarmupCosineScheduler(LEARNING_RATE, len(dpo_dataset), 50, LEARNING_RATE / 10)
trainer.train(dpo_dataset, 1, scheduler, DPO_MODEL)

epoch 1 step 100/460 loss 0.6913 lr 0.000001
epoch 1 step 200/460 loss 0.6926 lr 0.000001
epoch 1 step 300/460 loss 0.6761 lr 0.000000
epoch 1 step 400/460 loss 0.6767 lr 0.000000
epoch 1 saved DPO model to ../../../../data/tinyshakespeare-dpo.npz

推理¶

print(trainer.generate(dataset, prompt="ROMEO:"))

ROMEO:

HORTENSIO:
Villor, will I come: did a hear,
Not our drow supposessiance, that you married haste,
Like a joy a Jack for joy her.

First Servingman:
I do go of year.

PITONES:
Hereford thy country.

PRUTUS:
Call now.

CoRLIFFORD:

SICINIAN:

BRUTUS:

IS:
N

本章，我们用 DPO 数据集对 SFT 模型进行了直接偏好优化训练：让模型给接受答案打更高的分，给拒绝答案打更低的分；参考模型作为锚点，防止模型偏离太远。

训练损失在 0.68 到 0.72 之间波动，没有明显下降。当训练模型和参考模型给出的打分完全相同时，损失正好等于 $\ln 2 \approx 0.693$ ，所以这说明训练后的模型相对参考模型变化很小。我们的重点是理解 DPO 的流程，对这个很小的教学模型，不需要期待明显的效果。

至此，我们完成了模型微调，一步步让预训练模型学会和使用者对话：

SFT（监督微调）：让模型学会回答问题；
DPO（直接偏好优化）：让模型学会合理地回答问题。

课后练习¶

尝试把 DPO 损失函数的 beta 参数从 0.1 改成 0.5 或者 0.01，观察训练结果有什么变化？beta 参数用来控制什么？
尝试把学习率提高到 0.00001 或者 0.0001，观察训练损失的变化。