Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Pre-reg (FINAL): token-GRPO phase 2 — R0 smoke + R1 on the MolmoAct2 discrete pathway

2026-08-13 14:5xZ (frozen at commit 8548969, 14:55:52Z; stamp corrected from a mis-clocked 15:1xZ before first publication — no content change). The finalization the design memo §9.4 promised: constants frozen, HEAD + checkpoint pinned, launch follows the commit. Executed under the owner delegation (11:07Z “you make the decisions, ensure we make progress and GPU is always busy”; 11:18Z “don’t wait for my confirmations”) with the memo’s open forks resolved by the frozen rules: surface = §4 option B (the veto window, open since ~06:0xZ, passed unanswered; B was the recommendation), pathway = the molmoact2 discrete (AR) head (owner 10:02Z steering; the ar100 gate proved it success-capable, and arm B made grammar-masked decode the serving mode). Instrument items 1–4 are all landed and oracle-gated (418715c, 229d80f, a268046, fa739e9 — check.py 861 green); this run is the first thing the loop harness executes.

Plain words

Everything before this was preparation: a probe showed the policy’s sampled attempts spread out enough to rank, a gate eval showed the token pathway can actually succeed in our simulator, and three build sessions produced the machinery that records every sampled action token with its probability, replays them through the trainer, and verifies — bit for bit — that the trainer is scoring exactly the distribution the robot sampled from. This post freezes the first real training run: the robot re-tries fresh scenarios 8 times each with mild sampling noise, attempts that do better than their siblings get pushed up, worse ones get pushed down, and a battery of tripwires stops everything if the policy starts learning violence, collapsing its diversity, or getting worse at the held-out test. It is deliberately laddered — a 2-step smoke to measure the true pace and check the signal survives on this model, then a 15-step read — with hard budget gates and decision boundaries between rungs.

Pinned identities

  • Code: HEAD fa739e9 (loop harness sim/grpo_loop.py; rollout
    • replay a268046; GRPO step 229d80f; capture 418715c).
  • Checkpoint: allenai/MolmoAct2-SO100_101 (hub release), served by the parity-gated first-class port; official SO-101 shim (signs 1,-1,1,1,1,1, offsets 0,90,90,0,0,0), norm tag so100_so101_molmoact2; FAST artifact allenai/MolmoAct2-FAST-Tokenizer.
  • Anchor: the loaded checkpoint at step 0 (the KL reference; a CPU snapshot of the trainable tensors).
  • Frames: v3 (production default at finalization; sim100 amendment 5 — v4 shadows — remains an open owner ask and does NOT ride this run). Flipped mount, sim100 episode conventions, --episode-seconds 30, execute-horizon 30, workers 8.

Frozen algorithm constants (memo §2/§3/§4)

  • S=8 fresh spawn seeds/step from the dedicated stream (1000 + 8·step; disjoint from sim100 0–99, probe 0–14, held-out 200–219), G=8 grammar-masked sampled draws each at T=1.0 → 64 episodes/step, RNG keyed stable_sample_rng(run_seed=0, repo_id(draw), seed, replan, 0).
  • Reward/episode: progress_final_cm + 10·success − 2·[upright<0.9] − 5·[strikes>0]; advantages = within-group z-scores (ddof=0); groups with reward std < 0.05 cm dropped whole.
  • One on-policy gradient pass/step (μ=1): advantage-weighted clipped token-CE over the action block, clip-higher [0.8, 1.28], ratio under the SAME grammar-masked softmax the decode sampled (recorded per-token logprobs = π_old; the item-3 oracle bound 1e-5 on the CPU fixture; disk rows additionally carry the registered JPEG budget — the fixture measured ~1% of ratio at a fresh policy, R0 reads the real number). KL penalty OFF; k3 KL to the anchor recorded every step (one swapped reference forward, 32-row subsample).
  • Trainable surface (option B): the trunk TEXT stack (embeddings
    • transformer + lm_head) at lr 5e-6 flat, AdamW(0.9, 0.95, eps 1e-6, wd 0), grad-clip 1.0, fp32 text (TF32 matmul), vision frozen bf16. Registered fallback on §7 instability: option A (patch-only) as a NEW pre-reg, not an in-run swap.
  • Microbatch = 1 row; chunking is gradient-invariant (oracle-pinned), so raising it at a rung boundary is an execution note, not an amendment.
  • Held-out eval: seeds 200–219, greedy grammar-masked, at step 0 (pre-update baseline), every 5 steps, and at the endpoint; paired Δ composite reward, seeded 10k-bootstrap CI95.

Ladder + budget (measured-pace arithmetic)

Rollout pace measured on this exact serving stack (arm B: 100 episodes, 2,996 predicts, 63.1 min at workers 8 → 0.63 min/episode) → 64 episodes ≈ 40 min ≈ 0.67 GPU-h, ×1.0–1.3 for sampled streams and the fp32 text stack ≈ 0.7–0.9 GPU-h/step rollouts; gradient pass 1,920 batch-1 teacher-forced fwd+bwd ≈ 0.2–0.5 GPU-h; eval ~0.21 GPU-h per occurrence → ~1.0–1.4 GPU-h/step estimate — R0 exists to measure it.

rungsteps (cum)boundary readsbudget
R0 smoke1–2plumbing rc 0; measured GPU-h/step; ON-SURFACE signal: median group std ≥ 0.25 cm AND ≥ 8/16 groups non-degenerate, else STOP (the probe’s condition, transferred); step-1 mean_ratio ∈ [0.95, 1.05], clip_fraction < 0.2, else STOP (replay is not scoring the rollout distribution); KL telemetry sane → set the §7 KL numeric line; VRAM < 75 GiB≤ 3.5 GPU-h
R13–17memo §6 reads at the step-17 endpointre-priced at R0’s measured pace; stop before launch of R1 if projection > 22 GPU-h cum
R2 (conditional)+K stepsfull frozen readsK = what fits the 35 GPU-h total gate at measured pace (memo cap unchanged)

R1→R2 rule (memo §5 verbatim): extend iff (a) no tripwire fired and (b) held-out greedy composite at the endpoint is not worse than step 0 (paired CI not entirely below 0). Beyond R2 is a new pre-reg.

Frozen reads (memo §6, unchanged)

  1. Primary: paired Δ composite reward, held-out greedy, endpoint vs step 0 — IMPROVED iff CI95 entirely above 0; “phase 2 promising” = CI-above-0 on this or read 2.
  2. Knock-away rate under sampling: endpoint 5-step window below the transferred 10/120 baseline with binomial CI excluding it (recorded against R0’s own measured baseline too — this surface’s number lands at the R0 boundary).
  3. Success count (record + headline if > 0).
  4. Record-only: per-step median group std, non-degenerate fraction, KL-to-anchor curve, chosen-token NLL, tip rate, clip fraction, ratio extremes, per-seed traces.

Tripwires (§7, mechanized in the loop — exit 3 + heartbeat row)

Any reset strike in training rollouts; non-finite loss; median group std < 0.05 cm ×3 consecutive; knock-away rate > 2×(10/120) ×3 consecutive; held-out paired CI entirely below −1.0 cm. KL-to-anchor runaway: recorded, numeric line set at the R0 boundary.

Command (verbatim) + ops

R0:

MUJOCO_GL=egl fontaine/scripts/run_detached.sh fontaine-grpo-r0 \
  uv run python -m sim.grpo_loop \
  --checkpoint allenai/MolmoAct2-SO100_101 \
  --out-dir outputs/sim/grpo_phase2 --total-steps 2 \
  --eval-every 5 --save-every 1

(All other flags at their frozen defaults, which ARE the constants above.) R1 resumes the R0 checkpoint: --resume outputs/sim/grpo_phase2/step_0002.pt --total-steps 17. Babysit registry entry at launch (train-jsonl, outputs/sim/grpo_phase2/train.jsonl, probe eval_reward_mean, vram key vram_gib, vram gate 75 GiB); heartbeat carries reward/guard/ratio/KL facts per step. Training rows prune after each gradient pass (~0.3 GB/step transient). Rung-boundary checkpoints upload to fontaine-checkpoints if consumed by the next rung or a result claim.

Frozen at commit time; the launch immediately follows the push. Amendments only via numbered addenda below.


Addendum 1 (2026-08-13 16:1xZ — plumbing fix + relaunch, no constant changed). Launch 1 (14:58:55Z) crashed rc 1 at 15:51:26Z in the FIRST gradient step: grpo_objective_sums moved the caller’s advantages/rollout-logprob tensors to the training dtype but not the training device (cuda/cpu mix — invisible to the CPU oracles; the exact plumbing class R0 exists to catch). Everything before the step was healthy: step-0 baseline banked (held-out greedy composite 1.868, 2/20 successes), wave 0 complete (64 sampled episodes, ~35 min ≈ 0.58 GPU-h — inside the estimate band), 1,889 training rows written, mask verification passed, replay forward ran. Fix: normalize old_logprobs/advantages with .to(new_logprobs) (device+dtype) at one point in the surrogate — semantics unchanged, check.py 861 green. Relaunch rides the fix commit; the ~0.9 GPU-h of launch 1 counts against the R0 gate (honest accounting: R0 total budget may land ~3.1 of the 3.5 gate).*

Addendum 2 (2026-08-13 17:5xZ — memory fixes + relaunch, no frozen constant changed; R0 ops gate raised 3.5 → 5.5 GPU-h). Launch 2 (16:15:26Z) reproduced the step-0 baseline bit-identically (1.868, 2/20) and ran wave 0 + the FULL gradient accumulation cleanly, then OOM’d at 17:12:17Z inside the first optimizer.step() — Adam state init via the _foreach path materializes whole-surface temporaries. The measured fact that matters: 77 GiB PyTorch-allocated at the step → the option-B text stack is ~15 GB fp32 (~3.9B params — this checkpoint is a ~4B-class model, not the 2B the memo’s 69.2-GiB memory precedent was calibrated on; that precedent was measured on the OLD er60k surface and did not survive the 10:02Z retarget). Params+grads+2·Adam ≈ 62 GiB steady DOES fit; the two +P transients did not. Fixes (allocation-shape only, semantics oracle-pinned unchanged): (1) AdamW(foreach=False) — per-tensor step, no whole-surface temporary; (2) the anchor-KL swap now stages live weights to CPU for the reference forward instead of holding both copies on GPU; (3) PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True (launch 2 died with 1.36 GiB reserved-unallocated at a 96 MiB request). Projected steady peak ~68–70 GiB vs the 75 vram gate — tight; if launch 3 still OOMs, option B is measured-infeasible on one H100 for this model and the fallback discussion (§4 option A) goes in-channel as a new pre-reg, per the frozen rule. Gate accounting: launches 1+2 spent ~1.85 GPU-h on the two plumbing crashes; the R0 ops gate rises to 5.5 GPU-h to cover them — the 35 GPU-h ladder total is unchanged (R2 shrinks by whatever R0 overruns).*

Addendum 3 (2026-08-13 19:5xZ — worker-headroom fix + RESUME relaunch, no frozen constant changed). Launch 3 (17:56:31Z) delivered the run’s biggest milestone before dying: the first gradient step SURVIVED with every step-1 gate green (row 18:54:02Z: mean_ratio 1.00138, clip_fraction 0.132, median group std 4.17 cm, groups 8/8, approx_kl 0.0232, anchor_kl 0.0215, loss 0.0317, grad_norm 2.56, strikes 0, 4/64 sampled successes, vram_gib 70.93, 0.76 GPU-h/step measured) and step_0001.pt saved. Then wave 1’s first worker OOM’d at episode reset (18:57:55Z, _gather in the GPU photometric post): after the gradient pass the parent’s caching allocator retains the ~70 GiB activation peak as RESERVED segments (nvidia-smi showed 78–80 GiB; allocated was ~50), and the 8 spawned sim workers — own CUDA contexts + post tensors — could no longer fit. Wave 0 never saw this because Adam states didn’t exist yet. Fixes (allocation-shape only, semantics oracle-pinned unchanged): release_cached_vram() (torch.cuda.empty_cache()) before every wave and eval. Relaunch RESUMES step_0001.pt rather than re-running from scratch — deterministic keying makes the remaining work identical (step-2 wave seeds 1008–1015, baseline rides in the checkpoint), it saves ~1.1 GPU-h, and it exercises the exact resume path R1 depends on. One resume-correctness fix landed with it: the KL anchor is now snapshotted BEFORE the checkpoint restore (the old order would have silently rebased anchor_kl onto the resumed weights); new CPU oracle pins it (13 loop oracles). Watch item carried to the boundary: step-1 knockaway_frac 0.234 vs the 2×(10/120) = 0.167 tripwire line — ×3 consecutive fires it; R0’s own measured knockaway baseline lands at the boundary read as registered. Gate accounting: launch 3 spent ~1.0 GPU-h (plus ~1.0 h GPU idle after the 18:58Z crash — a babysit watcher bug, pgrep matching its own cmdline; fixed in-session). Crashes 1–3 total ~2.9 GPU-h; resume needs ~1.0 → R0 lands ~3.9 of the 5.5 ops gate. Ladder total 35 unchanged.*

R0 boundary read (2026-08-13 20:5xZ) — VERDICT: STOP

R0 completed rc 0 at 20:54:30Z (launch 4, the step_0001.pt resume — the R1 resume path validated in passing). R1 does not launch, by the frozen reads:

readvaluegateverdict
plumbingrc 0; ratio 1.00138, clip 0.132 (step 1)ratio ∈ [0.95, 1.05], clip < 0.2✓ green
pace0.76 / 0.68 GPU-h/step measuredR1 projection ~13 < 22 cum✓ green
VRAM76.53 GiB allocated peak at the steady-state step (nvidia-smi 79.5 of 81.6)< 75 GiB✗ FAIL
on-surface signalwave 1: median std 4.17 cm, 8/8 kept → wave 2: 0.0087 cm, 3/8 kept (5 groups with all 8 draws IDENTICAL); pooled 16-group median ~0.25, borderlinemedian ≥ 0.25 AND ≥ 8/16 nondeg✗ FAIL (trend decisive)
endpoint held-out−0.0 composite, 0/20 successes; paired Δ −1.868, CI95 [−4.41, −0.03]record; competence line −1.0CI entirely < 0 (no tripwire: CI hi −0.03 > −1.0)

The mechanism, from the recorded curves: one gradient step at lr 5e-6 on the ~3.9B fp32 text stack sharpened the policy into near-determinism — chosen-token NLL 0.766 → 0.329, anchor-KL 0.0215 → 0.0885 (≈4×/step), sampled diversity collapsed (T=1.0 draws within a group ending bit-identically), and held-out GREEDY competence dropped with it (1.868 → −0.0). Wave 0’s 4/64 sampled successes fed large positive advantages into a single update that overshot. No tripwire fired mechanically (each needs ×3 consecutive or CI < −1.0; a 2-step smoke ends first) — the boundary gates did exactly the job they were frozen for, at ~3.8 of the 5.5 GPU-h ops gate instead of R1’s ~13.

What R0 banked: the full plumbing chain works end-to-end (sampled waves → z-filter → chunked GRPO step → anchor-KL → paired eval → resume); the memory envelope on 1×H100 is measured-marginal for option B (steady-state peak 76.5 GiB allocated, ~2 GiB physical headroom); the pace book is real (0.7 GPU-h/step); and the FIRST gradient step at these constants is already destructive on this surface. Checkpoints step_0001/0002.pt stay on local disk for diagnosis (not uploaded: a STOP boundary consumes nothing; weights of a collapsed policy are a diagnostic artifact, not a bankable one).

Re-scope (next pre-reg, per the frozen fallback rule): option A (patch-only surface) — it dissolves the VRAM problem outright (trainable set orders of magnitude smaller) — plus explicit collapse mitigation for whichever surface runs: lower lr and/or fewer advantage-concentrated rows per step, with the R0 curves as the calibration. Goes in-channel as a NEW pre-registration before any launch.