summary: “用一个 16.9M 参数的迷你 GPT，把 tokenizer、预训练、SFT、奖励模型和 PPO/RLHF 这条大模型训练流水线完整跑一遍。”

手把手训练一个迷你大模型：从预训练到 RLHF

如果你第一次听到“训练大模型”，脑海里可能会自动浮现一排显卡、嗡嗡作响的机房、以及一张看起来不太友好的云账单。好消息是：我们今天不训练一个真正能上生产的 ChatGPT，也不挑战算力预算的物理极限。我们要做的是把大模型训练流程缩小成一个可以亲手跑通的“训练沙盘”。

这篇文章会带你从零训练一个迷你 GPT 模型。它会先学会把文本拆成 token，然后通过 next-token prediction 学会续写儿童故事，再经过监督微调变得更听指令，接着训练一个奖励模型判断“哪个回答更好”，最后用 PPO 做一轮教学版 RLHF。

0. 我们到底在训练什么？

所谓“迷你大模型”，不是参数量很大的模型，而是“大模型训练流程的迷你版”。它的价值不在于能力有多强，而在于流程完整、反馈快、容易观察。真正的大模型可能有几十亿、几百亿甚至更多参数，而这里的模型只有约 16.89M 参数。

我们使用的数据是 TinyStories。TinyStories 论文提出了一个有趣问题：小模型到底能不能写出连贯英文故事？作者构造了一个由 GPT-3.5 和 GPT-4 生成的短故事数据集，故事只使用典型 3 到 4 岁儿童能理解的简单词汇。论文显示，在这样受控、简单、结构清晰的数据上，远小于常见大模型的小语言模型也可以生成语法较好、情节连贯的故事。

数学上，语言模型学习的是一个条件概率分布：

$$p_{\theta }(x_1,x_2,\dots ,x_T)=\prod _{t=1}^T{p}_{\theta }(x_t\mid x_{<t})$$

这里的 (x_t) 是第 (t) 个 token，(\theta) 是模型参数。训练的本质就是调节 (\theta)，让模型对真实下一个 token 分配更高概率。

1. 安装环境

安装 PyTorch、Transformers、Datasets、Tokenizers、Accelerate、TRL 等库，python版本建议用3.10：

pip install -U torch torchvision torchaudio
pip install -U transformers datasets tokenizers accelerate trl evaluate rich tqdm

库	用途	在本文中的角色
torch	深度学习框架	模型训练、张量计算、GPU 加速
transformers	模型与训练工具	GPT2Config、GPT2LMHeadModel、Trainer
datasets	数据加载与缓存	加载 TinyStories、读取 text/jsonl 文件
tokenizers	高性能 tokenizer	从零训练 Byte-level BPE
trl	LLM 后训练工具	SFT、Reward Model、PPO/RLHF
accelerate	训练加速与设备管理	配合 Transformers/TRL 使用
tqdm / rich	进度条与输出美化	让训练过程不那么黑盒

Hugging Face Datasets 的 load_dataset 可以从 Hub 或本地文件加载数据，也支持手动指定 data_files。

如果你在 Colab 跑，建议选择 GPU Runtime。原 Notebook 的参考运行里，预训练 10,000 step 大约用了 44 分钟，SFT 2,000 step 大约 13 分钟，Reward Model 2,000 step 大约 13 分钟，PPO 教学跑法大约 11 分钟。不同 GPU、网络和库版本会有差异，这些数字只作为量级参考。

2. 准备 TinyStories数据集

我们先抽取 200,000 条训练故事和 5,000 条验证故事，写成普通文本文件：

from datasets import load_dataset
from pathlib import Path
from tqdm import tqdm

OUT = Path("data")
OUT.mkdir(exist_ok=True)

TRAIN_N = 200_000
VALID_N = 5_000

ds_train = load_dataset("roneneldan/TinyStories", split=f"train[:{TRAIN_N}]")
ds_valid = load_dataset("roneneldan/TinyStories", split=f"validation[:{VALID_N}]")

def write_txt(ds, path):
    with open(path, "w", encoding="utf-8") as f:
        for row in tqdm(ds, desc=f"writing {path}"):
            text = row["text"].strip()
            if text:
                f.write(text.replace("\r", " ") + "\n\n")

write_txt(ds_train, OUT / "tinystories_train.txt")
write_txt(ds_valid, OUT / "tinystories_valid.txt")

print("done")

Generating train split: 100%|██████████| 2119719/2119719 [..]
Generating validation split: 100%|██████████| 21990/21990 [..]
writing data/tinystories_train.txt: 100%|██████████| 200000/200000 [00:03<00:00, 62428.11it/s]
writing data/tinystories_valid.txt: 100%|██████████| 5000/5000 [00:00<00:00, 62167.12it/s]
done

这段代码做了三件朴素但重要的事。

第一，它固定了训练规模。 TinyStories 数据集在 Hugging Face 上约有 214 万行，其中训练集约 212 万行、验证集约 2.2 万行。我们只取子集，验证并跑通完整流程。

第二，它把数据写成 .txt。 预训练阶段最常见的数据格式就是一堆连续文本。语言模型并不天然知道“这是第 1 篇故事，这是第 2 篇故事”，它只看到 token 序列。我们在每个故事之间写入两个换行，相当于给故事之间留一点呼吸空间。

第三，它保留了验证集。 验证集可以让我们观察模型是否真的学到可泛化的语言模式。

3. 训练 tokenizer：决定“字词怎么切”

大模型训练的第一步不是训练模型，而是训练 tokenizer。原因很简单：神经网络不认识字符串，它只认识数字。Tokenizer 的任务是把文本变成 token ID，再在生成时把 token ID 还原成人能读的文本。

我们使用 Byte-level BPE：

from pathlib import Path
from tokenizers import ByteLevelBPETokenizer
from transformers import GPT2TokenizerFast

DATA = "data/tinystories_train.txt"
OUT = Path("artifacts/tokenizer")
OUT.mkdir(parents=True, exist_ok=True)

VOCAB_SIZE = 16_000

tokenizer = ByteLevelBPETokenizer()
tokenizer.train(
    files=[DATA],
    vocab_size=VOCAB_SIZE,
    min_frequency=2,
    special_tokens=["<s>", "<pad>", "</s>", "<unk>"],
)

tokenizer.save(str(OUT / "tokenizer.json"))

hf_tokenizer = GPT2TokenizerFast.from_pretrained(
    str(OUT),
    bos_token="<s>",
    pad_token="<pad>",
    eos_token="</s>",
    unk_token="<unk>",
)

hf_tokenizer.save_pretrained(str(OUT))

print("vocab size:", len(hf_tokenizer))
print("pad:", hf_tokenizer.pad_token_id, "eos:", hf_tokenizer.eos_token_id)

vocab size: 16000
pad: 1 eos: 2

3.1 BPE 在做什么？

BPE，全称 Byte Pair Encoding，可以理解为一种“频繁搭子合并法”。一开始，文本被拆成很小的单位。然后算法不断寻找最常一起出现的相邻单位，把它们合并成一个新 token。Hugging Face Tokenizers 的 quicktour 对 BPE 的描述也正是这个流程：从语料中的字符开始，反复合并最常见的 token pair，直到词表达到目标大小。Tokenizers quicktour

如果用数学写，某一步会选择：

$$(a^{\ast },b^{\ast })=\arg \underset{(a,b)}{\max }\mathrm{count}(a,b)$$

然后把所有相邻的 $a^{\ast },b^{\ast }$ 合并为新 token：

$$a^{\ast }\circ b^{\ast }\to c$$

反复执行后，常见片段会变成一个 token，罕见词会被拆成更小片段。比如英文里 story、little、Once 这类高频片段可能会被合并，而罕见拼写可以继续拆开。

3.2 为什么选 Byte-level BPE？

Byte-level BPE 的好处是“兜底能力”强。普通词级 tokenizer 遇到没见过的词，可能只能交给 <unk>。Byte-level 方法从字节层面出发，理论上任何 UTF-8 文本都能拆成字节序列，再通过 BPE 合并成常用片段。这对真实世界文本很重要，因为用户总会输入奇怪标点、拼写错误、混合语言和没见过的专有名词。

Hugging Face 的 Tokenizers 文档强调，该库支持训练新词表，并且由于底层 Rust 实现，训练和分词速度都很快。Tokenizers 文档 这也是为什么教学项目里直接用它，而不是手写 BPE。手写当然能学原理，但我们今天的主线是完整训练流水线，不是把每个齿轮都现场铸造一遍。

3.3 词表大小怎么选？

词表大小 (V) 是一个很有味道的取舍：

$$\text{embedding 参数量}\approx V\times d$$

这里 (d) 是隐藏维度。本文 (V=16000)，(d=384)，仅 token embedding 大约就是：

$$16000\times 384=6,144,000$$

词表越大，常见词越容易被一个 token 表示，序列长度可能变短；但 embedding 参数变多，低频 token 也更难学好。词表越小，参数更省，但一句话会被切得更碎，模型需要在更长序列上学习依赖。对 TinyStories 这种简单英文数据，16k 是一个比较舒服的教学选择。

4. 从零初始化 GPT：搭一个小型 decoder-only Transformer

Tokenizer 准备好后，我们就可以初始化 GPT 模型。我们使用 GPT2Config 和 GPT2LMHeadModel，也就是 GPT-2 风格的 decoder-only causal language model。Transformers 文档 说明 GPT2LMHeadModel 是在 GPT-2 Transformer 顶部加了语言建模头，用来输出词表上每个 token 的概率。

核心配置如下：

config = GPT2Config(
    vocab_size=len(tokenizer),
    n_positions=BLOCK_SIZE,
    n_ctx=BLOCK_SIZE,
    n_embd=384,
    n_layer=6,
    n_head=6,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
)

model = GPT2LMHeadModel(config)
print("parameters:", sum(p.numel() for p in model.parameters()) / 1e6, "M")

parameters: 16.889856 M

这 16.9M 参数大致来自三部分：

$$P_{\text{embed}}\approx Vd+L_{\text{ctx}}d$$

$$P_{\text{attn/block}}\approx 4d^2$$

$$P_{\text{mlp/block}}\approx 8d^2$$

其中 (V=16000)，(d=384)，上下文长度 (L_{\text{ctx}}=256)，层数为 6。GPT-2 这类模型通常使用多头自注意力和前馈网络堆叠。Transformer 架构最早由 Vaswani 等人在 Attention Is All You Need 中提出，核心思想是用 attention 机制建模序列中 token 之间的关系，而不是依赖 RNN 的逐步递归。

在 decoder-only 模型里，attention 必须带 causal mask。也就是说，第 (t) 个 token 只能看见 (1 \ldots t) 的内容，不能偷看未来答案。带 mask 的 scaled dot-product attention 可以写成：

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}+M\right)V$$

其中 mask (M) 通常满足：

$$M_{ij}=\{\begin{array}{ll}0,& j\le i\\ -\infty ,& j>i\end{array}$$

5. 预训练：让模型学会语言接龙

预训练目标是 causal language modeling，也就是 next-token prediction。Hugging Face 的 causal LM 文档明确说明：causal language model 会预测序列中的下一个 token，并且只能 attend 左侧上下文。Causal language modeling 文档

损失函数是标准交叉熵：

$$L_{\text{pretrain}}=-\frac{1}{T}\sum _{t=1}^T\log p_{\theta }(x_t\mid x_{<t})$$

困惑度 perplexity 则是：

$$\mathrm{PPL}=\exp (L_{\text{pretrain}})$$

直观地说，困惑度可以看成模型“平均有多纠结”。如果 $\mathrm{PPL}=100$，模型像是在 100 个差不多可能的下一个 token 里猜；如果 $\mathrm{PPL}=6.7$，它平均只像是在 6 到 7 个较可能 token 中选择。这个解释不严格等于分类选项数，但很适合作为直觉。

原 Notebook 的预训练代码如下：

import math
from pathlib import Path

import torch
from datasets import load_dataset
from transformers import (
    GPT2Config,
    GPT2LMHeadModel,
    GPT2TokenizerFast,
    DataCollatorForLanguageModeling,
    Trainer,
    TrainingArguments,
)

TOKENIZER_DIR = "artifacts/tokenizer"
OUT_DIR = "artifacts/pretrain"
BLOCK_SIZE = 256

tokenizer = GPT2TokenizerFast.from_pretrained(TOKENIZER_DIR)
tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token

config = GPT2Config(
    vocab_size=len(tokenizer),
    n_positions=BLOCK_SIZE,
    n_ctx=BLOCK_SIZE,
    n_embd=384,
    n_layer=6,
    n_head=6,
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
)

model = GPT2LMHeadModel(config)
print("parameters:", sum(p.numel() for p in model.parameters()) / 1e6, "M")

raw = load_dataset(
    "text",
    data_files={
        "train": "data/tinystories_train.txt",
        "validation": "data/tinystories_valid.txt",
    },
)

def tokenize_fn(examples):
    return tokenizer(examples["text"])

tokenized = raw.map(
    tokenize_fn,
    batched=True,
    remove_columns=["text"],
    num_proc=4,
)

def group_texts(examples):
    concatenated = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated["input_ids"])
    total_length = (total_length // BLOCK_SIZE) * BLOCK_SIZE

    result = {
        k: [t[i : i + BLOCK_SIZE] for i in range(0, total_length, BLOCK_SIZE)]
        for k, t in concatenated.items()
    }
    result["labels"] = result["input_ids"].copy()
    return result

lm_ds = tokenized.map(group_texts, batched=True, num_proc=4)

collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
)

args = TrainingArguments(
    output_dir=OUT_DIR,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    gradient_accumulation_steps=4,
    max_steps=10_000,
    eval_strategy="steps",
    eval_steps=500,
    save_steps=1_000,
    save_total_limit=2,
    logging_steps=20,
    learning_rate=5e-4,
    warmup_steps=300,
    weight_decay=0.1,
    lr_scheduler_type="cosine",
    fp16=torch.cuda.is_available(),
    tf32=False,
    gradient_checkpointing=True,
    report_to="none",
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=lm_ds["train"],
    eval_dataset=lm_ds["validation"],
    data_collator=collator,
)

trainer.train()
metrics = trainer.evaluate()

eval_loss = metrics["eval_loss"]
print("eval_loss:", eval_loss)
print("perplexity:", math.exp(eval_loss) if eval_loss < 20 else float("inf"))

trainer.save_model(OUT_DIR)
tokenizer.save_pretrained(OUT_DIR)

parameters: 16.889856 M
`loss_type=None` was set in the config but it is unrecognized. Using the default loss: `ForCausalLMLoss`.
[10000/10000 44:16, Epoch 1/2]
[456/456 00:09]
eval_loss: 1.904482364654541
perplexity: 6.7159303248114055
('artifacts/pretrain/tokenizer_config.json', 'artifacts/pretrain/tokenizer.json')

这里需要解释一下 group_texts。它把很多短文本 token 拼接起来，再切成固定长度 BLOCK_SIZE=256 的训练块。语言模型预训练通常不是“一篇故事一个样本”，而是“连续 token 流切块”。如果不切块，短故事长度差异会让 batch 很难高效组织；切成固定块后，GPU 更容易吃饱。

DataCollatorForLanguageModeling(mlm=False) 也很关键。Transformers 文档说明，当 mlm=False 时，labels 与 inputs 相同，padding token 会在 loss 中被忽略。DataCollator 文档 对 causal LM 来说，我们不需要像 BERT 那样随机 mask token；模型内部会做一位偏移，让当前位置的输出预测下一个 token。

此处预训练结果是：

eval_loss: 1.904482364654541
perplexity: 6.7159303248114055

对一个 16.9M 参数、只训练 10,000 step 的小模型来说，这个结果已经不错。更重要的是，loss 在下降，生成文本开始出现英文短句、角色、动作和简单故事结构，这说明预训练阶段成功了。

6. 生成测试：看看预训练模型会不会“说话”

预训练后，我们用同一个生成脚本检查模型输出。

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def generate_text(
    model_path,
    prompt,
    max_new_tokens=120,
    temperature=0.8,
    top_p=0.9,
):
    print(f"正在加载模型和分词器: {model_path} ...")
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(model_path)
    model.eval()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model.to(device)
    print(f"模型已加载至设备: {device}")

    inputs = tokenizer(prompt, return_tensors="pt").to(device)

    with torch.no_grad():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            top_p=top_p,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    return tokenizer.decode(out[0], skip_special_tokens=True)

print(generate_text(
    "artifacts/pretrain",
    "Once upon a time, there was a little rabbit",
))

输出示例：

正在加载模型和分词器: artifacts/pretrain ...
模型已加载至设备: cuda
正在生成文本...
========================================
Once upon a time, there was a little rabbit named Benny. Benny loved to hop around and play with his friends. One day, he found a big carrot in the forest. He was so happy and started hopping around.Suddenly, Benny's friend, a small squirrel named Nutty, came over to play. "What are you doing?" asked Sammy. "I'm hopping from my burrow," said Nutty. "That's very interesting! Do you want to play with me?"Nutty replied, "Yes please!"As they played, the sun began to set and it was time for Sammy to go home. "We have to leave now, Nutty
========================================

temperature 控制分布的尖锐程度。如果 logits 是 $z_i$，温度采样会使用：

$$p_i=\frac{\exp (z_i/\tau )}{{\sum }_j\exp (z_j/\tau )}$$

当 $\tau <1$ 时，模型更保守，更爱选高概率 token；当 $\tau >1$ 时，输出更发散，也更容易跑偏。top_p=0.9 是 nucleus sampling，它会从累计概率达到 0.9 的最小 token 集合中采样。repetition_penalty=1.1 则是轻轻提醒模型：不要把同一句话绕成一团。

原 Notebook 的预训练生成已经出现了“角色、动作、对话”的基本结构。不要期待这个模型回答数学题、写论文或当助理，因为它只在儿童故事上做了 next-token prediction。预训练模型会续写，但它还不一定懂“请按我的要求完成任务”。

这就引出了下一步：SFT。

7. 构造 SFT、Reward、PPO 数据：自制一套教学版偏好数据

真实 RLHF 需要人工示范数据和人工偏好数据。 InstructGPT 的训练流程就是先收集人类示范做 SFT，再收集人类对模型输出的排序训练 Reward Model，最后用 PPO 优化策略模型。InstructGPT 论文 和 OpenAI 介绍文章 都明确描述了这条三阶段路线。

但是我们不可能在教程里临时拉一个标注团队。我们使用 TinyStories 自动构造教学数据：

import json
import random
import re
from pathlib import Path
from datasets import load_dataset
from tqdm import tqdm

random.seed(42)

OUT = Path("data")
OUT.mkdir(exist_ok=True)

N = 20_000
ds = load_dataset("roneneldan/TinyStories", split=f"train[:{N}]")

def split_first_sentence(text):
    text = " ".join(text.strip().split())
    parts = re.split(r"(?<=[.!?])\s+", text, maxsplit=1)
    if len(parts) < 2:
        return None, None
    return parts[0], parts[1]

def make_bad_completion(good):
    sents = re.split(r"(?<=[.!?])\s+", good.strip())
    mode = random.choice(["repeat", "truncate", "shuffle", "unrelated"])

    if mode == "repeat":
        first = sents[0] if sents else good[:80]
        return (first + " ") * 5

    if mode == "truncate":
        return sents[0] if sents else good[:60]

    if mode == "shuffle":
        words = good.split()
        random.shuffle(words)
        return " ".join(words[: min(len(words), 80)])

    return "The story suddenly talks about computers, rockets, and numbers. It does not finish the original idea."

sft_rows = []
reward_rows = []
ppo_rows = []

for row in tqdm(ds):
    text = row["text"]
    first, rest = split_first_sentence(text)
    if not first or not rest:
        continue

    prompt = (
        "Write a short, kind, coherent children's story. "
        f"Begin with: {first}\nStory:"
    )

    chosen = rest.strip()
    rejected = make_bad_completion(chosen)

    if len(chosen.split()) < 20:
        continue

    sft_rows.append({
        "prompt": prompt,
        "completion": " " + chosen,
    })

    reward_rows.append({
        "prompt": prompt,
        "chosen": " " + chosen,
        "rejected": " " + rejected,
    })

    ppo_rows.append({
        "prompt": prompt,
    })

def write_jsonl(rows, path):
    with open(path, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")

write_jsonl(sft_rows, OUT / "sft_train.jsonl")
write_jsonl(reward_rows, OUT / "reward_train.jsonl")
write_jsonl(ppo_rows, OUT / "ppo_prompts.jsonl")

print("sft:", len(sft_rows))
print("reward:", len(reward_rows))
print("ppo:", len(ppo_rows))

100%|██████████| 20000/20000 [00:01<00:00, 12008.55it/s]
sft: 20000
reward: 20000
ppo: 20000

这一步的技巧是把每篇故事拆成两段：第一句作为开头，其余部分作为好答案。于是 prompt 像这样：

Write a short, kind, coherent children's story. Begin with: Lily found a small red ball.
Story:

而 completion 就是原故事后续。

Reward 数据则是三列：

{
  "prompt": "...",
  "chosen": " 好的故事后续",
  "rejected": " 故意构造的坏故事后续"
}

坏答案的构造方式有四类：

模式	做法	模型应该学到什么
repeat	把第一句重复多次	重复不是好故事
truncate	只保留第一句	太短、没发展不好
shuffle	打乱词序	语法和连贯性很重要
unrelated	突然跳到无关主题	回答要贴合 prompt

请注意：这不是真正的人类反馈数据，而是“规则构造的偏好数据”。 它能帮我们复现 RLHF 工程流程，但它学到的是这些规则，不是完整的人类审美。我们测试 reward accuracy 达到 0.998，听起来很高，但这不代表模型突然理解文学批评了，只代表“好故事”和“故意捣乱的坏故事”太容易分。

8. SFT：把续写模型变成听指令模型

预训练模型学会的是：

$$\text{给定前文，预测后文}$$

SFT 学的是：

$$\text{给定任务说明，生成符合任务的回答}$$

TRL 的 SFTTrainer 支持 language modeling 和 prompt-completion 数据格式，也支持 standard 和 conversational 格式。TRL SFT 文档 还解释了 SFT 的 loss：本质上仍然是 token-level negative log-likelihood，只是我们可以选择只在 completion 部分计算 loss。

对于 prompt-completion 数据，SFT loss 可以写成：

$$L_{\text{SFT}}=-\sum _{t\in \mathcal{C}}\log {\pi }_{\theta }(y_t\mid x,y_{<t})$$

这里 (x) 是 prompt，$\mathcal{C}$ 是 completion token 的位置集合。completion_only_loss=True 的含义就是：prompt 只是条件，不当作模型要学习复述的目标；模型真正要学的是 completion。

import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTConfig, SFTTrainer

BASE = "artifacts/pretrain"
OUT = "artifacts/sft"

tokenizer = AutoTokenizer.from_pretrained(BASE)
tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(BASE)

ds = load_dataset("json", data_files="data/sft_train.jsonl", split="train")
ds = ds.train_test_split(test_size=0.02, seed=42)

args = SFTConfig(
    output_dir=OUT,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=8,
    max_steps=2_000,
    learning_rate=1e-4,
    warmup_steps=100,
    logging_steps=20,
    eval_strategy="steps",
    eval_steps=200,
    save_steps=500,
    save_total_limit=2,
    max_length=256,
    completion_only_loss=True,
    fp16=torch.cuda.is_available(),
    tf32=False,
    gradient_checkpointing=True,
    report_to="none",
)

trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=ds["train"],
    eval_dataset=ds["test"],
    processing_class=tokenizer,
)

trainer.train()
trainer.save_model(OUT)
tokenizer.save_pretrained(OUT)

Loading weights: 100%|██████████| 76/76 [..]
Adding EOS to train dataset: 100%|██████████| 19600/19600 [..]
Tokenizing train dataset: 100%|██████████| 19600/19600 [..]
Adding EOS to eval dataset: 100%|██████████| 400/400 [..]
Tokenizing eval dataset: 100%|██████████| 400/400 [..]
[2000/2000 13:30, Epoch 3/4]
('artifacts/sft/tokenizer_config.json', 'artifacts/sft/tokenizer.json')

SFT 后测试：

print(generate_text(
    "artifacts/sft",
    "Write a short, kind, coherent children's story. Begin with: Lily found a small red ball.\nStory:",
))

正在加载模型和分词器: artifacts/sft ...
模型已加载至设备: cuda
正在生成文本...
========================================
Write a short, kind, coherent children's story. Begin with: Lily found a small red ball.
Story: It was big and round. She wanted to play with it. She asked her mom if she could have the ball. Her mom said yes, but she had to be careful. Lily took the ball and put it in her mouth. She smiled and hugged it. She thought the ball was pretty. She played with it every day. But then she saw something shiny on the ground. She picked it up and threw it at the ball. It went over a hill and landed on the ground. Lily did not see that the ball was too high. She slipped and fell down. She scraped her knee
========================================

SFT 输出已经明显会沿着 prompt 续写短故事。 它不再只是接着任意前文乱飘，而是更容易围绕“short, kind, coherent children’s story”这个指令组织文本。

9. Reward Model：训练一个偏好打分器

Reward Model 不负责生成文本，它负责打分。给定同一个 prompt 和两个 response，模型要让 chosen 的分数高于 rejected。

TRL 的 RewardTrainer 文档使用 Bradley-Terry 模型解释偏好训练。若 $y^{+}$ 是 preferred response，$y^{-}$ 是 rejected response，奖励函数为 $r_{\theta }(x,y)$，则：

$$P(y^{+}\succ y^{-}\mid x)=\sigma \left(r_{\theta }(x,y^{+})-r_{\theta }(x,y^{-})\right)$$

对应 loss 是：

$$L_{\text{RM}}=-{\mathbb{E}}_{(x,y^{+},y^{-})\sim D}\left[\log \sigma \left(r_{\theta }(x,y^{+})-r_{\theta }(x,y^{-})\right)\right]$$

如果好答案分数比坏答案高很多，sigmoid 接近 1，loss 接近 0；如果坏答案反而更高，loss 就会变大，模型会被迫反省。

训练代码：

import torch
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from trl import RewardConfig, RewardTrainer

SFT = "artifacts/sft"
OUT = "artifacts/reward"

tokenizer = AutoTokenizer.from_pretrained(SFT)
tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token

model = AutoModelForSequenceClassification.from_pretrained(
    SFT,
    num_labels=1,
    pad_token_id=tokenizer.pad_token_id,
)

ds = load_dataset("json", data_files="data/reward_train.jsonl", split="train")
ds = ds.train_test_split(test_size=0.05, seed=42)

args = RewardConfig(
    output_dir=OUT,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    gradient_accumulation_steps=8,
    max_steps=2_000,
    learning_rate=1e-4,
    warmup_steps=100,
    logging_steps=20,
    eval_strategy="steps",
    eval_steps=200,
    save_steps=500,
    save_total_limit=2,
    max_length=256,
    fp16=torch.cuda.is_available(),
    tf32=False,
    gradient_checkpointing=True,
    report_to="none",
)

trainer = RewardTrainer(
    model=model,
    args=args,
    train_dataset=ds["train"],
    eval_dataset=ds["test"],
    processing_class=tokenizer,
)

trainer.train()
trainer.save_model(OUT)
tokenizer.save_pretrained(OUT)

Loading weights: 100%|██████████| 76/76 [..]
GPT2ForSequenceClassification LOAD REPORT from: artifacts/sft
Key          | Status  |
-------------+---------+-
score.weight | MISSING |

Notes:
- MISSING : those params were newly initialized because missing from the checkpoint. Consider training on your downstream task.

Adding EOS to train dataset: 100%|██████████| 19000/19000 [..]
Tokenizing train dataset: 100%|██████████| 19000/19000 [..]
Filtering train >256 tokens: 100%|██████████| 19000/19000 [..]
Adding EOS to eval dataset: 100%|██████████| 1000/1000 [..]
Tokenizing eval dataset: 100%|██████████| 1000/1000 [..]
Filtering eval >256 tokens: 100%|██████████| 1000/1000 [..]
[2000/2000 13:07, Epoch 4/5]
('artifacts/reward/tokenizer_config.json', 'artifacts/reward/tokenizer.json')

我们可以用一个简单脚本验证 Reward Model 是否学到了偏好：

import json
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL = "artifacts/reward"
DATA = "data/reward_train.jsonl"

tokenizer = AutoTokenizer.from_pretrained(MODEL)
tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token

model = AutoModelForSequenceClassification.from_pretrained(MODEL)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()

correct = 0
total = 0

def score(text):
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=256,
        padding=True,
    ).to(device)

    with torch.no_grad():
        return model(**inputs).logits.squeeze().item()

with open(DATA, "r", encoding="utf-8") as f:
    for i, line in enumerate(f):
        if i >= 500:
            break

        row = json.loads(line)
        chosen_text = row["prompt"] + row["chosen"]
        rejected_text = row["prompt"] + row["rejected"]

        sc = score(chosen_text)
        sr = score(rejected_text)

        correct += int(sc > sr)
        total += 1

print("reward accuracy:", correct / total)

reward accuracy: 0.998

这个指标的含义是，在前 500 个样本里，Reward Model 有 99.8% 的时候给 chosen 更高分。再次强调，这里的数据是规则构造的，所以这个高分不宜解读为“模型学会了人类偏好”。它更像一场开卷小测：题目设计得清楚，答对率自然高。

10. PPO/RLHF：让模型朝偏好方向小步前进

PPO，全称 Proximal Policy Optimization，是 Schulman 等人在 2017 年提出的强化学习算法。PPO 论文 的核心思想是让策略更新不要一次迈太大步，从而在样本效率、稳定性和实现复杂度之间取得平衡。

在语言模型 RLHF 里，通常有四个角色：

角色	是否训练	作用
policy model	是	当前要优化的模型，初始来自 SFT
reference model	否	冻结的 SFT 模型，用来计算 KL，防止 policy 跑太远
reward model	否	冻结的偏好模型，给 prompt + response 打分
value model	是	估计价值函数，帮助计算 advantage

PPO 的 clipped objective 是：

$$r_t(\theta )=\frac{{\pi }_{\theta }(a_t\mid s_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t\mid s_t)}$$

$$L^{\text{CLIP}}(\theta )={\mathbb{E}}_t\left[\min \left(r_t(\theta ){\hat{A}}_t,\mathrm{clip}(r_t(\theta ),1-ϵ,1+ϵ){\hat{A}}_t\right)\right]$$

其中 ${\hat{A}}_t$ 是 advantage，表示“这个动作比预期好多少”。clip 的意义是给更新幅度装一个限速器：如果新策略相对旧策略变化太夸张，就把收益截断，避免模型因为奖励模型的一点偏好而猛拐弯。

在 RLHF 里，奖励通常不是只有 Reward Model 分数，还会加入 KL 惩罚：

$$R(x,y)=r_{\text{RM}}(x,y)-\beta D_{\text{KL}}\left({\pi }_{\theta }(\cdot \mid x)\Vert {\pi }_{\text{ref}}(\cdot \mid x)\right)$$

这个公式非常重要。Reward Model 像油门，KL 惩罚像方向盘和刹车。没有奖励，模型不知道往哪儿优化；没有 KL，模型可能为了讨好奖励模型而产生怪异输出，也就是 reward hacking。

TRL 的 PPOTrainer 文档也列出了 PPO 训练中常见指标，比如 objective/kl、objective/entropy、objective/non_score_reward、objective/rlhf_reward 等。TRL PPO 文档 其中 KL 正是衡量当前 policy 与 reference policy 距离的重要信号。

PPO 代码如下：

import torch
from datasets import load_dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    AutoModelForSequenceClassification,
)
from trl.experimental.ppo import PPOConfig, PPOTrainer

SFT = "artifacts/sft"
REWARD = "artifacts/reward"
OUT = "artifacts/ppo"

tokenizer = AutoTokenizer.from_pretrained(SFT, padding_side="left")
tokenizer.pad_token = tokenizer.pad_token or tokenizer.eos_token

dtype = torch.float16 if torch.cuda.is_available() else torch.float32

policy = AutoModelForCausalLM.from_pretrained(SFT, torch_dtype=dtype)
ref_policy = AutoModelForCausalLM.from_pretrained(SFT, torch_dtype=dtype)

reward_model = AutoModelForSequenceClassification.from_pretrained(
    REWARD,
    num_labels=1,
    torch_dtype=dtype,
)

value_model = AutoModelForSequenceClassification.from_pretrained(
    REWARD,
    num_labels=1,
    torch_dtype=dtype,
)

dataset = load_dataset("json", data_files="data/ppo_prompts.jsonl", split="train")

def tokenize(element):
    ids = tokenizer(element["prompt"], padding=False)["input_ids"]
    return {"input_ids": ids, "lengths": len(ids)}

dataset = dataset.map(
    tokenize,
    remove_columns=dataset.column_names,
)

dataset = dataset.filter(lambda x: 8 <= x["lengths"] <= 160)

args = PPOConfig(
    output_dir=OUT,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=16,
    total_episodes=2_000,
    learning_rate=3e-6,
    num_ppo_epochs=2,
    num_mini_batches=1,
    response_length=80,
    stop_token="eos",
    missing_eos_penalty=1.0,
    local_rollout_forward_batch_size=2,
    fp16=torch.cuda.is_available(),
    bf16=False,
    gradient_checkpointing=True,
    logging_steps=10,
    save_steps=500,
    save_total_limit=2,
    report_to="none",
)

trainer = PPOTrainer(
    args=args,
    processing_class=tokenizer,
    model=policy,
    ref_model=ref_policy,
    reward_model=reward_model,
    value_model=value_model,
    train_dataset=dataset,
)

trainer.train()
trainer.save_model(OUT)
tokenizer.save_pretrained(OUT)

TRLExperimentalWarning: You are importing from 'trl.experimental'. APIs here are unstable and may change or be removed without notice.
`torch_dtype` is deprecated! Use `dtype` instead!
Loading weights: 100%|██████████| 76/76 [..]
Loading weights: 100%|██████████| 76/76 [..]
Loading weights: 100%|██████████| 77/77 [..]
Loading weights: 100%|██████████| 77/77 [..]
Generating train split: 20000 examples [..]
Map: 100%|██████████| 20000/20000 [..]
Filter: 100%|██████████| 20000/20000 [..]
===training policy===
[63/63 11:22, Epoch 0/0.1]
('artifacts/ppo/tokenizer_config.json', 'artifacts/ppo/tokenizer.json')

这里有几个参数需要盯住：

参数	含义	为什么这样设
learning_rate=3e-6	PPO 阶段学习率	RL 更新更不稳定，所以比 SFT 小很多
total_episodes=2_000	采样并训练的 episode 数	教学版快速跑通，不追求极致
response_length=80	每次生成 response 长度	控制训练成本
missing_eos_penalty=1.0	没生成 EOS 时惩罚	鼓励模型学会结束
padding_side="left"	左 padding	自回归生成时更常用，便于对齐右侧有效 token

PPO 后同样可以测试：

print(generate_text(
    "artifacts/ppo",
    "Write a short, kind, coherent children's story. Begin with: Lily found a small red ball.\nStory:",
))

正在加载模型和分词器: artifacts/ppo ...
模型已加载至设备: cuda
正在生成文本...
========================================
Write a short, kind, coherent children's story. Begin with: Lily found a small red ball.
Story: She was happy and wanted to play with it. She put the ball on the floor and started to throw it at it. It went up and down, faster and faster. Lily laughed and said, "I caught a ball! It is so fast!" Suddenly, she saw her friend Max. Max had a toy car too. He took it and gave it back to Lily. Lily said, "Thank you, Max! You are my best friend." Max smiled and said, "You're welcome, Lily. I'm glad you like it." From that day on, Lily and Max played together
========================================

预训练教模型语言规律，SFT 教模型按指令回答，Reward Model 教模型比较好坏，PPO 教模型在偏好信号下小步更新。把这四步连起来，你就不再只是听说 RLHF，而是真的亲手摸到了它的形状。

一个 16.9M 参数的小模型不会替你改变世界，但它能帮你改变对大模型训练的感觉：从“神秘黑箱”变成“一条可以拆开看的流水线”。这已经很值了。

参考资料

1. TinyStories: How Small Can Language Models Be and Still Speak Coherent English?
2. roneneldan/TinyStories 数据集页
3. Attention Is All You Need
4. Hugging Face Tokenizers 文档
5. Hugging Face Tokenizers Quicktour: BPE training
6. Hugging Face Datasets: load_dataset
7. Hugging Face Transformers: Causal language modeling
8. Hugging Face Transformers: DataCollatorForLanguageModeling
9. Hugging Face Transformers: GPT2LMHeadModel
10. TRL SFTTrainer 文档
11. TRL RewardTrainer 文档
12. TRL PPOTrainer 文档
13. Proximal Policy Optimization Algorithms
14. Training language models to follow instructions with human feedback
15. OpenAI: Aligning language models to follow instructions