Lab: GRPO on a Maths or Code Task on One Machine

Validated on: written from the documentation cited above; not yet validated on hardware on any track. The versions, wall-clock times and pass@1 figures for each track will be recorded here when the validation pass is done.

Objective

By the end of this lab you will have run real reinforcement learning on your own machine: a model sampling groups of answers to problems it was not given solutions for, scored by a Python function you tested first, with the reward, KL and length curves written to a file you read rather than a dashboard you glance at. You will have pass@1 on held-out problems from before and after, and a run-log line that says exactly what produced the difference.

The model is Qwen3-1.7B on the 16 GB floor, whose card gives 1.7 billion parameters, 28 layers and an Apache-2.0 licence, or Qwen3-4B where memory allows, whose card gives 4.0 billion parameters, 3.6 billion outside the embeddings, 36 layers and the same licence. Both are instruction-tuned, which matters: GRPO can only amplify behaviour the model already samples, and a base model that never produces a marked final answer supplies no signal at all.

The task is multi-step arithmetic word problems, generated together with their answers by a script you can read. Generating them rather than downloading them means the reference answers cannot be wrong, and it means the course redistributes nothing. If you would rather train on a public set, make-tasks.py --source gsm8k builds the same files from GSM8K, whose dataset card gives the licence as MIT and whose problems it describes as needing “between 2 and 8 steps to solve”, with the final answer marked by ####.

⌂Single-machine path

Single machine on every track. A second machine is useful for one thing here: serving the rollout model on Tracks X and M, so that generation and training are not competing for the same accelerator. The lab works without it.

Requirements

Every track needs Part 11’s training libraries, the reward library from this part’s fifth lesson, about 10 GB of free disk, and roughly seventy-five minutes of which about twenty-five are attended. The memory floor is 16 GB and this part’s toolchain lesson has the arithmetic.

Be honest with yourself about the run length before you start. Unsloth’s reinforcement-learning guide, read on 2026-09-09, advises waiting “at least 300 steps for the reward to actually increase” and “at least 500 rows of data”, and warns to expect “a minimum of 12 hours” for a good result. This lab’s default is 120 steps on 240 problems, which is chosen to finish inside the hour and a quarter and to make the machinery visible. Treat a large gain from 120 steps as a surprise rather than an expectation, and use the unattended extension at the end if you want an effect big enough to measure comfortably.

Track S — NVIDIA DGX Spark

The primary path: TRL with vLLM in colocate mode, qwen3-1.7b for the matched baseline, eight rollouts per prompt, and a reference model loaded. The 128 GB pool holds all of it with the second memory diagram in the toolchain lesson to spare.

Check the installed vLLM against the range TRL documents as supported before you enable it; the course’s serving pin is outside that range, and in-process generation is the fallback that always works.

Track X — AMD Ryzen AI Max+ 395Partial

Training needs the ROCm build of PyTorch, which AMD's documentation does not qualify for gfx1151; vLLM's ROCm wheels for this chip have not been exercised by this course's validation pass. In-process generation on ROCm or the CPU completes the lab.

Use qwen3-1.7b and in-process generation, which needs nothing beyond Part 11’s environment. Confirm PyTorch sees the GPU before you start; if it does not, add --precision fp32 and finish on the CPU, recording that you did.

If you want a served rollout engine, the --rollouts llama-server path is available and carries the stale-weights caveat the toolchain lesson describes. Use it for a short run, not for the run you report.

Track M — Apple siliconPartial

No rollout engine such as vLLM runs on macOS, and mlx-lm's README read on 2026-09-09 documents no reinforcement-learning trainer. TRL runs on PyTorch's MPS backend in float32, which doubles the weight and activation footprint and is the slowest of the four tracks.

Use qwen3-1.7b at most, --precision fp32, four rollouts per prompt and a shorter completion cap. The reduced-path memory diagram below is written for exactly this configuration. Expect the run to take the whole budget and plan to leave it unattended.

A 0.6 billion parameter model will run faster and is a poor choice here for a reason worth understanding: it is below the size at which GRPO reliably has anything to amplify, and Unsloth’s guide advises “at least 1.5B in parameters” for this reason.

Track N — NVIDIA desktop or laptop

16 GB of VRAM runs qwen3-1.7b with beta at zero and eight rollouts; 24 GB runs it with a reference model, and qwen3-4b with four. vLLM in colocate mode is the documented single-card path, subject to the version constraint above.

On Windows, work inside WSL2 exactly as in Part 11.

Tasks

Working directory and terminal roles

Prepare the course execution workspace once before this procedure. It includes this part's scripts, data and shared Python helpers. In the client or training terminal, select this directory:

RunnableAll tracks

select this part’s execution directory

export LABS_ROOT="${LABS_ROOT:-$HOME/llm-course/labs}"
export LAB_DIR="$LABS_ROOT/part-14-preference-and-rl"
cd "$LAB_DIR"
pwd
test -f "make-tasks.py"

Expected result: pwd ends in part-14-preference-and-rl and the file check returns successfully. If it does not, finish workspace preparation before continuing. Activate the environment in the requirements for your track. Bare script and data filenames below are relative to this directory; paths to earlier experiments must point at the artefacts you actually retained.

Keep each foreground server in a separate terminal and send requests from this terminal. Reapply lesson-specific environment variables in each new shell. Stop at the first failed checkpoint and retain its output; the execution guide explains how to distinguish missing files, endpoint failures and capacity problems.

1. Collect the files

Four scripts, beside runlog.py from Part 11. train-grpo.py imports both rewards and runlog from its own directory.

RunnableAll tracks

make-tasks.py

"""Build a verifiable task set whose answers are correct by construction.

Purpose: the training and evaluation problems for the GRPO lab and the reality check.
    Every problem is generated from a template together with its answer, so the
    reference answer is not a label somebody typed: it is the number the generator
    computed before it wrote the question. Two families are produced. "arith" is
    multi-step arithmetic word problems, the family the lab trains on. "logic" is
    counting, calendar and ordering puzzles, the different-kind held-out family the
    reality check uses to ask whether anything generalised.
Platform: all (pure Python; --source gsm8k additionally needs the datasets package)
Minimum memory: 8 GB
Assumes: Python 3.10 or newer. For --source gsm8k, `datasets` installed and network
    access to the Hugging Face Hub; the GSM8K card gives the licence as MIT.

Usage: python3 make-tasks.py --out-dir tasks --seed 0
       python3 make-tasks.py --out-dir tasks --train 240 --heldout 60 --seed 0
       python3 make-tasks.py --out-dir tasks --source gsm8k --train 240 --heldout 60
"""

from __future__ import annotations

import argparse
import hashlib
import json
import random
from pathlib import Path
from typing import Callable

INSTRUCTION = (
    "Solve the problem. Put your working inside <think> and </think> tags, then end "
    "with a single line of the form 'Answer: N' where N is the final number and "
    "nothing follows it."
)

DAYS = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]
NAMES = ["Ada", "Bela", "Chi", "Dara", "Eero", "Fen", "Gita", "Hugo", "Ines", "Jonas",
         "Kira", "Liam", "Mira", "Nils", "Oona", "Pia", "Rafa", "Sena", "Tomas", "Ulla"]
GOODS = [("pencil", "pencils"), ("notebook", "notebooks"), ("cable", "cables"),
         ("mug", "mugs"), ("tile", "tiles"), ("filter", "filters"), ("bulb", "bulbs")]


# --------------------------------------------------------------------------------------
# Family: arithmetic word problems
# --------------------------------------------------------------------------------------

def arith_purchase(rng: random.Random) -> tuple[str, int]:
    boxes = rng.randint(3, 12)
    per_box = rng.randint(4, 15)
    used = rng.randint(2, boxes * per_box - 2)
    singular, plural = rng.choice(GOODS)
    who = rng.choice(NAMES)
    question = (f"{who} buys {boxes} boxes of {plural}, with {per_box} {plural} in each box. "
                f"{used} {plural if used != 1 else singular} are used during the week. "
                f"How many {plural} are left?")
    return question, boxes * per_box - used


def arith_rate(rng: random.Random) -> tuple[str, int]:
    per_hour = rng.randint(6, 30)
    hours = rng.randint(2, 9)
    extra = rng.randint(1, 40)
    who = rng.choice(NAMES)
    question = (f"A machine that {who} operates produces {per_hour} parts an hour. It runs for "
                f"{hours} hours, and {extra} parts from yesterday are added to the same crate. "
                f"How many parts are in the crate?")
    return question, per_hour * hours + extra


def arith_split(rng: random.Random) -> tuple[str, int]:
    people = rng.randint(3, 9)
    each = rng.randint(2, 12)
    leftover = rng.randint(1, people - 1)
    singular, plural = rng.choice(GOODS)
    question = (f"A crate of {plural} is shared between {people} people. Each person gets "
                f"{each} {plural}, and {leftover} are left over. How many {plural} were in "
                f"the crate?")
    return question, people * each + leftover


def arith_growth(rng: random.Random) -> tuple[str, int]:
    start = rng.randint(5, 40)
    gain = rng.randint(2, 15)
    days = rng.randint(3, 12)
    loss = rng.randint(1, 20)
    who = rng.choice(NAMES)
    question = (f"{who} starts a log with {start} entries and adds {gain} entries every day for "
                f"{days} days. At the end {loss} duplicate entries are deleted. How many entries "
                f"does the log hold?")
    return question, start + gain * days - loss


def arith_area_cost(rng: random.Random) -> tuple[str, int]:
    width = rng.randint(2, 15)
    length = rng.randint(2, 15)
    price = rng.randint(2, 20)
    spare = rng.randint(1, 10)
    question = (f"A floor is {width} metres by {length} metres. Tiles cost {price} units per "
                f"square metre, and {spare} extra units are spent on adhesive. What is the total "
                f"cost in units?")
    return question, width * length * price + spare


def arith_two_step_difference(rng: random.Random) -> tuple[str, int]:
    a_rate = rng.randint(5, 20)
    b_rate = rng.randint(1, a_rate - 1)
    minutes = rng.randint(4, 25)
    question = (f"Two pumps fill the same tank. The first adds {a_rate} litres a minute and the "
                f"second removes {b_rate} litres a minute. Both run for {minutes} minutes. How "
                f"many litres are in the tank at the end, if it started empty?")
    return question, (a_rate - b_rate) * minutes


ARITH_TEMPLATES: list[Callable[[random.Random], tuple[str, int]]] = [
    arith_purchase, arith_rate, arith_split, arith_growth, arith_area_cost,
    arith_two_step_difference,
]


# --------------------------------------------------------------------------------------
# Family: counting, calendar and ordering puzzles
# --------------------------------------------------------------------------------------

def logic_inclusion_exclusion(rng: random.Random) -> tuple[str, int]:
    both = rng.randint(2, 12)
    only_a = rng.randint(3, 25)
    only_b = rng.randint(3, 25)
    neither = rng.randint(0, 10)
    question = (f"In a group, {only_a + both} people read the newsletter, {only_b + both} people "
                f"attend the meeting, {both} do both, and {neither} do neither. How many people "
                f"are in the group?")
    return question, only_a + only_b + both + neither


def logic_calendar(rng: random.Random) -> tuple[str, int]:
    start = rng.randrange(7)
    ahead = rng.randint(9, 400)
    question = (f"A task starts on a {DAYS[start]}. Counting the start day as day 1, day {ahead} "
                f"falls on which day of the week? Answer with the position of that day in the "
                f"week where Monday is 1 and Sunday is 7.")
    return question, (start + ahead - 1) % 7 + 1


def logic_ordering(rng: random.Random) -> tuple[str, int]:
    people = rng.sample(NAMES, 5)
    positions = list(range(1, 6))
    rng.shuffle(positions)
    placed = dict(zip(people, positions))
    front = min(placed, key=placed.get)
    back = max(placed, key=placed.get)
    subject = rng.choice([p for p in people if p not in (front, back)])
    clues = []
    for person in people:
        if person == subject:
            continue
        ahead = placed[person] < placed[subject]
        clues.append(f"{person} is {'ahead of' if ahead else 'behind'} {subject}")
    rng.shuffle(clues)
    question = ("Five people stand in a queue, numbered 1 at the front to 5 at the back. "
                + "; ".join(clues) + f". What position is {subject} in?")
    return question, placed[subject]


def logic_handshakes(rng: random.Random) -> tuple[str, int]:
    people = rng.randint(6, 20)
    absent = rng.randint(1, 4)
    present = people - absent
    question = (f"{people} people are invited to a meeting and {absent} do not come. Everyone "
                f"who comes shakes hands with everyone else exactly once. How many handshakes "
                f"take place?")
    return question, present * (present - 1) // 2


def logic_comparison_chain(rng: random.Random) -> tuple[str, int]:
    base = rng.randint(10, 60)
    step_one = rng.randint(2, 20)
    step_two = rng.randint(2, 20)
    a, b, c = rng.sample(NAMES, 3)
    question = (f"{a} is {step_one} years older than {b}, and {b} is {step_two} years younger "
                f"than {c}. {c} is {base} years old. How old is {a}?")
    return question, base - step_two + step_one


def logic_parity(rng: random.Random) -> tuple[str, int]:
    total = rng.randint(20, 120)
    step = rng.choice([3, 4, 5, 6, 7])
    question = (f"The whole numbers from 1 to {total} are written down. How many of them are "
                f"multiples of {step}?")
    return question, total // step


LOGIC_TEMPLATES: list[Callable[[random.Random], tuple[str, int]]] = [
    logic_inclusion_exclusion, logic_calendar, logic_ordering, logic_handshakes,
    logic_comparison_chain, logic_parity,
]

FAMILIES = {"arith": ARITH_TEMPLATES, "logic": LOGIC_TEMPLATES}


# --------------------------------------------------------------------------------------
# Assembly
# --------------------------------------------------------------------------------------

def make_row(index: int, family: str, question: str, answer: int) -> dict:
    return {
        "id": f"{family}-{index:04d}",
        "family": family,
        "prompt": [{"role": "user", "content": f"{question}\n\n{INSTRUCTION}"}],
        "question": question,
        "answer": str(answer),
    }


def generate(family: str, count: int, rng: random.Random, seen: set[str], start_index: int) -> list[dict]:
    """Distinct problems from the family's templates, cycling so each is used equally."""
    templates = FAMILIES[family]
    rows: list[dict] = []
    attempts = 0
    while len(rows) < count:
        attempts += 1
        if attempts > count * 200:
            raise SystemExit(
                f"could not generate {count} distinct {family} problems; the templates repeat "
                f"themselves before that. Ask for fewer, or add a template."
            )
        template = templates[len(rows) % len(templates)]
        question, answer = template(rng)
        if question in seen:
            continue
        seen.add(question)
        rows.append(make_row(start_index + len(rows), family, question, answer))
    return rows


def from_gsm8k(count: int, split: str, offset: int) -> list[dict]:
    """A slice of GSM8K, whose card gives the licence as MIT and marks answers with ####."""
    try:
        from datasets import load_dataset  # noqa: PLC0415 - optional dependency
    except ImportError:
        raise SystemExit("--source gsm8k needs the datasets package: pip install datasets") from None
    data = load_dataset("openai/gsm8k", "main", split=split)
    rows = []
    for i in range(offset, min(offset + count, len(data))):
        item = data[i]
        answer = item["answer"].split("####")[-1].strip().replace(",", "")
        rows.append(make_row(i, "gsm8k", item["question"].strip(), int(answer)))
    return rows


def write_jsonl(path: Path, rows: list[dict]) -> str:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for row in rows:
            handle.write(json.dumps(row, ensure_ascii=False) + "\n")
    return hashlib.sha256(path.read_bytes()).hexdigest()


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--out-dir", default="tasks")
    parser.add_argument("--train", type=int, default=240, help="training problems, same family")
    parser.add_argument("--heldout", type=int, default=60,
                        help="held-out problems, per family; one set of the same kind and one of a different kind")
    parser.add_argument("--family", default="arith", choices=sorted(FAMILIES),
                        help="the family the lab trains on")
    parser.add_argument("--other-family", default="logic", choices=sorted(FAMILIES),
                        help="the different-kind family the reality check evaluates on")
    parser.add_argument("--source", default="generated", choices=["generated", "gsm8k"],
                        help="generated problems, or a slice of GSM8K (MIT licence, per its dataset card)")
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    if args.family == args.other_family:
        raise SystemExit("--family and --other-family must differ; the point is a different kind of task")

    out = Path(args.out_dir)
    rng = random.Random(args.seed)
    seen: set[str] = set()
    manifest: dict[str, object] = {"seed": args.seed, "source": args.source, "files": {}}

    if args.source == "gsm8k":
        train = from_gsm8k(args.train, "train", 0)
        same = from_gsm8k(args.heldout, "test", 0)
        print(f"GSM8K, licence MIT per the dataset card; {len(train)} train, {len(same)} held-out")
    else:
        train = generate(args.family, args.train, rng, seen, 0)
        same = generate(args.family, args.heldout, rng, seen, args.train)

    different = generate(args.other_family, args.heldout, rng, seen, 0)

    for name, rows in (("train", train), ("heldout-same", same), ("heldout-different", different)):
        digest = write_jsonl(out / f"{name}.jsonl", rows)
        manifest["files"][name] = {"rows": len(rows), "sha256": digest}
        print(f"{name:18s} {len(rows):5d} rows  sha256 {digest[:12]}...")

    (out / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
    print(f"\nwritten to {out}/  (train, heldout-same, heldout-different, manifest.json)")
    print("Read ten of them before you train on them; a generator with a bug produces a "
          "reward function that is confidently wrong.")

    sample = train[0]
    print(f"\nexample problem ({sample['family']}):\n  {sample['question']}\n  reference answer: {sample['answer']}")


if __name__ == "__main__":
    main()

Download make-tasks.py305 lines

RunnableAll tracks

train-grpo.py

"""Reinforcement learning with a verifiable reward, using TRL's GRPOTrainer.

Purpose: the course's reference GRPO run. Loads a prompt-only task set whose answers
    are correct by construction, samples a group of completions per prompt, scores
    them with the reward functions from rewards.py, and updates a LoRA adapter on the
    group-relative advantages. Writes the reward, KL and length curves to a CSV so
    they can be read rather than guessed at, and appends one run record to the lab
    notebook.
Platform: spark, nvidia (CUDA) and strix (ROCm) with in-process or vLLM rollouts;
    mac runs the same script on PyTorch's MPS backend in float32 with in-process
    rollouts and a small model. vLLM does not run on macOS, so --rollouts vllm-* is
    refused there; --rollouts llama-server is the served-model path for tracks
    without vLLM and carries the stale-weights caveat printed at start-up.
Minimum memory: 16 GB
Assumes: torch, transformers, trl, peft and datasets installed in the active
    environment; make-tasks.py has been run so that <tasks-dir>/train.jsonl exists;
    rewards.py and runlog.py sit next to this file.

Usage: python3 train-grpo.py --tasks-dir tasks --output-dir runs/grpo-qwen3-1.7b \
           --model Qwen/Qwen3-1.7B --labbook labbook.md
       python3 train-grpo.py --tasks-dir tasks --model Qwen/Qwen3-4B \
           --rollouts vllm-colocate --num-generations 8 --labbook labbook.md
       python3 train-grpo.py --tasks-dir tasks --model Qwen/Qwen3-1.7B \
           --rollouts llama-server --rollout-url http://127.0.0.1:8080 --labbook labbook.md
"""

from __future__ import annotations

import argparse
import csv
import json
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path
from typing import Any, Optional

import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoTokenizer
from trl import GRPOConfig, GRPOTrainer

import rewards as reward_lib
import runlog

DEFAULT_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]

# The metrics worth plotting, in the order the lab reads them.
CURVE_KEYS = ["reward", "reward_std", "kl", "completions/mean_length", "loss",
              "clip_ratio/region_mean", "learning_rate", "epoch"]


def pick_device() -> str:
    """CUDA (or ROCm, which reports as cuda), then MPS, then the CPU."""
    if torch.cuda.is_available():
        return "cuda"
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return "mps"
    return "cpu"


def use_bf16(device: str, requested: str) -> bool:
    if requested == "fp32":
        return False
    if requested == "bf16":
        return True
    return device == "cuda" and torch.cuda.is_bf16_supported()


# --------------------------------------------------------------------------------------
# Rollouts through an OpenAI-compatible llama-server
# --------------------------------------------------------------------------------------

STALE_WEIGHTS_WARNING = """\
--rollouts llama-server samples from a server holding a fixed GGUF file. TRL's vLLM
integration streams the updated weights into the rollout engine after every optimiser
step; llama.cpp's server README documents no equivalent endpoint, so after the first
update the completions come from an older policy than the one being trained. GRPO's
importance ratio corrects for a *small* mismatch, not an unbounded one. Use this path
for a short run whose purpose is to see the machinery work on a track without vLLM,
keep --num-iterations at 1, and record in the run log that the rollouts were stale.
The in-process path (--rollouts inproc) is always in sync and is the correct choice
whenever you can afford its speed.
"""


def post_json(url: str, payload: dict, api_key: Optional[str], timeout: int) -> dict:
    body = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(url, data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:400]
        raise RuntimeError(f"{url} returned HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"cannot reach {url}: {exc.reason}") from exc


def token_ids_and_logprobs(body: dict) -> tuple[list[int], list[float]]:
    """Pull the sampled tokens and their log-probabilities out of a llama-server reply.

    The server README documents `n_probs`: "If greater than 0, the response also
    contains the probabilities of top N tokens for each generated token", returned in
    `completion_probabilities`, where "Each item in the array has a nested array
    `top_logprobs`". The exact key names inside those items have changed between
    builds, so this reads defensively and fails with a message naming the field rather
    than training on log-probabilities it guessed at.
    """
    items = body.get("completion_probabilities")
    if not items:
        raise SystemExit(
            "the server returned no completion_probabilities. Send n_probs greater than 0 "
            "on the /completion endpoint, and check that your llama.cpp build returns it."
        )
    ids: list[int] = []
    logprobs: list[float] = []
    for item in items:
        chosen_id = item.get("id", item.get("tok"))
        candidates = item.get("top_logprobs") or []
        picked = None
        for candidate in candidates:
            if chosen_id is not None and candidate.get("id") == chosen_id:
                picked = candidate
                break
        if picked is None and candidates:
            picked = candidates[0]
        if picked is None:
            raise SystemExit("a completion_probabilities item has no top_logprobs entry")
        token_id = picked.get("id", chosen_id)
        logprob = picked.get("logprob")
        if token_id is None or logprob is None:
            raise SystemExit(
                "a top_logprobs entry is missing 'id' or 'logprob'; this build's response "
                "shape differs from the server README read for this lab"
            )
        ids.append(int(token_id))
        logprobs.append(float(logprob))
    return ids, logprobs


def make_llama_server_rollout(base_url: str, api_key: Optional[str], tokenizer,
                              num_generations: int, max_new_tokens: int,
                              temperature: float, top_p: float, seed: int, timeout: int):
    """A rollout function that samples from llama-server's native /completion endpoint.

    GRPOTrainer's `rollout_func` is documented as experimental: "It receives the list
    of prompts allocated to the current process and the trainer instance. It must
    return a dict with `prompt_ids`, `completion_ids`, and `logprobs` fields... This
    feature is experimental and may change or be removed at any time without prior
    notice." Because the calling convention may change, the wrapper below accepts the
    prompts positionally or by keyword and says so if it cannot find them.
    """
    endpoint = base_url.rstrip("/") + "/completion"
    state = {"warned_decode": False, "calls": 0}

    def one_completion(prompt_text: str, sample_index: int) -> tuple[list[int], list[float]]:
        payload = {
            "prompt": prompt_text,
            "n_predict": max_new_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "n_probs": 1,
            "cache_prompt": True,
            "seed": seed + sample_index + state["calls"] * num_generations,
        }
        body = post_json(endpoint, payload, api_key, timeout)
        ids, logprobs = token_ids_and_logprobs(body)
        text = body.get("content", "")
        if not state["warned_decode"]:
            decoded = tokenizer.decode(ids, skip_special_tokens=True)
            if decoded.strip() != text.strip():
                print("WARNING: the server's token ids do not decode to the text it returned "
                      "under the trainer's tokeniser. The GGUF file and the Hugging Face model "
                      "must be the same model with the same vocabulary.", file=sys.stderr)
            state["warned_decode"] = True
        return ids, logprobs

    def rollout(prompts: Any = None, trainer: Any = None, **call_kwargs: Any) -> dict:
        # TRL types rollout_func as Callable[[list[str], "GRPOTrainer"], dict], so the
        # prompts arrive first and the trainer second. The keyword fallback below exists
        # because the documentation marks the feature experimental.
        if prompts is None:
            prompts = call_kwargs.get("prompts")
        if prompts is None:
            raise SystemExit("rollout_func was called without a list of prompts; TRL's "
                             "experimental calling convention has changed")
        del trainer  # this path needs no trainer state; the server holds the weights

        prompt_ids: list[list[int]] = []
        completion_ids: list[list[int]] = []
        all_logprobs: list[list[float]] = []
        for prompt in prompts:
            text = prompt if isinstance(prompt, str) else tokenizer.apply_chat_template(
                prompt, tokenize=False, add_generation_prompt=True
            )
            ids = tokenizer(text, add_special_tokens=False)["input_ids"]
            for k in range(num_generations):
                completion, logprobs = one_completion(text, k)
                prompt_ids.append(list(ids))
                completion_ids.append(completion)
                all_logprobs.append(logprobs)
        state["calls"] += 1
        return {"prompt_ids": prompt_ids, "completion_ids": completion_ids, "logprobs": all_logprobs}

    return rollout


# --------------------------------------------------------------------------------------
# Curves
# --------------------------------------------------------------------------------------

def write_curves(history: list[dict], path: Path) -> dict[str, Any]:
    """One row per logged step, and the three numbers the lab asks you to compare."""
    rows = [row for row in history if "reward" in row or "loss" in row]
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["step"] + CURVE_KEYS)
        for row in rows:
            writer.writerow([row.get("step", "")] + [row.get(k, "") for k in CURVE_KEYS])

    rewarded = [row for row in rows if isinstance(row.get("reward"), (int, float))]
    lengths = [row["completions/mean_length"] for row in rows
               if isinstance(row.get("completions/mean_length"), (int, float))]
    kls = [row["kl"] for row in rows if isinstance(row.get("kl"), (int, float))]
    summary = {
        "logged_steps": len(rows),
        "first_reward": round(rewarded[0]["reward"], 4) if rewarded else None,
        "last_reward": round(rewarded[-1]["reward"], 4) if rewarded else None,
        "best_reward": round(max(r["reward"] for r in rewarded), 4) if rewarded else None,
        "first_mean_length": round(lengths[0], 1) if lengths else None,
        "last_mean_length": round(lengths[-1], 1) if lengths else None,
        "last_kl": round(kls[-1], 5) if kls else None,
        "max_kl": round(max(kls), 5) if kls else None,
    }
    return summary


# --------------------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--model", default="Qwen/Qwen3-1.7B",
                        help="base model repository id or local path")
    parser.add_argument("--adapter", default=None,
                        help="an existing LoRA adapter to continue training, e.g. a Part 13 fine-tune")
    parser.add_argument("--tasks-dir", default="tasks", help="directory holding train.jsonl from make-tasks.py")
    parser.add_argument("--output-dir", default="runs/grpo")
    parser.add_argument("--reward-kind", default="maths", choices=["maths", "exact", "code"])
    parser.add_argument("--no-format-reward", action="store_true")
    parser.add_argument("--no-length-reward", action="store_true")
    parser.add_argument("--fallback-to-last-number", action="store_true",
                        help="grade the last number in the completion when no answer marker is found")
    parser.add_argument("--num-generations", type=int, default=8,
                        help="rollouts per prompt, G in the lesson; the effective batch must divide by it")
    parser.add_argument("--max-completion-length", type=int, default=512)
    parser.add_argument("--target-tokens", type=int, default=256, help="soft budget for the length reward")
    parser.add_argument("--beta", type=float, default=0.04,
                        help="KL coefficient; TRL's default is 0.0, which skips the reference model entirely")
    parser.add_argument("--epsilon", type=float, default=0.2)
    parser.add_argument("--epsilon-high", type=float, default=None,
                        help="upper clipping bound; the DAPO paper recommends 0.28")
    parser.add_argument("--loss-type", default="dapo", choices=["grpo", "dapo", "dr_grpo"])
    parser.add_argument("--scale-rewards", default="group", choices=["group", "batch"],
                        help="divide the advantage by the group or the batch standard deviation; "
                             "to drop the term entirely, use --loss-type dr_grpo")
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top-p", type=float, default=1.0)
    parser.add_argument("--lr", type=float, default=1e-5)
    parser.add_argument("--batch-size", type=int, default=8, help="prompts per device per step")
    parser.add_argument("--grad-accum", type=int, default=1)
    parser.add_argument("--num-iterations", type=int, default=1, help="policy updates per batch of rollouts")
    parser.add_argument("--max-steps", type=int, default=200)
    parser.add_argument("--rank", type=int, default=16)
    parser.add_argument("--alpha", type=int, default=32)
    parser.add_argument("--target-modules", nargs="+", default=DEFAULT_TARGETS)
    parser.add_argument("--precision", choices=["auto", "bf16", "fp32"], default="auto")
    parser.add_argument("--gradient-checkpointing", action="store_true")
    parser.add_argument("--rollouts", default="inproc",
                        choices=["inproc", "vllm-colocate", "vllm-server", "llama-server"])
    parser.add_argument("--rollout-url", default="http://127.0.0.1:8080",
                        help="llama-server base URL for --rollouts llama-server")
    parser.add_argument("--rollout-api-key", default=None)
    parser.add_argument("--rollout-timeout", type=int, default=600)
    parser.add_argument("--report-to", default="none", choices=["none", "tensorboard"])
    parser.add_argument("--curves", default=None, help="write the reward, KL and length curves here as CSV")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--notes", default=None)
    args = parser.parse_args()

    device = pick_device()
    bf16 = use_bf16(device, args.precision)
    print(f"device: {device}   precision: {'bfloat16' if bf16 else 'float32'}   rollouts: {args.rollouts}")

    if args.rollouts.startswith("vllm") and device == "mps":
        raise SystemExit("vLLM's documented GPU path does not cover macOS; use --rollouts inproc "
                         "or --rollouts llama-server on Track M")

    effective_batch = args.batch_size * args.grad_accum
    if effective_batch % args.num_generations:
        raise SystemExit(
            f"the effective batch ({args.batch_size} x {args.grad_accum} = {effective_batch}) must be "
            f"divisible by --num-generations ({args.num_generations}); TRL requires it"
        )

    train_path = Path(args.tasks_dir) / "train.jsonl"
    if not train_path.is_file():
        raise SystemExit(f"{train_path} is missing; run make-tasks.py --out-dir {args.tasks_dir} first")
    dataset = load_dataset("json", data_files={"train": str(train_path)})["train"]
    print(f"train prompts: {len(dataset)}   families: {sorted(set(dataset['family']))}")

    tokenizer = AutoTokenizer.from_pretrained(args.adapter or args.model)
    if tokenizer.chat_template is None:
        raise SystemExit(f"{args.model} has no chat template; pick an instruction-tuned model")

    reward_funcs, reward_weights = reward_lib.build_reward_functions(
        kind=args.reward_kind,
        use_format=not args.no_format_reward,
        use_length=not args.no_length_reward,
        target_tokens=args.target_tokens,
        fallback_to_last_number=args.fallback_to_last_number,
    )
    print("rewards: " + ", ".join(f"{f.__name__} x{w}" for f, w in zip(reward_funcs, reward_weights)))

    config_kwargs: dict[str, Any] = dict(
        output_dir=args.output_dir,
        max_steps=args.max_steps,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=args.grad_accum,
        num_generations=args.num_generations,
        max_completion_length=args.max_completion_length,
        num_iterations=args.num_iterations,
        beta=args.beta,
        epsilon=args.epsilon,
        loss_type=args.loss_type,
        scale_rewards=args.scale_rewards,
        temperature=args.temperature,
        top_p=args.top_p,
        reward_weights=reward_weights,
        learning_rate=args.lr,
        lr_scheduler_type="constant_with_warmup",
        warmup_steps=5,
        bf16=bf16,
        gradient_checkpointing=args.gradient_checkpointing,
        logging_steps=1,
        save_strategy="steps",
        save_steps=max(25, args.max_steps // 4),
        save_total_limit=2,
        report_to=args.report_to,
        seed=args.seed,
        data_seed=args.seed,
        model_init_kwargs={"dtype": torch.bfloat16 if bf16 else torch.float32},
    )
    if args.epsilon_high is not None:
        config_kwargs["epsilon_high"] = args.epsilon_high
    if args.rollouts == "vllm-colocate":
        config_kwargs["use_vllm"] = True
        config_kwargs["vllm_mode"] = "colocate"
    elif args.rollouts == "vllm-server":
        config_kwargs["use_vllm"] = True
        config_kwargs["vllm_mode"] = "server"

    trainer_kwargs: dict[str, Any] = {}
    if args.rollouts == "llama-server":
        print(STALE_WEIGHTS_WARNING, file=sys.stderr)
        trainer_kwargs["rollout_func"] = make_llama_server_rollout(
            base_url=args.rollout_url,
            api_key=args.rollout_api_key,
            tokenizer=tokenizer,
            num_generations=args.num_generations,
            max_new_tokens=args.max_completion_length,
            temperature=args.temperature,
            top_p=args.top_p,
            seed=args.seed,
            timeout=args.rollout_timeout,
        )

    if args.adapter:
        from peft import AutoPeftModelForCausalLM  # noqa: PLC0415 - only needed on this branch
        model = AutoPeftModelForCausalLM.from_pretrained(args.adapter, is_trainable=True)
        config_kwargs.pop("model_init_kwargs", None)
        print(f"continuing the adapter in {args.adapter}")
    else:
        model = args.model
        trainer_kwargs["peft_config"] = LoraConfig(
            r=args.rank, lora_alpha=args.alpha, lora_dropout=0.0,
            target_modules=args.target_modules, bias="none", task_type="CAUSAL_LM",
        )

    trainer = GRPOTrainer(
        model=model,
        args=GRPOConfig(**config_kwargs),
        reward_funcs=reward_funcs,
        train_dataset=dataset,
        processing_class=tokenizer,
        **trainer_kwargs,
    )

    started = time.time()
    trainer.train()
    elapsed = time.time() - started

    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)

    curves_path = Path(args.curves or Path(args.output_dir) / "curves.csv")
    summary = write_curves(trainer.state.log_history, curves_path)
    summary["seconds"] = round(elapsed, 1)
    print(json.dumps(summary, indent=2))
    print(f"curves written to {curves_path}")
    print(f"adapter saved to {args.output_dir}")

    if args.labbook:
        record = runlog.record(
            labbook=args.labbook,
            lab="part-14/train-grpo",
            model=args.adapter or args.model,
            dataset={
                "path": str(train_path),
                "sha256": runlog.file_sha256(train_path),
                "train_prompts": len(dataset),
            },
            hyperparameters={
                "method": "grpo-lora",
                "rollouts": args.rollouts,
                "rollouts_in_sync": args.rollouts != "llama-server",
                "num_generations": args.num_generations,
                "max_completion_length": args.max_completion_length,
                "beta": args.beta,
                "epsilon": args.epsilon,
                "epsilon_high": args.epsilon_high,
                "loss_type": args.loss_type,
                "scale_rewards": args.scale_rewards,
                "temperature": args.temperature,
                "top_p": args.top_p,
                "learning_rate": args.lr,
                "batch_size": args.batch_size,
                "grad_accum": args.grad_accum,
                "num_iterations": args.num_iterations,
                "max_steps": args.max_steps,
                "rank": args.rank,
                "alpha": args.alpha,
                "reward_functions": [f.__name__ for f in reward_funcs],
                "reward_weights": reward_weights,
                "precision": "bfloat16" if bf16 else "float32",
            },
            seed=args.seed,
            losses=summary,
            scores={},
            config_path=__file__,
            notes=args.notes,
        )
        print(f"recorded run {record['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()

Download train-grpo.py465 lines

RunnableAll tracks

eval-pass-at-1.py

"""Measure pass@1 on a held-out task set, before and after a reinforcement-learning run.

Purpose: the number the GRPO lab and the reality check are decided on. Sends every
    held-out problem to the model, grades the answer with the same verifier the reward
    used, and reports the share solved with its standard error and a breakdown by task
    family. Running it once on the starting model and once on the trained one is the
    whole measurement; running it on a family the model never trained on is what turns
    it into a claim about generalisation.
Platform: all. The model is loaded locally through transformers and PEFT by default, so
    no server is needed; --base-url sends the same prompts to any OpenAI-compatible
    endpoint instead, which is the faster path when you already have one running.
Minimum memory: 16 GB for a 1B to 4B model loaded locally; very little against a server.
Assumes: Python 3.10 or newer; torch, transformers and peft for the local path;
    make-tasks.py has written the held-out files; rewards.py and runlog.py sit next to
    this file.

Usage: python3 eval-pass-at-1.py --tasks tasks/heldout-same.jsonl --model Qwen/Qwen3-1.7B \
           --label before --labbook labbook.md
       python3 eval-pass-at-1.py --tasks tasks/heldout-same.jsonl --model Qwen/Qwen3-1.7B \
           --adapter runs/grpo-qwen3-1.7b --label after --labbook labbook.md
       python3 eval-pass-at-1.py --tasks tasks/heldout-different.jsonl \
           --base-url http://127.0.0.1:8080/v1 --served-model local --attempts 4 --temperature 0.7
"""

from __future__ import annotations

import argparse
import json
import math
import time
import urllib.error
import urllib.request
from collections import defaultdict
from pathlib import Path
from typing import Optional

import rewards as reward_lib
import runlog


def post_json(url: str, payload: dict, api_key: Optional[str], timeout: int) -> dict:
    body = json.dumps(payload).encode("utf-8")
    headers = {"Content-Type": "application/json"}
    if api_key:
        headers["Authorization"] = f"Bearer {api_key}"
    request = urllib.request.Request(url, data=body, headers=headers, method="POST")
    try:
        with urllib.request.urlopen(request, timeout=timeout) as response:
            return json.loads(response.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", "replace")[:400]
        raise RuntimeError(f"{url} returned HTTP {exc.code}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"cannot reach {url}: {exc.reason}") from exc


class ServedModel:
    def __init__(self, base_url: str, model: str, api_key: Optional[str], timeout: int,
                 temperature: float, max_tokens: int):
        self.endpoint = base_url.rstrip("/") + "/chat/completions"
        self.model, self.api_key, self.timeout = model, api_key, timeout
        self.temperature, self.max_tokens = temperature, max_tokens

    def answer(self, messages: list[dict], seed: int) -> str:
        payload = {"model": self.model, "messages": messages, "temperature": self.temperature,
                   "max_tokens": self.max_tokens, "seed": seed}
        body = post_json(self.endpoint, payload, self.api_key, self.timeout)
        return (body["choices"][0]["message"]["content"] or "").strip()


class LocalModel:
    def __init__(self, model_id: str, adapter: Optional[str], temperature: float, max_tokens: int):
        import torch  # noqa: PLC0415
        from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: PLC0415

        self.torch = torch
        self.temperature, self.max_tokens = temperature, max_tokens
        if adapter:
            from peft import AutoPeftModelForCausalLM  # noqa: PLC0415
            self.model = AutoPeftModelForCausalLM.from_pretrained(adapter)
            self.tokenizer = AutoTokenizer.from_pretrained(adapter)
        else:
            self.model = AutoModelForCausalLM.from_pretrained(model_id)
            self.tokenizer = AutoTokenizer.from_pretrained(model_id)
        self.model.eval()
        print(f"evaluating {adapter or model_id} on {self.model.device}")

    def answer(self, messages: list[dict], seed: int) -> str:
        self.torch.manual_seed(seed)
        inputs = self.tokenizer.apply_chat_template(
            messages, tokenize=True, add_generation_prompt=True, return_dict=True, return_tensors="pt"
        ).to(self.model.device)
        sampling = {"do_sample": False} if self.temperature <= 0 else {
            "do_sample": True, "temperature": self.temperature, "top_p": 0.95}
        with self.torch.no_grad():
            out = self.model.generate(**inputs, max_new_tokens=self.max_tokens, **sampling)
        generated = out[0][inputs["input_ids"].shape[-1]:]
        return self.tokenizer.decode(generated, skip_special_tokens=True).strip()


def standard_error(successes: float, trials: int) -> float:
    """The binomial standard error of a proportion, so a difference can be read honestly."""
    if trials <= 0:
        return float("nan")
    p = successes / trials
    return math.sqrt(max(0.0, p * (1 - p)) / trials)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--tasks", required=True, help="a JSONL file from make-tasks.py")
    parser.add_argument("--model", default="Qwen/Qwen3-1.7B")
    parser.add_argument("--adapter", default=None)
    parser.add_argument("--base-url", default=None, help="evaluate a served model instead")
    parser.add_argument("--served-model", default=None)
    parser.add_argument("--api-key", default=None)
    parser.add_argument("--attempts", type=int, default=1,
                        help="samples per problem; 1 at temperature 0 is deterministic pass@1, "
                             "more at a higher temperature estimates it with an error bar")
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--max-tokens", type=int, default=512)
    parser.add_argument("--limit", type=int, default=None, help="evaluate only the first N problems")
    parser.add_argument("--tolerance", type=float, default=1e-6)
    parser.add_argument("--fallback-to-last-number", action="store_true")
    parser.add_argument("--label", default="run", help="a name for this side of the comparison")
    parser.add_argument("--out", default=None, help="write every answer here as JSON")
    parser.add_argument("--timeout", type=int, default=300)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--labbook", default=None)
    parser.add_argument("--notes", default=None)
    args = parser.parse_args()

    if args.attempts > 1 and args.temperature <= 0:
        raise SystemExit("--attempts above 1 with temperature 0 samples the same answer every "
                         "time; raise --temperature or drop back to one attempt")

    rows = [json.loads(line) for line in Path(args.tasks).read_text(encoding="utf-8").splitlines() if line.strip()]
    if args.limit:
        rows = rows[: args.limit]
    print(f"{len(rows)} problems from {args.tasks}, {args.attempts} attempt(s) each")

    if args.base_url:
        if not args.served_model:
            raise SystemExit("--base-url needs --served-model")
        model = ServedModel(args.base_url, args.served_model, args.api_key, args.timeout,
                            args.temperature, args.max_tokens)
    else:
        model = LocalModel(args.model, args.adapter, args.temperature, args.max_tokens)

    grade = reward_lib.numeric_reward(args.tolerance, args.fallback_to_last_number)
    per_family: dict[str, list[float]] = defaultdict(list)
    answers = []
    started = time.time()

    for i, row in enumerate(rows, start=1):
        scores = []
        for attempt in range(args.attempts):
            text = model.answer(row["prompt"], args.seed + attempt)
            correct = grade(completions=[text], answer=[row["answer"]])[0]
            scores.append(correct)
            answers.append({"id": row["id"], "family": row["family"], "attempt": attempt,
                            "answer": text, "reference": row["answer"], "correct": bool(correct)})
        share = sum(scores) / len(scores)
        per_family[row["family"]].append(share)
        print(f"  [{i}/{len(rows)}] {row['id']:14s} {share:.2f}")

    elapsed = time.time() - started
    overall = [s for values in per_family.values() for s in values]
    scores = {
        "label": args.label,
        "tasks": args.tasks,
        "problems": len(rows),
        "attempts_each": args.attempts,
        "pass_at_1": round(sum(overall) / len(overall), 4) if overall else None,
        "standard_error": round(standard_error(sum(overall), len(overall)), 4) if overall else None,
        "by_family": {family: round(sum(v) / len(v), 4) for family, v in sorted(per_family.items())},
        "seconds": round(elapsed, 1),
    }
    print("\n" + json.dumps(scores, indent=2))
    print("\nThe standard error is the one on this sample of problems. Two runs whose "
          "intervals overlap have not been shown to differ.")

    if args.out:
        Path(args.out).write_text(json.dumps({"scores": scores, "answers": answers}, indent=2),
                                  encoding="utf-8")
        print(f"answers written to {args.out}")

    if args.labbook:
        record = runlog.record(
            labbook=args.labbook,
            lab="part-14/eval-pass-at-1",
            model=args.adapter or args.served_model or args.model,
            dataset={"path": args.tasks, "sha256": runlog.file_sha256(args.tasks),
                     "problems": len(rows)},
            hyperparameters={"attempts": args.attempts, "temperature": args.temperature,
                             "max_tokens": args.max_tokens, "tolerance": args.tolerance,
                             "fallback_to_last_number": args.fallback_to_last_number,
                             "label": args.label},
            seed=args.seed,
            losses={},
            scores=scores,
            config_path=__file__,
            notes=args.notes,
        )
        print(f"recorded run {record['run_id']} in {args.labbook}")


if __name__ == "__main__":
    main()

Download eval-pass-at-1.py209 lines

The reward library and its tests are the two files from this part’s fifth lesson. If you have not downloaded them, do that now; nothing below runs without rewards.py.

2. Build the task set, and read ten of them

RunnableAll tracks

generate the problems and their answers

python3 make-tasks.py --out-dir tasks --train 240 --heldout 60 --seed 0

Output — what you should see

train                240 rows  sha256 91efdbfe1984...
heldout-same          60 rows  sha256 6d8692fd3e82...
heldout-different     60 rows  sha256 c22061c2525f...

written to tasks/  (train, heldout-same, heldout-different, manifest.json)
Read ten of them before you train on them; a generator with a bug produces a
reward function that is confidently wrong.

example problem (arith):
  Rafa buys 9 boxes of cables, with 10 cables in each box. 7 cables are used during the week. How many cables are left?
  reference answer: 83

Three files, and the difference between them is the whole experiment. train.jsonl is what the run optimises on. heldout-same.jsonl is more problems of the same kind, never trained on, which is where the honest before-and-after comparison happens. heldout-different.jsonl is a second family of counting, calendar and ordering puzzles, which the reality check at the end of this part uses to ask whether anything transferred.

RunnableAll tracks

read ten problems and check the arithmetic yourself

head -n 10 tasks/train.jsonl | python3 -c "
import sys, json
for line in sys.stdin:
    row = json.loads(line)
    print(row['id'], '|', row['question'], '=>', row['answer'])
"

Do the arithmetic on three of them with a pencil. The generator computes the answer before it writes the question, so a mismatch means a bug in a template, and a bug in a template means every reward this lab computes is wrong in the same direction.

3. Test the reward before you optimise it

RunnableAll tracks

the reward's own tests, then its worked examples

python3 test-rewards.py
python3 rewards.py --demo

Thirty-six tests in about five seconds, then the table of six completions from the fifth lesson, two of which score as well as a correct answer. You are about to spend most of an hour optimising this function, and this is the last cheap moment to change it.

4. Measure pass@1 before anything is trained

RunnableAll tracks

pass@1 on held-out problems, starting model

python3 eval-pass-at-1.py \
  --tasks tasks/heldout-same.jsonl \
  --model Qwen/Qwen3-1.7B \
  --limit 40 \
  --label before \
  --out eval-before.json \
  --labbook labbook.md

Output — what you should see

40 problems from tasks/heldout-same.jsonl, 1 attempt(s) each
evaluating Qwen/Qwen3-1.7B on cuda:0
  [1/40] arith-0240   1.00
  ...
{
  "label": "before",
  "problems": 40,
  "pass_at_1": 0.xxx,
  "standard_error": 0.0xx,
  "by_family": { "arith": 0.xxx }
}

Write the two numbers down now. The standard error is the one on this sample of forty problems, and it is what stops you calling a small difference a result later.

✕Common mistake

Do not skip this step because you can run it afterwards. After training you will have a trained model, an untrained model and an hour of enthusiasm, and it is remarkably easy to evaluate the wrong one. Measuring first, with a run-log line, removes the possibility.

5. Do the arithmetic before the run

Apply the toolchain lesson’s five terms to your own settings. On the primary path the first diagram in that lesson is the budget. Track M and any 16 GB machine that needs headroom take the reduced path below.

The reduced path: Qwen3-1.7B in float32 on PyTorch MPS, 4 rollouts of 256 completion tokens, beta zero

Policy weights, float32

6.8 GB

Adapter, gradients and Adam states

0.3 GB

Rollout key-value cache, 4 sequences

0.4 GB

Logits and activations for the update

1.5 GB

Reserved for the operating system

2 GB

Free

5 GB

Total

16 GB

Estimate from arithmetic, not a measurement. Weights are twice the BF16 figure from the course model reference because MPS runs this lab in float32; the key-value figure is 4 sequences of 448 tokens at twice the model's documented 112 KiB per token, for the same reason; the logits term is 4 by 256 by Qwen3's 151,936-entry vocabulary at 4 bytes, about 0.6 GB, plus an allowance for the remaining activations. Beta is zero here, so no reference model is loaded.

Write the peak you predict in the notebook now, so the number the run reports has something to be compared against.

6. Run it

Track S — NVIDIA DGX Spark

RunnableTrack S · DGX Spark

GRPO with vLLM rollouts in colocate mode

python3 train-grpo.py \
  --tasks-dir tasks \
  --model Qwen/Qwen3-1.7B \
  --output-dir runs/grpo-qwen3-1.7b \
  --rollouts vllm-colocate \
  --num-generations 8 \
  --batch-size 8 \
  --grad-accum 4 \
  --max-completion-length 512 \
  --beta 0.04 \
  --lr 1e-5 \
  --max-steps 120 \
  --curves curves.csv \
  --labbook labbook.md

Track X — AMD Ryzen AI Max+ 395Partial

ROCm build required for the GPU path; in-process generation on ROCm or the CPU completes the lab.

RunnableTrack X · Ryzen AI Max+

GRPO with in-process rollouts

python3 train-grpo.py \
  --tasks-dir tasks \
  --model Qwen/Qwen3-1.7B \
  --output-dir runs/grpo-qwen3-1.7b \
  --rollouts inproc \
  --num-generations 8 \
  --batch-size 8 \
  --grad-accum 4 \
  --max-completion-length 384 \
  --beta 0.04 \
  --lr 1e-5 \
  --max-steps 120 \
  --curves curves.csv \
  --labbook labbook.md

Track M — Apple siliconPartial

MPS in float32, in-process rollouts only; no vLLM on macOS.

RunnableTrack M · Apple silicon

GRPO on MPS, the reduced path

python3 train-grpo.py \
  --tasks-dir tasks \
  --model Qwen/Qwen3-1.7B \
  --output-dir runs/grpo-qwen3-1.7b \
  --rollouts inproc \
  --precision fp32 \
  --num-generations 4 \
  --batch-size 4 \
  --grad-accum 4 \
  --max-completion-length 256 \
  --beta 0.0 \
  --lr 1e-5 \
  --max-steps 60 \
  --curves curves.csv \
  --labbook labbook.md

Track N — NVIDIA desktop or laptop

RunnableTrack N · NVIDIA GPU

GRPO with vLLM rollouts in colocate mode

python3 train-grpo.py \
  --tasks-dir tasks \
  --model Qwen/Qwen3-1.7B \
  --output-dir runs/grpo-qwen3-1.7b \
  --rollouts vllm-colocate \
  --num-generations 8 \
  --batch-size 8 \
  --grad-accum 4 \
  --max-completion-length 512 \
  --beta 0.04 \
  --lr 1e-5 \
  --max-steps 120 \
  --curves curves.csv \
  --labbook labbook.md

The first lines of output are the ones to check.

Output — what you should see

device: cuda   precision: bfloat16   rollouts: vllm-colocate
train prompts: 240   families: ['arith']
rewards: numeric_reward x1.0, format_reward x1.0, length_reward x1.0

If the effective batch does not divide by the rollout count the script stops before loading anything and says so, because TRL requires that the effective batch size be evenly divisible by num_generations and the failure otherwise arrives several minutes later.

✱Hardware note

The optional served-rollout path for Tracks X and M starts a llama.cpp server and points the trainer at it with the experimental rollout_func hook. It prints a warning at start-up explaining that llama.cpp’s server documents no way to receive updated weights, so the rollouts drift away from the policy as training proceeds, and it records rollouts_in_sync: false in the run log.

RunnableAll tracks

serve the same model for rollouts, then train against it

llama-server \
  --model ~/models/Qwen3-1.7B-Q8_0.gguf \
  --alias rollout \
  --ctx-size 4096 \
  --host 127.0.0.1 \
  --port 8080

RunnableAll tracks

GRPO with served rollouts, short runs only

python3 train-grpo.py \
  --tasks-dir tasks \
  --model Qwen/Qwen3-1.7B \
  --output-dir runs/grpo-served \
  --rollouts llama-server \
  --rollout-url http://127.0.0.1:8080 \
  --num-generations 4 \
  --batch-size 4 \
  --grad-accum 4 \
  --max-steps 20 \
  --labbook labbook.md

7. Read the curves, all three together

RunnableAll tracks

the reward, KL and length curves

column -s, -t curves.csv | head -30

Output — what you should see

step  reward  reward_std  kl       completions/mean_length  loss     ...
1     0.4xx   0.4xx       0.0000   2xx.x                    0.00xx
2     0.4xx   0.4xx       0.0001   2xx.x                    0.00xx
...

Four columns, read as a group rather than one at a time.

reward is the weighted sum across your reward functions. It should rise, slowly and noisily. A flat line at zero means no rollout ever scored, which is the failure the previous lesson explains: no correct answers means no spread, which means no advantage and no gradient.

reward_std is the spread within the batch. It is the health check nobody looks at. If it falls towards zero while the reward is still low, the groups have collapsed into agreement on a wrong answer and there is nothing left to learn from.

kl is the divergence from the reference, logged only when beta is above zero. It should rise from zero and level off. A KL that climbs without limit while the reward is flat means the policy is wandering, and the answer is a larger beta or a smaller learning rate.

completions/mean_length is the one to be suspicious of. If it climbs and the reward climbs with it, the honest first hypothesis is not that the model is reasoning more carefully: the Dr. GRPO paper identifies an “optimization bias in Group Relative Policy Optimization (GRPO), which artificially increases response length (especially for incorrect outputs) during training”. Try --loss-type dr_grpo and see whether the reward gain survives without the length gain.

=Expected result

On a 120-step run at these settings, expect the reward to move a little and noisily, the KL to be small and rising, and the mean length to change by more than you would like. That is what a short run looks like. The result this lab is after is that you can read the four columns and say what each one is telling you, not that the model became a mathematician in an hour.

8. Read the completions, not just the curves

RunnableAll tracks

generate from the trained adapter and look at it

python3 eval-pass-at-1.py \
  --tasks tasks/heldout-same.jsonl \
  --model Qwen/Qwen3-1.7B \
  --adapter runs/grpo-qwen3-1.7b \
  --limit 8 \
  --label spotcheck \
  --out spotcheck.json
python3 -c "
import json
for row in json.load(open('spotcheck.json'))['answers'][:4]:
    print('---', row['id'], 'correct' if row['correct'] else 'wrong')
    print(row['answer'][:600])
"

You are looking for the four hacks from the reward lesson: sprays of numbers, immaculate format around a wrong answer, padding, and answers that stop before the working is finished. A reward curve cannot show you any of them.

9. Measure pass@1 after

RunnableAll tracks

the same held-out problems, the trained adapter

python3 eval-pass-at-1.py \
  --tasks tasks/heldout-same.jsonl \
  --model Qwen/Qwen3-1.7B \
  --adapter runs/grpo-qwen3-1.7b \
  --limit 40 \
  --label after \
  --out eval-after.json \
  --labbook labbook.md

Compare the two pass_at_1 figures with their standard errors. Forty problems is a small sample, and two runs whose intervals overlap have not been shown to differ. Writing “the intervals overlap” in the notebook is a result; writing “it improved” when they overlap is not.

10. Optionally, run it properly

The lab’s defaults are sized for the hour. If you want a change large enough to survive the reality check, the run to leave overnight is longer, on more problems, with a bigger group:

RunnableAll tracks

the unattended run

python3 make-tasks.py --out-dir tasks-large --train 600 --heldout 120 --seed 1
python3 train-grpo.py \
  --tasks-dir tasks-large \
  --model Qwen/Qwen3-1.7B \
  --output-dir runs/grpo-long \
  --rollouts vllm-colocate \
  --num-generations 8 \
  --batch-size 8 \
  --grad-accum 8 \
  --max-steps 400 \
  --curves curves-long.csv \
  --labbook labbook.md

11. Record the run

RunnableAll tracks

the last three run records

tail -n 3 labbook.md | python3 -m json.tool

Check six fields: dataset.sha256 matches the generator’s output, hyperparameters.rollouts names the path you actually used, hyperparameters.rollouts_in_sync is true unless you deliberately used the served path, hyperparameters.num_generations and beta are what you meant, and losses.first_reward and losses.last_reward are both present. Add the peak memory by hand.

Complete the verifier and policy checks before scaling

Keep rewards.py, test-rewards.py and runlog.py with the training script. Run the reward tests before downloading or loading the policy. Inspect a correct answer, a wrong answer and a malformed answer through the verifier; stop if the verdict differs from the documented contract.

Choose the policy size once and use it for both the before and after evaluation. The commands use the 1.7B policy on every track. If you choose the optional larger policy, change both baseline and training commands and use a new output directory. Comparing different bases would not isolate reinforcement learning.

Start with the in-process path when diagnosing correctness. If using a separate rollout engine, verify how the trainer refreshes its policy weights; an unchanged external checkpoint is not an on-policy sampler. During training, inspect reward distributions and actual completions, including all-fail groups and suspicious high-reward answers. Evaluate on both held-out task families with the same sample policy as the baseline. Save verifier revision, generated tasks, model lineage, curves and per-task results. A reward increase with unchanged or worse held-out success is a result to report and investigate, not evidence that the model learned general reasoning.

Validation

You are done when all of the following are true:

• make-tasks.py wrote three files and a manifest, and you checked three answers by hand;
• test-rewards.py reported no failures;
• a before run-log line exists with a pass@1 and a standard error, recorded before training;
• the training script printed a device, a precision, a rollout path and the reward functions with their weights;
• curves.csv has one row per logged step with reward, reward_std, kl and completions/mean_length columns;
• the adapter directory exists and contains adapter_config.json;
• an after run-log line exists on the same held-out file with the same number of problems;
• you have read at least four full completions from the trained model;
• you can say whether the two pass@1 intervals overlap.

Expected outcome

Three run-log lines, four curves and a sentence you are willing to defend. The table below is what to record per track; the validation pass will fill it in from the course’s own machines.

Pending validationWhat to record from this lab, per track

Track	Model	Rollouts	Wall clock	Peak memory, GB	pass@1 before → after	Mean length first → last
S: DGX Spark, 128 GB	Qwen3-1.7B	vLLM colocate	to be measured	to be measured	to be measured	to be measured
X: Ryzen AI Max+ 395	Qwen3-1.7B	in-process	to be measured	to be measured	to be measured	to be measured
M: Apple silicon	Qwen3-1.7B, float32	in-process	to be measured	to be measured	to be measured	to be measured
N: NVIDIA desktop or laptop	Qwen3-1.7B	vLLM colocate	to be measured	to be measured	to be measured	to be measured

the four platform tracks, one machine each · TRL GRPOTrainer with a PEFT LoRA adapter; vLLM in colocate mode where supported trl 1.12.0, transformers 5.16.1, peft 0.20.0 · Qwen3-1.7B and Qwen3-4B with rank-16 LoRA adapters, BF16, float32 on Track M · 1,024 tokens of context · 2026-09-09

Not yet run on hardware on any track. Until the validation pass fills these in, the table records what to measure. A pass@1 change smaller than the standard error printed by eval-pass-at-1.py is not a finding, and the reality check at the end of this part is where that gets tested properly.

Troubleshooting

The script stops with a message about the effective batch. TRL requires the effective batch size to be evenly divisible by num_generations. Adjust --batch-size or --grad-accum until it is; the script checks before loading the model so you find out in a second rather than in five minutes.

Reward is flat at zero from the first step. The model never produces a correct, findable answer. Check three things in order: run eval-pass-at-1.py on the training file and confirm pass@1 is not zero; read four completions and see whether the answer is present but unmarked, in which case add --fallback-to-last-number for the first run; and confirm you are using an instruction-tuned model rather than a base one.

Reward moves and reward_std collapses towards zero. The groups agree with each other. If the reward is high, the model has learned the task well enough that the batch has stopped teaching, and the answer is harder problems. If the reward is low, they agree on something wrong, and the answer is a higher temperature or a wider group.

Out of memory. In order: set --beta 0.0 to drop the reference model, lower --max-completion-length, lower --num-generations with --batch-size to match, then add --gradient-checkpointing. The first is the largest single saving.

vLLM fails to initialise. Check its version against the range TRL documents as supported, which this course’s serving pin sits outside. In-process generation with --rollouts inproc finishes the lab on every track.

Completion length grows steadily and the reward barely moves. The documented optimisation bias. Try --loss-type dr_grpo, which removes the length normalisation, or raise the weight on the length penalty by lowering --target-tokens.

Every completion is cut off mid-sentence. --max-completion-length is too small for the problem and the format. A truncated answer scores zero for a reason that is not about reasoning, so raise it or shorten the requested working.

The served rollout path fails with a message about completion_probabilities. Your llama.cpp build’s response shape differs from the server README this lab was written against. The script names the missing field rather than training on log-probabilities it guessed at; fall back to --rollouts inproc.

Track M is extremely slow. Expected. float32 on MPS with in-process generation is the slowest of the four paths, and there is no vLLM to add. Cut --max-steps, cut --num-generations to 4, and treat the run as a demonstration of the machinery.

Cleanup

Keep labbook.md, tasks/, curves.csv, both evaluation files and the adapter: the reality check at the end of this part uses all of them.

RunnableAll tracks

reclaim the disk, keeping what the reality check needs

rm -rf runs/grpo-qwen3-1.7b/checkpoint-* runs/grpo-served

⚠Warning

That deletes the intermediate checkpoints and the short served-rollout run. It is safe because the final adapter is kept and the run log records the settings and the dataset hash that produced it. Do not extend the pattern to runs/grpo-* without reading what is in there first.

What you learned

• A reinforcement-learning step is sampling plus a comparison. You watched G rollouts per prompt turn into advantages with no critic anywhere, on a machine you own.
• The reward is the specification, and it is code you tested. Thirty-six tests ran in five seconds before an hour of optimisation, which is the cheapest ratio in this part.
• Four curves, read together. Reward, its spread, the KL and the mean length each answer a different question, and only the group of them describes the run.
• Length is the first thing to suspect. The objective has a documented bias towards longer answers, so a length increase is a property of the algorithm until you have shown otherwise.
• A small sample gives a small conclusion. Forty held-out problems produce a standard error, and the standard error is what decides whether your before-and-after difference means anything.
• The rollout engine is a component with its own memory and its own correctness. Colocate mode holds the weights twice; a server that cannot receive weights samples from a policy you have left behind.

Record in the notebook: the predicted and actual peak memory; the wall clock and seconds per step; the first and last reward and the best; the first and last mean completion length; the KL at the end; pass@1 before and after with both standard errors; whether the intervals overlap; and one sentence about what the completions looked like that the curves did not show.

Check your understanding

Question 1. The reward curve is flat at zero from the first step. What is the mechanism, and what is the first thing to check?

• The learning rate is too low; raise it
• Every rollout in every group scores the same, so every advantage is zero and there is no gradient; check whether the model ever produces a correct, findable answer at all
• The KL penalty is too strong; lower beta
• The dataset is too small; generate more problems

Show the answer and why

Answer: Every rollout in every group scores the same, so every advantage is zero and there is no gradient; check whether the model ever produces a correct, findable answer at all

GRPO learns from the spread inside a group. All-zero rewards give a spread of zero. The two usual causes are a task the model cannot do and an answer it produces but the extractor cannot find, and reading four completions distinguishes them in a minute.

Question 2. Why does the lab measure pass@1 before training rather than only afterwards?

• To warm up the GPU
• Because the comparison is the result, and measuring first with a run-log line removes both the temptation and the possibility of evaluating the wrong model afterwards
• Because the model changes during evaluation
• To check the dataset is valid

Show the answer and why

Answer: Because the comparison is the result, and measuring first with a run-log line removes both the temptation and the possibility of evaluating the wrong model afterwards

A before-and-after measurement taken in that order, each with its own log line, is the difference between an experiment and a story. It also gives you the standard error that decides whether the difference means anything.

Question 3. Mean completion length rises steadily through the run while the reward rises slightly. Which reading is best supported?

• The model is learning to reason in more steps
• The objective has a documented bias that increases response length, especially for incorrect outputs, so length growth is expected from the algorithm and is not evidence about reasoning
• The temperature is drifting upward
• The reward function is broken

Show the answer and why

Answer: The objective has a documented bias that increases response length, especially for incorrect outputs, so length growth is expected from the algorithm and is not evidence about reasoning

The Dr. GRPO paper identifies this as an optimisation bias in the objective itself. The way to separate the two explanations is to rerun with the unbiased loss and see whether the accuracy gain survives without the length gain.

Question 4. Which of these reduce the memory a GRPO run needs? Select all that apply.

• Setting --beta 0.0
• Lowering --max-completion-length
• Lowering --num-generations and --batch-size together
• Raising --lr

Show the answer and why

Answer: Setting --beta 0.0, Lowering --max-completion-length, Lowering --num-generations and --batch-size together

Beta at zero means no reference model is loaded, which is a whole copy of the weights. The other two shrink the rollout cache and the logits tensor, which are the terms that scale with your settings. The learning rate changes the step, not the footprint.

Question 5. The lab warns against the llama-server rollout path for the run you report. Why?

• It is slower than in-process generation
• llama.cpp's server documents no way to receive updated weights during training, so the rollouts keep coming from the checkpoint it loaded while the policy moves away from it
• It cannot return token identifiers
• It only supports one model at a time

Show the answer and why

Answer: llama.cpp's server documents no way to receive updated weights during training, so the rollouts keep coming from the checkpoint it loaded while the policy moves away from it

TRL streams weights into the vLLM server after every optimiser step for exactly this reason. Without that path the mismatch between the sampling policy and the trained policy grows step by step, which is why the script prints a warning and records the run as out of sync.

Check my answersClear