Forge

A minimal, readable GRPO fine-tuning engine for HuggingFace causal language models.

GRPO (Group Relative Policy Optimization) is the reinforcement-learning algorithm behind DeepSeek-R1 and DeepSeekMath. It fine-tunes an LLM's policy using RL, but - unlike PPO - without training a separate critic (value) network. Forge is a from-scratch, single-GPU-friendly reference implementation: no RL framework, no distributed training abstraction, no config indirection. Just the algorithm, in ~400 lines you can read in one sitting.

Think of it as the nanoGPT of RL fine-tuning.

Illustrative reward curve - simulated, not from a real training run. See Training on GSM8K Math for what a real run looks like.

Package name on PyPI: forge-rl (import as from forge import ... - the top-level name forge was already taken).

---

Table of Contents

• Why GRPO over PPO
• Install
• Quickstart
• Architecture
• The GRPO Algorithm, In Detail
• Walkthrough: One Training Step
• Module Reference
• Built-in Reward Functions
• Custom Reward Functions
• GRPOConfig
• Data Pipeline
• Logging & Checkpoints
• Training on GSM8K Math
• Known Limitations
• How This Connects to AgentLens and Verdict
• Citation

---

Why GRPO over PPO

PPO needs a critic (value) network - a second model, usually the same size as the policy, trained alongside it to estimate expected future reward. That doubles GPU memory and adds a whole extra set of hyperparameters and failure modes (value loss spikes, GAE tuning, etc.).

GRPO's insight: instead of a learned value baseline, sample G responses per prompt (a group) and use the group's own mean and standard deviation as the baseline. A response is "good" if it scored above the group average, "bad" if below - normalized into a per-response advantage with no extra network required.

	PPO	GRPO
Critic network	Required (~same size as policy)	None
Baseline for advantage	Learned value function	Group mean/std of sampled rewards
Extra GPU memory	+1 model	+1 frozen copy of the policy (for the KL reference - cheaper, no gradients/optimizer state)
Samples per prompt	1 (typically)	G (group_size, default 8)
Rollout cost	Lower	Higher (G× generations per prompt)

The tradeoff: GRPO trades generation compute (you sample G completions per prompt) for architectural simplicity (no critic to train or debug).

Install

pip install forge-rl

Core dependencies are just torch>=2.1.0 and transformers>=4.40.0. Optional extras:

pip install forge-rl[demo]    # datasets + accelerate, for GSM8K training
pip install forge-rl[wandb]   # Weights & Biases logging
pip install forge-rl[all]     # everything

To install from source instead:

git clone https://github.com/madhumithakolkar/forge.git
cd forge
pip install -e ".[dev]"

Quickstart

from transformers import AutoModelForCausalLM, AutoTokenizer
from forge import GRPOTrainer, GRPOConfig, ExactMatchReward

model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM2-135M")
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M")
tokenizer.pad_token = tokenizer.eos_token

dataset = [
    {"prompt": "What is 2 + 2? Answer:", "reference": "4"},
    {"prompt": "What is 10 - 3? Answer:", "reference": "7"},
]

trainer = GRPOTrainer(
    model=model,
    tokenizer=tokenizer,
    reward_fn=ExactMatchReward(),
    config=GRPOConfig(group_size=4, num_train_epochs=1),
)
trainer.train(dataset)

This runs on CPU with a small enough model (it'll just be slow). For anything beyond a smoke test, use a GPU.

Architecture

Forge has five moving parts. GRPOTrainer (forge/core/trainer.py) is the orchestrator; it owns the policy model, a frozen reference copy, the reward function, and the optimizer. Nothing else in the library holds state.

Key structural decisions:

• The reference model is a one-time frozen deepcopy, made in GRPOTrainer.__init__, with requires_grad = False on every parameter. It never receives gradients or an optimizer state - it costs extra VRAM for weights only, not activations/gradients/Adam moments. This is cheaper than PPO's critic, which needs a full trainable second model.
• Advantages are computed once per prompt, from the raw reward scores of that prompt's G sampled responses - never across prompts, never across a whole batch. This is what "group relative" means: the baseline is local to each group.
• Gradient accumulation happens at the sample level, not response level: one train_step() call processes an entire group (G responses) for one prompt and produces a single scalar loss; the optimizer only steps every gradient_accumulation_steps prompts.

The GRPO Algorithm, In Detail

This is the literal loss Forge computes, from forge/core/grpo.py:

L_GRPO = -E[ min(ratio · A, clip(ratio, 1-ε, 1+ε) · A) ] + β · KL(π_θ ‖ π_ref)

where:
  ratio = exp(logp_θ(o|q) - logp_ref(o|q))     - policy vs. frozen reference, in prob space
  A     = normalized group advantage
  ε     = clip_epsilon   (default 0.2)
  β     = kl_beta         (default 0.01)

Walking through each piece as it's actually implemented:

1. Group advantage normalization - compute_group_advantages(rewards, group_size)

rewards = rewards.view(-1, group_size)           # (batch, G)
mean = rewards.mean(dim=1, keepdim=True)
std = rewards.std(dim=1, keepdim=True) + 1e-8     # epsilon guards div-by-zero
advantages = (rewards - mean) / std

For a single prompt's group of G=8 responses with rewards [1, 1, 0, 0, 1, 0, 0, 1] (mean 0.5, std ≈0.53), the two responses scoring 1 get a positive advantage (≈+0.94), the 0s get a negative one (≈−0.94). If every response in a group scores identically - e.g. all 0 because the model always fails, or all 1 because the task is trivial - every advantage collapses to 0 (the +1e-8 prevents this becoming a division by zero, though the signal is correctly zero either way: there's nothing to differentiate). This is the mechanism that replaces PPO's learned baseline - no critic needed, because the group itself supplies the baseline.

2. Per-token log-probabilities, masked to the response - compute_token_log_probs(...)

Given input_ids for prompt + response concatenated, and a response_mask that is 1 only over the response tokens:

logits = model(input_ids, attention_mask).logits            # (batch, seq_len, vocab)
log_probs = log_softmax(logits[:, :-1, :], dim=-1)           # predict token t+1 from position t
target_ids = input_ids[:, 1:]
selected = gather(log_probs, target_ids)                     # per-token logp of the actual next token
masked = selected * response_mask_shifted                    # zero out prompt-token positions
return masked.sum(dim=1) / token_counts                      # mean logp over response tokens only

This is standard next-token teacher-forcing, but explicitly masked so the prompt tokens contribute nothing to the log-probability used in the ratio - the model is never rewarded or penalized for how likely it found tokens it didn't generate. The result is a single scalar mean log-probability per response, computed once for the trainable policy (requires_grad active) and once for the frozen reference (inside torch.no_grad()).

3. Clipped surrogate objective + KL penalty - grpo_loss(...)

log_ratio = policy_log_probs - ref_log_probs
ratio = exp(log_ratio)
clipped_ratio = clamp(ratio, 1 - ε, 1 + ε)
policy_loss = -min(ratio * A, clipped_ratio * A).mean()

kl = (ratio - log_ratio - 1).mean()     # k3 estimator, always >= 0
loss = policy_loss + β * kl

The min(...) clip is the same PPO-style trust-region trick: if the policy would move too far from the reference in a direction that increases the objective, the gradient is clipped off. It protects against a single very-high-advantage response causing a destructively large update. The KL term is a second, separate regularizer that pulls the policy back toward the reference distribution regardless of sign - Forge uses the k3 KL estimator (exp(x) - x - 1 where x = log_ratio), which is non-negative and lower-variance than the naive log_ratio estimator.

Walkthrough: One Training Step

Tracing GRPOTrainer.train_step() for a single prompt end to end:

1. Tokenize the prompt. prompt_enc = tokenizer(prompt, return_tensors="pt").

2. Sample a group. generate_group() repeats the prompt group_size times along the batch dimension and calls model.generate(..., do_sample=True, temperature=config.temperature) - one forward-generation call produces all G responses at once, each a stochastically different completion of the same prompt.

3. Score the group. The reward function is called once with G copies of the prompt and the G decoded responses: reward_fn([prompt]*G, responses, references=[reference]*G). Returns a RewardOutput with G float scores.

4. Compute advantages. compute_group_advantages(rewards, group_size) - see above. This is the only place the reward scale matters; scores can be any range (0/1, −1..1, continuous) since they get normalized away.

5. Compute log-probs, per response. For each of the G responses: re-encode prompt + response together, build the response_mask, then run it through both the trainable policy (gradients on) and the frozen reference (torch.no_grad()). This is a Python loop over G items - not batched - so each response can have a different length without padding complications.

6. Compute the loss. grpo_loss(policy_log_probs, ref_log_probs, advantages, clip_epsilon, kl_beta) - one scalar for the whole group.

7. Back in train(): (loss / gradient_accumulation_steps).backward() accumulates gradients; every gradient_accumulation_steps prompts, clip_grad_norm_ is applied, the AdamW optimizer steps, the cosine-with-warmup LR scheduler steps, and global_step increments. Logging (logging_steps) and checkpointing (save_steps) are both gated on global_step, not on the raw sample index - so with small datasets or high gradient_accumulation_steps, make sure logging_steps/save_steps are low enough to actually fire (total optimizer steps ≈ len(dataset) * num_train_epochs / gradient_accumulation_steps).

Module Reference

forge/
  __init__.py                 Public API: GRPOTrainer, GRPOConfig, all reward classes
  core/
    grpo.py                   GRPOConfig dataclass + pure-tensor GRPO math (no model/IO deps)
    trainer.py                GRPOTrainer - owns model, ref_model, optimizer, the training loop
    reward.py                 BaseReward ABC + 4 built-in reward functions
  data/
    gsm8k.py                  GSM8K loader (HF datasets) + a hardcoded 20-sample mock set
  logging/
    logger.py                 ForgeLogger - console + JSONL, optional wandb passthrough
tests/
  test_grpo.py                Pure tensor-math tests: advantage normalization, loss shape/values, KL behavior
  test_reward.py               Reward function correctness
  test_data.py                 Prompt formatting / answer extraction
examples/
  train_math.py                Full GSM8K run on SmolLM2-1.7B (needs GPU)
  custom_reward.py             Template for writing your own reward function

grpo.py deliberately has zero dependency on trainer.py, the tokenizer, or generation - it only operates on tensors (rewards, logits, log_probs). That's what makes it directly unit testable with tests/test_grpo.py on CPU, no model or GPU required: compute_group_advantages and grpo_loss are tested with hand-constructed tensors, not a real forward pass.

Built-in Reward Functions

All rewards implement BaseReward.__call__(self, prompts: list[str], responses: list[str], **kwargs) -> RewardOutput, returning one float score per response. Scores can be any scale - GRPO's group normalization makes the absolute scale irrelevant, only the relative ordering within a group matters.

Reward	Signature-relevant kwargs	Behavior
ExactMatchReward(references=None)	references: list[str] (at init or call time)	1.0 if the reference string appears (case-insensitive) in the response, else 0.0
FormatReward(format_fn)	-	1.0 if format_fn(response) is truthy, else 0.0. format_fn is any str -> bool callable
LengthPenaltyReward(optimal_length=200, min_length=50, max_length=500)	-	1.0 at the optimal word count, decaying linearly toward 0.0 outside [min_length, max_length]
CompositeReward(rewards_and_weights)	forwards matching kwargs to each sub-reward	Weighted average of component scores; weights are normalized to sum to 1 automatically

CompositeReward inspects each sub-reward's __call__ signature and only forwards the kwargs it actually accepts (e.g. references goes to ExactMatchReward but not FormatReward), so you can freely mix reward types that do and don't need reference answers in the same composite.

reward_fn = CompositeReward([
    (ExactMatchReward(), 0.7),
    (FormatReward(lambda r: "answer:" in r.lower()), 0.2),
    (LengthPenaltyReward(optimal_length=150), 0.1),
])

Custom Reward Functions

from forge.core.reward import BaseReward, RewardOutput

class MyReward(BaseReward):
    def __call__(self, prompts, responses) -> RewardOutput:
        scores = [1.0 if "correct" in r else 0.0 for r in responses]
        return RewardOutput(scores=scores)

RewardOutput is a plain dataclass: scores: list[float], plus an optional metadata: dict for anything you want to inspect later (not currently logged by ForgeLogger, but available on the object returned from reward_fn(...)). See examples/custom_reward.py for a fuller template.

GRPOConfig

All fields, from forge/core/grpo.py:

Field	Default	Description
group_size	8	Number of responses sampled per prompt (G). Directly multiplies generation cost per step.
clip_epsilon	0.2	PPO-style clip range [1-ε, 1+ε] on the policy/reference probability ratio
kl_beta	0.01	Weight of the KL penalty term against the frozen reference policy
temperature	0.9	Sampling temperature used during generate_group()
max_new_tokens	512	Max tokens generated per response
learning_rate	5e-6	AdamW learning rate
gradient_accumulation_steps	4	Prompts accumulated before an optimizer step
max_grad_norm	1.0	Gradient clipping norm, applied right before optimizer.step()
num_train_epochs	3	Passes over the dataset
save_steps	100	Save a checkpoint every N optimizer steps
logging_steps	10	Log metrics every N optimizer steps
output_dir	"forge_output"	Root directory for checkpoints ({output_dir}/{step}/) and training_log.jsonl
use_wandb	False	Mirror logged metrics to Weights & Biases
wandb_project	"forge"	W&B project name
wandb_run_name	"grpo_run"	W&B run name
seed	42	torch.manual_seed set once at the start of train()

Note save_steps/logging_steps count optimizer steps, not dataset samples - with gradient_accumulation_steps=4 on a 20-sample dataset for 1 epoch, you get only 5 optimizer steps total, so logging_steps=10 would never fire. Scale these down for small datasets or smoke tests.

Data Pipeline

forge/data/gsm8k.py provides two loaders, both returning list[{"prompt": str, "reference": str}] - the shape GRPOTrainer.train() expects:

• load_gsm8k(split="train", max_samples=500) - streams from openai/gsm8k via HuggingFace datasets (requires the demo extra). Parses the #### <number> suffix GSM8K uses to mark the ground-truth final answer.
• load_gsm8k_mock(n=20) - 20 hardcoded, hand-verified grade-school problems with correct answers, no network or dataset download required. This is what backs the CPU smoke test and can be used to sanity-check a new reward function or config before spending GPU time.

Both funnel every question through the same format_prompt():

You are a math reasoning assistant. Solve the problem step by step.
At the end of your solution, write the final answer on its own line in the format: Answer: <number>

Problem: <question>

Solution:

extract_answer() is a standalone utility for pulling the final numeric answer back out of a generated solution (used by reward functions or eval code, not by the loaders themselves) - it looks for an Answer: line first, falling back to the last number that appears anywhere in the text.

Logging & Checkpoints

ForgeLogger (forge/logging/logger.py) writes to three places, all driven off GRPOConfig:

1. Console - one line per logged step: step {N} | epoch {E} | loss {L} | reward {R} | kl {K}.
2. JSONL file - {output_dir}/training_log.jsonl, one JSON object per logged step (everything in metrics except the raw responses list, plus step/epoch). Appended, not overwritten, so resuming a run's logs concatenates.
3. Weights & Biases (optional) - only if use_wandb=True and wandb is importable; otherwise it prints a warning and silently falls back to console + file only, rather than raising.

Checkpoints are full model.save_pretrained() + tokenizer.save_pretrained() dumps (not just adapter weights - Forge doesn't currently do LoRA/PEFT) written to {output_dir}/{step}/, taken every save_steps optimizer steps and once more, unconditionally, at the end of training under {output_dir}/final/.

Training on GSM8K Math

See examples/train_math.py for a full training run on GSM8K using SmolLM2-1.7B-Instruct. Requires a GPU with 16GB+ VRAM (a T4 works if you drop group_size to 4). Estimated cost on a RunPod A100: ~$1.50 for a 3-epoch run over 500 samples (~2 hours).

What to expect from a real run:

• Reward starts around 0.1–0.2 (a base model rarely gets grade-school math right zero-shot)
• After ~200 steps: reward climbs to roughly 0.4–0.6
• By the end of training: reward reaches roughly 0.6–0.8 on the training set

Reward curves are logged to {output_dir}/training_log.jsonl as the run progresses - pipe that file into the same style of plot as scripts/plot_reward_curve.py to reproduce the hero image above with real numbers.

For a quick, GPU-free correctness check of the training loop itself (not real learning), swap in HuggingFaceTB/SmolLM2-135M, set USE_MOCK_DATA = True, and shrink group_size/max_new_tokens

• it will run to completion on CPU in a couple of minutes; the reward will just be noise, which is expected from an untrained tiny model.

Known Limitations

• No batching across prompts. Each train_step() handles exactly one prompt's group; there's no cross-prompt batch dimension, so throughput on a single GPU is bounded by group_size × sequence length, not by a configurable batch size.
• No LoRA/PEFT support. Checkpoints save full model weights.
• No distributed training. Single process, single device - next(model.parameters()).device determines where everything runs.
• The per-response log-prob loop is Python-level, not vectorized (see step 5 in the walkthrough above), trading some throughput for simplicity and avoiding padding/masking bugs across variable-length responses.

These are intentional scope cuts, not oversights - Forge optimizes for being readable end to end over being maximally fast. Contributions that add capability without compromising that are welcome.

How This Connects to AgentLens and Verdict

Forge fine-tunes the model. Verdict evaluates whether its outputs are good. AgentLens evaluates whether its agent behaviour (tool calls, trajectory) is correct. Together they form a complete RL fine-tuning and evaluation stack.

Citation

If you use Forge, please cite the original GRPO paper:

DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models (2024)