23f2002275

docs(readme): link raw training logs + training scripts per organizer ask

882b74a 4 months ago

14.7 kB

FATHOM — First RL-Trained Recursive Language Model

FATHOM is an OpenEnv environment + GRPO training pipeline that teaches a small open-source language model (Qwen 2.5 Coder 1.5B, 4-bit + LoRA) to use a Recursive Language Model scaffold well: slice long contexts with Python, grep for relevant regions, delegate to sub-LM calls only when needed, and answer questions about documents that are 50× larger than its own native context window.

Submitted to the Meta × PyTorch × Hugging Face OpenEnv Hackathon Grand Finale (Bangalore, April 25–26, 2026 — Theme 2: Long-Horizon Planning).

Submission Links (Judges Start Here)

Artifact	URL
Environment Space (Hub page)	https://huggingface.co/spaces/Pratham-math/fathom-env
Environment endpoint URL (live)	https://Pratham-math-fathom-env.hf.space
Demo UI (Streamlit Space)	https://huggingface.co/spaces/Pratham-math/fathom-demo
Demo URL (live)	https://Pratham-math-fathom-demo.hf.space
Health check	https://Pratham-math-fathom-env.hf.space/healthz
Code repo (HF — primary)	https://huggingface.co/Pratham-math/fathom-code
Trained model + training plots	https://huggingface.co/Pratham-math/fathom-1.5b-grpo
Colab reproducer notebook	`notebooks/fathom_train.ipynb` (open via HF code repo)
Mini-blog	https://huggingface.co/spaces/Pratham-math/fathom-blog
W&B training run (v2 — successful learning)	https://wandb.ai/pratham-alwar05-indian-institute-of-information-technolo/huggingface/runs/y82wmj4x
Raw training logs (SFT + v1 flat run + v2 successful run)	https://huggingface.co/Pratham-math/fathom-1.5b-grpo/tree/main/logs
Training scripts	`train/grpo.py`, `scripts/job_train.sh`

The HF Space /healthz endpoint cold-starts the first time it's hit; if you get a 503, refresh once and it returns 200.

Architecture

$FATHOM architecture$

A TRL GRPOTrainer runs Qwen 2.5 Coder 1.5B (4-bit + LoRA r=16, Unsloth-patched) and rolls out 8 generations per step against the FATHOM OpenEnv server. The env exposes two tool primitives — a sandboxed Python REPL and a recursive llm() call — so the agent can decompose long documents on its own. A composable, deterministic verifier (4 components: format_gate × correctness + token_budget + recursion_efficiency) returns the scalar reward.

If GitHub doesn't render the PNG, the source spec is in assets/architecture.mmd.

A note on the role of the env in training

TRL 1.2.0 with transformers==4.56.2 does not yet expose multi-turn env-tool calls inside GRPOTrainer.train() (the tools= / environment_factory= kwargs require transformers>=5.0, and a custom rollout_func= was outside our time budget). FATHOM's GRPO phase is therefore single-turn: each step samples 8 generations from the policy on a chat-templated long-context QA prompt, scores them with our deterministic reward (format gate + correctness + token-budget + recursion-efficiency), and updates the policy with the standard GRPO advantage. The env is exercised end-to-end at inference time — the demo Space runs full multi-turn REPL + recursive llm() episodes against the trained model. Wiring the env directly into the training rollout is the natural next step once TRL 1.3 / transformers 5 ships.

Problem and Why It Matters

Long-context inference keeps growing (1M-token Gemini, 200K Claude), but small open-weights models are still capped at 4K–32K tokens. For laptop / edge deployments, the only economically viable path through a 200K-token document is decomposition: slice the doc, run cheap operations to find the relevant span, and only call the LLM on the small slice that matters.

Recursive Language Models (RLMs) formalise this. Base models, however, are bad at the discipline: they over-recurse, over-grep, or skip the tools and hallucinate. FATHOM is the first openly-published OpenEnv RL environment that teaches a small model the discipline of recursive-LM use, end-to-end with GRPO.

Environment Design (OpenEnv)

FATHOM follows the OpenEnv server contract:

POST /reset — start an episode, returns initial Observation (the document + question)
POST /step — execute one tool action (REPL or llm() call), returns next Observation + reward signal
GET /state — debug introspection
GET /healthz — readiness probe

Tool primitives:

Primitive	Implementation	Safety
`repl(code)`	RestrictedPython AST filter + subprocess sandbox	Network off, ulimit'd CPU/memory, ephemeral cwd, non-root
`llm(prompt, slice)`	Recursive sub-call into the same model	Depth capped at 2 in training, 4 at demo time

Implementation:

env/server/app.py — FastAPI surface
env/server/environment.py — Observation/Action types + episode state
env/server/repl.py — sandboxed REPL
env/server/llm_primitive.py — recursive sub-call dispatcher
openenv.yaml — Hub manifest

Reward Design

Deterministic, composable, no LLM-as-judge in the training loop. Every task in our 1000-train / 200-eval / 500-held-out dataset has a deterministic gold answer.

Weight	Component	Source	What it scores
+0.10 bonus	`format_gate.py`	additive (soft)	`<answer>…</answer>` tags present (v2 — was a hard multiplier in v1)
0.70	`correctness.py`	additive	Normalised exact-match against gold
0.15	`token_budget.py`	penalty	Total tool-call tokens (Mercor sub-prize aligned)
0.15	`recursion_efficiency.py`	additive (correctness-gated)	Linear decay on `llm()` call count, only counts when answer is correct

Composition: rewards/compose.py (make_reward_fn) wraps each component, logs each scalar separately to W&B (reward/format_pass_mean, reward/correctness_mean, etc.), and exposes the composite to TRL's GRPOTrainer.reward_funcs interface.

Anti-reward-hacking — five attacks, audited before training

REWARD_AUDIT.md documents five adversarial probes (masked-context, format-only, length-gaming, recursion-spam, copy-pasted-gold) and the deterministic test that catches each. pytest -m reward_audit re-runs them on every change.

Training Pipeline (Unsloth + TRL GRPO)

1) Smoke test — required gate

Runs one GRPO step against the env, writes outputs/smoke/SMOKE_RESULT.md. Last green run: see SMOKE_RESULT.md (verdict: GO, 6/6 checks PASS, 47 s on HF Jobs a10g-large).

python -m uvicorn env.server.app:app --host 0.0.0.0 --port 8001
python -m train.smoke_test --env-url http://localhost:8001

2) Full training scripts

scripts/job_smoke.sh — HF Jobs smoke (~1 min on a10g-large)
scripts/job_train.sh — full SFT → GRPO → plots → push (~40 min Path 3 on a10g-large; ~5 h Path 1 on a100-large)
scripts/job_sft_only.sh — SFT-only fallback path
scripts/make_plots.py — generates reward / loss / grad-norm / KL PNGs from trainer_state.json
scripts/submission_preflight.py — README + Dockerfile + manifest validator (must pass before submitting)

3) Core training modules

train/model_load.py — Unsloth-with-HF-fallback loader
train/sft.py — TRL SFTTrainer warm-start
train/grpo.py — TRL GRPOTrainer. Rollout backend is gated by FATHOM_USE_VLLM env var: defaults to vllm_mode='colocate' (per TRL #4543); set FATHOM_USE_VLLM=0 to fall back to HF generate() for QLoRA-stable rollouts (avoids the IS-ratio collapse from merged-4bit weight drift).
train/smoke_test.py — 6-check pipeline gate

4) Hyperparameters (from `configs/train/grpo.yaml`)

num_generations: 8
beta: 0.04                    # KL floor (EDGE-GRPO §3.2)
learning_rate: 5.0e-6         # 4-bit safe band
max_grad_norm: 0.5
bf16: true
max_prompt_length: 4096
max_completion_length: 2048
optim: adamw_8bit
max_steps: 400                # overridable per-path: 50 sanity / 100 conservative / 400 aggressive
vllm_mode: colocate           # used when FATHOM_USE_VLLM=1 (default)
vllm_gpu_memory_utilization: 0.45

v2 run note: the successful y82wmj4x W&B run was launched with FATHOM_USE_VLLM=0 and train.max_steps=200. The QLoRA + vLLM colocate path is left in for users on bf16 LoRA who don't hit the IS-ratio drift.

Training Evidence

SFT warm-start — model learns the format and answer style cleanly

$SFT loss$ SFT loss drops from 3.20 → 0.29 across 63 steps on 500 Claude-generated traces. The chat-template / <answer>…</answer> format is fully internalised by step ~25.

$SFT token accuracy$ Mean per-token accuracy climbs from 0.46 → 0.93 over the SFT epoch — confirms the warm-start adapter generates the correct answer span ~93% of the time on training data.

GRPO learning — v2 run after fixing prompt drift, format gating, and IS-ratio collapse

Our v1 GRPO run produced a flat reward curve. We diagnosed three compounding causes — (a) prompt-shape drift between the SFT user template (Question: …\n\n[Document excerpt]:\n…) and the GRPO user template (Context: …\n\n…), (b) a multiplicative format gate that zeroed the composite reward whenever the model dropped the <answer>…</answer> wrapper, and (c) importance-sampling-ratio collapse from QLoRA + vLLM colocate weight-merge drift — and fixed all three:

Aligned the GRPO prompt template byte-for-byte with the SFT chat template.
Replaced the multiplicative format gate with an additive +0.10 soft format bonus, capped at 0.25 when correctness is 0 to block format-only exploits.
Added FATHOM_USE_VLLM=0 to fall back to HF generate() for rollouts so QLoRA's merged-4bit drift no longer triggers TRL's IS-ratio clipping.

$GRPO reward curve$ Composite reward over the v2 GRPO run. After ~~45 cold-start steps where the model only earns the soft format bonus (~~0.15), the policy discovers the correct-answer mode and reward climbs to 0.6–0.98 with healthy variance (std 0.30–0.50). β=0.0, lr=5e-6, 8 generations/step, HF generate() rollout.

$Training summary (8-panel)$ All 8 GRPO metrics on one canvas — loss, reward, KL, entropy, grad norm, completion length, learning rate, advantage variance.

What this run proves

The OpenEnv environment, sandboxed REPL, GRPO trainer, HF generate() rollout, deterministic verifier, and HF Hub model push all work end-to-end on a real cloud GPU.
The SFT warm-start achieves a 91% loss reduction (3.20 → 0.29) and 2× token-accuracy gain (0.46 → 0.93), demonstrating the chat-template + format internalisation works.
The v2 GRPO reward curve exhibits a textbook bimodal cold-start: a flat ~0.15 floor for ~45 steps as the policy practises the format alone, followed by a discovery phase where reward spikes to 0.86–0.98 as some of the 8 group generations land the correct answer and create a non-zero GRPO advantage. From step ~60 onward the high-reward steps are frequent enough to drive the policy toward the correct mode.
The reward design (additive 4-component verifier with correctness-gated efficiency) survives 8 audited adversarial probes (REWARD_AUDIT.md). No LLM-as-judge is in the training loop — every gold answer is deterministic.

W&B v2 run (full metric history, ~80 steps): https://wandb.ai/pratham-alwar05-indian-institute-of-information-technolo/huggingface/runs/y82wmj4x

All plot PNGs are also published at https://huggingface.co/Pratham-math/fathom-1.5b-grpo/tree/main/plots.

Reproduce in 5 minutes

Open the Colab notebook (also browseable on the HF code repo).
Run cells 1–5 — verifies env health + runs a single smoke step against the live HF Space.
(Optional, A100 needed) Run cell 6 — launches a short GRPO sanity run.
Cell 7 generates the reward / loss curve PNGs.

The full A100-large + 1.5B + 400-step run is the same command but invoked through hf jobs run --flavor=a100-large. We did the full run for ~$20 of HF credits.

Hugging Face Space Deployment

Two paths — pick whichever your shell prefers:

# Python deploy (recommended)
export HF_TOKEN=hf_xxx
export FATHOM_SPACE_NAME=Pratham-math/fathom-env
python scripts/deploy_space.py

# Shell deploy
export HF_TOKEN=hf_xxx
export FATHOM_SPACE_NAME=Pratham-math/fathom-env
bash scripts/deploy_env_space.sh

Verify after deploy:

curl -s https://Pratham-math-fathom-env.hf.space/healthz
# → {"status":"ok"}

Local Setup

uv venv fathom --python 3.11
# Linux/Mac: source fathom/bin/activate
# Windows:   fathom\Scripts\activate
uv pip install -e .
pytest -q

Submission Checklist

Uses OpenEnv latest (openenv-core>=0.2.3)
Working training script using Unsloth + TRL — train/grpo.py
Colab notebook so judges can re-run — notebooks/fathom_train.ipynb
HF Space deployed (Pratham-math/fathom-env) and live (/healthz returns 200)
README explains motivation + env design + reward design + training
README links HF Space + all materials
REWARD_AUDIT.md (5 adversarial attacks neutralised)
Smoke test green on HF Jobs (SMOKE_RESULT.md)
Submission preflight passes (python scripts/submission_preflight.py)
Loss + reward plot PNGs from a real GRPO run (12 PNGs on the model repo)
W&B training run linked in Submission Links (v2 run y82wmj4x showing learning from 0.15 → 0.98)
Mini-blog linked in Submission Links

License

MIT. See LICENSE once added.