# Optional: GRPO on a real model with Hugging Face TRL

**Optional, needs a GPU, not run by this lab's tests.** This takes the algorithm you built on TinyLM and runs it on a small open model (Qwen2.5-0.5B-Instruct) on grade-school math (GSM8K), using the `GRPOTrainer` from Hugging Face TRL. A free Colab T4 can run a short version. A 24 GB GPU is more comfortable.

Checked against the TRL documentation in October 2026 (<https://huggingface.co/docs/trl/main/en/grpo_trainer>). TRL changes quickly: if something here fails, the docs page is the source of truth.

## Install

```bash
pip install -U trl datasets transformers accelerate
```

## The script

```python
# train_gsm8k_grpo.py
import re
from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer

# GSM8K answers end with "#### <number>". Keep the number as a column called "answer":
# TRL passes every dataset column except "prompt" to the reward functions as a keyword argument.
def to_example(ex):
    return {
        "prompt": [
            {"role": "system", "content": "Solve the problem step by step. "
             "Put the final answer, a single number, after 'Answer:' on the last line."},
            {"role": "user", "content": ex["question"]},
        ],
        "answer": ex["answer"].split("####")[-1].strip().replace(",", ""),
    }

train = load_dataset("openai/gsm8k", "main", split="train").map(to_example)

def extract(text):
    """The verifier: the last number after 'Answer:', or None."""
    m = re.findall(r"Answer:\s*\$?(-?[\d,]*\.?\d+)", text)
    return m[-1].replace(",", "") if m else None

def correctness_reward(completions, answer, **kwargs):
    # Conversational format: each completion is a list of messages.
    texts = [c[0]["content"] for c in completions]
    out = []
    for t, a in zip(texts, answer):
        pred = extract(t)
        try:
            out.append(1.0 if pred is not None and float(pred) == float(a) else 0.0)
        except ValueError:
            out.append(0.0)
    return out

config = GRPOConfig(
    output_dir="qwen0.5b-gsm8k-grpo",
    num_generations=8,              # G
    max_completion_length=256,
    learning_rate=1e-6,
    beta=0.04,                      # DeepSeekMath's KL coefficient (TRL's default is 0.0)
    loss_type="grpo",               # per-response length averaging, as in DeepSeekMath (default "dapo")
    scale_rewards="group",          # divide by the group std (default); "none" = Dr. GRPO style
    per_device_train_batch_size=8,  # completions per device per step; keep it a multiple of G
    max_steps=200,
    logging_steps=5,
    # TRL defaults to bf16=True, which a Colab T4 does not support. On a T4 uncomment:
    # bf16=False, fp16=True,
)

trainer = GRPOTrainer(
    model="Qwen/Qwen2.5-0.5B-Instruct",
    reward_funcs=correctness_reward,
    args=config,
    train_dataset=train,
)
trainer.train()
```

Run it with `python train_gsm8k_grpo.py` on one GPU, or `accelerate launch train_gsm8k_grpo.py` on several.

## Mapping to your lab

| Your `grpo.py` | TRL |
|---|---|
| `G` | `num_generations` |
| `beta` (0.04 in the lab) | `beta` (default 0.0: no reference model is loaded at all) |
| `clip_eps` | `epsilon` (default 0.2) |
| `inner_steps` | `num_iterations` (default 1, so the ratio is 1 and the clip is inactive) |
| `scale_std=True / False` | `scale_rewards="group"` / `"none"` |
| per-response mean, then mean over responses | `loss_type="grpo"` (the default `"dapo"` averages over all tokens instead; see Chapter 16) |
| plain `k3_kl` | set `use_bias_correction_kl=False` (the default `True` adds DeepSeek-V3.2's importance-weight correction; it only matters when `beta > 0`) |
| `frac_zero_signal` | logged as `frac_reward_zero_std` |
| `entropy`, `kl`, `clip_frac` | logged as `entropy`, `kl` (only if `beta > 0`), `clip_ratio/...` |

## What to look at

- `reward` should rise over the first hundred steps. A 0.5B instruct model already solves a good fraction of GSM8K, so expect improvement, not a transformation.
- `frac_reward_zero_std`: the chapter's zero-signal groups. Watch it grow as easy problems saturate.
- `completions/mean_length`: does reasoning get longer? Compare `loss_type="grpo"` with `"dr_grpo"`.
- Format failures early on (no `Answer:` line) earn 0. You can add a second reward function for the format, as R1-Zero did, by passing `reward_funcs=[correctness_reward, format_reward]`. Then try to break your format reward (Exercise 15.6).

To go faster on a bigger GPU, TRL can generate with vLLM (`use_vllm=True`; see the docs for `vllm_mode`).