Post-training a Small Language Model with GRPO¶

Group Relative Policy Optimization (GRPO) samples several completions for the same prompt, scores them, and uses their relative rewards to update the policy. Here we use Qwen2-0.5B-Instruct, short math prompts, LoRA, and verifiable answer rewards. The default data is reproducible elementary arithmetic so a small model can reach a final answer within the output budget; the original NuminaMath dataset remains optional.

Unlike DPO, the training responses are generated online. Unlike the PPO workflow with a value head, this GRPO example obtains a baseline from each group's rewards and does not learn a critic. A correct final answer does not certify the quality of every reasoning step.

1. Kaggle setup¶

In [1]:
import os
import sys
import json
import subprocess
import importlib
from pathlib import Path
from importlib.metadata import version, PackageNotFoundError

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["HF_HUB_DISABLE_XET"] = "1"

# These are the GRPO training dependencies only. Keep Kaggle's CUDA PyTorch.
PINS = {
    "transformers": "4.51.3", "trl": "0.18.2", "datasets": "3.5.0",
    "accelerate": "1.6.0", "peft": "0.15.2",
    "huggingface-hub": "0.36.2", "tokenizers": "0.21.4",
    "math-verify": "0.7.0", "antlr4-python3-runtime": "4.13.2",
    "latex2sympy2-extended": "1.10.1",
}
ENV_DIR = Path("/kaggle/working/.grpo_env")
SETUP_LOG = Path("/kaggle/working/grpo_setup.log")
fingerprint = {"revision": "private-env-v3-no-ensurepip", "python": sys.version.split()[0],
               "executable": sys.executable, "pins": PINS}
marker = ENV_DIR / "grpo_setup_complete.json"
ready = marker.exists() and json.loads(marker.read_text()) == fingerprint

# Imported modules cannot safely be switched to a different version in place.
guard_names = ("transformers", "trl", "datasets", "accelerate", "peft",
               "huggingface_hub", "tokenizers", "fsspec", "dill", "multiprocess",
               "math_verify", "antlr4", "latex2sympy2_extended")
loaded = [name for name in guard_names if name in sys.modules]
if loaded and not globals().get("GRPO_PRIVATE_ENV_ACTIVE", False):
    raise RuntimeError("Start a fresh Kaggle session and run this cell first. Already loaded: " + str(loaded))

def setup_command(command):
    with SETUP_LOG.open("a") as log:
        result = subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, text=True)
    if result.returncode:
        print(SETUP_LOG.read_text()[-16000:])
        raise RuntimeError(f"GRPO setup failed (exit {result.returncode}). See {SETUP_LOG}.")

if not ready:
    SETUP_LOG.parent.mkdir(parents=True, exist_ok=True)
    SETUP_LOG.write_text("GRPO private environment setup\n")
    print("Preparing the private GRPO environment; this may take a few minutes.", flush=True)
    # Kaggle may not support ensurepip. The host pip can manage a pip-free venv.
    if tuple(int(part) for part in version("pip").split(".")[:2]) < (22, 3):
        raise RuntimeError("This setup needs host pip >= 22.3 for the --python option.")
    setup_command([sys.executable, "-m", "venv", "--without-pip",
                   "--system-site-packages", str(ENV_DIR)])
    env_python = str(ENV_DIR / "bin/python")
    # Ignore unrelated shared-app conflict summaries only. Pip failures and the
    # trainer's own dependency/import checks remain visible and fail loudly.
    setup_command([sys.executable, "-m", "pip", "--python", env_python,
                   "install", "--disable-pip-version-check",
                   "--no-warn-conflicts", *[f"{k}=={v}" for k, v in PINS.items()]])
else:
    env_python = str(ENV_DIR / "bin/python")

env_site = subprocess.check_output(
    [env_python, "-c", "import sysconfig; print(sysconfig.get_paths()['purelib'])"], text=True
).strip()
if env_site in sys.path:
    sys.path.remove(env_site)
sys.path.insert(0, env_site)
importlib.invalidate_caches()

def installed_version(name):
    try:
        return version(name)
    except PackageNotFoundError:
        return None

actual = {k: installed_version(k) for k in PINS}
assert actual == PINS, f"Unexpected GRPO package versions: {actual}"

# This text-only notebook does not need optional vision/audio integrations.
# Blocking these optional imports locally avoids binary mismatches without
# removing torchvision, torchaudio, or librosa from the shared environment.
for optional_name in ("torchvision", "torchaudio", "librosa"):
    if optional_name not in sys.modules:
        sys.modules[optional_name] = None

from trl import GRPOConfig, GRPOTrainer
from math_verify import parse, verify, LatexExtractionConfig
assert verify(parse(r"\boxed{5}"), parse(r"\boxed{5}")), "Math verifier check failed."
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import transformers
import datasets
import accelerate
import peft

GRPO_PRIVATE_ENV_ACTIVE = True
marker.write_text(json.dumps(fingerprint, indent=2))
print(actual)
print("PyTorch:", torch.__version__)
print("Transformers loaded from:", transformers.__file__)
print("GRPO setup checks passed. Continue to the next cell.")
print("Installation log:", SETUP_LOG)

Preparing the private GRPO environment; this may take a few minutes.
{'transformers': '4.51.3', 'trl': '0.18.2', 'datasets': '3.5.0', 'accelerate': '1.6.0', 'peft': '0.15.2', 'huggingface-hub': '0.36.2', 'tokenizers': '0.21.4', 'math-verify': '0.7.0', 'antlr4-python3-runtime': '4.13.2', 'latex2sympy2-extended': '1.10.1'}
PyTorch: 2.11.0+cu128
Transformers loaded from: /kaggle/working/.grpo_env/lib/python3.13/site-packages/transformers/__init__.py
GRPO setup checks passed. Continue to the next cell.
Installation log: /kaggle/working/grpo_setup.log
In [2]:
import json
import math
import random
from pathlib import Path
import torch
from datasets import Dataset, load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed

SEED = 42
set_seed(SEED)
assert torch.cuda.is_available(), "Enable a Kaggle GPU accelerator, then start a fresh session."
USE_BF16 = torch.cuda.is_bf16_supported()
USE_FP16 = not USE_BF16
print("PyTorch:", torch.__version__)
print("GPU:", torch.cuda.get_device_name(0))
print("Visible GPUs:", torch.cuda.device_count())
print("Training mixed precision:", "bf16" if USE_BF16 else "fp16")

PyTorch: 2.11.0+cu128
GPU: Tesla T4
Visible GPUs: 1
Training mixed precision: bf16

2. Load the starting policy and tokenizer¶

We use a 0.5B instruction model to keep the demonstration small. LoRA trains a small number of adapter parameters. The frozen base with adapters disabled serves as the reference policy for the KL penalty.

In [3]:
from peft import LoraConfig
from trl import GRPOConfig, GRPOTrainer
from math_verify import LatexExtractionConfig, ExprExtractionConfig, parse, verify
from transformers import AutoConfig, GenerationConfig

MODEL_ID = "Qwen/Qwen2-0.5B-Instruct"
DATASET_ID = "AI-MO/NuminaMath-TIR"
DATA_SOURCE = "arithmetic"  # Optional: "numinamath" for the original, harder problems.
N_TRAIN = 128
N_EVAL = 12
MAX_STEPS = 64
GROUP_SIZE = 4
MAX_PROMPT_LENGTH = 256
MAX_COMPLETION_LENGTH = 192
OUTPUT_DIR = Path("/kaggle/working/grpo_qwen_optimized")
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

# Screen training prompts only. Held-out questions never enter this selection.
SCREEN_MAX_PROMPTS = 24
TARGET_MIXED_PROMPTS = 8
FORMAT_WEIGHT = 0.1
SAVE_RUNTIME_ARTIFACTS = False  # True saves to this runtime; download separately to keep them.

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "left"
model_config = AutoConfig.from_pretrained(MODEL_ID)
# Our short contexts use full causal attention. Do not advertise an inactive
# sliding window to the SDPA implementation in this Transformers version.
model_config.use_sliding_window = False
model_config.sliding_window = None
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, config=model_config, torch_dtype=torch.float32, attn_implementation="sdpa",
).to("cuda")
model.config.pad_token_id = tokenizer.pad_token_id
model.config.use_cache = False
model.generation_config = GenerationConfig(
    bos_token_id=tokenizer.bos_token_id, eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
)

tokenizer_config.json: 0.00B [00:00, ?B/s]
vocab.json: 0.00B [00:00, ?B/s]
merges.txt: 0.00B [00:00, ?B/s]
tokenizer.json: 0.00B [00:00, ?B/s]
config.json:   0%|          | 0.00/659 [00:00<?, ?B/s]
model.safetensors:   0%|          | 0.00/988M [00:00<?, ?B/s]
generation_config.json:   0%|          | 0.00/242 [00:00<?, ?B/s]

3. Prepare math prompts and verifiable targets¶

The default dataset contains 128 mixed-difficulty arithmetic candidate training questions and 12 fixed held-out questions. No model generates the labels: ordinary integer arithmetic supplies the reference answers. Each row has prompt and solution; only the prompt is passed to the policy.

Set DATA_SOURCE="numinamath" to stream the original math dataset instead. We retain only short prompts with an extractable, parseable final reference answer. Short competition-math questions can still be too hard for a 0.5B model. We do not force incorrect responses or manufacture a positive training result.

The training pool mixes carrying/borrowing, signed subtraction, two-digit multiplication, two-step expressions, and short word problems. The bounded pre-training screening pass prefers prompts where actual sampled completions include both correct and incorrect answers. This is an explicit training-data selection step; it does not fabricate responses and never screens the held-out questions. The smaller selected set is reused during training and can overfit.

Two fixed worked examples encourage brief calculations and signed answers. They are included in both the before and after prompts; their prompting effect must not be attributed to GRPO training.

In [4]:
import re

SYSTEM_PROMPT = (
    "Solve the problem with at most two short arithmetic steps inside <think>...</think>. "
    "Put only the final number inside <answer>...</answer>. Preserve minus signs."
)

def last_boxed(text):
    # Balanced braces support nested answers such as boxed fractions.
    marker = chr(92) + "boxed{"
    start = text.rfind(marker)
    if start < 0:
        return None
    content_start = start + len(marker)
    depth = 1
    for i in range(content_start, len(text)):
        depth += int(text[i] == "{") - int(text[i] == "}")
        if depth == 0:
            return marker + text[content_start:i] + "}"
    return None

def normalize_answer_span(span):
    span = span.replace(chr(0x2212), "-")
    boxed = last_boxed(span)
    if boxed:
        return boxed
    # Do not interpret the left side of an equation as the final answer.
    span = re.sub(r"<[^>]*>", " ", span)
    span = " ".join(span.split()).strip()
    tail = span.rsplit("=", 1)[-1].strip()
    for prefix in ("The final answer is ", "The answer is ", "Answer: ", "Final answer: "):
        if tail.lower().startswith(prefix.lower()):
            tail = tail[len(prefix):].strip()
            break
    tail = tail.rstrip(" .")
    # Allow units after a declared numeric answer, e.g. '275 pencils'.
    number = r"([-+]?(?:[0-9][0-9,]*(?:[.][0-9]+)?|[.][0-9]+)(?:/[-+]?[0-9]+)?)"
    match = re.fullmatch(number + r"(?: +[A-Za-z][A-Za-z ]*)?", tail)
    return match.group(1).replace(",", "") if match else None

def extract_answer(text):
    text = text.replace(chr(0x2212), "-")
    # Use the final declared answer region, not the first intermediate answer.
    closing = text.rfind("</answer>")
    if closing >= 0:
        opening = text.rfind("<answer>", 0, closing)
        if opening >= 0:
            return normalize_answer_span(text[opening + len("<answer>"):closing])
    boxed = last_boxed(text)
    if boxed:
        return boxed
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    if lines and re.fullmatch(r"[-+]?[0-9]+(?:[.][0-9]+)?(?:/[0-9]+)?", lines[-1]):
        return lines[-1]
    if lines and re.match(r"(?i)(?:final answer|answer|result)[ ]*(?:is|:|=)", lines[-1]):
        return normalize_answer_span(lines[-1])
    return None

def make_conversation(example):
    return {
        "prompt": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "What is (12 + 7) * 3?"},
            {"role": "assistant", "content": "<think>12 + 7 = 19; 19 * 3 = 57.</think><answer>57</answer>"},
            {"role": "user", "content": "What is 14 - 23?"},
            {"role": "assistant", "content": "<think>23 - 14 = 9; 14 is smaller, so the result is -9.</think><answer>-9</answer>"},
            {"role": "user", "content": example["problem"]},
        ],
        "solution": extract_answer(example["solution"]),
    }

def parse_math(text):
    # Disable extraction fallback: an unparsed string must not count as a valid answer.
    return parse(text, extraction_config=[LatexExtractionConfig(), ExprExtractionConfig()],
                 extraction_mode="first_match", fallback_mode="no_fallback")

def load_short_math(split, count):
    rows = []
    seen = set()
    stream = load_dataset(DATASET_ID, split=split, streaming=True)
    for index, example in enumerate(stream):
        row = make_conversation(example)
        prompt_text = tokenizer.apply_chat_template(row["prompt"], tokenize=False, add_generation_prompt=True)
        problem = example["problem"].strip()
        if len(tokenizer.encode(prompt_text)) <= MAX_PROMPT_LENGTH and problem not in seen:
            try:
                valid_gold = bool(row["solution"] and parse_math(row["solution"]))
            except Exception:
                valid_gold = False
            if valid_gold:
                rows.append(row)
                seen.add(problem)
        if len(rows) >= count or index >= 4999:
            break
    assert len(rows) == count, f"Only {len(rows)} eligible rows in {split}; increase the scan or context limit."
    return Dataset.from_list(rows)

def arithmetic_data():
    rng = random.Random(SEED)
    # Fixed BEFORE/AFTER questions. Do not select them according to model outputs.
    held_out = [
        ("What is 27 * 34?", 27 * 34),
        ("What is 387 + 568?", 387 + 568),
        ("What is (17 + 26) * 8?", (17 + 26) * 8),
        ("What is 27 - 64?", 27 - 64),
        ("What is 144 / 12 + 7 * 3?", 144 // 12 + 7 * 3),
        ("A shop has 18 boxes of 24 pencils. It sells 157 pencils. How many remain?", 18 * 24 - 157),
        ("What is 46 * 17?", 46 * 17),
        ("What is 623 - 278?", 623 - 278),
        ("What is 39 - 82?", 39 - 82),
        ("What is 23 * 7 - 48?", 23 * 7 - 48),
        ("What is (63 - 28) * 6?", (63 - 28) * 6),
        ("There are 14 bags of 18 marbles. Give away 95 marbles. How many remain?", 14 * 18 - 95),
    ]
    seen = {question for question, _ in held_out} | {"What is (12 + 7) * 3?", "What is 14 - 23?"}
    rows = []
    kinds = ["add", "subtract", "negative", "multiply", "two_step", "word"]
    while len(rows) < N_TRAIN:
        kind = kinds[len(rows) % len(kinds)]
        if kind == "add":
            a, b = rng.randint(80, 700), rng.randint(80, 700)
            question, answer = f"What is {a} + {b}?", a + b
        elif kind == "subtract":
            a, b = rng.randint(200, 900), rng.randint(100, 700)
            question, answer = f"What is {a} - {b}?", a - b
        elif kind == "negative":
            a, b = rng.randint(10, 49), rng.randint(50, 99)
            question, answer = f"What is {a} - {b}?", a - b
        elif kind == "multiply":
            a, b = rng.randint(12, 39), rng.randint(11, 35)
            question, answer = f"What is {a} * {b}?", a * b
        elif kind == "two_step":
            a, b, d = rng.randint(10, 39), rng.randint(10, 39), rng.randint(3, 12)
            question, answer = f"What is ({a} + {b}) * {d}?", (a + b) * d
        else:
            a, b, sold = rng.randint(8, 24), rng.randint(12, 29), rng.randint(20, 90)
            question = f"A shop has {a} boxes of {b} pencils. It sells {sold} pencils. How many remain?"
            answer = a * b - sold
        if question not in seen:
            seen.add(question)
            row = make_conversation({"problem": question, "solution": str(answer)})
            row["task_type"] = kind
            rows.append(row)
    # Interleave task types so the bounded screening pass sees a mixture.
    test = [make_conversation({"problem": question, "solution": str(answer)})
            for question, answer in held_out[:N_EVAL]]
    return Dataset.from_list(rows), Dataset.from_list(test)

if DATA_SOURCE == "arithmetic":
    train_dataset, test_dataset = arithmetic_data()
elif DATA_SOURCE == "numinamath":
    train_dataset = load_short_math("train", N_TRAIN)
    test_dataset = load_short_math("test", N_EVAL)
else:
    raise ValueError("DATA_SOURCE must be arithmetic or numinamath.")
train_questions = {r["prompt"][-1]["content"].strip() for r in train_dataset}
test_dataset = test_dataset.filter(lambda r: r["prompt"][-1]["content"].strip() not in train_questions)
assert len(test_dataset) > 0, "No disjoint evaluation questions remain."
for row in list(train_dataset) + list(test_dataset):
    text = tokenizer.apply_chat_template(row["prompt"], tokenize=False, add_generation_prompt=True)
    assert len(tokenizer.encode(text)) <= MAX_PROMPT_LENGTH, "Increase MAX_PROMPT_LENGTH; do not truncate the question."
print(train_dataset)
print(json.dumps(train_dataset[0], indent=2, ensure_ascii=False))

Filter:   0%|          | 0/12 [00:00<?, ? examples/s]
Dataset({
    features: ['prompt', 'solution', 'task_type'],
    num_rows: 128
})
{
  "prompt": [
    {
      "content": "Solve the problem with at most two short arithmetic steps inside <think>...</think>. Put only the final number inside <answer>...</answer>. Preserve minus signs.",
      "role": "system"
    },
    {
      "content": "What is (12 + 7) * 3?",
      "role": "user"
    },
    {
      "content": "<think>12 + 7 = 19; 19 * 3 = 57.</think><answer>57</answer>",
      "role": "assistant"
    },
    {
      "content": "What is 14 - 23?",
      "role": "user"
    },
    {
      "content": "<think>23 - 14 = 9; 14 is smaller, so the result is -9.</think><answer>-9</answer>",
      "role": "assistant"
    },
    {
      "content": "What is 194 + 105?",
      "role": "user"
    }
  ],
  "solution": "299",
  "task_type": "add"
}

4. Define and check the reward functions¶

We keep two distinct signals:

  1. Format reward: up to 1, with 0.25 each for opening <think>, closing it with nonempty content, opening <answer>, and closing it with nonempty content. Partial credit makes this signal less sparse. It rewards structure, not the truth of an explanation.
  2. Correctness reward: 1 when an extractable final answer matches the reference, otherwise 0. The final declared <answer> region, the last balanced LaTeX box, or a bare number on the final line can supply the answer. Numeric equations are scored using the rightmost result, not an intermediate left-hand expression; units after a numeric answer are allowed. Correctness does not depend on earning the format reward. Unparseable references or answers receive 0.

The combined reward is $0.1r_{format}+r_{correct}$. Correctness is the larger component. We verify only the final answer, not every statement in the explanation. Unit checks below use hand-written examples and are not model generations.

In [5]:
import re

FORMAT_PATTERN = re.compile(
    r"<think>(?P<think>.*?)</think>\s*<answer>(?P<answer>.*?)</answer>",
    flags=re.DOTALL,
)

def completion_text(completion):
    if isinstance(completion, str):
        return completion
    return "".join(message.get("content", "") for message in completion)

def format_reward(completions, **kwargs):
    rewards = []
    for completion in completions:
        text = completion_text(completion)
        think = re.search(r"<think>(.*?)</think>", text, flags=re.DOTALL)
        answer = re.search(r"<answer>(.*?)</answer>", text, flags=re.DOTALL)
        parts = ["<think>" in text, bool(think and think.group(1).strip()),
                 "<answer>" in text, bool(answer and answer.group(1).strip())]
        score = sum(parts) / 4.0
        full = FORMAT_PATTERN.fullmatch(text.strip())
        tags_once = all(text.count(tag) == 1 for tag in ("<think>", "</think>", "<answer>", "</answer>"))
        if score == 1.0 and not (full and tags_once):
            score = 0.75  # Out-of-order tags or extra text are only partial format compliance.
        rewards.append(score)
    return rewards

def accuracy_reward(completions, solution, **kwargs):
    if len(completions) != len(solution):
        raise ValueError("Each generated completion must have one aligned reference solution.")
    rewards = []
    for completion, gold_text in zip(completions, solution):
        answer = extract_answer(completion_text(completion))
        try:
            gold = parse_math(gold_text)
            predicted = parse_math(answer) if answer else []
            # math_verify expects the reference (gold) first, then the prediction.
            rewards.append(float(bool(gold and predicted) and verify(gold, predicted)))
        except Exception:
            rewards.append(0.0)
    return rewards

reward_examples = [
    [{"role": "assistant", "content": "<think>Two plus three.\nThe sum is five.</think><answer>\\boxed{5}</answer>"}],
    [{"role": "assistant", "content": "<think>Adding.</think><answer>\\boxed{4}</answer>"}],
    [{"role": "assistant", "content": "5"}],
    [{"role": "assistant", "content": "<think>Adding.</think><answer>nonsense</answer>"}],
]
gold_examples = [r"\boxed{5}"] * len(reward_examples)
assert format_reward(reward_examples) == [1.0, 1.0, 0.0, 1.0]
assert accuracy_reward(reward_examples, gold_examples) == [1.0, 0.0, 1.0, 0.0]
assert format_reward(["<think>Starting a calculation"]) == [0.25]
assert format_reward(["<answer>5</answer><think>Add.</think>"]) == [0.75]
assert accuracy_reward(["An intermediate value is 5, but the final answer is unknown."], ["5"]) == [0.0]
assert accuracy_reward(["Working." + chr(10) + "5"], ["5"]) == [1.0]
assert accuracy_reward(["<answer>4</answer>"], ["5"]) == [0.0]
assert accuracy_reward(reward_examples[:1], ["not a parseable gold answer"]) == [0.0]
fraction_example = [[{"role": "assistant", "content": r"<think>Equivalent fractions.</think><answer>\boxed{\frac{2}{4}}</answer>"}]]
assert accuracy_reward(fraction_example, [r"\boxed{\frac{1}{2}}"]) == [1.0]
print("Reward checks passed.")


def describe_reward(text, reference):
    candidate = extract_answer(text)
    gold = parse_math(reference)
    expected = str(gold[0]) if gold else str(reference)
    correct = accuracy_reward([text], [reference])[0]
    if candidate is None:
        status = "NO FINAL ANSWER EXTRACTED"
    else:
        try:
            parsed = parse_math(candidate)
        except Exception:
            parsed = []
        status = "CORRECT" if correct else ("WRONG ANSWER" if parsed else "UNPARSEABLE ANSWER")
    return {"extracted": candidate, "expected": expected, "status": status,
            "format_reward": format_reward([text])[0], "correctness_reward": correct}

# These are verifier regression cases, not generated model responses.
assert extract_answer("<answer>(17 + 26) * 8 = 408</answer>") == "408"
assert accuracy_reward(["<answer>(17 + 26) * 8 = 408</answer>"], ["344"]) == [0.0]
assert accuracy_reward(["<answer>27 - 64 = -37</answer>"], ["-37"]) == [1.0]
pencil_case = "<think>Calculate.</think><answer>18 * 24 = 432</answer> After selling 157, 432 - 157 = 275 pencils.</answer>"
assert extract_answer(pencil_case) == "275"
assert accuracy_reward([pencil_case], ["275"]) == [1.0]
assert format_reward([pencil_case])[0] < 1.0
assert extract_answer("<answer>1,000 pencils</answer>") == "1000"
assert accuracy_reward(["<answer>" + chr(0x2212) + "37</answer>"], ["-37"]) == [1.0]
assert describe_reward("No final answer.", "5")["status"] == "NO FINAL ANSWER EXTRACTED"
assert describe_reward("<answer>4</answer>", "5")["status"] == "WRONG ANSWER"
assert describe_reward("<answer>5</answer>", "5")["status"] == "CORRECT"
print("Final-answer extraction and status checks passed.")

Reward checks passed.
Final-answer extraction and status checks passed.

5. From group rewards to a learning signal¶

For one prompt, sample $G$ completions from the current policy. Each gets a combined reward $r_i$. The group-relative advantage is

$$A_i = \frac{r_i-\operatorname{mean}(r_1,\ldots,r_G)}{\operatorname{std}(r_1,\ldots,r_G)+\epsilon}.$$

The trainer uses this advantage in a clipped policy-ratio objective, with a KL penalty to the starting reference. No learned value model is needed. We explicitly select the original grpo loss aggregation and group reward scaling for this teaching example.

When all rewards in a group are equal, this relative-reward signal is zero. The next cell uses illustrative rewards to make that visible. KL regularization is a separate term.

In [6]:
for rewards in [torch.tensor([1.2, 0.2, 1.2, 0.0]), torch.zeros(GROUP_SIZE)]:
    advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-4)
    print("Rewards:", rewards.tolist())
    print("Advantages:", [round(a, 3) for a in advantages.tolist()])

Rewards: [1.2000000476837158, 0.20000000298023224, 1.2000000476837158, 0.0]
Advantages: [0.859, -0.703, 0.859, -1.015]
Rewards: [0.0, 0.0, 0.0, 0.0]
Advantages: [0.0, 0.0, 0.0, 0.0]
In [7]:
from contextlib import nullcontext
from transformers import GenerationConfig

@torch.inference_mode()
def generate_response(model, tokenizer, messages, max_new_tokens=192):
    model.eval()
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    device = next(model.parameters()).device
    inputs = tokenizer(text, return_tensors="pt", return_token_type_ids=False).to(device)
    # A fresh config prevents Qwen's sampling defaults leaking into greedy decoding.
    config = GenerationConfig(
        do_sample=False, max_new_tokens=max_new_tokens, use_cache=True,
        bos_token_id=tokenizer.bos_token_id, eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
    )
    precision = (torch.autocast("cuda", dtype=torch.bfloat16 if USE_BF16 else torch.float16)
                 if device.type == "cuda" else nullcontext())
    with precision:
        output = model.generate(**inputs, generation_config=config, use_model_defaults=False)
    return tokenizer.decode(output[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)

@torch.inference_mode()
def sample_group(model, tokenizer, messages, group_size=4):
    model.eval()
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    device = next(model.parameters()).device
    inputs = tokenizer(text, return_tensors="pt", return_token_type_ids=False).to(device)
    config = GenerationConfig(
        do_sample=True, temperature=0.9, top_p=0.95, top_k=50,
        num_return_sequences=group_size, max_new_tokens=MAX_COMPLETION_LENGTH,
        use_cache=True, bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id,
    )
    precision = (torch.autocast("cuda", dtype=torch.bfloat16 if USE_BF16 else torch.float16)
                 if device.type == "cuda" else nullcontext())
    with precision:
        output = model.generate(**inputs, generation_config=config, use_model_defaults=False)
    return tokenizer.batch_decode(output[:, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
In [8]:
# Store the fixed held-out outputs before any optimization.
eval_rows = list(test_dataset)
baseline_questions = [row["prompt"][-1]["content"] for row in eval_rows]
before_responses = [generate_response(model, tokenizer, row["prompt"], MAX_COMPLETION_LENGTH)
                    for row in eval_rows]
print(f"Stored {len(eval_rows)} fixed held-out responses.")

# This is training-data screening, not part of the GRPO loss itself.
# Only the training pool is screened. No held-out question is used here.
assert SCREEN_MAX_PROMPTS >= 1 and TARGET_MIXED_PROMPTS >= 1 and GROUP_SIZE >= 2
candidate_train_dataset = train_dataset
mixed_indices, varying_indices, screen_records = [], [], []
for index, row in enumerate(candidate_train_dataset):
    if index >= SCREEN_MAX_PROMPTS or len(mixed_indices) >= TARGET_MIXED_PROMPTS:
        break
    set_seed(SEED + 1000 + index)
    texts = sample_group(model, tokenizer, row["prompt"], GROUP_SIZE)
    completions = [[{"role": "assistant", "content": output}] for output in texts]
    correctness = accuracy_reward(completions, [row["solution"]] * GROUP_SIZE)
    formats = format_reward(completions)
    rewards = torch.tensor([FORMAT_WEIGHT * f + a for f, a in zip(formats, correctness)])
    mixed = 0 < sum(correctness) < GROUP_SIZE
    if mixed:
        mixed_indices.append(index)
    if rewards.std().item() > 0:
        varying_indices.append(index)
    screen_records.append({"index": index, "question": row["prompt"][-1]["content"],
                           "solution": row["solution"], "outputs": texts,
                           "correctness": correctness, "format": formats,
                           "rewards": rewards.tolist(), "mixed_correctness": mixed})
    print(f"Probe {index + 1}: {sum(correctness):.0f}/{GROUP_SIZE} correct | mixed={mixed}", flush=True)

if mixed_indices:
    train_dataset = candidate_train_dataset.select(mixed_indices)
    print(f"Training on {len(train_dataset)} prompts with observed mixed correctness rewards.")
elif varying_indices:
    train_dataset = candidate_train_dataset.select(varying_indices)
    print("No mixed-correctness group was observed. Selected groups vary only in format/combined reward.")
else:
    print("No varying-reward group was observed in the bounded screening pass. Keeping the original training pool; inspect the training diagnostics before interpreting the run.")

# Show a complete actual group, including why each reward was assigned.
record = next((r for r in screen_records if r["mixed_correctness"]), screen_records[0])
rewards = torch.tensor(record["rewards"])
advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-4)
print("\nTRAINING GROUP:", record["question"])
for i, output in enumerate(record["outputs"]):
    details = describe_reward(output, record["solution"])
    print(f"\nCompletion {i + 1}:\n{output}")
    print(f"Extracted={details['extracted']} | expected={details['expected']} | {details['status']}")
    print(f"reward={rewards[i]:.2f} | advantage={advantages[i]:+.3f}")
(OUTPUT_DIR / "screening_groups.json").write_text(json.dumps(screen_records, indent=2))
set_seed(SEED)

Stored 12 fixed held-out responses.
Probe 1: 0/4 correct | mixed=False
Probe 2: 0/4 correct | mixed=False
Probe 3: 3/4 correct | mixed=True
Probe 4: 3/4 correct | mixed=True
Probe 5: 0/4 correct | mixed=False
Probe 6: 0/4 correct | mixed=False
Probe 7: 1/4 correct | mixed=True
Probe 8: 1/4 correct | mixed=True
Probe 9: 3/4 correct | mixed=True
Probe 10: 1/4 correct | mixed=True
Probe 11: 2/4 correct | mixed=True
Probe 12: 1/4 correct | mixed=True
Training on 8 prompts with observed mixed correctness rewards.

TRAINING GROUP: What is 24 - 58?

Completion 1:
<think>58 - 24 = 34; 34 is smaller, so the result is -34.</think><answer>-34</answer>
Extracted=-34 | expected=-34 | CORRECT
reward=1.10 | advantage=+0.500

Completion 2:
<think>58 - 24 = 34; 24 is smaller, so the result is -34.</think><answer>-34</answer>
Extracted=-34 | expected=-34 | CORRECT
reward=1.10 | advantage=+0.500

Completion 3:
<think>58 - 24 = 34; 34 is larger, so the result is 34.</think><answer>34</answer>
Extracted=34 | expected=-34 | WRONG ANSWER
reward=0.10 | advantage=-1.500

Completion 4:
<think>58 - 24 = 34; 34 is smaller, so the result is -34.</think><answer>-34</answer>
Extracted=-34 | expected=-34 | CORRECT
reward=1.10 | advantage=+0.500

6. Configure LoRA and GRPO training¶

We sample groups of four completions, use one GPU, and perform 64 optimizer steps. The generation batch contains a complete group. Two update iterations reuse each sampled group; the clipped ratio limits changes relative to the policy that sampled it.

The KL coefficient is set explicitly. The PEFT policy uses its frozen base, with adapters disabled, as the reference. No separately trained reward model is loaded: the Python reward functions supply the scores.

Training may be slower than DPO because it generates responses online. Watch the reward components, group reward variation, completion length, and KL rather than expecting a visible reasoning breakthrough from this small run.

The printed scalar loss can be near zero even with nonzero gradients: group advantages are centered, so positive and negative terms may cancel. Equal rewards for every completion are a different case: they produce zero relative-reward advantage. We inspect reward variation, gradient norms, and actual adapter parameter changes rather than judging by loss alone. Gradient checkpointing is unnecessary for this short 0.5B LoRA run and is disabled to avoid conflicting with rollout caching.

Compared with the minimal version, this run uses 64 optimizer steps, a 2e-5 learning rate, rank-16 LoRA on all attention projections, and a 0.02 KL coefficient. These are proposed teaching settings, not a proven optimal configuration. Two microbatches of two completions make one four-completion group per optimizer step. Screening adds generation work before training.

In [9]:
lora_config = LoraConfig(
    task_type="CAUSAL_LM", r=16, lora_alpha=32,
    lora_dropout=0.0, bias="none", target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
)
training_args = GRPOConfig(
    output_dir=str(OUTPUT_DIR),
    per_device_train_batch_size=2,
    gradient_accumulation_steps=2,
    generation_batch_size=GROUP_SIZE,
    num_generations=GROUP_SIZE, num_iterations=2,
    max_steps=MAX_STEPS, learning_rate=2e-5, warmup_steps=4,
    lr_scheduler_type="constant", optim="adamw_torch", max_grad_norm=1.0,
    max_prompt_length=MAX_PROMPT_LENGTH,
    max_completion_length=MAX_COMPLETION_LENGTH,
    temperature=0.9, top_p=0.95,
    beta=0.02, epsilon=0.2, loss_type="grpo", scale_rewards=True,
    reward_weights=[FORMAT_WEIGHT, 1.0],
    gradient_checkpointing=False,
    bf16=USE_BF16, fp16=USE_FP16,
    use_vllm=False, remove_unused_columns=False,
    eval_strategy="no", save_strategy="no", logging_steps=1,
    report_to="none", push_to_hub=False, seed=SEED,
    dataloader_num_workers=0, disable_dropout=True, label_names=[],
)
assert training_args.generation_batch_size % GROUP_SIZE == 0
trainer = GRPOTrainer(
    model=model, processing_class=tokenizer, args=training_args,
    reward_funcs=[format_reward, accuracy_reward],
    train_dataset=train_dataset, peft_config=lora_config,
)
assert trainer.args.n_gpu == 1, "This teaching configuration expects one visible GPU."
trainer.model.print_trainable_parameters()
adapter_before = {name: p.detach().float().cpu().clone()
                  for name, p in trainer.model.named_parameters() if p.requires_grad}
train_result = trainer.train()
assert math.isfinite(train_result.training_loss), "Training returned a non-finite loss."
adapter_delta_sq = sum((p.detach().float().cpu() - adapter_before[name]).square().sum().item()
                       for name, p in trainer.model.named_parameters() if name in adapter_before)
print("Adapter parameter L2 change:", math.sqrt(adapter_delta_sq))
log_rows = [row for row in trainer.state.log_history if "loss" in row]
for row in log_rows:
    print({key: row[key] for key in ("step", "loss", "grad_norm", "reward_std", "frac_reward_zero_std", "kl") if key in row})
if adapter_delta_sq == 0:
    print("No adapter update was measured. Check reward variation and gradient norms; do not present this as learned improvement.")

trainable params: 2,162,688 || all params: 496,195,456 || trainable%: 0.4359
[64/64 02:46, Epoch 4/4]
Step Training Loss
1 0.000000
2 -0.002600
3 0.000000
4 -0.001200
5 0.000000
6 0.000100
7 0.000000
8 0.000000
9 0.000000
10 -0.003900
11 0.000100
12 -0.002400
13 0.000000
14 0.000000
15 0.000100
16 -0.004000
17 0.000000
18 -0.001700
19 0.000100
20 -0.003500
21 0.000100
22 -0.000700
23 0.000100
24 0.000100
25 0.000000
26 -0.000500
27 0.000200
28 0.000300
29 0.000000
30 -0.000100
31 0.000100
32 -0.000500
33 0.000100
34 -0.003500
35 0.000100
36 -0.002300
37 0.000100
38 0.000100
39 0.000100
40 -0.000500
41 0.000300
42 -0.001300
43 0.000100
44 0.000100
45 0.000200
46 0.000500
47 0.000200
48 0.000200
49 0.000300
50 -0.002800
51 0.000100
52 0.000100
53 0.000300
54 -0.000200
55 0.000300
56 0.000300
57 0.000300
58 0.000300
59 0.000200
60 0.000200
61 0.000400
62 0.000400
63 0.000100
64 -0.005000 

Adapter parameter L2 change: 0.4325107437525202
{'step': 1, 'loss': 0.0, 'grad_norm': 1.0351554155349731, 'reward_std': 0.5, 'frac_reward_zero_std': 0.0, 'kl': 0.0}
{'step': 2, 'loss': -0.0026, 'grad_norm': 0.9860057830810547, 'kl': 0.000444301258539781}
{'step': 3, 'loss': 0.0, 'grad_norm': 0.699398398399353, 'reward_std': 0.5773502588272095, 'frac_reward_zero_std': 0.0, 'kl': 0.0001506053413322661}
{'step': 4, 'loss': -0.0012, 'grad_norm': 0.6876218318939209, 'kl': 0.00010133934847544879}
{'step': 5, 'loss': 0.0, 'grad_norm': 0.5320498943328857, 'reward_std': 0.5, 'frac_reward_zero_std': 0.0, 'kl': 0.0003567295061657205}
{'step': 6, 'loss': 0.0001, 'grad_norm': 0.5782124996185303, 'kl': 9.250781386072049e-05}
{'step': 7, 'loss': 0.0, 'grad_norm': 0.0038216696120798588, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.000678150710882619}
{'step': 8, 'loss': 0.0, 'grad_norm': 0.003442243207246065, 'kl': 0.0006137032178230584}
{'step': 9, 'loss': 0.0, 'grad_norm': 1.585867166519165, 'reward_std': 0.5, 'frac_reward_zero_std': 0.0, 'kl': 0.0018451823270879686}
{'step': 10, 'loss': -0.0039, 'grad_norm': 1.561866044998169, 'kl': 0.0023824720410630107}
{'step': 11, 'loss': 0.0001, 'grad_norm': 1.0387073755264282, 'reward_std': 0.5, 'frac_reward_zero_std': 0.0, 'kl': 0.002675531664863229}
{'step': 12, 'loss': -0.0024, 'grad_norm': 1.0646030902862549, 'kl': 0.003047882695682347}
{'step': 13, 'loss': 0.0, 'grad_norm': 0.004058769438415766, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.0012033652456011623}
{'step': 14, 'loss': 0.0, 'grad_norm': 0.0040434543043375015, 'kl': 0.0012179385812487453}
{'step': 15, 'loss': 0.0001, 'grad_norm': 1.0272185802459717, 'reward_std': 0.5, 'frac_reward_zero_std': 0.0, 'kl': 0.0047800319734960794}
{'step': 16, 'loss': -0.004, 'grad_norm': 0.9862939119338989, 'kl': 0.006136302719824016}
{'step': 17, 'loss': 0.0, 'grad_norm': 1.1033082008361816, 'reward_std': 0.5773502588272095, 'frac_reward_zero_std': 0.0, 'kl': 0.0011763959046220407}
{'step': 18, 'loss': -0.0017, 'grad_norm': 1.0587873458862305, 'kl': 0.0017763163923518732}
{'step': 19, 'loss': 0.0001, 'grad_norm': 1.358246922492981, 'reward_std': 0.5773502588272095, 'frac_reward_zero_std': 0.0, 'kl': 0.003959024092182517}
{'step': 20, 'loss': -0.0035, 'grad_norm': 0.5897223949432373, 'kl': 0.0034150180872529745}
{'step': 21, 'loss': 0.0001, 'grad_norm': 2.217958450317383, 'reward_std': 0.550000011920929, 'frac_reward_zero_std': 0.0, 'kl': 0.0027764136902987957}
{'step': 22, 'loss': -0.0007, 'grad_norm': 2.198162078857422, 'kl': 0.0026544641004875302}
{'step': 23, 'loss': 0.0001, 'grad_norm': 0.003966786433011293, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.0025462907506152987}
{'step': 24, 'loss': 0.0001, 'grad_norm': 0.003906985279172659, 'kl': 0.002560557797551155}
{'step': 25, 'loss': 0.0, 'grad_norm': 0.5425760746002197, 'reward_std': 0.5773502588272095, 'frac_reward_zero_std': 0.0, 'kl': 0.0018025299068540335}
{'step': 26, 'loss': -0.0005, 'grad_norm': 0.5642411708831787, 'kl': 0.0021219695918262005}
{'step': 27, 'loss': 0.0002, 'grad_norm': 0.740838885307312, 'reward_std': 0.5773502588272095, 'frac_reward_zero_std': 0.0, 'kl': 0.009380503091961145}
{'step': 28, 'loss': 0.0003, 'grad_norm': 0.7368197441101074, 'kl': 0.009383933618664742}
{'step': 29, 'loss': 0.0, 'grad_norm': 1.1298210620880127, 'reward_std': 0.5, 'frac_reward_zero_std': 0.0, 'kl': 0.00177583412732929}
{'step': 30, 'loss': -0.0001, 'grad_norm': 1.1210988759994507, 'kl': 0.0016945380484685302}
{'step': 31, 'loss': 0.0001, 'grad_norm': 0.9170358777046204, 'reward_std': 0.5773502588272095, 'frac_reward_zero_std': 0.0, 'kl': 0.005144597264006734}
{'step': 32, 'loss': -0.0005, 'grad_norm': 0.984567403793335, 'kl': 0.005621601128950715}
{'step': 33, 'loss': 0.0001, 'grad_norm': 2.0503695011138916, 'reward_std': 0.550000011920929, 'frac_reward_zero_std': 0.0, 'kl': 0.006754715461283922}
{'step': 34, 'loss': -0.0035, 'grad_norm': 2.0486481189727783, 'kl': 0.008028591983020306}
{'step': 35, 'loss': 0.0001, 'grad_norm': 1.0030733346939087, 'reward_std': 0.5773502588272095, 'frac_reward_zero_std': 0.0, 'kl': 0.004371508257463574}
{'step': 36, 'loss': -0.0023, 'grad_norm': 1.0071489810943604, 'kl': 0.004760880139656365}
{'step': 37, 'loss': 0.0001, 'grad_norm': 0.011982212774455547, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.003609871375374496}
{'step': 38, 'loss': 0.0001, 'grad_norm': 0.018087197095155716, 'kl': 0.004922280320897698}
{'step': 39, 'loss': 0.0001, 'grad_norm': 0.6724298596382141, 'reward_std': 0.5, 'frac_reward_zero_std': 0.0, 'kl': 0.004765854449942708}
{'step': 40, 'loss': -0.0005, 'grad_norm': 0.6781410574913025, 'kl': 0.0050692601362243295}
{'step': 41, 'loss': 0.0003, 'grad_norm': 0.9643133282661438, 'reward_std': 0.5773502588272095, 'frac_reward_zero_std': 0.0, 'kl': 0.014563204254955053}
{'step': 42, 'loss': -0.0013, 'grad_norm': 0.8823528289794922, 'kl': 0.01641524862498045}
{'step': 43, 'loss': 0.0001, 'grad_norm': 0.0017313749995082617, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.0030840933322906494}
{'step': 44, 'loss': 0.0001, 'grad_norm': 0.0017373405862599611, 'kl': 0.0030671562999486923}
{'step': 45, 'loss': 0.0002, 'grad_norm': 1.3954052925109863, 'reward_std': 0.5773502588272095, 'frac_reward_zero_std': 0.0, 'kl': 0.01147513696923852}
{'step': 46, 'loss': 0.0005, 'grad_norm': 1.3935962915420532, 'kl': 0.011951538734138012}
{'step': 47, 'loss': 0.0002, 'grad_norm': 0.01714283600449562, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.009299900382757187}
{'step': 48, 'loss': 0.0002, 'grad_norm': 0.0197350662201643, 'kl': 0.010163025930523872}
{'step': 49, 'loss': 0.0003, 'grad_norm': 0.9702491760253906, 'reward_std': 0.5773502588272095, 'frac_reward_zero_std': 0.0, 'kl': 0.013018805999308825}
{'step': 50, 'loss': -0.0028, 'grad_norm': 0.9800103902816772, 'kl': 0.013623438309878111}
{'step': 51, 'loss': 0.0001, 'grad_norm': 0.0019908789545297623, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.004270221572369337}
{'step': 52, 'loss': 0.0001, 'grad_norm': 0.0028420889284461737, 'kl': 0.004688122775405645}
{'step': 53, 'loss': 0.0003, 'grad_norm': 1.3191063404083252, 'reward_std': 0.051538821309804916, 'frac_reward_zero_std': 0.0, 'kl': 0.014647746924310923}
{'step': 54, 'loss': -0.0002, 'grad_norm': 1.3181734085083008, 'kl': 0.013659677002578974}
{'step': 55, 'loss': 0.0003, 'grad_norm': 0.0026868758723139763, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.013021013233810663}
{'step': 56, 'loss': 0.0003, 'grad_norm': 0.0026746320072561502, 'kl': 0.013159879948943853}
{'step': 57, 'loss': 0.0003, 'grad_norm': 0.00761838536709547, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.014744222164154053}
{'step': 58, 'loss': 0.0003, 'grad_norm': 0.008670816197991371, 'kl': 0.015929841436445713}
{'step': 59, 'loss': 0.0002, 'grad_norm': 0.009598137810826302, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.008403334999457002}
{'step': 60, 'loss': 0.0002, 'grad_norm': 0.009757611900568008, 'kl': 0.008458983153104782}
{'step': 61, 'loss': 0.0004, 'grad_norm': 0.011399062350392342, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'kl': 0.02068170439451933}
{'step': 62, 'loss': 0.0004, 'grad_norm': 0.011508246883749962, 'kl': 0.020991764962673187}
{'step': 63, 'loss': 0.0001, 'grad_norm': 2.848334789276123, 'reward_std': 0.5, 'frac_reward_zero_std': 0.0, 'kl': 0.007139152963645756}
{'step': 64, 'loss': -0.005, 'grad_norm': 1.2507139444351196, 'kl': 0.006105866516008973}

7. Inspect multiple responses from the actual trained policy¶

The policy below is the actual in-memory model trained above. We do not load another author's published checkpoint. The saved adapter must be loaded together with the original base model; a standalone copy of base weights is not included.

We display every held-out question with its reference answer and both generated responses, using the same decoding settings before and after training. Format and correctness rewards show how the scoring functions treat each response; the goal here is to illustrate the method. A response may change or remain identical.

Adapter saving is optional and disabled by default. Before/after answers are shown for all 12 fixed questions. Each answer shows the extracted value, reference value, and the reason for its correctness reward. A reward of zero means no correctness reward, not an execution error.

In [10]:
trainer.model.gradient_checkpointing_disable()
assert baseline_questions == [row["prompt"][-1]["content"] for row in eval_rows], "The before and after question sets must match."
after_responses = [generate_response(trainer.model, tokenizer, r["prompt"], MAX_COMPLETION_LENGTH) for r in eval_rows]
references = [r["solution"] for r in eval_rows]
def score_responses(responses):
    completions = [[{"role": "assistant", "content": text}] for text in responses]
    return {"mean_format_reward": sum(format_reward(completions)) / len(completions),
            "correct_rate": sum(accuracy_reward(completions, references)) / len(completions)}

summary = {"n_held_out": len(eval_rows), "before": score_responses(before_responses),
           "after": score_responses(after_responses)}
comparisons = [{**r, "before": b, "after": a}
               for r, b, a in zip(eval_rows, before_responses, after_responses)]
(OUTPUT_DIR / "held_out_comparison.json").write_text(json.dumps(comparisons, indent=2))
(OUTPUT_DIR / "held_out_metrics.json").write_text(json.dumps(summary, indent=2))
for i, (row, before, after) in enumerate(zip(eval_rows, before_responses, after_responses), 1):
    print("\n" + "=" * 80)
    print(f"Example {i}/{len(eval_rows)} | {row['prompt'][-1]['content']}")
    for label, response in (("BEFORE GRPO", before), ("AFTER GRPO", after)):
        details = describe_reward(response, row["solution"])
        print(f"\n{label}:\n{response}")
        print(f"Extracted answer: {details['extracted']}")
        print(f"Reference answer: {details['expected']}")
        print(f"Result: {details['status']}")
        print(f"Learning signal: correctness={details['correctness_reward']:.0f}, format={details['format_reward']:.2f}")
    print("Same generated text:", before == after)

if SAVE_RUNTIME_ARTIFACTS:
    trainer.save_model(str(OUTPUT_DIR / "final_adapter"))
    tokenizer.save_pretrained(OUTPUT_DIR / "final_adapter")
    trainer.save_state()
    (OUTPUT_DIR / "training_log.json").write_text(json.dumps(trainer.state.log_history, indent=2))
    train_dataset.to_json(str(OUTPUT_DIR / "selected_train_prompts.jsonl"))
    print("Saved to runtime only:", OUTPUT_DIR)
else:
    print("Adapter saving is disabled. The trained policy is still available in trainer.model for this session.")

================================================================================
Example 1/12 | What is 27 * 34?

BEFORE GRPO:
<think>27 * 34 = 888; 888 is larger, so the result is -888.</think><answer>-888</answer>
Extracted answer: -888
Reference answer: 918
Result: WRONG ANSWER
Learning signal: correctness=0, format=1.00

AFTER GRPO:
<think>27 * 34 = 888; 888 is larger, so the result is 888.</think><answer>888</answer>
Extracted answer: 888
Reference answer: 918
Result: WRONG ANSWER
Learning signal: correctness=0, format=1.00
Same generated text: False

================================================================================
Example 2/12 | What is 387 + 568?

BEFORE GRPO:
<think>387 + 568 = 955; 955 is larger, so the result is -955.</think><answer>-955</answer>
Extracted answer: -955
Reference answer: 955
Result: WRONG ANSWER
Learning signal: correctness=0, format=1.00

AFTER GRPO:
<think>387 + 568 = 955; 955 is larger, so the result is 955.</think><answer>955</answer>
Extracted answer: 955
Reference answer: 955
Result: CORRECT
Learning signal: correctness=1, format=1.00
Same generated text: False

================================================================================
Example 3/12 | What is (17 + 26) * 8?

BEFORE GRPO:
<think>17 + 26 = 43; 43 * 8 = 344.</think><answer>344</answer>
Extracted answer: 344
Reference answer: 344
Result: CORRECT
Learning signal: correctness=1, format=1.00

AFTER GRPO:
<think>17 + 26 = 43; 43 * 8 = 344.</think><answer>344</answer>
Extracted answer: 344
Reference answer: 344
Result: CORRECT
Learning signal: correctness=1, format=1.00
Same generated text: True

================================================================================
Example 4/12 | What is 27 - 64?

BEFORE GRPO:
<think>64 - 27 = 37; 37 is smaller, so the result is -37.</think><answer>-37</answer>
Extracted answer: -37
Reference answer: -37
Result: CORRECT
Learning signal: correctness=1, format=1.00

AFTER GRPO:
<think>27 - 64 = -37; -37 is smaller, so the result is -37.</think><answer>-37</answer>
Extracted answer: -37
Reference answer: -37
Result: CORRECT
Learning signal: correctness=1, format=1.00
Same generated text: False

================================================================================
Example 5/12 | What is 144 / 12 + 7 * 3?

BEFORE GRPO:
<think>144 / 12 = 12; 12 + 7 * 3 = 21; 21 + 144 = 165.</think><answer>165</answer>
Extracted answer: 165
Reference answer: 33
Result: WRONG ANSWER
Learning signal: correctness=0, format=1.00

AFTER GRPO:
<think>144 / 12 = 12; 12 + 7 * 3 = 12 + 21 = 33.</think><answer>33</answer>
Extracted answer: 33
Reference answer: 33
Result: CORRECT
Learning signal: correctness=1, format=1.00
Same generated text: False

================================================================================
Example 6/12 | A shop has 18 boxes of 24 pencils. It sells 157 pencils. How many remain?

BEFORE GRPO:
Think: 18 * 24 = 432 pencils. 157 - 432 = -375. 375 is smaller, so the result is -375.
Extracted answer: None
Reference answer: 275
Result: NO FINAL ANSWER EXTRACTED
Learning signal: correctness=0, format=0.00

AFTER GRPO:
<think>18 * 24 = 432 pencils; 432 - 157 = 275 pencils remain.</think><answer>275</answer>
Extracted answer: 275
Reference answer: 275
Result: CORRECT
Learning signal: correctness=1, format=1.00
Same generated text: False

================================================================================
Example 7/12 | What is 46 * 17?

BEFORE GRPO:
<think>46 * 17 = 742; 742 is larger, so the result is -742.</think><answer>-742</answer>
Extracted answer: -742
Reference answer: 782
Result: WRONG ANSWER
Learning signal: correctness=0, format=1.00

AFTER GRPO:
<think>46 * 17 = 742; 742 is larger, so the result is 742.</think><answer>742</answer>
Extracted answer: 742
Reference answer: 782
Result: WRONG ANSWER
Learning signal: correctness=0, format=1.00
Same generated text: False

================================================================================
Example 8/12 | What is 623 - 278?

BEFORE GRPO:
<think>623 - 278 = 345; 345 is smaller, so the result is -345.</think><answer>-345</answer>
Extracted answer: -345
Reference answer: 345
Result: WRONG ANSWER
Learning signal: correctness=0, format=1.00

AFTER GRPO:
<think>623 - 278 = 345; 345 is smaller, so the result is -345.</think><answer>-345</answer>
Extracted answer: -345
Reference answer: 345
Result: WRONG ANSWER
Learning signal: correctness=0, format=1.00
Same generated text: True

================================================================================
Example 9/12 | What is 39 - 82?

BEFORE GRPO:
<think>82 - 39 = 43; 43 is smaller, so the result is -43.</think><answer>-43</answer>
Extracted answer: -43
Reference answer: -43
Result: CORRECT
Learning signal: correctness=1, format=1.00

AFTER GRPO:
<think>39 - 82 = -43; -43 is smaller, so the result is -43.</think><answer>-43</answer>
Extracted answer: -43
Reference answer: -43
Result: CORRECT
Learning signal: correctness=1, format=1.00
Same generated text: False

================================================================================
Example 10/12 | What is 23 * 7 - 48?

BEFORE GRPO:
<think>23 * 7 = 161; 161 - 48 = 113.</think><answer>113</answer>
Extracted answer: 113
Reference answer: 113
Result: CORRECT
Learning signal: correctness=1, format=1.00

AFTER GRPO:
<think>23 * 7 = 161; 161 - 48 = 113.</think><answer>113</answer>
Extracted answer: 113
Reference answer: 113
Result: CORRECT
Learning signal: correctness=1, format=1.00
Same generated text: True

================================================================================
Example 11/12 | What is (63 - 28) * 6?

BEFORE GRPO:
<think>63 - 28 = 35; 35 * 6 = 210.</think><answer>210</answer>
Extracted answer: 210
Reference answer: 210
Result: CORRECT
Learning signal: correctness=1, format=1.00

AFTER GRPO:
<think>63 - 28 = 35; 35 * 6 = 210.</think><answer>210</answer>
Extracted answer: 210
Reference answer: 210
Result: CORRECT
Learning signal: correctness=1, format=1.00
Same generated text: True

================================================================================
Example 12/12 | There are 14 bags of 18 marbles. Give away 95 marbles. How many remain?

BEFORE GRPO:
Think: 14 bags * 18 marbles/bag = 252 marbles. 252 - 95 = 157 marbles remain.
Extracted answer: None
Reference answer: 157
Result: NO FINAL ANSWER EXTRACTED
Learning signal: correctness=0, format=0.00

AFTER GRPO:
<think>14 bags * 18 marbles/bag = 252 marbles; 252 - 95 = 157 marbles remain.</think><answer>157</answer>
Extracted answer: 157
Reference answer: 157
Result: CORRECT
Learning signal: correctness=1, format=1.00
Same generated text: False
Adapter saving is disabled. The trained policy is still available in trainer.model for this session.

Takeaways¶

  • GRPO consumes prompts and generates response groups during training.
  • A reward function or reward model scores the responses; this example uses verifiable final answers plus a smaller format reward.
  • Group-relative advantages replace the learned critic used in the PPO demonstration.
  • DPO and GRPO need different data at training time. They can also appear at different stages of a larger post-training pipeline.

Credits and references¶

Adapted from Sergio Paniego, Post training an LLM for reasoning with GRPO in TRL, Hugging Face Open-Source AI Cookbook. Original article: https://huggingface.co/learn/cookbook/fine_tuning_llm_grpo_trl . Original notebook: https://github.com/huggingface/cookbook/blob/main/notebooks/en/fine_tuning_llm_grpo_trl.ipynb . Author: https://github.com/sergiopaniego . The source project is licensed under Apache 2.0: https://github.com/huggingface/cookbook/blob/main/LICENSE .

Method: Zhihong Shao et al., DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models (2024): https://arxiv.org/abs/2402.03300 .

Implementation reference (version used here): https://huggingface.co/docs/trl/v0.18.2/grpo_trainer . Policy: Qwen Team, Qwen/Qwen2-0.5B-Instruct. Optional original dataset: Project Numina and AI-MO, AI-MO/NuminaMath-TIR. The default elementary-arithmetic questions are generated locally using a fixed seed and integer arithmetic. Mathematical equivalence checking: Hugging Face Math-Verify contributors.

This revision changes dependency setup, GPU precision, subset size, reward validation, group batching, local saving, and evaluation of the actual trained adapter for a compact Kaggle teaching run. Original example outputs and training screenshots have been removed.

This optimized teaching revision adds mixed-difficulty arithmetic, training-only prompt screening, explicit final-answer diagnostics, and a longer but still small LoRA run. No performance improvement is asserted in advance.