In this tutorial, we build an end-to-end post-training pipeline for a compact instruction-tuned language model using AllenAI’s Open Instruct framework. We move through three major training stages: Supervised Fine-Tuning, Direct Preference Optimization, and Reinforcement Learning with Verifiable Rewards using GRPO, while adapting the original multi-GPU Tulu 3 stack to fit within a 16 GB runtime. We clone the Open Instruct repository, selectively load its native loss and utility functions, configure LoRA adapters, prepare GSM8K data for each training stage, and use deterministic verifiers to evaluate generated mathematical answers. Throughout the workflow, we preserve the core optimization logic of Open Instruct while replacing distributed components such as vLLM, Ray actors, DeepSpeed, and asynchronous rollout queues with lightweight Hugging Face and PyTorch implementations suitable for 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
We install the required lightweight dependencies, clone the Open Instruct repository, and configure the Colab environment for stable execution. We detect the available GPU precision mode and select either FP16 or BF16 autocasting based on the hardware capabilities. We also extract the original DPO, GRPO, masking, and log-probability functions directly from the repository without importing its full distributed training stack.
@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()}")
We define a centralized configuration class that controls the model, dataset sizes, learning rates, batch settings, and optimization parameters for every training stage. We initialize the Open Instruct tokenizer while preserving the model’s chat template and ensuring that padding and end-of-sequence tokens remain correctly separated. We then tokenize a sample conversation and visualize which assistant tokens contribute to the supervised training loss.
gsm = load_dataset("openai/gsm8k", "main")
SYS = "You are a careful math assistant. Reason step by step, then finish with 'The answer is N.'"
def gsm_answer(a):
return a.split("####")[-1].strip().replace(",", "")
def gsm_solution(a):
body = a.split("####")[0].strip()
body = re.sub(r"<<.*?>>", "", body)
return f"{body}nThe answer is {gsm_answer(a)}."
def as_messages(row):
return [{"role": "system", "content": SYS},
{"role": "user", "content": row["question"]},
{"role": "assistant", "content": gsm_solution(row["answer"])}]
train_rows = [gsm["train"][i] for i in range(cfg.n_sft + cfg.n_dpo)]
eval_rows = [gsm["test"][i] for i in range(cfg.n_eval)]
def to_lists(row):
for k in ("input_ids", "labels", "attention_mask"):
row[k] = row[k].tolist()
return row
sft_ds = Dataset.from_list([{"messages": as_messages(r)} for r in train_rows[: cfg.n_sft]])
sft_ds = sft_ds.map(lambda r: to_lists(sft_tulu_tokenize_and_truncate_v1(r, tok, cfg.max_seq_len)),
remove_columns=["messages"], desc="sft tokenize")
sft_ds = sft_ds.filter(sft_tulu_filter_v1, fn_kwargs={"tokenizer": tok}, desc="drop all-masked")
def make_pair(r):
gold = gsm_answer(r["answer"])
bad = (str(int(float(gold)) + random.choice([-10, -3, -1, 1, 2, 7]))
if gold.replace('.', '', 1).lstrip('-').isdigit() else gold + "0")
prompt = [{"role": "system", "content": SYS}, {"role": "user", "content": r["question"]}]
good_txt = gsm_solution(r["answer"])
bad_txt = good_txt.rsplit("The answer is", 1)[0] + f"The answer is {bad}."
return {"chosen": prompt + [{"role": "assistant", "content": good_txt}],
"rejected": prompt + [{"role": "assistant", "content": bad_txt}]}
dpo_ds = Dataset.from_list([make_pair(r) for r in train_rows[cfg.n_sft:]])
dpo_ds = dpo_ds.map(
lambda r: {k: (v.tolist() if torch.is_tensor(v) else v) for k, v in
preference_tulu_tokenize_and_truncate_v1_2(r, tok, cfg.max_seq_len).items()},
remove_columns=["chosen", "rejected"], desc="dpo tokenize")
rlvr_rows = [{"messages": as_messages(r)[:2], "ground_truth": gsm_answer(r["answer"]), "dataset": "gsm8k"}
for r in train_rows[: cfg.n_sft]]
rlvr_ds = Dataset.from_list(rlvr_rows).map(lambda r: rlvr_tokenize_v1(r, tok),
remove_columns=["messages"], desc="rlvr tokenize")
print(f"nsft={len(sft_ds)} dpo={len(dpo_ds)} rlvr={len(rlvr_ds)}")
VERIFIERS = {"gsm8k": GSM8KVerifier(), "math": MathVerifier(), "ifeval_old": IFEvalVerifierOld()}
print("n[verifier smoke test]")
print(" gsm8k :", VERIFIERS["gsm8k"]([], "9 + 3 = 12. The answer is 12.", "12").score)
print(" gsm8k :", VERIFIERS["gsm8k"]([], "The answer is 11.", "12").score)
print(" math :", VERIFIERS["math"]([], r"hence boxed{0.5}", r"frac{1}{2}").score)
print(" ifeval:", VERIFIERS["ifeval_old"]([], "one two three four five six seven",
json.dumps({"func_name": "validate_word_constraint",
"N": 6, "quantifier": "at least"})).score)
def verify_batch(responses, ground_truths, sources, tokenized=None):
out = []
for i, (resp, gt, src) in enumerate(zip(responses, ground_truths, sources)):
v = VERIFIERS.get(src, VERIFIERS["gsm8k"])
out.append(v(tokenized[i] if tokenized else [], resp, gt).score * v.weight)
return np.array(out, dtype=np.float32)
We load GSM8K and transform its questions and solutions into a consistent conversational format for SFT, DPO, and RLVR training. We create supervised examples, preference pairs with deliberately incorrect final answers, and verifier-ready prompts with structured ground-truth labels. We also initialize Open Instruct’s GSM8K, mathematical, and instruction-following verifiers and use them to score generated responses deterministically.
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)
We load the Qwen instruction model, apply LoRA adapters to its attention and feed-forward projection layers, and restrict optimization to the trainable adapter parameters. We configure mixed-precision execution, gradient scaling, gradient clipping, learning-rate scheduling, and temporary KV-cache activation for generation. We then evaluate the untrained baseline on GSM8K using greedy decoding and verifier-based answer accuracy.
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)
We construct a padded SFT DataLoader and train the LoRA adapters on tokenized GSM8K conversations using gradient accumulation. We optimize the model with cross-entropy loss calculated only over the unmasked assistant response tokens. We track the training loss and learning rate throughout the stage and evaluate the updated model after supervised fine-tuning.
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)
We batch the chosen and rejected responses separately and calculate their length-normalized sequence log probabilities with Open Instruct’s native utilities. We compare the active LoRA policy against the frozen base reference policy and optimize the model using the repository’s DPO loss. We monitor preference accuracy, reward margins, and training loss before measuring the model’s post-DPO verifier performance.
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(0)
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[:, :-1], old_lp[sl][:, :-1], ref_lp[sl][:, :-1]
m_, a_ = mask[sl][:, 1:], adv[sl][:, 1:]
ratio = torch.exp((new_lp_ - old_lp_).clamp(-20, 20))
pg, clipfrac, kl = compute_grpo_loss(new_lp_, ratio, a_, ref_lp_, grpo_cfg,
torch.ones_like(ratio))
loss = masked_mean(pg + cfg.grpo_kl_beta * kl, m_) / n_chunks
scaler.scale(loss).backward()
with torch.no_grad():
stats["pg"] += masked_mean(pg.detach(), m_).item() / n_chunks
stats["kl"] += masked_mean(kl.detach(), m_).item() / n_chunks
stats["clip"] += masked_mean(clipfrac.detach(), m_).item() / n_chunks
del new_lp, ratio, pg, kl
step_opt(opt, sched, scaler)
if DEV == "cuda":
torch.cuda.empty_cache()
print(f" grpo iter {it_i+1}/{cfg.grpo_iters} ep{ep+1} reward {scores.mean():.3f} "
f"(solved {int(scores.sum())}/{len(scores)}) pg {stats['pg']:+.4f} "
f"kl {stats['kl']:.4f} clipfrac {stats['clip']:.3f}")
print("n sample rollout ->", textwrap.shorten(texts[0].replace("n", " "), 220))
rlvr_acc = evaluate("after-rlvr", eval_rows)
print("n" + "=" * 90)
print(f"{'stage':<14}{'verifier acc':>14}")
for name, val in [("base", f"{base_acc:.3f}"), ("sft", f"{sft_acc:.3f}"),
("dpo", f"{dpo_acc:.3f}"), ("rlvr", f"{rlvr_acc:.3f}")]:
print(f"{name:<14}{val:>14}")
print("=" * 90)
OUT = "/content/tulu-mini" if os.path.isdir("/content") else "./tulu-mini"
merged = model.merge_and_unload()
merged.save_pretrained(OUT); tok.save_pretrained(OUT)
print(f"merged checkpoint -> {OUT} (equivalent to `python open_instruct/merge_lora.py`)")
We generate multiple sampled responses for each prompt, score them with deterministic verifiers, and calculate group-relative advantages from their reward distributions. We optimize the policy with Open Instruct’s GRPO and DAPO-style clipping logic while applying response masks, importance ratios, and KL regularization against the reference model. We finally compare accuracy across the baseline, SFT, DPO, and RLVR stages before merging the LoRA adapters and saving the completed checkpoint.
In conclusion, we implemented a practical miniature version of the Tulu 3 post-training stack and observed how each training stage changes model performance on verifier-scored mathematical reasoning tasks. We first established a baseline, improved instruction-following through supervised fine-tuning, refined response preferences through length-normalized DPO, and finally optimized verified task rewards using group-relative advantages and the repository’s GRPO loss implementation. We also used LoRA to maintain an accessible reference policy, apply response masking and KL regularization during reinforcement learning, compare accuracy across all training stages, and export a merged checkpoint for later inference or evaluation.
Check out the FULL CODES here. Also, feel free to follow us on Twitter and don’t forget to join our 150k+ML SubReddit and Subscribe to our Newsletter. Wait! are you on telegram? now you can join us on telegram as well.
Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us
The post AllenAI Open Instruct Tulu 3 Post-Training with SFT, DPO, RLVR, GRPO, and Verifier-Based Evaluation appeared first on MarkTechPost.