| # 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: |
|
|
| ```bash |
| 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"]`). |
|
|
