上一章，我们构建了 SFT 训练数据：512 个（提示词，回答）数据对。本章，我们将用这些数据对预训练模型进行监督微调。
监督微调¶
监督微调的核心思想：只计算回答部分的损失。
预训练时，模型对每个位置都计算损失，因为每个位置都在学习下一个词元。但到了监督微调阶段，我们只关心模型是否学会了生成合适的答案，不需要它去学习怎样生成提示词。
掩码¶
解决办法就是给每个样本加一个掩码（Mask）：
提示词部分：掩码为 0，不计算损失；
回答部分：掩码为 1，正常计算损失。
这样，提示词部分不会产生损失，模型只从回答部分学习。
策略¶
学习率：监督微调通常采用比预训练阶段小一个数量级的学习率。因为模型已经训练好了，我们只想微调它生成文本的方向，而不是重新学习语言知识；
轮数：SFT 数据量小，训练太久容易过拟合，所以轮数通常很少。本章只训练 1 个轮次。
import pickle
import numpy as np
from src.core import *
from src.gpt import *np.random.seed(42)数据集¶
SFT 数据集¶
SFT 数据集 SFTDataset 和预训练的数据集不同，要处理（提示词，回答）数据对：
拼接：把提示词和回答拼接到一起，成为一个词元序列；
截断：如果超出上下文尺寸，从提示词部分的头部截断，保留完整的回答部分；
掩码：给提示词部分添加掩码 0，给回答部分添加掩码 1。结果返回
(特征, 标签, 掩码)。
此外，各个样本的长度不同，所以批大小固定为 1，不需要填充。
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)
# 加载函数：读取 SFT 样本文件
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])损失函数¶
SFT 交叉熵损失函数¶
SFT 损失函数 SFTLoss 是带掩码的交叉熵损失函数：
计算所有位置的交叉熵；
乘以掩码，只保留回答部分的损失；
除以回答部分的长度（掩码中 1 的个数），求取有效值的平均。
反向传播时，提示词位置的梯度同样为 0。
class SFTLoss(Loss):
def loss(self, p: Tensor, y: Tensor, mask):
# 内置 Softmax
exp = np.exp(p.data - np.max(p.data, axis=-1, keepdims=True))
softmax = exp / np.sum(exp, axis=-1, keepdims=True)
# 展平
flat_softmax = softmax.reshape(-1, softmax.shape[-1])
flat_y = y.data.reshape(-1).astype(np.int64)
# 掩码
flat_mask = mask.data.reshape(-1)
n = max(np.sum(flat_mask), 1.0)
rows = np.arange(len(flat_y))
# 创建损失值张量
log = np.log(np.clip(flat_softmax[rows, flat_y], 1e-10, 1))
# 使用掩码
ce = Tensor(0 - np.sum(log * flat_mask) / n)
# 损失值张量的梯度（闭包）函数
def gradient_fn():
flat_grad = flat_softmax.copy()
flat_grad[rows, flat_y] -= 1
# 使用掩码：提示词部分掩码为 0，梯度也是 0
flat_grad *= flat_mask[:, None]
p.grad += ce.grad * flat_grad.reshape(softmax.shape) / n
# 设置损失值张量梯度（闭包）函数和父节点列表
return ce.attach(gradient_fn, parents=[p])训练器¶
SFT 训练器¶
SFT 训练器 SFTTrainer 继承自 GPT 训练器 GPTTrainer，有以下主要区别：
train（训练函数）：从数据集接收掩码，作为参数传给损失函数；test（测试函数）：监督微调阶段不需要测试，直接返回None。
SFT 只做 1 轮训练，结束后直接保存模型参数。
class SFTTrainer(GPTTrainer):
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)
# 传递掩码
feature, label, mask = dataset[i]
prediction = self.model(feature)
loss = self.loss_fn(prediction, label, mask)
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 SFT model to {filename}")
# 测试函数：微调阶段不做测试
def test(self, dataset):
return NoneLEARNING_RATE = 0.0001模型参数¶
其他模型参数需要保持和预训练阶段一致。
BATCH_SIZE = 16
CONTEXT_SIZE = 256
EMBEDDING_SIZE = 384
HEADS = 6
BLOCKS = 6DATA_FILE = "../../../../data/tinyshakespeare.txt"
MODEL_FILE = "../../../../data/tinyshakespeare-gpt.npz"
dataset = CharDataset(DATA_FILE, BATCH_SIZE, CONTEXT_SIZE)
model = GPT(dataset.vocab_size, CONTEXT_SIZE, EMBEDDING_SIZE, HEADS, BLOCKS)
loss_fn = SFTLoss()
optimizer = AdamWOptimizer(model.parameters, lr=LEARNING_RATE)
trainer = SFTTrainer(model, loss_fn, optimizer)
trainer.load(MODEL_FILE)加载 SFT 数据¶
SFT 样本按 9:1 划分训练集和测试集。训练集有 460 个样本，批大小为 1，所以一轮训练共 460 步；测试集的 52 个样本在本章不使用。
SFT_SAMPLES = "../../../../data/sft-samples.pkl"
sft_dataset = SFTDataset(SFT_SAMPLES, CONTEXT_SIZE)训练¶
使用预热余弦衰减学习率调度器：前 50 步线性预热到 LEARNING_RATE，之后按余弦曲线衰减到 LEARNING_RATE / 10。
SFT_MODEL = "../../../../data/tinyshakespeare-sft.npz"
scheduler = WarmupCosineScheduler(LEARNING_RATE, len(sft_dataset), 50, LEARNING_RATE / 10)
trainer.train(sft_dataset, 1, scheduler, SFT_MODEL)epoch 1 step 100/460 loss 1.5225 lr 0.000097
epoch 1 step 200/460 loss 1.4181 lr 0.000074
epoch 1 step 300/460 loss 1.3249 lr 0.000040
epoch 1 step 400/460 loss 1.3159 lr 0.000015
epoch 1 saved SFT model to ../../../../data/tinyshakespeare-sft.npz
推理¶
让微调后的模型从提示词 ROMEO: 开始生成。
print(trainer.generate(dataset, prompt="ROMEO:"))ROMEO:
He theople, that low with many thrieve weak:
While crampudged you I know the valous in your blood,
I must my patience, and my Edward's life!
TRANIO:
ANORTEN:
POLIXENES:
I know not mine a deep a shado, your whom
PETRUCHIO:
You say you know! but know re
本章，我们用 SFT 数据集对预训练模型进行了监督微调：只计算回答部分的损失，让模型学习“提示词，回答”的对话格式。从生成的结果可以看到，模型保持了“角色：台词”轮流发言的格式，但是内容依然不通顺。
但是 SFT 也有自己的局限：它只告诉模型怎样回答问题，却不能告诉模型什么是更好的答案。比如两个答案都正确，但是一个简洁明了、另一个冗长啰嗦，SFT 无法告诉模型哪个更好。
下一章，我们将引入直接偏好优化（Direct Preference Optimization，DPO），让模型学会从多个答案中选出更好的一个。
课后练习¶
试试把 SFT 的学习率提高，比如和预训练一样，观察训练效果如何？思考为什么微调需要更小的学习率？