121 lines
4.6 KiB
Markdown
121 lines
4.6 KiB
Markdown
# M5.5 — verl training loop
|
||
|
||
| Field | Value |
|
||
|---|---|
|
||
| Phase | M5 — Post-training |
|
||
| Size | L — 3+ days |
|
||
| Status | ⬜ Not started |
|
||
| Flags | — |
|
||
| Spec | inlined below |
|
||
| Blocks | M5.3, M5.4 |
|
||
|
||
## Goal
|
||
|
||
Train the LoRA that makes the update and exit gates better than prompting alone.
|
||
|
||
## Facts (inlined — no spec read needed)
|
||
|
||
Python, using verl (`github.com/volcengine/verl`), which is what the paper used.
|
||
Lives outside the Rust workspace; the JSONL corpus is the only interface.
|
||
|
||
Paper hyperparameters, Table 3 — start here rather than guessing:
|
||
|
||
```
|
||
chunk size 5000
|
||
max prompt length 8192
|
||
max response length 2048
|
||
clip ratio 0.20
|
||
learning rate 1e-6
|
||
temperature (train) 1.0 top_p 1.0
|
||
temperature (val) 1.0 top_p 0.7
|
||
train batch size 128
|
||
rollout N 16
|
||
mini batch size 128
|
||
LR warmup steps 20
|
||
```
|
||
|
||
Advantage, paper §3.2.2 — the part most likely to be implemented wrong:
|
||
|
||
```
|
||
Â_traj = r_traj_g − mean over the GROUP of trajectories
|
||
Â_turn = r_update_{g,t} − mean over turns AT STEP t across groups
|
||
 = α·Â_traj + (1−α)·Â_turn α = 0.9
|
||
```
|
||
|
||
Two distinct baselines. `Â_turn` is normalised across groups **at the same `t`**,
|
||
and the group size at step `t` can differ from the trajectory group size, because
|
||
trajectories that exited early have fewer turns.
|
||
|
||
α=0.9 is the paper's default and the ablation (Figure 8) explains why: at α=1
|
||
there is no update-gate reward and accuracy on evidence-free chunks collapses —
|
||
the model updates indiscriminately, which is exactly the failure this whole
|
||
system exists to avoid.
|
||
|
||
Expect instability. The paper's own limitations section says the extra rewards
|
||
"reduce training stability, requiring a smaller off-policy degree and longer
|
||
convergence time."
|
||
|
||
## Steps
|
||
|
||
1. `training/` directory, Python, `uv`-managed. System Python is 3.9.6; this
|
||
needs 3.11+.
|
||
2. Corpus loader for M5.3's format.
|
||
3. Configure verl for LoRA on Qwen2.5-3B-Instruct, rank 16–32.
|
||
4. Implement the three rewards and the two-baseline advantage, α configurable.
|
||
5. Log per step: update accuracy split by evidence-present and evidence-free,
|
||
exact-exit ratio, format correctness, mean response length, validation reward.
|
||
6. Hold out a validation split by **project**, not by trajectory — same-project
|
||
trajectories share vocabulary and leak.
|
||
7. Export the adapter, version it, publish where M5.4 can mount it.
|
||
|
||
## Acceptance
|
||
|
||
- Training runs to convergence on the validation reward.
|
||
- Both advantage terms are computed with their own baselines.
|
||
- Update accuracy on evidence-free chunks does not collapse.
|
||
- The adapter loads in M5.4's server.
|
||
|
||
## Verify
|
||
|
||
**Harness:** pytest for reward and advantage maths; a short training run for the
|
||
loop itself.
|
||
|
||
**Integration test** — `training/tests/test_rewards.py`:
|
||
1. `a1_r_update_signs` — matching label +1, mismatching −1.
|
||
2. `a2_r_exit_bands` — exact 0, early −0.75, late −0.5.
|
||
3. `a3_r_format_strict` — any unparsed turn zeroes the whole trajectory.
|
||
4. `a4_traj_baseline` — `Â_traj` uses the group mean; hand-computed fixture.
|
||
5. `a5_turn_baseline_at_step_t` — `Â_turn` normalises across groups at the same
|
||
`t`; a fixture with unequal trajectory lengths must not misalign. This is the
|
||
assertion that catches the most likely implementation error.
|
||
6. `a6_alpha_mix` — α=1 yields pure trajectory advantage; α=0 pure turn.
|
||
7. `a7_alpha_1_degenerates` — train 50 steps at α=1 on a fixture; assert
|
||
evidence-free accuracy drops relative to α=0.9, reproducing the paper's
|
||
Figure 8b.
|
||
8. `a8_validation_split_by_project` — assert no project appears in both splits.
|
||
9. `a9_adapter_loads` — export, mount in M5.4, assert a completion returns 200.
|
||
|
||
**Command:** `uv run pytest training/tests -v`
|
||
|
||
**False pass:**
|
||
- Normalising `Â_turn` over the whole batch rather than per step `t`. It trains,
|
||
loss goes down, and the turn-level signal is diluted into noise — assertion 5
|
||
is the only thing that catches it.
|
||
- Splitting validation by trajectory. Same-project trajectories share phrasing
|
||
and file paths, so validation reward looks excellent and generalisation is
|
||
untested.
|
||
- Skipping assertion 7 as "too slow". It is the only end-to-end evidence that the
|
||
update reward is wired to anything.
|
||
|
||
## Traps
|
||
|
||
- Tuning α before the rewards are verified. Every α is wrong if `r_update`'s sign
|
||
is flipped, and the symptom looks identical.
|
||
- Training on labels whose κ was never measured. The policy learns the labeler,
|
||
and there is no held-out signal that would reveal it — M5.2 exists for this and
|
||
M5.3 refuses to export without it.
|
||
|
||
---
|
||
|
||
Background: [DESIGN.md](../DESIGN.md) — P6 · paper §3.2.2, Table 3, Fig 8
|