5.1 reasoning training on a 22M model

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 — 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

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, §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.