code/docs/RL_LEAN.md

Lean RL path (lane R): single-turn, on-policy, LoRA, Qwen3-8B-family provers

Code: align/rl/lean_rl.py (config, rollout -> batch, losses, gradient step), align/rl/lean_rl_run.py (loop, atomic checkpoints, PEFT export, ledger, backends), scripts/p2/rl_lean.py (CLI), configs/p2_rl_lean.toml. Tests: tests/test_lean_rl_loss.py, tests/test_lean_rl_run.py, tests/test_rl_audit_u04_fixes.py. The multi-turn env (align/rl/envs/lean_multiturn.py) is untouched and stays opt-in.

One iteration k (policy = ckpt/iter_k/adapter)

1. gen: extras/vllm_rollout/lean_gen.py --lora-adapter ckpt/iter_k/adapter --top-p 1.0 --record-logprobs (G samples per prompt; fallback rollout.adapter_mode="merged_export" serves a bf16-merged checkpoint instead, lossy).
2. verify: scripts/p2/verify_generations.py (workers configurable). reward = accepted; timeout, crash, protocol verdicts are MASKED (dropped from the batch and from the group baseline), never scored 0.
3. batch: group-centered advantages over valid samples, dead groups dropped (tolerance DEAD_ATOL), loss modes cispo | dr_grpo | ei (ei = same pipeline, negatives dropped; keeps all-solved positives unless batch.ei_keep_all_solved=false). Loss normalisation = explicit global response-token count over micro-batches and ranks.
4. update: LoRA only (FSDP2-sharded frozen base), AdamW, grad clip, optional FP32 head (loss.fp32_head).
5. commit ckpt/iter_{k+1} atomically (adapter, optimizer, RNG, prompt-stream position, ledger, PEFT dir), keep 2.

Three logprob bases are stamped in every ledger row: behaviour (sampler, un-truncated: top_p == 1.0 is enforced in config and per record), old (trainer at rollout weights; telemetry only), current (differentiable). The importance ratio is pi_current / pi_behaviour (audit U04-3), so the CISPO clamp is a trust region against the policy that actually produced the tokens (checkpoint lag and engine mismatch included). old_vs_behaviour_abs_logratio in the ledger is the engine-mismatch figure; if it is large with fp32_head=false, rerun the probe with fp32_head=true.

K1 contrast on the SAME cached rollouts: run arm A, then arm B with --set 'replay_from="<A run_dir>"' --set 'loss.mode="ei"' (no regeneration, no re-verification; iteration 0 is exactly on-policy for both arms; later iterations replay the SOURCE arm's rollouts, which are off-policy for arm B by construction).

Budget ledger (<run>/ledger.jsonl, also in each checkpoint)

Per iteration: generated_tokens, prompt_tokens, verifier_seconds (sum of per-proof verify_s), verifier_wall_s, generate_wall_s, train_tokens (response tokens trained), train_padded_tokens, train_wall_s, wall_s, stage_attempts (re-runs after a kill), cumulative sums, batch telemetry (masked reasons, group filter counts), loss metrics, and on GPU cuda_peak_alloc_gb / cuda_peak_reserved_gb / step_wall_s. iters/iter_k/stage_times.json also has train_proc_wall_s (the whole training child: model load + step + commit). Wall of a killed attempt is not recoverable; only the attempt count is.

Phase 0b GPU probe (NOT RUN; needs the GPUs idle; each child takes the heavy-job lock itself)

Prerequisite: branch phase0/genharness merged (lean_gen.py --lora-adapter/--top-p/--temperature), then:

cd /apps/development/rlhf/align            # after merging phase0/rl-lean + phase0/genharness
POOL=<P2 train-pool jsonl, an in-band slice, 8 problems is enough>
RUN=/apps/development/rlhf/align/runs/p2_rl_probe
.venv/bin/python scripts/p2/rl_lean.py run configs/p2_rl_lean.toml --nproc 4 --until 1 \
  --set "run_dir=\"$RUN\"" --set "problems=\"$POOL\"" \
  --set iterations=1 --set prompts_per_iter=8 --set batch.group_size=8 --set rollout.cap_tokens=16384 \
  --set 'loss.mode="cispo"' --set loss.fp32_head=false --set verify_workers=8
# FP32-head variant on the SAME rollouts (no regeneration): memory + parity
.venv/bin/python scripts/p2/rl_lean.py run configs/p2_rl_lean.toml --nproc 4 --until 1 \
  --set "run_dir=\"${RUN}_fp32\"" --set "problems=\"$POOL\"" --set "replay_from=\"$RUN\"" \
  --set iterations=1 --set prompts_per_iter=8 --set batch.group_size=8 --set rollout.cap_tokens=16384 \
  --set loss.fp32_head=true
python3 -c "import json;[print(json.dumps({k:e[k] for k in ('generated_tokens','generate_wall_s','verifier_wall_s','train_wall_s')}|{'step':e['metrics'].get('steps',[{}])[-1]}, indent=1)) for e in map(json.loads,open('$RUN/ledger.jsonl'))]"
cat $RUN/iters/iter_0000/stage_times.json     # train_proc_wall_s = load + step + commit; gen/verify walls = handoff budget

Read off: cuda_peak_alloc_gb / cuda_peak_reserved_gb (must fit 23 GB; if short at 16k drop the FP32 head and watch old_vs_behaviour_abs_logratio), step_wall_s, train_proc_wall_s - train_wall_s (model load + commit), generate_wall_s (engine start included), then export fidelity: serve ckpt/iter_1/adapter un-merged and compare to a merged export (--set export_merged=true, ledger[...]["merged_export"]["cosine"]).