code/grpo_train.py

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