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 TinyGPT 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 — 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) 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
-
tinygpt grpo --base flagship-22m.tinygpt --task gsm8k --steps 2000 --out grpo.loraruns 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.shpasses 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, §3 — the GRPO loss and the verifiable-reward formulation. Cite, don’t redocument.
- Tulu 3, Lambert et al. 2024 — 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.