llm-fine-tuning-grpo-countdown

GRPO Fine-Tuning: Teach a Small Model to Solve Math Puzzles

Take a tiny model that mostly fails a number puzzle, reward it only when its answer is correct, and watch GRPO make it solve far more of them. No labeled solutions, no answer key to copy. Just problems and a verifier.

On a 0.5B model, in about 15 minutes on a single T4:

	Base model	After GRPO
Solve rate (held-out)	10.0%	36.7%

That is +26.7 percentage points, and it holds up under scrutiny (see How We Know It's Real). Train for a few more epochs and it climbs past 43%. This is the cleanest kind of RL demo: an objective task, an unhackable reward, and a real, measured jump.

The Task: Countdown

Given a few numbers and a target, write an arithmetic expression that uses each number exactly once (with + - * / and parentheses) and equals the target:

Numbers: [3, 7, 9]. Target: 20.
  ->  3 * 9 - 7        (27 - 7 = 20)  correct

It is an ideal task for reinforcement learning on a small model:

• Verifiable. Just evaluate the expression and check. There is no sandbox to run (a safe AST evaluator handles the arithmetic) and no human judgment.
• Unhackable reward. Every puzzle has different numbers and a different target, so there is no constant or template you can spam to farm reward. The only way to score is to actually solve it.
• Homogeneous. It is one skill applied to endless instances, so what the model learns on training puzzles transfers to unseen ones.
• Always solvable. Each puzzle is generated by folding its numbers into the target, so a valid answer is guaranteed to exist.

Setup

cd tutorials/llm-fine-tuning-grpo-countdown
uv venv .venv --python 3.11
source .venv/bin/activate
uv pip install -r requirements.txt

Optionally set a HuggingFace token (only needed for gated models):

echo "HF_TOKEN=hf_your_token_here" > .env

Run

# The result above (0.5B, single T4, ~15 min)
flyte run workflow.py pipeline \
  --model_name "Qwen/Qwen2.5-0.5B-Instruct" \
  --n_numbers 3 --max_num 9 --epochs 2

# Train longer for a bigger jump (~43% solved)
flyte run workflow.py pipeline \
  --model_name "Qwen/Qwen2.5-0.5B-Instruct" --n_numbers 3 --epochs 5

# Harder puzzles (4 numbers) or a bigger base model
flyte run workflow.py pipeline --n_numbers 4 --max_num 12
flyte run workflow.py pipeline --model_name "Qwen/Qwen2.5-1.5B-Instruct"

It uses LoRA by default (--method full for full fine-tuning) and trains comfortably on a single T4 (16 GB).

What to Watch

• Training report. The solve rate climbing live as GRPO reinforces correct answers.
• Eval report. Base versus GRPO solve rate, the bucket breakdown (fixed, both, regressed), and side-by-side equations you can verify by hand.

Deploy the Model

Once you have a trained model, serve it as an API with a web UI. It is the fun payoff of a workshop: let people throw puzzles at the model they just watched learn.

1. Serve it (FastAPI on Flyte):

# deploy the model from the latest pipeline run
python serve.py

# or a specific run (grab the run name from the Flyte UI)
python serve.py --run-name <run-name>

serve.py mounts the model directory the pipeline returns (via RunOutput) and exposes a /solve endpoint that runs the model and verifies its answer with the same safe evaluator used for the training reward:

curl -X POST https://<your-app-url>/solve \
  -H "Content-Type: application/json" \
  -d '{"numbers": [3, 7, 9], "target": 20}'

{ "expression": "3 * 9 - 7", "value": 20.0, "correct": true, "raw_output": "3 * 9 - 7" }

2. Add the web UI (Gradio):

# auto-discovers the deployed serve.py endpoint
python app_gradio.py

# or point at a specific server
SERVER_URL=https://<your-app-url> python app_gradio.py

Enter numbers and a target, hit Solve, and the UI shows the model's expression with a pass or fail verified by the server. Both apps scale to zero (replicas=(0, 1)), so they cost nothing when idle.

File	What it does
serve.py	FastAPI endpoint that loads the GRPO model; /solve generates and verifies
app_gradio.py	Gradio UI that calls the /solve endpoint

Parameters

Flag	Default	Description
--model_name	Qwen/Qwen2.5-1.5B-Instruct	instruct model to fine-tune
--method	lora	lora or full
--epochs	3	training epochs
--lr	1e-5	learning rate
--batch_size	8	completions per step (must divide by --num_generations)
--num_generations	8	attempts (rollouts) the model makes per puzzle; GRPO compares them against each other (the "group")
--max_completion_length	320	max tokens per answer
--beta	0.005	strength of the KL leash to the base model (see note below); higher keeps it closer, 0 removes it
--n_numbers	3	how many numbers are in each puzzle (task difficulty, not a training setting; 4 is much harder than 3)
--max_num	9	largest number used in a puzzle
--max_train_samples	300	training puzzles
--num_eval_examples	60	held-out puzzles compared before and after

For the fastest clean demo, use --model_name Qwen/Qwen2.5-0.5B-Instruct --n_numbers 3 --epochs 2. Bigger models and 4-number puzzles are harder and slower but scale further.

A note on beta (the KL penalty)

KL stands for Kullback-Leibler divergence, a standard measure of how different two probability distributions are. A language model's output is a probability distribution over the next token, and the KL penalty measures how far the model you are training has drifted from where it started (the frozen base model), then charges that drift as a cost during training.

Why it is there: RL optimizes the reward and nothing else, so left unchecked the model will contort itself into whatever scores highest, even if that means degenerate or incoherent text. The KL penalty is a leash back to a competent base model, so the model improves at the task without forgetting how to write normally. beta sets the leash length: higher keeps it closer to the base (safer, learns more slowly), lower lets it change more (faster, but risks drifting into nonsense), and 0 removes the leash entirely.

---

How It Works

Do We Give It a Dataset to Learn From?

This is the question that trips everyone up, and the answer is the whole point of GRPO:

No. We never show the model a single solved example, or even a correct answer to copy.

	Supervised fine-tuning (SFT)	GRPO (what we do here)
What you must provide	Thousands of problem -> worked solution demonstrations	Just problems plus a reward function
Where the signal comes from	Copying the human-written answers	The model's own attempts, scored by the reward
Needs labeled solutions?	Yes	No

The "dataset" here is just a list of problems (Numbers: [3,7,9]. Target: 20.). No solution travels with the problem into training. The model generates several attempts at each puzzle, and a reward function checks each one: does this expression use the given numbers once and equal the target? Correct attempts get reinforced. The reward function is the teacher, not a labeled dataset.

This is the core idea of RLVR (Reinforcement Learning with Verifiable Rewards): if you can check an answer with a program, you do not need to demonstrate it. Correctness is cheap to verify but expensive to demonstrate, which is exactly the trade RLVR exploits.

What about evaluation? Surely that needs a dataset. It needs a held-out set of problems (this pipeline generates one and de-duplicates it from the training problems, so they are genuinely unseen), but it still needs no labeled solutions. The same verifier that trains the model also grades it. The only ground truth is the target, which is part of the problem itself. The general rule: RLVR needs a verifiable target per problem, for both the reward and the eval, but never the reasoning that reaches it.

How GRPO Learns to Solve

For each puzzle, the model generates a group of attempts. The reward scores each one, and GRPO reinforces the attempts that beat the group's average:

Puzzle: "Numbers: [3, 7, 9]. Target: 20."

  Attempt 1: "3 * 9 - 7"   ->  20  correct   reward 1.0   positive advantage
  Attempt 2: "3 + 7 + 9"   ->  19  wrong     reward 0.0   negative advantage
  Attempt 3: "9 * 3 - 7"   ->  20  correct   reward 1.0   positive advantage
  Attempt 4: "7 * 3 - 9"   ->  12  wrong     reward 0.0   negative advantage

The policy shifts toward the kind of expressions that hit the target.
No solution was ever shown. The model discovers what works from its own tries.

GRPO stands for Group Relative Policy Optimization. Advantages are computed within each group of attempts, which is why the group needs a mix of right and wrong answers to produce a learning signal. That is also why the task has to sit in the model's "learnable zone," meaning it is solvable often enough that some attempts succeed.

GRPO vs PPO vs DPO: Why GRPO Here

GRPO is one of a few ways to train a model from feedback instead of from labeled answers. The three you will hear about most are PPO, DPO, and GRPO. They differ mostly in what extra pieces they need and whether the model gets to explore.

	PPO	DPO	GRPO (this tutorial)
How it learns	Model generates a response, a reward model scores it, a separate value model estimates a baseline, policy is nudged toward higher reward	No generation. Directly optimizes the policy on pairs of (preferred answer, rejected answer)	Model generates several responses per prompt, each is scored, and the group's average is the baseline for which to reinforce
Extra models needed	A reward model and a value/critic model, plus a reference model	Just a reference model	Just a reference model (the reward is a program, no critic)
Data it needs	Prompts plus a trained reward model	A dataset of preference pairs (chosen vs rejected)	Prompts plus a way to score an answer (a verifier or reward model)
Does the model explore?	Yes (on-policy generation)	No (learns from a fixed set of pairs)	Yes (on-policy generation)
Cost and complexity	Highest (juggles up to four models)	Lowest (no generation loop, no reward or value model)	Medium (generation loop, but no critic)
Best for	Large-scale RLHF with a learned reward model	Cheap alignment when you already have preference data	Tasks where an answer can be checked (math, code, puzzles)

The short version:

• PPO is the classic RLHF method (it trained the first ChatGPT). It is powerful but heavy: it needs a learned reward model to score responses and a separate value model to estimate a baseline, so you are training several models at once.
• DPO skips reinforcement learning almost entirely. You hand it pairs of answers labeled better and worse, and it directly tunes the model to prefer the better one. It is simple and stable, but it only learns from the pairs you give it. The model never explores or discovers new solutions, and you need preference data.
• GRPO keeps PPO's idea of the model exploring by generating its own attempts, but throws out the expensive value model. Instead of a learned baseline, it compares each attempt to the average of the group. That makes it much lighter than PPO, and it pairs perfectly with a verifiable reward, where a short program (not a trained reward model, not human labels) says whether an answer is correct.

For Countdown that last point is the whole reason to pick GRPO. We can check an answer instantly with a few lines of code, so we do not need human preference pairs (DPO) or a trained reward model plus a critic (PPO). We just let the model try, check each try, and reinforce what worked.

Designing the Reward

In GRPO, the reward function is the task definition. The model has no other notion of "good." It optimizes exactly the number you return, and it will find the laziest path to a high number whether or not that path is what you intended. So most of the work is reward design, not RL. Ours is deliberately simple:

def reward(completion, numbers, target):
    expr = extract_answer(completion)          # pull out the expression
    if not uses_each_number_once(expr, numbers):
        return 0.0                             # must use the given numbers, each once
    return 1.0 if evaluates_to(expr, target) else 0.0

Result	Reward
Uses each given number once and equals the target	1.0
Anything else	0.0

Three design decisions, and why each one:

1. Binary, not graded. We do not give partial credit for being "close to the target." A graded reward is the classic reward-hacking trap: if a near-miss scores something, the model drifts toward a safe average guess instead of solving. Binary means the only way to earn reward is a genuinely correct answer.

2. The "uses each number once" check matters as much as the arithmetic. Without it, the reward is trivially hackable. The model could ignore the numbers and just emit the target (20), reuse a number, or invent new ones. The constraint is the puzzle, so the reward has to enforce it. This is the general rule: before you ship a reward, ask "what is the laziest output that scores well?" If that output is not what you want, the reward is wrong.

3. One reward, focused on the outcome. We reward correctness and nothing else. It is tempting to add extra terms (rewards for a certain output format, for showing work, and so on), but every term you add is another thing the model will optimize and possibly game, and it splits the model's attention away from actually solving. Add sub-rewards only when the model genuinely cannot learn without them.

How We Know It's Real

RL results are notoriously easy to fake. A model can hack the reward, or the eval can flatter it, and a naive before/after number looks great while the model learned nothing. This tutorial's eval is built to catch that, and the result survives all of it:

1. The reward cannot be gamed by a constant. Every puzzle has different numbers and a different target, so there is no single answer that scores well across the set.

2. Held-out problems. The 60 eval puzzles are de-duplicated from training, so they are genuinely unseen.

3. Lenient answer extraction, applied to both models. The grader accepts a bare expression (3 * 9 - 7) and does not require any special format, so the base model gets full credit for correct answers. The gain is not "GRPO learned the output format," it is real solving.

4. The bucket breakdown. The eval report does not cherry-pick wins. It categorizes every problem:

   Bucket	Count	Meaning
   GRPO fixed it	18	base wrong, GRPO right
   Both solved	4	already solvable
   Both failed	36	still too hard
   GRPO regressed	2	base right, GRPO wrong

   18 fixed versus 2 broken is the signature of broad improvement. A model that had merely collapsed to one lucky template would show a weak fixed-to-regressed ratio. This does not.

Read the actual equations in the eval report and you will see varied, correct expressions. The model is solving, not pattern-matching.

What's in the Pipeline

+--------------+     +------------------+     +------------+
| Prepare Data | --> |  GRPO Training   | --> |  Evaluate  |
|  (CPU task)  |     |   (GPU task)     |     | (GPU task) |
+--------------+     +------------------+     +------------+
 Generate solvable    Reward = correct        Base vs GRPO solve
 puzzles (no answers) equation (LoRA/full)    rate + bucket breakdown

1. Prepare data. Generate solvable Countdown puzzles (numbers and target only).
2. Train with GRPO. Generate attempts, reward the correct ones, reinforce.
3. Evaluate. Run held-out puzzles through the base and trained models; report solve rate, the bucket breakdown, and side-by-side equations.

A cached download_model task fetches the base model once and reuses it across runs, so repeat runs and the deploy step do not re-download from HuggingFace.

Where the Problems Come From

They are generated, so we never need a label and can guarantee solvability. Each puzzle folds a few random numbers into a target with random + - *:

nums = [3, 7, 9]
# 3 * 7 = 21, 21 - 9 = 12  ->  target 12  (a solution is guaranteed: the one we built it from)

The model is free to find a different valid expression. Train and eval puzzles are de-duplicated so held-out puzzles are truly unseen. No human solves anything.

Choosing a Task and Model for GRPO

The biggest predictor of whether GRPO will work is not the algorithm. It is whether your task and your model fit together. Countdown was chosen because it scores well on four axes. Use them to judge your own problem before you spend a GPU:

1. Learnable-zone density. For what fraction of problems does the base model succeed sometimes but not always? GRPO's advantage is computed within a group of attempts, so if every attempt scores 0 (too hard) or every attempt scores 1 (too easy), that problem contributes zero gradient. RL only sharpens what the model can already do occasionally.
2. Transfer. Does improving on training instances raise held-out performance, or is each problem its own island?
3. Homogeneity. One skill across endless instances, versus many distinct micro-skills.
4. Verifiability. Can a cheap program score it? This is what makes math, code, and games like Countdown such good targets.

Countdown passes all four: the base model already solves some puzzles by luck (dense learnable zone), the skill of combining numbers transfers to new puzzles, it is one skill over endless instances, and a short function verifies every answer.

Where the model comes in

The first two axes are not fixed. A bigger model moves them:

• Learnable-zone density rises with capability. A stronger base solves more problems sometimes, pulling them out of the "always fails" dead zone and into the range where GRPO has signal.
• Transfer improves with scale too. Larger models are more sample-efficient, so they generalize a skill from far fewer examples instead of memorizing individual instances.

This is why a task that looks hopeless on a tiny model can work well with a larger base and more data, and why production RLVR systems run on large, capable models. A verifiable task with weak learnable-zone density at 0.5B often becomes an easy win at 7B or larger, sometimes with no other change than the base model. GRPO is the method; the base model is the fuel. Small models are used here for cost and speed, not because they are optimal.

The one-line takeaway: match the task to what RL can do (sharpen a repeatable, verifiable skill that transfers), pick a base model capable enough to have a learnable zone, give it an unhackable reward, and read the actual outputs before you trust the number.

Expanding the Pipeline

This tutorial runs the whole loop (generate, score, update) in a single GPU task on purpose, so it stays easy to follow. In a real system there are two natural ways to scale it out, and Flyte is a good fit for both.

Fan out the rollouts (scale one training). In GRPO the wall-clock is dominated by generation, the model producing its attempts, not by the gradient step. Generation is also embarrassingly parallel, so it is the thing to spread out. The catch is that every step's attempts must come from the current policy, so whatever generates them needs the latest weights. Two levers here:

• Swap the generation backend for vLLM (continuous batching, paged attention), which speeds up rollouts on a single GPU with no distribution at all.
• For larger scale, stand up a warm pool of inference workers with a flyte.ReusePolicy so the model stays loaded across steps, have each worker generate a shard of the rollouts, and sync the policy to them each step. With LoRA the sync is cheap because the adapter delta is tiny (tens of MB), not the whole model. Flyte keeps the pool warm and orchestrates it; an RL engine (TRL's vLLM mode or veRL) owns the tight generate, sync, update loop.

Fan out independent runs (sweep configs). Because each training run is self-contained, you can launch many at once with flyte.map or asyncio.gather to sweep a hyperparameter (beta, n_numbers, seed) and compare results, or keep the best model. This is parallel experiments rather than distributed training of one model, so it costs N times the GPUs but finishes in roughly the wall-clock of one run.

What Flyte gives you is the substrate: warm worker pools, fan-out, caching, and the secure verifier. What it does not replace is the RL-engine internals such as per-step weight sync and off-policy correction. At small scale (a 0.5B on one T4) the single-GPU loop here is simplest and fast enough; the pool approach earns its complexity once the model is larger and generation is expensive.

Scaling Up to Reasoning

This exact setup is one step away from the DeepSeek-R1-Zero result reproduced by TinyZero. Add a <think>...</think><answer>...</answer> format and use a larger model, and something remarkable happens: the model discovers that reasoning step by step before answering raises its hit rate, so its responses grow longer and more structured on their own. Reasoning emerges from nothing but a correctness reward.

On a small (0.5B to 1.5B) model that emergence is weak, since the model tends to find a terse answer rather than reason out loud, so this tutorial keeps to the solid, measurable deliverable of improving solve rate. If you have a bigger GPU, add the reasoning format on a 3B to 7B base and watch the response-length curve grow. Same recipe, more capable base, visible reasoning.

Further Reading

• DeepSeekMath introduces GRPO (group-relative advantages, no value network).
• DeepSeek-R1 shows RL inducing reasoning from a base model with no reasoning demonstrations (R1-Zero).
• TinyZero reproduces the R1-Zero result on a small model using exactly this Countdown task.