tinymind-native-8b-remote-handoff/bundle/train/grpo_trainer.py

"""
GRPO Trainer — Group Relative Policy Optimization

สำหรับ reasoning ที่ verifiable (math, logic, coding):
- สร้าง N คำตอบจาก policy model
- ให้ reward จาก verifier (ถูก=+1, ผิด=-1)
- normalize reward ภายใน group → เป็น advantage
- optimize ด้วย policy gradient ที่ clip แบบ PPO

อ้างอิง: DeepSeekMath GRPO + DeepSeek-R1 recipe
"""

from __future__ import annotations

import json
import math
import re
import time
from pathlib import Path
from typing import Callable

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from tokenizers import Tokenizer

from model.config import OmegaConfig, small_config
from model.architecture import OmegaModel
from model.reasoning import build_cot_prompt, extract_thinking


# ─── Reward Functions ─────────────────────────────────────────────────────────

def math_reward(question: str, generated: str, ground_truth: str) -> float:
    """Reward สำหรับ math: ตรวจคำตอบตัวเลขว่าตรงกับ ground truth ไหม"""
    _, answer = extract_thinking(generated)
    answer = answer.strip()

    def extract_number(text: str) -> float | None:
        nums = re.findall(r"-?\d+(?:\.\d+)?", text.replace(",", ""))
        return float(nums[-1]) if nums else None

    pred = extract_number(answer)
    truth = extract_number(ground_truth)
    if pred is None or truth is None:
        return -0.5
    if abs(pred - truth) < 1e-6:
        return 1.0
    if abs(pred - truth) / (abs(truth) + 1e-9) < 0.01:
        return 0.5
    return -1.0


def format_reward(generated: str) -> float:
    """Reward สำหรับ format: มี <think> และ <answer> ไหม"""
    has_think  = bool(re.search(r"<think>[\s\S]+</think>", generated, re.IGNORECASE))
    has_answer = bool(re.search(r"<answer>[\s\S]+</answer>", generated, re.IGNORECASE))
    if has_think and has_answer:
        return 0.2
    if has_answer:
        return 0.0
    return -0.2


def combined_reward(question: str, generated: str, ground_truth: str) -> float:
    r1 = math_reward(question, generated, ground_truth)
    r2 = format_reward(generated)
    return r1 + r2


# ─── GRPO Dataset ─────────────────────────────────────────────────────────────

class ReasoningDataset(Dataset):
    """Dataset สำหรับ GRPO — แต่ละ record มี question + ground_truth answer"""

    def __init__(self, path: str, tokenizer: Tokenizer, max_prompt_len: int = 512):
        self.tokenizer = tokenizer
        self.max_prompt_len = max_prompt_len
        self.records: list[dict] = []
        with open(path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                rec = json.loads(line)
                if rec.get("question") and rec.get("answer"):
                    self.records.append(rec)

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, idx: int) -> dict:
        rec = self.records[idx]
        lang = rec.get("lang", "th")
        prompt = build_cot_prompt(rec["question"], lang=lang)
        enc = self.tokenizer.encode(prompt)
        prompt_ids = enc.ids[: self.max_prompt_len]
        return {
            "question":    rec["question"],
            "ground_truth": rec["answer"],
            "prompt_ids":  prompt_ids,
            "lang":        lang,
        }


def grpo_collate(batch: list[dict], pad_id: int = 0) -> dict:
    max_len = max(len(b["prompt_ids"]) for b in batch)
    padded = [b["prompt_ids"] + [pad_id] * (max_len - len(b["prompt_ids"])) for b in batch]
    return {
        "questions":     [b["question"]     for b in batch],
        "ground_truths": [b["ground_truth"] for b in batch],
        "prompt_ids":    torch.tensor(padded, dtype=torch.long),
        "prompt_lens":   torch.tensor([len(b["prompt_ids"]) for b in batch], dtype=torch.long),
    }


# ─── GRPO Core ────────────────────────────────────────────────────────────────

@torch.no_grad()
def sample_group(
    model: OmegaModel,
    prompt_ids: torch.Tensor,
    n_samples: int,
    max_new_tokens: int,
    temperature: float,
    tokenizer: Tokenizer,
    device: torch.device,
) -> list[str]:
    """Sample N completions สำหรับ 1 prompt"""
    completions: list[str] = []
    model.eval()
    for _ in range(n_samples):
        generated = model.generate(
            prompt_ids.to(device),
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=0.9,
        )
        new_tokens = generated[0, prompt_ids.shape[1]:].tolist()
        text = tokenizer.decode(new_tokens)
        completions.append(text)
    return completions


def compute_group_advantages(rewards: list[float]) -> list[float]:
    """Normalize rewards ภายใน group → advantages"""
    mean = sum(rewards) / len(rewards)
    std  = math.sqrt(sum((r - mean) ** 2 for r in rewards) / len(rewards) + 1e-8)
    return [(r - mean) / std for r in rewards]


def grpo_policy_loss(
    model: OmegaModel,
    prompt_ids: torch.Tensor,
    completion_ids: torch.Tensor,
    advantages: torch.Tensor,
    old_log_probs: torch.Tensor,
    clip_eps: float = 0.2,
) -> torch.Tensor:
    """PPO-clip style policy loss สำหรับ GRPO"""
    full_ids = torch.cat([prompt_ids, completion_ids], dim=1)  # [1, T_prompt+T_comp]
    out = model(full_ids)
    logits = out["logits"][:, :-1, :]  # [1, T-1, V]
    target = full_ids[:, 1:]            # [1, T-1]

    log_probs = F.log_softmax(logits, dim=-1)
    token_lp = log_probs.gather(2, target.unsqueeze(-1)).squeeze(-1)  # [1, T-1]

    # เก็บเฉพาะ completion part
    p_len = prompt_ids.shape[1]
    comp_lp = token_lp[:, p_len - 1:p_len - 1 + completion_ids.shape[1]]  # [1, T_comp]

    # mask padding
    comp_mask = (completion_ids != 0).float()
    seq_lp = (comp_lp * comp_mask).sum(dim=-1) / comp_mask.sum(dim=-1).clamp(min=1)

    # PPO ratio
    ratio = torch.exp(seq_lp - old_log_probs)
    clipped = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps)
    policy_loss = -torch.min(ratio * advantages, clipped * advantages).mean()
    return policy_loss


# ─── GRPO Trainer ─────────────────────────────────────────────────────────────

GRPO_CFG = {
    "data_path":      "data/filtered/reasoning_qa.jsonl",
    "tokenizer_path": "data/tokenizer/tokenizer.json",
    "ref_checkpoint": "checkpoints/omega_best.pt",
    "out_dir":        "checkpoints",
    "n_samples":      4,          # G — completions per prompt
    "max_new_tokens": 256,
    "temperature":    0.8,
    "max_prompt_len": 512,
    "lr":             1e-6,
    "clip_eps":       0.2,
    "max_steps":      3_000,
    "save_every":     300,
    "dtype":          "bfloat16",
    "grad_clip":      1.0,
}


class GRPOTrainer:
    def __init__(self, cfg: dict = GRPO_CFG, model_cfg: OmegaConfig | None = None):
        self.cfg = cfg
        self.model_cfg = model_cfg or small_config()
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.dtype = torch.bfloat16 if cfg["dtype"] == "bfloat16" else torch.float16
        self.step = 0

    def setup(self):
        print(f"GRPO Trainer | device={self.device} | G={self.cfg['n_samples']}")
        tok_path = self.cfg["tokenizer_path"]
        if not Path(tok_path).exists():
            raise FileNotFoundError(f"Tokenizer not found: {tok_path}")
        self.tokenizer = Tokenizer.from_file(tok_path)

        ds = ReasoningDataset(self.cfg["data_path"], self.tokenizer, self.cfg["max_prompt_len"])
        self.loader = DataLoader(
            ds, batch_size=1, shuffle=True,
            collate_fn=lambda b: grpo_collate(b, pad_id=self.model_cfg.pad_token_id),
            num_workers=0,
        )
        print(f"Reasoning dataset: {len(ds):,} problems")

        ref_path = self.cfg.get("ref_checkpoint")
        if ref_path and Path(ref_path).exists():
            ckpt = torch.load(ref_path, map_location=self.device, weights_only=False)
            saved_cfg: OmegaConfig = ckpt["model_cfg"]
            self.model = OmegaModel(saved_cfg).to(self.device)
            self.model.load_state_dict(ckpt["model_state"])
            self.model_cfg = saved_cfg
        else:
            self.model = OmegaModel(self.model_cfg).to(self.device)

        self.optimizer = torch.optim.AdamW(
            self.model.parameters(), lr=self.cfg["lr"], betas=(0.9, 0.95)
        )

    def train_step(self, batch: dict) -> float:
        question     = batch["questions"][0]
        ground_truth = batch["ground_truths"][0]
        prompt_ids   = batch["prompt_ids"][:1]  # [1, T_prompt]

        # 1. Sample G completions
        completions = sample_group(
            self.model, prompt_ids,
            n_samples=self.cfg["n_samples"],
            max_new_tokens=self.cfg["max_new_tokens"],
            temperature=self.cfg["temperature"],
            tokenizer=self.tokenizer,
            device=self.device,
        )

        # 2. Compute rewards
        rewards = [combined_reward(question, c, ground_truth) for c in completions]

        # 3. Normalize → advantages
        advantages = compute_group_advantages(rewards)

        if all(r == rewards[0] for r in rewards):
            return 0.0  # ทุก completion ได้ reward เท่ากัน — skip

        # 4. Policy gradient update (per completion)
        total_loss = 0.0
        self.model.train()

        for comp_text, adv in zip(completions, advantages):
            enc = self.tokenizer.encode(comp_text)
            comp_ids = torch.tensor([enc.ids[:self.cfg["max_new_tokens"]]], dtype=torch.long)

            # old log prob (no grad)
            with torch.no_grad():
                full_ids = torch.cat([prompt_ids.to(self.device),
                                      comp_ids.to(self.device)], dim=1)
                out_ref = self.model(full_ids)
                lp_ref = F.log_softmax(out_ref["logits"][:, :-1, :], dim=-1)
                target = full_ids[:, 1:]
                old_lp = lp_ref.gather(2, target.unsqueeze(-1)).squeeze(-1)
                p_len = prompt_ids.shape[1]
                old_seq_lp = old_lp[:, p_len - 1:p_len - 1 + comp_ids.shape[1]].mean()

            adv_t = torch.tensor([adv], device=self.device, dtype=self.dtype)

            with torch.amp.autocast(device_type=self.device.type, dtype=self.dtype,
                                    enabled=self.device.type == "cuda"):
                loss = grpo_policy_loss(
                    self.model,
                    prompt_ids.to(self.device),
                    comp_ids.to(self.device),
                    adv_t,
                    old_seq_lp,
                    clip_eps=self.cfg["clip_eps"],
                ) / self.cfg["n_samples"]

            loss.backward()
            total_loss += loss.item() * self.cfg["n_samples"]

        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg["grad_clip"])
        self.optimizer.step()
        self.optimizer.zero_grad()
        return total_loss

    def save(self, tag: str = "grpo_latest"):
        path = Path(self.cfg["out_dir"]) / f"omega_{tag}.pt"
        torch.save({
            "step": self.step,
            "model_state": self.model.state_dict(),
            "model_cfg": self.model_cfg,
        }, path)
        print(f"  Saved → {path}")

    def train(self):
        self.setup()
        data_iter = iter(self.loader)
        t0 = time.time()
        running_loss = 0.0

        print(f"GRPO training for {self.cfg['max_steps']:,} steps\n")
        while self.step < self.cfg["max_steps"]:
            try:
                batch = next(data_iter)
            except StopIteration:
                data_iter = iter(self.loader)
                batch = next(data_iter)

            running_loss += self.train_step(batch)
            self.step += 1

            if self.step % 10 == 0:
                dt = time.time() - t0
                print(f"step {self.step:5d} | grpo_loss {running_loss/10:.4f} | {dt:.1f}s")
                running_loss = 0.0
                t0 = time.time()

            if self.step % self.cfg["save_every"] == 0:
                self.save(f"grpo_step{self.step}")

        self.save("grpo_final")
        print("GRPO training complete!")


if __name__ == "__main__":
    trainer = GRPOTrainer()
    trainer.train()