Fine-tune a model with GRPO for mathematical reasoning

Post-training a language model with reinforcement learning used to mean juggling a policy network, a reward model, and a separate critic network—three moving parts that demand substantial GPU memory and careful synchronisation. Group Relative Policy Optimization (GRPO), introduced in the DeepSeekMath paper and later used at the core of DeepSeek-R1's training pipeline, eliminates the critic entirely. Instead of learning a value baseline, GRPO samples a group of completions for each prompt, computes their relative rewards within that group, and uses those group statistics as the advantage signal. The result is a reinforcement learning loop that fits inside a single GPU session without sacrificing the quality gains that make RL post-training worthwhile for domain-specific reasoning.

Who needs this? Any team that has a capable base model but finds that supervised fine-tuning alone leaves structured reasoning tasks—multi-step mathematics, formal proofs, constrained code generation—still falling short. GRPO is particularly effective here because the reward signal can be decomposed into format correctness (did the model emit the expected reasoning scaffold?) and answer correctness (is the final answer mathematically equivalent to the ground truth?), two objectives that combine naturally into a multi-reward training loop. If you have been watching the broader trend of test-time compute scaling and want to push a smaller model further on reasoning benchmarks without building a full RLHF pipeline, GRPO is the most practical entry point available today.

The configuration used here targets a single GPU with at least 16 GB of VRAM. The guide uses Qwen/Qwen2-0.5B-Instruct with LoRA adapters, keeping the trainable parameter count low enough to fit on an A10G or a cloud T4. With num_generations=4 and max_completion_length=64, one epoch over 5 % of NuminaMath-TIR is a practical demo workload. Scaling to larger models or more generations multiplies memory proportionally. This guide is adapted from the Hugging Face Cookbook's fine_tuning_llm_grpo_trl.ipynb, released under the Apache-2.0 licence.

---

Prerequisites

Hardware: A single NVIDIA GPU with ≥ 16 GB VRAM (A10G, A100, or equivalent). CPU-only execution is not practical for training.

Software versions tested by the notebook maintainers:

• transformers == 4.47.1
• trl == 0.14.0
• datasets == 3.2.0
• peft == 0.14.0
• accelerate == 1.2.1
• math_verify == 0.3.3

Accounts: A Hugging Face account with a write-access token. You will push the trained adapter to the Hub at the end of training, so the token must have repository-creation permission.

---

Step 1: Install dependencies and authenticate

Install the three packages that are not part of a standard Hugging Face environment, then authenticate so the Hub push at the end of training works without interruption.

!pip install  -U -q trl peft math_verify
# Tested with transformers==4.47.1, trl==0.14.0, datasets==3.2.0, peft==0.14.0, accelerate==1.2.1, math_verify==0.3.3

from huggingface_hub import notebook_login

notebook_login()

math_verify is the library that parses LaTeX expressions from model completions and compares them symbolically against ground-truth solutions. It is not bundled with TRL, so the explicit install is mandatory.

---

Step 2: Load and format the dataset

NuminaMath-TIR is a mathematics reasoning dataset that pairs problem statements with detailed tool-integrated reasoning chains and final answers. The 5 % slice used here is enough to demonstrate the training loop without requiring multi-hour runtimes.

from datasets import load_dataset

dataset_id = 'AI-MO/NuminaMath-TIR'
train_dataset, test_dataset = load_dataset(dataset_id, split=['train[:5%]', 'test[:5%]'])

print(train_dataset)

print(train_dataset[0])

Inspect the raw sample before transforming it. The key columns are problem (the question text) and solution (the ground-truth answer, expressed in LaTeX). The messages column contains the original reasoning chain but is not used directly by GRPO—you are training the model to generate that reasoning, not to copy it.

Next, wrap each problem in the system prompt that instructs the model to structure its output with explicit <think> and <answer> tags. This prompt mirrors the one used in DeepSeek-R1's cold-start stage and is what the format reward function will verify.

SYSTEM_PROMPT = (
    "A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant "
    "first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning "
    "process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., "
    "<think> reasoning process here </think><answer> answer here </answer>"
)

def make_conversation(example):
    return {
        "prompt": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": example["problem"]},
        ],
    }

train_dataset = train_dataset.map(make_conversation)
test_dataset = test_dataset.map(make_conversation)

print(train_dataset[0]['prompt'])

train_dataset = train_dataset.remove_columns(['messages', 'problem'])
print(train_dataset)

After the column removal, the training dataset contains exactly two columns: prompt (a list of role/content dictionaries) and solution (the LaTeX answer string). Keeping solution in the dataset is critical—GRPOTrainer forwards it to the reward functions as a keyword argument.

---

Step 3: Load the model and attach LoRA adapters

GRPO requires the policy model and a frozen reference copy of it. TRL manages the reference model internally, but it doubles the GPU memory footprint of the base weights. LoRA cuts the trainable parameter count sharply and keeps the reference copy cheap, because only the adapter weights change—the frozen base weights are shared.

import torch
from transformers import AutoModelForCausalLM

model_id = "Qwen/Qwen2-0.5B-Instruct"
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    device_map="auto",
)

from peft import LoraConfig, get_peft_model

lora_config = LoraConfig(
    task_type="CAUSAL_LM",
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    target_modules=["q_proj", "v_proj"],
)

model = get_peft_model(model, lora_config)

model.print_trainable_parameters()

With r=8 applied only to the query and value projection matrices, the trainable parameter count is typically under 1 % of total model size. The print_trainable_parameters() call confirms the exact count and is worth checking before committing to a long training run.

---

Step 4: Define the reward functions

GRPO's training signal comes entirely from reward functions—there is no labelled "correct completion" to supervise against. Two reward functions are composed here, each returning a list of floats aligned with the batch of completions.

Format reward — checks that the completion matches the required XML-like tag structure. A completion that does not contain properly ordered <think> and <answer> blocks scores zero regardless of whether the answer is correct.

import re
def format_reward(completions, **kwargs):
    """Reward function that checks if the completion has a specific format."""
    pattern = r"^<think>.*?</think>\s*<answer>.*?</answer>$"
    completion_contents = [completion[0]["content"] for completion in completions]
    matches = [re.match(pattern, content) for content in completion_contents]
    rewards_list = [1.0 if match else 0.0 for match in matches]
    return [1.0 if match else 0.0 for match in matches]

Accuracy reward — parses the model's <answer> block and the ground-truth solution using math_verify, then checks symbolic equivalence. If the ground-truth cannot be parsed (e.g. it is a free-text answer), the function awards a score of 1.0 to avoid penalising the model for unparseable references.

from math_verify import LatexExtractionConfig, parse, verify
def accuracy_reward(completions, **kwargs):
    """Reward function that checks if the completion is the same as the ground truth."""
    solutions = kwargs['solution']
    completion_contents = [completion[0]["content"] for completion in completions]
    rewards = []
    for content, solution in zip(completion_contents, solutions):
        gold_parsed = parse(solution, extraction_mode="first_match", extraction_config=[LatexExtractionConfig()])
        answer_parsed = parse(content, extraction_mode="first_match", extraction_config=[LatexExtractionConfig()])
        if len(gold_parsed) != 0:
            try:
                rewards.append(float(verify(answer_parsed, gold_parsed)))
            except Exception:
                rewards.append(0.0)
        else:
            rewards.append(1.0)
    return rewards

Both functions receive completions as a list-of-lists (one inner list per sample, each containing a single assistant turn dictionary). The **kwargs mechanism is how GRPOTrainer passes dataset columns—here solution—into reward functions at runtime.

---

Step 5: Configure and launch training

GRPOConfig extends TrainingArguments with GRPO-specific parameters. The table below explains the non-obvious ones before you read the configuration block.

from trl import GRPOConfig

# Configure training arguments using GRPOConfig
training_args = GRPOConfig(
    output_dir="Qwen2-0.5B-GRPO-test",
    learning_rate=1e-5,
    remove_unused_columns=False, # to access the solution column in accuracy_reward
    gradient_accumulation_steps=16,
    num_train_epochs=1,
    bf16=True,

    # Parameters that control de data preprocessing
    max_completion_length=64, # default: 256
    num_generations=4, # default: 8

    # Parameters related to reporting and saving
    report_to=["tensorboard"],
    logging_steps=10,
    push_to_hub=True,
    save_strategy="steps",
    save_steps=10,
)

Pass both reward functions as a list to GRPOTrainer. The trainer sums their outputs to form the combined reward signal used for advantage estimation.

from trl import GRPOTrainer

trainer = GRPOTrainer(
    model=model,
    reward_funcs=[format_reward, accuracy_reward],
    args=training_args,
    train_dataset=train_dataset
)

trainer.train()

trainer.save_model(training_args.output_dir)
trainer.push_to_hub(dataset_name=dataset_id)

---

Step 6: Evaluate the trained model

Load the pushed adapter and run inference on the held-out test set to verify that the model now structures its responses correctly.

from transformers import AutoTokenizer

model_id = "sergiopaniego/Qwen2-0.5B-GRPO"
trained_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype="auto",
    device_map="auto",
)
trained_tokenizer = AutoTokenizer.from_pretrained(model_id)

print(test_dataset['prompt'][0])

import time

def generate_with_reasoning(prompt):
  # Build the prompt from the dataset
  prompt = " ".join(entry['content'] for entry in prompt)

  # Tokenize and move to the same device as the model
  inputs = trained_tokenizer(prompt, return_tensors="pt").to(trained_model.device)

  # Generate text without gradients
  start_time = time.time()
  with torch.no_grad():
      output_ids = trained_model.generate(**inputs, max_length=500)
  end_time = time.time()

  # Decode and extract model response
  generated_text = trained_tokenizer.decode(output_ids[0], skip_special_tokens=True)

  # Get inference time
  inference_duration = end_time - start_time

  # Get number of generated tokens
  num_input_tokens = inputs['input_ids'].shape[1]
  num_generated_tokens = output_ids.shape[1] - num_input_tokens

  return generated_text, inference_duration, num_generated_tokens

prompt = test_dataset['prompt'][0]
generated_text, inference_duration, num_generated_tokens = generate_with_reasoning(prompt)
print(generated_text)

print(f"Inference time: {inference_duration:.2f} seconds")
print(f"Generated tokens: {num_generated_tokens}")

prompt_text = " ".join(entry['content'] for entry in prompt)
response_text = generated_text[len(prompt_text):].strip()
print(response_text)

The final print isolates only the model's response by slicing away the prompt prefix. A well-trained model produces a <think> block containing intermediate reasoning steps followed by an <answer> block containing the final value. If neither tag appears, the format reward has not taken hold—see the failure modes section below.

---

What to watch out for

remove_unused_columns=False is not optional. The Hugging Face Trainer base class removes dataset columns that are not named arguments of the model's forward method. Because solution is not a model input, it gets silently dropped before reaching the reward functions unless you explicitly disable this behaviour. The accuracy reward will then receive an empty kwargs['solution'] and fail—sometimes silently, sometimes with a cryptic index error deep inside math_verify.

max_completion_length=64 is very short for real mathematical reasoning. The tutorial sets it deliberately low to fit a demo on constrained hardware. At 64 tokens a model cannot develop multi-step reasoning chains; it will learn to produce the correct tag format but may produce truncated or trivially short answers. For production use, budget at least 256–512 tokens and be prepared to upgrade your GPU accordingly.

Reward hacking on the format signal. A model can achieve a format reward of 1.0 by emitting <think></think><answer>x</answer> with an empty think block. Watch the format_reward and accuracy_reward curves separately in TensorBoard: if format reward saturates early but accuracy reward stays near zero, the model has learned to game the tag check without actually reasoning. Tightening the regex—for example, requiring a minimum number of characters inside <think>—is one mitigation.

LoRA target modules matter. The configuration targets only q_proj and v_proj. Adding k_proj, o_proj, or the MLP layers increases adapter expressiveness but also increases memory and convergence time. For tasks requiring more representational change than format learning, expanding the target modules is worth the cost.

num_generations must be ≥ 2. GRPO computes advantages relative to the group mean. With only one completion there is no variance to normalise against, and the advantage is always zero. The default in TRL is 8; this guide uses 4 to halve the sampling cost. Going below 2 produces a degenerate training signal.

BF16 requires Ampere or newer. On V100 or T4 instances, bf16=True will either raise an error or silently fall back to FP32, doubling memory usage. If you are on an older instance, set bf16=False and fp16=True instead, though FP16 is more prone to gradient overflow during RL training.

Dataset slice affects reward coverage. Using 5 % of NuminaMath-TIR means the model sees a narrow problem distribution. If your target domain is geometry or number theory specifically, verify that the 5 % slice is not dominated by algebra before concluding that GRPO is failing on your problem type.

---

Where to go next

The training loop here is intentionally minimal—one epoch, a 0.5 B parameter model, 5 % of the dataset. The natural extensions are: increasing max_completion_length to allow genuine chain-of-thought development, scaling to a larger base model such as Qwen2-7B, and adding further reward functions such as step-count penalties or unit-consistency checks. The Hugging Face Open-R1 project applies the same principles at full DeepSeek-R1 scale and is the most authoritative public reference for what GRPO can achieve with sufficient compute. For broader context on how reinforcement learning has evolved as a post-training technique and the debates around its alignment properties, the work traced back through RLHF's origins provides useful background. The full source notebook, including additional commentary and visualisations, is available at github.com/huggingface/cookbook under the Apache-2.0 licence.