---
title: "5.1 reasoning training on a 22M model"
sidebar:
  order: 501
---

<!-- PRD metadata (moved from frontmatter for Blume schema compatibility) -->
**status:** not-started  
**owner:** unassigned  
**created:** 2026-06-13  
**parent_plan:** docs/PLAN.md §3 Tier 5 (5.1)  
**related_prds:** 5.2-testtime-compute-scaling.md (sibling research arc)

# PRD — Reasoning training on a 22M-class model

## Goal

Train a 22M-class model with reasoning-style supervision (GRPO / DAPO
on verifiable rewards; SFT-on-traces as the warm-up), publish the
*negative result* if (as expected) chain-of-thought capability does
not emerge at this scale. Publishable shape: "smallest model where
RLVR demonstrably fails to elicit CoT; here's the scaling-curve
intercept."

Tier-5 research arc. Estimated 5–7 days end-to-end. The deliverable
is a paper-shaped artifact + reproducible code + a single
scaling-curve data point — not a polished UX feature.

## Why now

- The platform has the substrate: SFT, DPO, distillation, all
  shipped. RLVR-style GRPO/DAPO is the next training paradigm and
  we have no point on the curve.
- Publishable negative results are valuable: "small enough that CoT
  doesn't emerge" is information the field needs to bound the
  emergence threshold (currently rough; DeepSeek-R1 distilled down
  to 7B, but no one's published cleanly at 22M).
- The 22M scale is a known posttrainllm trained-size — flagship-huge-v5
  is 221M; one of the smaller pretrained configs is ~22M. Both are
  in the bench history.

## Scope — in

- **Phase 1 (SFT-on-traces):** finetune a 22M base on the
  DeepSeek-R1-distill traces (or a similar open trace corpus).
  This is the warm-up; baseline numbers before RL.
- **Phase 2 (GRPO):** implement Group Relative Policy Optimization
  from [DeepSeek-R1, Jan 2025](https://arxiv.org/abs/2501.12948) —
  K rollouts per prompt, group-relative advantage normalization,
  KL to a reference policy. Verifiable rewards from math + simple
  code tasks (a small GSM8K-class subset).
- **New module:** `Sources/TinyGPT/Grpo.swift` — the RL loop. Reuses
  the shipped sampler, the eval harnesses for reward computation,
  the existing optimizer.
- **Eval:** GSM8K accuracy + a per-step CoT-length curve (does the
  model learn to think longer when it gets reward signal?). Both
  before and after Phase 2.
- **Artifact:** `docs/research/reasoning-22m-results.md` —
  paper-shaped report: hypothesis, recipe, numbers, scaling-curve
  point, interpretation.
- **Decision log:** one entry noting GRPO is the "mental model"
  implementation; DAPO ([Mar 2025](https://arxiv.org/abs/2503.14476))
  is the production refinement we'd port if results justify.

## Scope — out

- **Multi-turn reasoning tasks** (e.g. agentic browse). V1 sticks
  to single-shot math + code where the reward is unambiguous.
- **Process reward models** (per-step reward via a learned PRM).
  Verifier-based outcome rewards only.
- **Cross-scale ablation** (run 22M, 50M, 100M to find the
  emergence threshold). V1 is one data point — the scaling-curve
  arc lives in 5.2.

## Files to touch

| File | Change |
|---|---|
| `Sources/TinyGPT/Grpo.swift` | new — RL loop |
| `Sources/TinyGPTModel/AdvantageNormalizer.swift` | new — group-relative advantage compute |
| `Sources/TinyGPT/TinyGPT.swift` | `case "grpo"` |
| `Sources/TinyGPT/SftOnTraces.swift` | new (or extend `SFT.swift`) — load DeepSeek-R1-distill JSONL, sequence-pack |
| `docs/research/reasoning-22m-results.md` | new — the paper-shaped artifact |
| `evals/grpo-smoke.sh` | new — 20-step GRPO on a 10-prompt GSM8K subset; assert reward signal is non-degenerate |
| `docs/PLAN.md` | flip 5.1 ⬜ → ✅ on ship (regardless of positive/negative result) |

## Don't touch

- Existing SFT/DPO paths — GRPO is sibling, not a replacement.
- The OpenAI surface — GRPO is a training-side feature, no serve
  surface needed.

## Acceptance criteria

- [ ] `posttrainllm grpo --base flagship-22m.tinygpt --task gsm8k
  --steps 2000 --out grpo.lora` runs end-to-end on M5 Pro in < 24h.
- [ ] Phase-1 SFT improves GSM8K above the base by ≥ 1pp (small
  but positive — confirms the trace data is being absorbed).
- [ ] Phase-2 GRPO either (a) lifts GSM8K further by ≥ 2pp, or
  (b) plateaus, and the report documents the plateau as the negative
  result with the per-step reward curve.
- [ ] CoT-length-over-step curve plotted in the results doc;
  shows whether the model learns to think longer under RL pressure.
- [ ] `grpo-smoke.sh` passes in CI — small fixture, no GPU model
  needed beyond what CI already has.

## Reference patterns

- `Sources/TinyGPT/Dpo.swift` — the closest training shape; preference
  learning with KL-to-ref. GRPO is "DPO but K samples per prompt
  with task-defined reward instead of preference pairs."
- [DeepSeek-R1 paper, Jan 2025](https://arxiv.org/abs/2501.12948),
  §3 — the GRPO loss and the verifiable-reward formulation. Cite,
  don't redocument.
- [Tulu 3, Lambert et al. 2024](https://arxiv.org/abs/2411.15124) —
  open RLVR recipe at larger scale; informs the SFT-warm-up shape.

## Open questions

- Reward shape for math: exact-match on the final numeric answer
  vs partial credit on intermediate steps. **Recommendation:**
  exact-match for V1 — the simplest verifiable signal.
- Whether to publish even if both phases plateau. **Recommendation:**
  yes — "we tried at 22M and it didn't work" is the publishable
  finding. Tier 5's whole point is paper-shaped artifacts, including
  negatives.