docs: M4 — GRPO results (infra + memory/rollout walls + capability-wall negative result)
Implementation log (docs/18) + Phase-3 row (evolution.md): the clipped_pg_loss op + gates, the actor-learner loop, the easy-task SFT baseline (held-out 18.7%, plateaus → no generalization), the two systems walls the design doc flagged (two 1B models OOM the 32GB box → β=0; naive rollout fragments the allocator → cached temperature sampling, rollout still the long pole), and the result: format holds, held-out 20.0% (+1.3pp, statistically flat) — the same wall as DPO. Closes the SFT→KV-cache→DPO→GRPO post-training arc with honest limits. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -466,3 +466,59 @@ verifiable reward* online (sample → check → reinforce what is genuinely corr
|
||||
fixed-pair proxy — though GRPO faces the same 8%-correct sparsity, so whether it moves the metric
|
||||
is M4's open question. Gate met for M3 = the infra is correct (op grad-checks, log2-at-init,
|
||||
margin/acc rise); the correctness flatness is the reported finding, not a bug.
|
||||
|
||||
### M4 — GRPO (online RL, critic-free, landed; infra + two honest systems walls)
|
||||
|
||||
The centerpiece: generation INSIDE the training loop. Infra built and gated; the run surfaces
|
||||
two concrete systems findings (the memory long-pole + the rollout long-pole, both flagged in the
|
||||
design doc's Risks) and the same capability wall as M3.
|
||||
|
||||
**Task made learnable first (per the aligned decision "easier task → then M4"):** the v12 SFT
|
||||
model scores ~8% on the hard task *and* on easy problems — it learned format, not arithmetic. So
|
||||
the easy task (operands ≤20, ops `+ − ×`) was re-SFT'd from the v12 base → **held-out 18.7%**
|
||||
(100% format), a baseline with reward variance for GRPO. Note: even easy arithmetic plateaus at
|
||||
~19% held-out (250 vs 600 SFT steps identical) — a 1B web-text model does not generalize the
|
||||
add/sub algorithm from ~550 examples; it memorizes train (982 total problems, 550 seen).
|
||||
|
||||
**New op (`xtrain-autodiff`, reuses the CE kernel + one new primitive):**
|
||||
- `clipped_pg_loss(logits, target, logp_old, logp_ref, A, ε, β)` — per completion token
|
||||
`ρ_t = exp(logπθ_t − logp_old_t)`, `L = −mean min(ρA, clip(ρ,1±ε)A) + β·mean KL` (k3), masked
|
||||
to completion tokens. Backward reuses `(probs − onehot)` + `scale_rows` (a new ~5-line per-row
|
||||
scale kernel — the per-token coefficient varies, which CE-backward's single scalar can't
|
||||
express). **Gate:** grad-check the active PG path + the A=0 (KL-only) path; degenerate value
|
||||
checks ε→∞ ⇒ vanilla PG, β=0 ⇒ no KL.
|
||||
|
||||
**Loop (`train_grpo`):** per step — sample B prompts, roll out G completions each, score (reward
|
||||
0/1), group-relative advantage `A=(r−mean)/(std+ε)` (no critic; all-correct/all-wrong groups
|
||||
skipped — zero advantage), capture `logπθ_old`/`logπref` per token, K inner clipped-PG epochs.
|
||||
Rollout uses the M2 KV-cache engine with **temperature sampling** (added in M4): single-row
|
||||
`[1,vocab]` logits per step vs the naive sampler's `[seq,vocab]`.
|
||||
|
||||
**Systems wall #1 — memory (the design doc's "two/three resident models"):** KL-leash GRPO needs
|
||||
policy + frozen reference, two 1.05B fp32-master models + AdamW m/v ≈ 21 GB fixed + training
|
||||
activations → unreliably OOMs on a 32 GB 5090 (fragmentation tips it over). To get a completing
|
||||
run, `β=0` (pure PG) drops the reference model (−4.2 GB). So the *principled* KL-leash version is
|
||||
memory-bound at this model size on this hardware — a real, reported constraint, not a bug.
|
||||
|
||||
**Systems wall #2 — rollout (the design doc's "rollout is the long pole"):** the naive sampler's
|
||||
growing `[seq,vocab]` allocations fragment the caching allocator over a long rollout → OOM. The
|
||||
cached temperature rollout (single-row logits) is lighter; but single-sequence cached decode is
|
||||
slow (the M2a host-round-trip), so rollout still dominates wall-clock (~16 s/step at G=6·B=6).
|
||||
Batched ragged decode (M2b) is the real fix and is deferred to where it is load-bearing.
|
||||
|
||||
**Result (easy task, β=0, G=6·B=6, 40 steps, lr 5e-7; 150 held-out, vs SFT 28/150 = 18.7%):**
|
||||
mean rollout reward fluctuates ~0.58–0.81 (noisy, inflated by train-set overlap in the sampled
|
||||
problems); **format stays 100/100** (no collapse even without the KL leash, at this gentle lr);
|
||||
**held-out 30/150 = 20.0%** — `+1.3 pp`, within the ~3% std-error of 150 prompts, i.e.
|
||||
**statistically flat**, the same wall as M3 DPO.
|
||||
|
||||
**The consistent M3+M4 lesson:** on a task where the base model lacks the underlying capability,
|
||||
**neither offline preference optimization (DPO) nor online RL (GRPO) moves held-out correctness**
|
||||
— each optimizes its objective (margin / reward) on the *training distribution* it can reach
|
||||
(here inflated by memorization), but cannot install a *generalizable* algorithm the model never
|
||||
had. RL reinforces what the model already does; it does not teach arithmetic. Gate met for M4 =
|
||||
the infra is correct (PG/KL grad-checks + degenerate checks, the loop runs, reward signal + KL
|
||||
leash wired, format held); the held-out flatness + the two memory/throughput walls are the
|
||||
reported findings. The honest end-state of the post-training arc: **a complete, correctness-gated
|
||||
SFT → KV-cache → DPO → GRPO stack** — the infrastructure learned in full, with measured, honest
|
||||
limits on what alignment can do for a capability the base model lacks.
|
||||
|
||||
Reference in New Issue
Block a user