Metadata-Version: 2.4
Name: forge-rl
Version: 0.1.0
Summary: Minimal GRPO fine-tuning engine for language models
License: MIT
Project-URL: Homepage, https://github.com/madhumithakolkar/forge
Keywords: grpo,rlhf,rl,fine-tuning,llm,language-model,reinforcement-learning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1.0
Requires-Dist: transformers>=4.40.0
Provides-Extra: demo
Requires-Dist: datasets>=2.19.0; extra == "demo"
Requires-Dist: accelerate>=0.29.0; extra == "demo"
Provides-Extra: wandb
Requires-Dist: wandb>=0.16.0; extra == "wandb"
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: black; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Provides-Extra: all
Requires-Dist: datasets>=2.19.0; extra == "all"
Requires-Dist: accelerate>=0.29.0; extra == "all"
Requires-Dist: wandb>=0.16.0; extra == "all"
Dynamic: license-file

# 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.

![Reward Curve](assets/reward_curve.png)
*Illustrative reward curve — simulated, not from a real training run. See [Training on GSM8K Math](#training-on-gsm8k-math) for what a real run looks like.*

Planned package name on PyPI: `forge-rl` (import as `from forge import ...` — the top-level name `forge` was already taken). Not yet published; see [Install](#install).

---

## Table of Contents

- [Why GRPO over PPO](#why-grpo-over-ppo)
- [Install](#install)
- [Quickstart](#quickstart)
- [Architecture](#architecture)
- [The GRPO Algorithm, In Detail](#the-grpo-algorithm-in-detail)
- [Walkthrough: One Training Step](#walkthrough-one-training-step)
- [Module Reference](#module-reference)
- [Built-in Reward Functions](#built-in-reward-functions)
- [Custom Reward Functions](#custom-reward-functions)
- [GRPOConfig](#grpoconfig)
- [Data Pipeline](#data-pipeline)
- [Logging & Checkpoints](#logging--checkpoints)
- [Training on GSM8K Math](#training-on-gsm8k-math)
- [Known Limitations](#known-limitations)
- [How This Connects to AgentLens and Verdict](#how-this-connects-to-agentlens-and-verdict)
- [Citation](#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

> **Not yet published to PyPI.** Install from source for now:

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

Once released, the plan is:

```bash
pip install forge-rl
```

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

```bash
pip install -e ".[demo]"    # datasets + accelerate, for GSM8K training
pip install -e ".[wandb]"   # Weights & Biases logging
pip install -e ".[all]"     # everything
```

## Quickstart

```python
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.

```mermaid
flowchart TB
    subgraph Input
        DS["Dataset<br/>[{prompt, reference}, ...]"]
    end

    subgraph Trainer["GRPOTrainer.train_step()"]
        direction TB
        GEN["generate_group()<br/>sample G responses per prompt<br/>(policy model, no grad)"]
        RM["reward_fn(prompts, responses)<br/>BaseReward subclass"]
        ADV["compute_group_advantages()<br/>normalize within group"]
        ENC["encode_prompt_response()<br/>build response_mask per response"]
        PLP["compute_token_log_probs()<br/>on policy model (grad on)"]
        RLP["compute_token_log_probs()<br/>on frozen ref_model (no grad)"]
        LOSS["grpo_loss()<br/>clipped surrogate + KL penalty"]

        GEN --> RM --> ADV
        GEN --> ENC
        ENC --> PLP
        ENC --> RLP
        PLP --> LOSS
        RLP --> LOSS
        ADV --> LOSS
    end

    subgraph Update["Optimizer step (every gradient_accumulation_steps)"]
        BW["loss.backward()"]
        CLIP["clip_grad_norm_"]
        STEP["AdamW.step() + cosine LR scheduler"]
    end

    subgraph Sidecar
        POLICY["model (policy)<br/>trainable"]
        REF["ref_model<br/>frozen deepcopy, made once at init"]
        LOG["ForgeLogger<br/>console + JSONL + optional wandb"]
        CKPT["save_checkpoint()<br/>every save_steps + final"]
    end

    DS --> GEN
    POLICY -.weights.-> GEN
    POLICY -.weights.-> PLP
    REF -.weights.-> RLP
    LOSS --> BW --> CLIP --> STEP
    STEP -.updates.-> POLICY
    LOSS -.metrics.-> LOG
    STEP -.periodically.-> CKPT
```

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)`

```python
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
`0`s 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:

```python
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(...)`

```python
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`).

```mermaid
sequenceDiagram
    participant T as train()
    participant S as train_step()
    participant P as policy model
    participant R as ref_model (frozen)
    participant Rw as reward_fn

    T->>S: prompt, references
    S->>P: generate_group() [G samples, no_grad]
    P-->>S: G responses
    S->>Rw: reward_fn(prompts×G, responses, references×G)
    Rw-->>S: G reward scores
    S->>S: compute_group_advantages()
    loop for each of G responses
        S->>P: compute_token_log_probs (grad on)
        S->>R: compute_token_log_probs (no_grad)
    end
    S->>S: grpo_loss(policy_lp, ref_lp, advantages)
    S-->>T: loss, metrics
    T->>T: (loss / grad_accum).backward()
    alt every gradient_accumulation_steps
        T->>P: clip_grad_norm_ + optimizer.step() + scheduler.step()
    end
    alt every logging_steps
        T->>T: logger.log(step, epoch, metrics)
    end
    alt every save_steps
        T->>T: save_checkpoint(step)
    end
```

## 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.

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

## Custom Reward Functions

```python
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](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)
