""" GRPO (Group Relative Policy Optimization) with a verifiable reward, for the Nawah reasoning models. Why hand-rolled: TRL is not installed and cannot be added cleanly here — .venv-lfm2 runs transformers 5.15 while current TRL targets 4.x, and .venv-vllm is pinned to transformers 4.51.3 and must not be upgraded (see HANDOFF.md). At 51.8M parameters the algorithm is small enough that a direct implementation is less risk than a third environment. The method, briefly: sample G completions per prompt, score each with a mechanical verifier, normalise the rewards WITHIN the group to get advantages, and do a policy-gradient step. No value network — the group mean is the baseline, which is the whole point of GRPO. A_i = (r_i - mean(r)) / (std(r) + eps) loss = -mean_i( A_i * mean_t log pi(token_t) ) + beta * KL_k3(pi || pi_ref) Sampling is on-policy and each batch takes exactly one gradient step, so there is no importance ratio and no PPO clipping to get wrong — the ratio would be identically 1. ⚠️ The measured precondition: a group where every sample scores the same has zero advantage and contributes NO gradient. On v5, 40.4% of synth problems are never solved in 8 tries, so those groups are dead weight — they are skipped and counted, not silently averaged in. The exploitable headroom is pass@k - pass@1, measured at +27.1 points (32.5% -> 59.6%) by passk_diag.py. """ import json import os import sys from pathlib import Path import torch import torch.nn.functional as F from transformers import AutoModelForCausalLM, AutoTokenizer sys.path.insert(0, ".") from eval_reasoning import numbers, parse MODEL = os.environ.get("MODEL", "./Nawah-Reasoning-v5") OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "./Nawah-Reasoning-v6-grpo") TRAIN_FILE = os.environ.get("TRAIN_FILE", "data_synth_sft/train.jsonl") GROUP = int(os.environ.get("GROUP", 8)) # completions per prompt PROMPTS_PER = int(os.environ.get("PROMPTS_PER", 8)) # prompts per optimiser step STEPS = int(os.environ.get("STEPS", 500)) LR = float(os.environ.get("LR", 1e-6)) # RL wants far less than SFT's 3e-4 BETA = float(os.environ.get("BETA", 0.02)) # KL to the frozen reference TEMP = float(os.environ.get("TEMP", 1.0)) MAX_NEW = int(os.environ.get("MAX_NEW", 320)) MAX_GRAD = 1.0 SAVE_EVERY = int(os.environ.get("SAVE_EVERY", 100)) LOG_EVERY = int(os.environ.get("LOG_EVERY", 5)) SEED = int(os.environ.get("SEED", 42)) FORMAT_PENALTY = float(os.environ.get("FORMAT_PENALTY", 0.1)) def reward(completion: str, ref: float) -> float: """Verifiable reward: does the stated final answer equal the reference number? Deliberately NOT a partial-credit score. The failure this is meant to fix is a model that reasons plausibly and lands on the wrong number, so rewarding anything short of the right number would reinforce exactly that. """ _, ans, well_formed = parse(completion) ns = numbers(ans or "") correct = bool(ns) and ref is not None and ns[-1] == ref r = 1.0 if correct else 0.0 if not well_formed: r -= FORMAT_PENALTY return r def completion_logprobs(model, ids, attn, prompt_len): """-> (mean log pi over completion tokens, per-token log pi, completion mask). Prompt tokens are masked out: the prompt is not an action the policy chose, so including it would add a term with no gradient meaning and would dilute the per-sequence mean. """ out = model(input_ids=ids, attention_mask=attn) logits = out.logits[:, :-1] # position t predicts token t+1 targets = ids[:, 1:] logp = torch.log_softmax(logits.float(), dim=-1) tok_logp = logp.gather(-1, targets.unsqueeze(-1)).squeeze(-1) mask = attn[:, 1:].clone().float() mask[:, : prompt_len - 1] = 0.0 # only the generated part is the "action" mean_lp = (tok_logp * mask).sum(1) / mask.sum(1).clamp(min=1) return mean_lp, tok_logp, mask def k3_kl(pol_tok, ref_tok, mask): """Schulman's k3 estimator of KL(pi || pi_ref): exp(d) - d - 1 where d = log pi_ref - log pi. Unbiased and always non-negative, unlike the naive (log pi - log pi_ref) difference. Computed per token and averaged over the completion, which is what GRPO regularises. """ d = (ref_tok - pol_tok).clamp(-20, 20) per_tok = torch.exp(d) - d - 1.0 return (per_tok * mask).sum(1) / mask.sum(1).clamp(min=1) def main(): torch.manual_seed(SEED) tok = AutoTokenizer.from_pretrained(MODEL) policy = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32).cuda() ref = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32).cuda().eval() for p in ref.parameters(): p.requires_grad_(False) policy.config.use_cache = True rows = [json.loads(l) for l in open(TRAIN_FILE, encoding="utf-8")] refs = [] for r in rows: ns = numbers(r["answer"]) if ns: refs.append((r["instruction"], ns[-1])) print(f"[*] {len(refs):,} prompts with a parseable reference answer", flush=True) opt = torch.optim.AdamW(policy.parameters(), lr=LR, weight_decay=0.0) rng = torch.Generator().manual_seed(SEED) Path(OUTPUT_DIR).mkdir(exist_ok=True) history, dead_groups, seen = [], 0, 0 for step in range(1, STEPS + 1): idx = torch.randint(0, len(refs), (PROMPTS_PER,), generator=rng).tolist() batch = [refs[i] for i in idx] step_rewards, step_solved, losses = [], [], [] opt.zero_grad(set_to_none=True) used_groups = 0 for question, ref_num in batch: prompt = tok.apply_chat_template([{"role": "user", "content": question}], tokenize=False, add_generation_prompt=True) enc = tok(prompt, return_tensors="pt").to("cuda") plen = enc["input_ids"].shape[1] policy.eval() with torch.no_grad(): gen = policy.generate(**enc, max_new_tokens=MAX_NEW, do_sample=True, temperature=TEMP, top_p=0.95, num_return_sequences=GROUP, pad_token_id=tok.pad_token_id) policy.train() texts = tok.batch_decode(gen[:, plen:], skip_special_tokens=True) r = torch.tensor([reward(t, ref_num) for t in texts], dtype=torch.float32) step_rewards.append(r.mean().item()) step_solved.append(float((r > 0.5).any())) seen += 1 # A group with no reward spread carries no learning signal — skip it rather than # letting a zero-advantage group dilute the batch. if r.std() < 1e-6: dead_groups += 1 continue adv = ((r - r.mean()) / (r.std() + 1e-6)).cuda() attn = (gen != tok.pad_token_id).long() attn[:, :plen] = 1 mean_lp, pol_tok, mask = completion_logprobs(policy, gen, attn, plen) with torch.no_grad(): _, ref_tok, _ = completion_logprobs(ref, gen, attn, plen) pg = -(adv * mean_lp).mean() kl = k3_kl(pol_tok, ref_tok, mask).mean() loss = (pg + BETA * kl) / PROMPTS_PER loss.backward() losses.append(pg.item()) used_groups += 1 if used_groups: torch.nn.utils.clip_grad_norm_(policy.parameters(), MAX_GRAD) opt.step() history.append({"step": step, "mean_reward": sum(step_rewards) / len(step_rewards), "any_solved": sum(step_solved) / len(step_solved), "pg_loss": (sum(losses) / len(losses)) if losses else None, "live_groups": used_groups}) if step % LOG_EVERY == 0: h = history[-1] print(f"[{step:>4}/{STEPS}] reward {h['mean_reward']:.3f} " f"pass@{GROUP} {h['any_solved']:.2f} live {used_groups}/{PROMPTS_PER} " f"dead so far {100*dead_groups/max(seen,1):.0f}%", flush=True) if step % SAVE_EVERY == 0 or step == STEPS: policy.config.use_cache = True policy.save_pretrained(OUTPUT_DIR) tok.save_pretrained(OUTPUT_DIR) Path(OUTPUT_DIR, "grpo_history.json").write_text( json.dumps(history, ensure_ascii=False, indent=2), encoding="utf-8") print(f" [+] saved at step {step}", flush=True) print(f"[+] done -> {OUTPUT_DIR} ({100*dead_groups/max(seen,1):.1f}% of groups had no signal)") if __name__ == "__main__": main()