"""Lab 15: GRPO on TinyLM, from group sampling to the training loop. Your file. Five tasks, in order; each is marked TODO. Run the tests for one task at a time: python -m pytest -q -k task1 Chapter: https://learn.xhelio.ai/rl/learn/15-rlvr-grpo/ Everything below "given: helpers and the training loop" is complete. Read train_grpo once: it is the outer loop that calls your grpo_step. """ import random import time from dataclasses import dataclass import torch import torch.nn.functional as F from tinylm import ( accuracy, decode_responses, load_or_train_sft, make_batch, reward, sample, sequence_logprobs, ) # ============================================================== Task 1 def sample_group(model, prompts, answers, G, temperature=1.0): """Sample a group of G responses for every prompt and score them. Args: model: the policy (a TinyLM). prompts: list of P prompt strings, e.g. ["37+48=", "5+5="]. answers: list of P answer strings, e.g. ["85", "10"]. G: group size (responses per prompt). temperature: sampling temperature. Returns a dict: "out": the dict returned by tinylm.sample for the P*G repeated prompts. Rows are grouped contiguously: rows 0..G-1 answer prompt 0, rows G..2G-1 answer prompt 1, and so on. "responses": list of P*G decoded response strings. "rewards": float tensor of shape (P*G,), the verifiable reward tinylm.reward(response, answer) of each row (1.0 or 0.0). """ # TODO (Task 1): repeat each prompt (and answer) G times, consecutively, call # tinylm.sample, decode with tinylm.decode_responses, score with tinylm.reward. raise NotImplementedError("Task 1: sample_group") # ============================================================== Task 2 def group_advantages(rewards, G, scale_std=True, eps=1e-4): """Group-relative advantages, one number per response. Args: rewards: (P*G,) tensor, grouped contiguously as in sample_group. G: group size. scale_std: if True, divide by the group's standard deviation (GRPO); if False, only subtract the group mean. eps: added to the std so all-equal groups do not divide by zero. Returns: (P*G,) tensor with A_i = (r_i - mean(group)) / (std(group) + eps) where mean and std are taken over the G responses of the same prompt. Use the population std (divide by G, i.e. torch.std(..., unbiased=False)). A group whose rewards are all equal gets advantage 0 for every response. """ # TODO (Task 2): reshape to (P, G), subtract the row mean, optionally divide by # (row std + eps), reshape back to (P*G,). raise NotImplementedError("Task 2: group_advantages") # ============================================================== Task 3 def k3_kl(logp, ref_logp): """Per-token k3 estimator of KL(pi_theta || pi_ref). Args: logp: (B, T) log pi_theta(y_t | ...) of the sampled tokens. ref_logp: (B, T) log pi_ref(y_t | ...) of the same tokens. Returns: (B, T) tensor k3 = pi_ref/pi_theta - log(pi_ref/pi_theta) - 1, computed from log-probs: with d = ref_logp - logp, k3 = exp(d) - d - 1. Every entry is >= 0. Gradients must flow through logp (not through ref_logp, which should come from a frozen model anyway). """ # TODO (Task 3): two lines. raise NotImplementedError("Task 3: k3_kl") # ============================================================== Task 4 def grpo_loss(logp, old_logp, ref_logp, adv, mask, clip_eps=0.2, beta=0.04): """The GRPO loss (DeepSeekMath's objective, negated so we can minimize it). Args: logp: (B, T) log-probs of the sampled tokens under the current policy (with grad). old_logp: (B, T) log-probs under the policy that sampled them (no grad). ref_logp: (B, T) log-probs under the frozen reference model (no grad). adv: (B,) one advantage per response, shared by all its tokens. mask: (B, T) 1.0 for real response tokens, 0.0 for padding. clip_eps: PPO clip range epsilon. beta: KL coefficient. Per token: rho = exp(logp - old_logp) obj = min(rho * A, clip(rho, 1 - eps, 1 + eps) * A) - beta * k3 Aggregate: average obj over each response's own tokens (divide by |y_i|), then average over the B responses. loss = -that. Returns (loss, stats) where stats is a dict of floats: "kl": mean k3 over all real tokens of the batch, "clip_frac": fraction of real tokens where the clipped term was the smaller one (i.e. where the clip changed the objective), "ratio_max": the largest rho over real tokens. """ # TODO (Task 4): ratio, unclipped and clipped surrogate, min, minus beta * k3, # per-response masked mean, mean over responses, negate. Then the stats # (compute them under torch.no_grad()). raise NotImplementedError("Task 4: grpo_loss") # ============================================================== Task 5 def grpo_step(policy, ref, opt, prompts, answers, cfg): """One GRPO iteration: sample groups, score, compute advantages, then update. Steps: 1. groups = sample_group(policy, prompts, answers, cfg.G, cfg.temperature) 2. adv = group_advantages(groups["rewards"], cfg.G, cfg.scale_std) 3. With torch.no_grad(): old_logp = sequence_logprobs(policy, out) and ref_logp = sequence_logprobs(ref, out). These stay fixed for the whole iteration. 4. Repeat cfg.inner_steps times: logp = sequence_logprobs(policy, out); loss, s = grpo_loss(logp, old_logp, ref_logp, adv, mask, cfg.clip_eps, cfg.beta); opt.zero_grad(); loss.backward(); clip grad norm to cfg.max_grad_norm; opt.step(). Returns a dict of floats: "reward" (mean reward of the samples), "loss" (last inner step), "kl", "clip_frac" (from the last inner step), "frac_zero_signal" (fraction of groups whose rewards are all equal), "entropy" (mean per-token entropy of the policy at the sampled positions, from `token_entropy`). """ # TODO (Task 5): follow the five steps in the docstring. raise NotImplementedError("Task 5: grpo_step") # ============================================================== given: helpers and the training loop @torch.no_grad() def token_entropy(model, out): """Mean entropy (nats) of the model's next-token distribution at the response positions.""" ids, L = out["ids"], out["prompt_len"] logits = model(ids[:, :-1])[:, L - 1:] logp = F.log_softmax(logits, -1) ent = -(logp.exp() * logp).sum(-1) m = out["mask"] return float((ent * m).sum() / m.sum().clamp(min=1.0)) @dataclass class GRPOConfig: iters: int = 300 # outer iterations (each samples a fresh batch) prompts_per_iter: int = 16 # P G: int = 8 # group size inner_steps: int = 2 # gradient steps per batch (mu in the DeepSeekMath paper) lr: float = 3e-4 clip_eps: float = 0.2 beta: float = 0.04 temperature: float = 1.0 scale_std: bool = True max_grad_norm: float = 1.0 eval_every: int = 10 eval_n: int = 500 seed: int = 0 def train_grpo(cfg, sft_path=None, verbose=True): """Run GRPO from the cached SFT checkpoint. Returns (policy, history). The SFT model is trained once (about 5 s) and cached in tinylm_sft.pt next to tinylm.py. History holds per-iteration stats and the periodic evaluations. """ import os import tinylm if sft_path is None: sft_path = os.path.join(os.path.dirname(os.path.abspath(tinylm.__file__)), "tinylm_sft.pt") policy = load_or_train_sft(sft_path) ref = load_or_train_sft(sft_path) # an identical, frozen copy for p in ref.parameters(): p.requires_grad_(False) torch.manual_seed(cfg.seed) opt = torch.optim.Adam(policy.parameters(), lr=cfg.lr) rng = random.Random(cfg.seed) hist = {k: [] for k in ["reward", "loss", "kl", "clip_frac", "frac_zero_signal", "entropy"]} hist.update(eval_iter=[], acc_sampled=[], acc_greedy=[]) def evaluate(it): hist["eval_iter"].append(it) hist["acc_sampled"].append(accuracy(policy, n=cfg.eval_n, greedy=False)) hist["acc_greedy"].append(accuracy(policy, n=cfg.eval_n, greedy=True)) t0 = time.time() evaluate(0) for it in range(1, cfg.iters + 1): prompts, answers = make_batch(cfg.prompts_per_iter, rng) s = grpo_step(policy, ref, opt, prompts, answers, cfg) for k in hist: if k in s: hist[k].append(s[k]) if it % cfg.eval_every == 0 or it == cfg.iters: evaluate(it) if verbose: print(f"iter {it:4d} reward {s['reward']:.2f} sampled acc {hist['acc_sampled'][-1]:.3f} " f"greedy acc {hist['acc_greedy'][-1]:.3f} kl {s['kl']:.4f} " f"entropy {s['entropy']:.3f} zero-signal {s['frac_zero_signal']:.2f} " f"({time.time() - t0:.0f}s)") return policy, hist