File size: 5,059 Bytes
d909bd8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
# 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"]`).