用 AllenAI Open Instruct 构建 RLVR/GRPO 后训练流水线

在本教程中，我们使用 AllenAI 的 Open Instruct 框架为紧凑型指令微调语言模型构建端到端的后训练流程。我们将依次经历三个主要训练阶段：监督微调（Supervised Fine-Tuning）、直接偏好优化（Direct Preference Optimization）以及使用 GRPO 的基于可验证奖励的强化学习（Reinforcement Learning with Verifiable Rewards），同时将原始的多 GPU Tulu 3 技术栈适配到 16 GB 运行时环境中。我们克隆 Open Instruct 仓库，有选择性地加载其原生损失函数和工具函数，配置 LoRA 适配器，为每个训练阶段准备 GSM8K 数据，并使用确定性验证器（deterministic verifier）来评估生成的数学答案。在整个工作流程中，我们保留 Open Instruct 的核心优化逻辑，同时用轻量级的 Hugging Face 和 PyTorch 实现替代分布式组件（如 vLLM、Ray actors、DeepSpeed 和异步 rollout 队列），使其适合在 Colab 环境中运行。

import os, sys, subprocess, textwrap, json, math, random, re, ast, types, dataclasses, gc, contextlib
REPO_URL = "https://github.com/allenai/open-instruct.git"
REPO_DIR = "/content/open-instruct" if os.path.isdir("/content") else "./open-instruct"
PIP_PKGS = [
   "peft", "accelerate",
   "ray", "wandb", "beaker-py",
   "langdetect==1.0.9", "immutabledict==1.2.0", "nltk",
   "absl-py", "sympy", "antlr4-python3-runtime==4.11",
   "tiktoken",
]
def sh(*args):
   print("$", " ".join(args))
   subprocess.run(args, check=False)
def setup():
   sh(sys.executable, "-m", "pip", "install", "-q", *PIP_PKGS)
   if not os.path.isdir(REPO_DIR):
       sh("git", "clone", "--depth", "1", REPO_URL, REPO_DIR)
   if REPO_DIR not in sys.path:
       sys.path.insert(0, REPO_DIR)
   os.environ.setdefault("WANDB_MODE", "disabled")
   os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
   os.environ.setdefault("RAY_DISABLE_IMPORT_WARNING", "1")
setup()
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from datasets import load_dataset, Dataset
from transformers import AutoModelForCausalLM, DataCollatorForSeq2Seq, get_cosine_schedule_with_warmup
from peft import LoraConfig, get_peft_model
DEV = "cuda" if torch.cuda.is_available() else "cpu"
try:
   _bf16 = DEV == "cuda" and torch.cuda.is_bf16_supported(including_emulation=False)
except TypeError:
   _bf16 = DEV == "cuda" and torch.cuda.get_device_properties(0).major >= 8
AMP_DTYPE = torch.bfloat16 if _bf16 else torch.float16
USE_SCALER = AMP_DTYPE is torch.float16
print(f"device={DEV}  autocast dtype={AMP_DTYPE}  gpu={torch.cuda.get_device_name(0) if DEV=='cuda' else '-'}")
def oi_load(relpath, names, ns=None):
   src = open(os.path.join(REPO_DIR, relpath)).read()
   tree = ast.parse(src)
   found = {n.name: n for n in tree.body
            if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and n.name in names}
   missing = set(names) - set(found)
   if missing:
       raise KeyError(f"{relpath}: could not find {missing} (upstream may have renamed them)")
   ns = {} if ns is None else dict(ns)
   ns.update({"torch": torch, "F": F, "np": np, "enum": __import__("enum"),
              "dataclasses": dataclasses, "math": math, "os": os})
   future = ast.parse("from __future__ import annotations").body
   mod = ast.Module(body=future + [found[n] for n in names], type_ignores=[])
   exec(compile(ast.fix_missing_locations(mod), f"<open_instruct:{relpath}>", "exec"), ns)
   return {n: ns[n] for n in names}
_dpo  = oi_load("open_instruct/dpo_utils.py", ["dpo_loss", "_get_batch_logps"])
_pf   = oi_load("open_instruct/padding_free_collator.py", ["calculate_per_token_logps"])
_rl   = oi_load("open_instruct/rl_utils.py", ["masked_mean"])
_mu   = oi_load("open_instruct/model_utils.py", ["estimate_kl"])
_grpo = oi_load("open_instruct/grpo_utils.py", ["GRPOLossType", "compute_grpo_loss"],
               ns={"model_utils": types.SimpleNamespace(**_mu)})
dpo_loss           = _dpo["dpo_loss"]
get_batch_logps    = _dpo["_get_batch_logps"]
per_token_logps_fn = _pf["calculate_per_token_logps"]
masked_mean        = _rl["masked_mean"]
compute_grpo_loss  = _grpo["compute_grpo_loss"]
GRPOLossType       = _grpo["GRPOLossType"]
print("lifted from repo:", [f.__name__ for f in (dpo_loss, get_batch_logps, per_token_logps_fn,
                                                masked_mean, compute_grpo_loss)])
from open_instruct.dataset_transformation import (
   CHAT_TEMPLATES, TokenizerConfig,
   sft_tulu_tokenize_and_truncate_v1, sft_tulu_filter_v1,
   preference_tulu_tokenize_and_truncate_v1_2,
   rlvr_tokenize_v1, visualize_token_role,
)
from open_instruct.ground_truth_utils import GSM8KVerifier, MathVerifier, IFEvalVerifierOld

我们安装所需的轻量级依赖项，克隆 Open Instruct 仓库，并配置 Colab 环境以确保稳定执行。我们检测可用的 GPU 精度模式，根据硬件能力选择 FP16 或 BF16 自动转换。同时，我们直接从仓库中提取原始的 DPO、GRPO、掩码和对数概率函数，而无需引入其完整的分布式训练栈。

@dataclasses.dataclass
class CFG:
   model: str = "Qwen/Qwen2.5-0.5B-Instruct"
   max_seq_len: int = 640
   seed: int = 42
   n_sft: int = 192
   sft_steps: int = 40
   sft_micro_bs: int = 2
   sft_accum: int = 4
   sft_lr: float = 1e-4
   n_dpo: int = 96
   dpo_steps: int = 24
   dpo_micro_bs: int = 1
   dpo_accum: int = 4
   dpo_lr: float = 5e-5
   dpo_beta: float = 0.1
   dpo_norm: bool = True
   grpo_iters: int = 6
   prompts_per_iter: int = 4
   samples_per_prompt: int = 4
   grpo_micro_bs: int = 1
   grpo_inner_epochs: int = 2
   grpo_lr: float = 2e-5
   grpo_temperature: float = 1.0
   grpo_max_new: int = 200
   grpo_kl_beta: float = 0.02
   clip_lower: float = 0.2
   clip_higher: float = 0.272
   kl_estimator: int = 2
   adv_norm: str = "centered"
   n_eval: int = 24
cfg = CFG()
random.seed(cfg.seed); np.random.seed(cfg.seed); torch.manual_seed(cfg.seed)
tc = TokenizerConfig(tokenizer_name_or_path=cfg.model, chat_template_name=None, use_fast=True)
tok = tc.tokenizer
print(f"\navailable CHAT_TEMPLATES: {list(CHAT_TEMPLATES)[:12]} ... ({len(CHAT_TEMPLATES)} total)")
print(f"pad={tok.pad_token!r}({tok.pad_token_id})  eos={tok.eos_token!r}({tok.eos_token_id})")
_demo = {"messages": [
   {"role": "user", "content": "What is 12 * 3?"},
   {"role": "assistant", "content": "12 * 3 = 36. The answer is 36."},
   {"role": "user", "content": "And minus 6?"},
   {"role": "assistant", "content": "36 - 6 = 30. The answer is 30."},
]}
_enc = sft_tulu_tokenize_and_truncate_v1(dict(_demo), tok, cfg.max_seq_len)
print("\n[SFT label masking — colour 0 = masked out of the loss, colour 1 = trained on]")
visualize_token_role(_enc["input_ids"].tolist(), (_enc["labels"] != -100).long().tolist(), tok)
print(f"trainable tokens: {(_enc['labels'] != -100).sum().item()}/{_enc['labels'].numel()}")

我们定义了一个集中式的配置类，用于控制模型、数据集大小、学习率、批量设置以及每个训练阶段的优化参数。我们在保留模型聊天模板的同时初始化 Open Instruct 分词器，并确保填充标记和序列结束标记正确分离。随后，我们对一个示例对话进行分词，并可视化哪些助手标记对监督训练损失有贡献。

我们加载 GSM8K 并将其问题和解答转换为一致的对话格式，用于 SFT、DPO 和 RLVR 训练。我们创建监督示例、带故意错误最终答案的偏好对，以及带有结构化 ground-truth 标签的 verifier 就绪提示。我们还初始化 Open Instruct 的 GSM8K、数学和指令遵循验证器，并用它们对生成的响应进行确定性评分。

model = AutoModelForCausalLM.from_pretrained(cfg.model, dtype=torch.float32).to(DEV)
model.config.use_cache = False
if len(tok) > model.get_input_embeddings().weight.shape[0]:
   model.resize_token_embeddings(len(tok))
def _patch_peft_torchao():
   import importlib
   for mod in ("peft.import_utils", "peft.tuners.lora.torchao",
               "peft.tuners.lora.model", "peft.tuners.lora.layer"):
       try:
           m = importlib.import_module(mod)
       except Exception:
           continue
       if hasattr(m, "is_torchao_available"):
           m.is_torchao_available = lambda: False
_patch_peft_torchao()
model = get_peft_model(model, LoraConfig(
   r=32, lora_alpha=64, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
model.print_trainable_parameters()
TRAINABLE = [p for p in model.parameters() if p.requires_grad]
@contextlib.contextmanager
def with_cache():
   old = model.config.use_cache
   model.config.use_cache = True
   try:
       yield
   finally:
       model.config.use_cache = old
def amp():
   return torch.autocast(device_type="cuda", dtype=AMP_DTYPE) if DEV == "cuda" \
       else torch.autocast(device_type="cpu", enabled=False)
def new_opt(lr, steps):
   opt = torch.optim.AdamW(TRAINABLE, lr=lr, weight_decay=0.0, betas=(0.9, 0.999))
   sched = get_cosine_schedule_with_warmup(opt, int(0.05 * steps) + 1, steps)
   scaler = torch.amp.GradScaler("cuda", enabled=USE_SCALER)
   return opt, sched, scaler
def step_opt(opt, sched, scaler):
   scaler.unscale_(opt)
   torch.nn.utils.clip_grad_norm_(TRAINABLE, 1.0)
   scaler.step(opt); scaler.update(); sched.step(); opt.zero_grad(set_to_none=True)
@torch.no_grad()
def evaluate(tag, rows, max_new=256):
   model.eval()
   tok.padding_side = "left"
   correct, bs = 0.0, 4
   for i in range(0, len(rows), bs):
       chunk = rows[i:i + bs]
       prompts = [tok.apply_chat_template(
           [{"role": "system", "content": SYS}, {"role": "user", "content": r["question"]}],
           add_generation_prompt=True, tokenize=False) for r in chunk]
       enc = tok(prompts, return_tensors="pt", padding=True, add_special_tokens=False).to(DEV)
       with amp(), with_cache():
           out = model.generate(**enc, max_new_tokens=max_new, do_sample=False,
                                pad_token_id=tok.pad_token_id)
       texts = tok.batch_decode(out[:, enc["input_ids"].shape[1]:], skip_special_tokens=True)
       correct += verify_batch(texts, [gsm_answer(r["answer"]) for r in chunk],
                               ["gsm8k"] * len(chunk)).sum()
   acc = correct / len(rows)
   print(f"  [eval:{tag}] verifier accuracy = {acc:.3f}  ({int(correct)}/{len(rows)})")
   model.train(); tok.padding_side = "right"
   return acc
print("\n" + "=" * 90); print("BASELINE"); print("=" * 90)
base_acc = evaluate("base", eval_rows)

我们加载 Qwen 指令模型，向其注意力和前馈投影层应用 LoRA 适配器，并将优化限制为可训练的适配器参数。我们配置混合精度执行、梯度缩放、梯度裁剪、学习率调度，以及用于生成的临时 KV-cache 激活。然后使用贪心解码和基于验证器的答案准确率在 GSM8K 上评估未训练基线。

print("\n" + "=" * 90); print("STAGE 1 — SFT"); print("=" * 90)
sft_collate = DataCollatorForSeq2Seq(tokenizer=tok, padding="longest", label_pad_token_id=-100)
sft_dl = DataLoader(sft_ds, batch_size=cfg.sft_micro_bs, shuffle=True, collate_fn=sft_collate, drop_last=True)
opt, sched, scaler = new_opt(cfg.sft_lr, cfg.sft_steps)
model.train(); it, step, run = iter(sft_dl), 0, 0.0
while step < cfg.sft_steps:
   for _ in range(cfg.sft_accum):
       try:
           batch = next(it)
       except StopIteration:
           it = iter(sft_dl); batch = next(it)
       batch = {k: v.to(DEV) for k, v in batch.items()}
       with amp():
           loss = model(**batch).loss / cfg.sft_accum
       scaler.scale(loss).backward()
       run += loss.item()
   step_opt(opt, sched, scaler); step += 1
   if step % 10 == 0 or step == 1:
       print(f"  sft step {step:>3}/{cfg.sft_steps}  loss {run:.4f}  lr {sched.get_last_lr()[0]:.2e}")
   run = 0.0
sft_acc = evaluate("after-sft", eval_rows)

我们构建一个带填充的 SFT DataLoader，并使用梯度累积在分词后的 GSM8K 对话上训练 LoRA 适配器。我们使用仅在未屏蔽的 assistant 响应 token 上计算的交叉熵损失来优化模型。我们在整个阶段中追踪训练损失和学习率，并在监督微调后评估更新后的模型。

print("\n" + "=" * 90); print("STAGE 2 — DPO (dpo_norm)"); print("=" * 90)
def pad_side(seqs, pad, maxlen):
   return torch.tensor([s + [pad] * (maxlen - len(s)) for s in seqs], dtype=torch.long)
def dpo_collate(features):
   out = {}
   for pfx in ("chosen", "rejected"):
       L = max(len(f[f"{pfx}_input_ids"]) for f in features)
       out[f"{pfx}_input_ids"] = pad_side([f[f"{pfx}_input_ids"] for f in features], tok.pad_token_id, L)
       out[f"{pfx}_labels"] = pad_side([f[f"{pfx}_labels"] for f in features], -100, L)
       out[f"{pfx}_attention_mask"] = pad_side([f[f"{pfx}_attention_mask"] for f in features], 0, L)
   return out
def seq_logps(input_ids, attn, labels):
   with amp():
       logits = model(input_ids=input_ids, attention_mask=attn).logits
   ptl = per_token_logps_fn(logits, labels)
   return get_batch_logps(ptl, labels, average_log_prob=cfg.dpo_norm)
dpo_dl = DataLoader(dpo_ds, batch_size=cfg.dpo_micro_bs, shuffle=True, collate_fn=dpo_collate, drop_last=True)
opt, sched, scaler = new_opt(cfg.dpo_lr, cfg.dpo_steps)
it, step = iter(dpo_dl), 0
while step < cfg.dpo_steps:
   agg = {"loss": 0.0, "acc": 0.0, "margin": 0.0}
   for _ in range(cfg.dpo_accum):
       try:
           b = next(it)
       except StopIteration:
           it = iter(dpo_dl); b = next(it)
       b = {k: v.to(DEV) for k, v in b.items()}
       with torch.no_grad(), model.disable_adapter():
           ref_c = seq_logps(b["chosen_input_ids"], b["chosen_attention_mask"], b["chosen_labels"])
           ref_r = seq_logps(b["rejected_input_ids"], b["rejected_attention_mask"], b["rejected_labels"])
       pol_c = seq_logps(b["chosen_input_ids"], b["chosen_attention_mask"], b["chosen_labels"])
       pol_r = seq_logps(b["rejected_input_ids"], b["rejected_attention_mask"], b["rejected_labels"])
       losses, r_c, r_r = dpo_loss(pol_c, pol_r, ref_c, ref_r, beta=cfg.dpo_beta, label_smoothing=0.0)
       loss = losses.mean() / cfg.dpo_accum
       scaler.scale(loss).backward()
       agg["loss"] += loss.item()
       agg["acc"] += (r_c > r_r).float().mean().item() / cfg.dpo_accum
       agg["margin"] += (r_c - r_r).mean().item() / cfg.dpo_accum
   step_opt(opt, sched, scaler); step += 1
   if step % 8 == 0 or step == 1:
       print(f"  dpo step {step:>3}/{cfg.dpo_steps}  loss {agg['loss']:.4f} "
             f"reward_acc {agg['acc']:.2f}  margin {agg['margin']:+.3f}")
dpo_acc = evaluate("after-dpo", eval_rows)

我们将 chosen 和 rejected 响应分别组成批次，使用 Open Instruct 原生工具计算其长度归一化的序列对数概率。我们将活跃的 LoRA 策略与冻结的基础参考策略进行比较，并使用仓库的 DPO loss 优化模型。我们监控偏好准确率、奖励边际和训练 loss，然后评估模型在 DPO 之后的验证器性能。

print("\n" + "=" * 90); print("STAGE 3 — RLVR / GRPO"); print("=" * 90)
grpo_cfg = types.SimpleNamespace(loss_fn=GRPOLossType.dapo, clip_lower=cfg.clip_lower,
                                clip_higher=cfg.clip_higher, kl_estimator=cfg.kl_estimator)
_gen_eos = getattr(getattr(model, "generation_config", None), "eos_token_id", None)
_terms = {tok.eos_token_id, tok.pad_token_id}
_terms |= set(_gen_eos) if isinstance(_gen_eos, (list, tuple)) else {_gen_eos}
TERMINATORS = torch.tensor(sorted(t for t in _terms if t is not None), device=DEV)
def token_logps(seq, attn, temperature, grad=True):
   pos = (attn.cumsum(-1) - 1).clamp(min=0)
   ctx = torch.enable_grad() if grad else torch.no_grad()
   with ctx, amp():
       logits = model(input_ids=seq, attention_mask=attn, position_ids=pos).logits
   return per_token_logps_fn(logits / temperature, seq)
def rollout(batch_rows):
   G = cfg.samples_per_prompt
   ids = [r["input_ids_prompt"] for r in batch_rows]
   P = max(len(x) for x in ids)
   pin = torch.tensor([[tok.pad_token_id] * (P - len(x)) + x for x in ids], device=DEV)
   pmask = torch.tensor([[0] * (P - len(x)) + [1] * len(x) for x in ids], device=DEV)
   model.eval()
   with torch.no_grad(), amp(), with_cache():
       seq = model.generate(input_ids=pin, attention_mask=pmask, do_sample=True,
                            temperature=cfg.grpo_temperature, top_p=1.0, top_k=0,
                            max_new_tokens=cfg.grpo_max_new, num_return_sequences=G,
                            pad_token_id=tok.pad_token_id)
   model.train()
   resp = seq[:, P:]
   is_term = torch.isin(resp, TERMINATORS)
   first = torch.where(is_term.any(1), is_term.float().argmax(1),
                       torch.full((resp.shape[0],), resp.shape[1] - 1, device=DEV))
   idx = torch.arange(resp.shape[1], device=DEV).unsqueeze(1)
   resp_mask = (idx <= first.unsqueeze(1)).long()
   full_mask = torch.cat([torch.zeros(seq.shape[0], P, dtype=torch.long, device=DEV), resp_mask], 1)
   attn = torch.cat([pmask.repeat_interleave(G, 0), resp_mask], 1)
   texts = tok.batch_decode(resp, skip_special_tokens=True)
   gts = [r["ground_truth"] for r in batch_rows for _ in range(G)]
   srcs = [r["dataset"] for r in batch_rows for _ in range(G)]
   scores = verify_batch(texts, gts, srcs)
   per_prompt = scores.reshape(-1, G)
   mean_g = np.repeat(per_prompt.mean(-1), G, 0)
   if cfg.adv_norm == "standard":
       adv = (scores - mean_g) / (np.repeat(per_prompt.std(-1), G, 0) + 1e-8)
   else:
       adv = scores - mean_g
   adv_t = torch.tensor(adv, device=DEV, dtype=torch.float32).unsqueeze(1).expand_as(full_mask.float())
   return seq, attn, full_mask, adv_t, scores, texts
opt, sched, scaler = new_opt(cfg.grpo_lr, cfg.grpo_iters * cfg.grpo_inner_epochs)
order = list(range(len(rlvr_ds))); random.shuffle(order)
for it_i in range(cfg.grpo_iters):
   rows = [rlvr_ds[order[(it_i * cfg.prompts_per_iter + j) % len(rlvr_ds)]]
           for j in range(cfg.prompts_per_iter)]
   seq, attn, mask, adv, scores, texts = rollout(rows)
   with torch.no_grad():
       old_lp = torch.cat([token_logps(seq[i:i + cfg.grpo_micro_bs], attn[i:i + cfg.grpo_micro_bs],
                                       cfg.grpo_temperature, grad=False)
                           for i in range(0, seq.shape[0], cfg.grpo_micro_bs)])
       with model.disable_adapter():
           ref_lp = torch.cat([token_logps(seq[i:i + cfg.grpo_micro_bs], attn[i:i + cfg.grpo_micro_bs],
                                           cfg.grpo_temperature, grad=False)
                               for i in range(0, seq.shape[0], cfg.grpo_micro_bs)])
   n_chunks = math.ceil(seq.shape[0] / cfg.grpo_micro_bs)
   for ep in range(cfg.grpo_inner_epochs):
       stats = {"pg": 0.0, "kl": 0.0, "clip": 0.0}
       for i in range(0, seq.shape[0], cfg.grpo_micro_bs):
           sl = slice(i, i + cfg.grpo_micro_bs)
           new_lp = token_logps(seq[sl], attn[sl], cfg.grpo_temperature, grad=True)
           new_lp_, old_lp_, ref_lp_ = new_lp[

我们为每个提示生成多个采样响应，使用确定性验证器对其进行评分，并根据奖励分布计算组相对优势。我们使用 Open Instruct 的 GRPO 和 DAPO 风格裁剪逻辑优化策略，同时应用响应掩码、重要性比率和针对参考模型的 KL 正则化。最后，我们比较基线、SFT、DPO 和 RLVR 各阶段的准确率，然后合并 LoRA 适配器并保存完成的检查点。

总之，我们实现了一个实用版的 Tulu 3 后训练栈微型版本，并观察了每个训练阶段如何改变模型在验证器评分数学推理任务上的表现。我们首先建立了基线，通过监督微调改进指令跟随，通过长度归一化 DPO 细化响应偏好，最后使用组相对优势和仓库的 GRPO loss 实现优化已验证任务奖励。我们还使用 LoRA 来维护一个可访问的参考策略，在强化学习期间应用响应掩码和 KL 正则化，比较所有训练阶段的准确率，并导出一个合并的检查点以供后续推理或评估。