code/grpo_train.py

"""GRPO RL phase for BLT-Reasoner.

Architecture-specific differences from Abstract-CoT GRPO (`experiments/abstract_cot/grpo_train.py`):

* **z is deterministic given x and current policy weights.** No discrete
  sampling in the latent. The only stochasticity is in y. So per-prompt we
  compute z once and sample K different y's from `π(y | x, z)`. This is much
  cheaper than the abstract-vocab version which had to sample z too.

* **Two model instances on the GPU**: policy (trainable, init = SFT ckpt) and
  reference (frozen, identical init). KL is policy↔reference per y token. We
  cannot use `peft.disable_adapter()` as the reference, because the SFT ckpt's
  LoRA + projector + head IS the reference — disabling adapters would compare
  against vanilla Qwen, which is not what we want.

* **InfoNCE optionally retained as a low-weight aux loss** during GRPO so the
  z geometry doesn't drift while the policy learns to use rewards. This is
  the BLT-specific safety net.

Reward = math-verifier on the extracted final number after `####`.
Advantage = group-normalized within K rollouts per prompt.

Usage:
    python -m experiments.blt_reasoner.grpo_train --config configs/grpo_from_step12000.json
"""
from __future__ import annotations

import argparse
import copy
import json
import math
import os
import random
import re
import sys
import time
from pathlib import Path
from typing import List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

from .data import GSM8KDataset, collate_batch
from .losses import InfoNCEHead, encode_answer_for_infonce, infonce_loss
from .model import (
    NEG,
    BLTConfig,
    LatentProjector,
    build_base,
    build_blt_mask,
    forward_with_latent,
    generate_with_latent,
)
from .train import _load_state_from_ckpt


GSM8K_NUM = re.compile(r"####\s*(-?\d+(?:\.\d+)?)")
ANY_NUM = re.compile(r"-?\d+(?:\.\d+)?")


def parse_pred(text: str) -> Optional[str]:
    m = GSM8K_NUM.search(text)
    if m:
        return m.group(1)
    nums = ANY_NUM.findall(text)
    return nums[-1] if nums else None


def reward_for(decoded: str, gold: str, *, length_pen: float = 0.0, n_tokens: int = 0) -> float:
    """Math-verifier reward.

    +1.0   exact-match the gold number
    -0.5   parseable but wrong, or no number found
    -length_pen * (n_tokens / 192)  small length penalty to discourage rambling
    """
    pred = parse_pred(decoded)
    base = -0.5
    if pred is not None:
        try:
            if abs(float(pred) - float(gold)) < 1e-4:
                base = 1.0
        except ValueError:
            pass
    return base - length_pen * (n_tokens / 192.0)


def set_seed(s: int):
    random.seed(s)
    torch.manual_seed(s)
    torch.cuda.manual_seed_all(s)


def per_token_logp(logits: torch.Tensor, target_ids: torch.Tensor, target_mask: torch.Tensor) -> torch.Tensor:
    """logits [B, L, V] (already shifted so logits[:, t] predicts target[:, t]);
    target_ids [B, L]; target_mask [B, L]. Returns [B, L] per-token logp.
    Masked positions get zero (not NaN).
    """
    logp = F.log_softmax(logits.float(), dim=-1)
    gathered = logp.gather(-1, target_ids.unsqueeze(-1)).squeeze(-1)
    return gathered * target_mask.float()


def forward_logp_for_rollout(model, projector, x_ids, x_attn, y_ids, y_mask, K, *, block_y_to_x=True):
    """Run forward_with_latent and return per-y-token logps under the given model.

    `y_ids` is the rollout's sampled answer. We use forward_with_latent to
    recompute z under the current params and grab logits at y positions.
    """
    logits_y, z, _ = forward_with_latent(
        model, x_ids, x_attn, y_ids, projector, K,
        block_y_to_x=block_y_to_x, return_z=True,
    )
    # logits_y[:, t, :] predicts y_ids[:, t]. So logp at position t is on target y_ids[:, t].
    return per_token_logp(logits_y, y_ids, y_mask), z


def grpo_loss(policy_logp, ref_logp, advantages, y_mask, beta, kl_clamp):
    """
    policy_logp, ref_logp : [B*K, L_y]   per-token logps (masked to 0 at pads)
    advantages            : [B*K]
    y_mask                : [B*K, L_y]   1 where real token
    Returns (loss, info_dict).
    """
    # Policy gradient: -A * sum_t logp_policy(y_t | ...)
    seq_logp = policy_logp.sum(dim=-1)                          # [B*K]
    pg = -(advantages * seq_logp).mean()

    # Per-token KL via k3 estimator: clamp log-ratio then (exp - 1 - log_ratio).
    log_ratio = (policy_logp - ref_logp).clamp(min=-kl_clamp, max=kl_clamp)
    kl_per_tok = (log_ratio.exp() - 1.0 - log_ratio)             # [B*K, L_y]
    n_tok = y_mask.sum().clamp(min=1.0)
    kl_term = (kl_per_tok * y_mask.float()).sum() / n_tok

    loss = pg + beta * kl_term
    return loss, {"pg": pg.detach(), "kl": kl_term.detach()}


def make_optimizer(model, projector, head, cfg):
    try:
        from bitsandbytes.optim import PagedAdamW8bit as AdamW
        use_8bit = True
    except Exception:
        from torch.optim import AdamW
        use_8bit = False
    groups = [
        {"params": [p for p in model.parameters() if p.requires_grad], "lr": cfg["lr_lora"]},
        {"params": list(projector.parameters()), "lr": cfg["lr_proj"]},
        {"params": list(head.parameters()),      "lr": cfg["lr_head"]},
    ]
    return AdamW(groups, weight_decay=cfg["weight_decay"]), use_8bit


def sample_rollouts(
    model, tokenizer, projector,
    x_ids, x_attn, *,
    K_latents, group_size, max_new_tokens, temperature, block_y_to_x,
    eos_id, pad_id,
):
    """Sample `group_size` rollouts per prompt by expanding the batch K-fold.

    Returns (y_ids, y_mask): both [B*K, L_y_gen]
    """
    B = x_ids.size(0)
    x_ids_rep = x_ids.repeat_interleave(group_size, dim=0)
    x_attn_rep = x_attn.repeat_interleave(group_size, dim=0)

    gen = generate_with_latent(
        model, tokenizer, projector,
        x_ids=x_ids_rep, x_attn=x_attn_rep, K=K_latents,
        block_y_to_x=block_y_to_x,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        eos_token_id=eos_id,
    )                                          # [B*K, L_gen]
    y_mask = (gen != pad_id).long()
    return gen, y_mask


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", required=True)
    args = parser.parse_args()
    with open(args.config) as f:
        cfg = json.load(f)
    set_seed(cfg["seed"])

    out_dir = Path(cfg["output_dir"])
    out_dir.mkdir(parents=True, exist_ok=True)
    log_f = open(out_dir / "grpo.log", "a", buffering=1)
    met_f = open(out_dir / "metrics.jsonl", "a", buffering=1)

    def log(m):
        line = f"[{time.strftime('%H:%M:%S')}] {m}"
        print(line, flush=True)
        log_f.write(line + "\n")

    log(f"config={json.dumps(cfg)}")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # ---- 1) Build policy + reference (both init from same SFT ckpt) ----------
    bcfg = BLTConfig(
        base_model=cfg["base_model"],
        use_lora=cfg["use_lora"],
        lora_r=cfg["lora_r"], lora_alpha=cfg["lora_alpha"],
        lora_dropout=cfg["lora_dropout"],
        lora_target_modules=tuple(cfg["lora_target_modules"]),
        K_latents=cfg["K_latents"], block_y_to_x=cfg["block_y_to_x"],
        proj_init_scale=cfg["proj_init_scale"],
        dtype=cfg["dtype"], attn_impl=cfg["attn_impl"],
    )

    log("building policy model …")
    policy_model, tokenizer = build_base(bcfg)
    policy_model.to(device)

    inner_pol = policy_model.get_base_model() if hasattr(policy_model, "get_base_model") else policy_model
    d_model = inner_pol.config.hidden_size
    model_dtype = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[cfg["dtype"]]
    pol_projector = LatentProjector(
        d_model, init_scale=cfg["proj_init_scale"],
        use_mlp=cfg.get("proj_mlp", False),
        hidden_mult=cfg.get("proj_hidden_mult", 4),
    ).to(device=device, dtype=model_dtype)
    pol_head = InfoNCEHead(d_z=d_model, d_y=d_model, d_out=256).to(device)

    log(f"loading policy from SFT ckpt {cfg['warmup_ckpt']}")
    _load_state_from_ckpt(policy_model, pol_projector, pol_head, cfg["warmup_ckpt"], device)

    log("building reference model (frozen copy of policy init) …")
    ref_model, _ = build_base(bcfg)
    ref_model.to(device)
    ref_projector = LatentProjector(
        d_model, init_scale=cfg["proj_init_scale"],
        use_mlp=cfg.get("proj_mlp", False),
        hidden_mult=cfg.get("proj_hidden_mult", 4),
    ).to(device=device, dtype=model_dtype)
    ref_head = InfoNCEHead(d_z=d_model, d_y=d_model, d_out=256).to(device)
    _load_state_from_ckpt(ref_model, ref_projector, ref_head, cfg["warmup_ckpt"], device)
    for p in ref_model.parameters():
        p.requires_grad_(False)
    for p in ref_projector.parameters():
        p.requires_grad_(False)
    for p in ref_head.parameters():
        p.requires_grad_(False)
    ref_model.eval()
    ref_projector.eval()
    ref_head.eval()

    # ---- 2) Optimizer over POLICY only ----------
    opt, used_8bit = make_optimizer(policy_model, pol_projector, pol_head, cfg)
    log(f"optimizer 8bit={used_8bit}")

    n_trainable = sum(p.numel() for p in policy_model.parameters() if p.requires_grad) \
                  + sum(p.numel() for p in pol_projector.parameters()) \
                  + sum(p.numel() for p in pol_head.parameters())
    log(f"policy trainable params = {n_trainable/1e6:.2f}M")

    # ---- 3) Data: GSM8K-train problems for RL ----------
    train_ds = GSM8KDataset(split="train", max_examples=cfg.get("rl_data_size"))
    log(f"GSM8K RL pool size = {len(train_ds)}")
    loader = DataLoader(
        train_ds, batch_size=cfg["per_prompt_batch"], shuffle=True, drop_last=True,
        collate_fn=lambda b: collate_batch(
            b, tokenizer,
            max_prompt_len=cfg["max_prompt_len"],
            max_answer_len=cfg["max_answer_len"],
        ),
    )

    # ---- 4) GRPO loop ----------
    K = cfg["K_latents"]
    G = cfg["group_size"]
    beta = cfg["beta"]
    kl_clamp = cfg.get("kl_clamp", 20.0)
    eos_id = tokenizer.eos_token_id
    pad_id = tokenizer.pad_token_id

    step = 0
    t0 = time.time()
    reward_hist: List[float] = []

    while step < cfg["max_steps"]:
        for batch in loader:
            if step >= cfg["max_steps"]:
                break

            x_ids = batch.x_ids.to(device)
            x_attn = batch.x_attn.to(device)
            # golds: parse "#### NUMBER" from batch.final_strs ("#### N")
            golds = [s.replace("#### ", "").strip() for s in batch.final_strs]
            B = x_ids.size(0)

            # 4a) Sample rollouts under POLICY
            policy_model.eval()                      # no grad needed for sampling
            with torch.no_grad():
                y_ids, y_mask = sample_rollouts(
                    policy_model, tokenizer, pol_projector,
                    x_ids, x_attn,
                    K_latents=K, group_size=G,
                    max_new_tokens=cfg["max_new_tokens"],
                    temperature=cfg["sample_temperature"],
                    block_y_to_x=cfg["block_y_to_x"],
                    eos_id=eos_id, pad_id=pad_id,
                )
            policy_model.train()
            # Decode for verifier reward
            decoded = tokenizer.batch_decode(y_ids, skip_special_tokens=True)
            rewards = []
            for i in range(B):
                for g in range(G):
                    idx = i * G + g
                    n_tok = int(y_mask[idx].sum().item())
                    rewards.append(reward_for(
                        decoded[idx], golds[i],
                        length_pen=cfg.get("length_penalty_coef", 0.0),
                        n_tokens=n_tok,
                    ))
            rewards = torch.tensor(rewards, device=device, dtype=torch.float32)   # [B*K]
            reward_hist.append(rewards.mean().item())

            # 4b) Group-normalized advantages
            rew = rewards.view(B, G)
            adv = (rew - rew.mean(dim=1, keepdim=True)) / (rew.std(dim=1, keepdim=True) + 1e-6)
            adv = adv.view(-1)                                # [B*K]

            # Skip steps with no reward variance (all-correct or all-wrong group) — no signal.
            n_groups_with_signal = int((rew.std(dim=1) > 1e-6).sum().item())
            if n_groups_with_signal == 0:
                opt.zero_grad(set_to_none=True)
                step += 1
                if step % cfg["log_every"] == 0:
                    log(f"step={step} reward_mean={rewards.mean().item():.3f} "
                        f"n_signal=0 (skipped) elapsed={time.time()-t0:.0f}s")
                    met_f.write(json.dumps({"step": step, "reward_mean": rewards.mean().item(),
                                             "n_groups_with_signal": 0,
                                             "skipped": 1,
                                             "elapsed_s": time.time()-t0}) + "\n")
                continue

            # 4c) Compute per-token logps under policy (with grad) and reference (no grad)
            # Re-expand x for B*K. (Cheap to repeat tensors; we re-forward through model regardless.)
            x_ids_r = x_ids.repeat_interleave(G, dim=0)
            x_attn_r = x_attn.repeat_interleave(G, dim=0)

            policy_logp, _z_pol = forward_logp_for_rollout(
                policy_model, pol_projector, x_ids_r, x_attn_r, y_ids, y_mask, K,
                block_y_to_x=cfg["block_y_to_x"],
            )
            with torch.no_grad():
                ref_logp, _ = forward_logp_for_rollout(
                    ref_model, ref_projector, x_ids_r, x_attn_r, y_ids, y_mask, K,
                    block_y_to_x=cfg["block_y_to_x"],
                )

            loss, info = grpo_loss(policy_logp, ref_logp, adv, y_mask, beta, kl_clamp)
            (loss / cfg["grad_accum"]).backward()

            if (step + 1) % cfg["grad_accum"] == 0:
                torch.nn.utils.clip_grad_norm_(
                    [p for p in policy_model.parameters() if p.requires_grad]
                    + list(pol_projector.parameters()) + list(pol_head.parameters()),
                    cfg["max_grad_norm"],
                )
                opt.step()
                opt.zero_grad(set_to_none=True)

            step += 1
            if step % cfg["log_every"] == 0:
                tail = reward_hist[-min(50, len(reward_hist)):]
                log(f"step={step} reward_mean={rewards.mean().item():.3f} "
                    f"reward_avg50={sum(tail)/len(tail):.3f} "
                    f"adv|abs|={adv.abs().mean().item():.3f} "
                    f"pg={info['pg'].item():.4f} kl={info['kl'].item():.4f} "
                    f"n_signal={n_groups_with_signal}/{B} "
                    f"elapsed={time.time()-t0:.0f}s")
                met_f.write(json.dumps({
                    "step": step,
                    "reward_mean": rewards.mean().item(),
                    "reward_avg50": sum(tail)/len(tail),
                    "adv_abs_mean": adv.abs().mean().item(),
                    "pg": float(info["pg"].item()),
                    "kl": float(info["kl"].item()),
                    "n_groups_with_signal": n_groups_with_signal,
                    "skipped": 0,
                    "elapsed_s": time.time() - t0,
                }) + "\n")

            if cfg["save_every"] > 0 and step % cfg["save_every"] == 0:
                save_dir = out_dir / f"grpo-step{step}"
                save_dir.mkdir(exist_ok=True)
                policy_model.save_pretrained(save_dir / "model")
                torch.save(pol_projector.state_dict(), save_dir / "projector.pt")
                torch.save(pol_head.state_dict(), save_dir / "head.pt")
                log(f"[save] {save_dir}")

    # Final save
    save_dir = out_dir / "final"
    save_dir.mkdir(exist_ok=True)
    policy_model.save_pretrained(save_dir / "model")
    torch.save(pol_projector.state_dict(), save_dir / "projector.pt")
    torch.save(pol_head.state_dict(), save_dir / "head.pt")
    log(f"[done] final saved at {save_dir}")
    log_f.close()
    met_f.close()


if __name__ == "__main__":
    main()