Постобучение на Tulu 3 с AllenAI Open Instruct чрез SFT, DPO, RLVR, GRPO и оценяване на базата на верификатор

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

@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()}")

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)

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)

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)

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)

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`)")