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