File size: 13,526 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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
# Dependency policy

`align` runs on **torch, triton, numpy, safetensors**. That is the whole list. It is
enforced by `tests/test_dependency_firewall.py` (AST scan of every module, lazy
imports included), not by convention.

## What was removed, and what replaced it

| Removed | Was used for | Replacement in `align` |
|---|---|---|
| `transformers` | tokenizer, `MistralConfig`, HF-layout load for the golden regression | `align/tokenizer/` (pure-Python BPE from `tokenizer.json`, parity-certified), `align/model/hf_config.py` (own `config.json` read/write), golden regression loads through `align.model` |
| `omegaconf` | dataclass+YAML config with struct mode | `align/config.py`: dataclass + stdlib `tomllib`, struct mode (unknown keys raise), dotted CLI overrides. Configs are `.toml` |
| `datasets` | HF streaming | `align/data/hub.py`: HTTP streaming of JSONL/parquet row groups from the HF hub via `urllib`; parquet ingest is a one-time step behind the `ingest` extra |
| `pyarrow` | parquet | only `align/data/ingest_parquet.py` (whitelisted, lazy, one-time ingest into the repo's own shard format) |
| `math_verify` + `sympy` | answer equivalence | `align/verify/`: boxed extraction, normalisation, exact rational/decimal/interval/set equivalence, small expression evaluator. Ceiling-tested |
| `wandb` | metrics | `align/utils/tracking.py`: JSONL `history.jsonl` (mtime is the liveness signal), optional stdout |
| `bitsandbytes` | 8-bit AdamW (`optimizer.name = "adamw8bit"`, named by every shipped full-FT config) | `align/train/adamw8bit.py`: pure-torch block-wise 8-bit AdamW (Dettmers et al. 2022, arXiv:2110.02861), codebook and quantize/dequantize in `align/train/quant8.py`. Dynamic 256-entry maps, block 2048, uint8 codes + fp32 absmax per block, fp32 state below 4096 elements, no master copy (stochastic rounding into bf16 params), DTensor/FSDP2 local-shard stepping, global pre-clip grad norm, bit-exact resume. 2 + 8/block = 2.004 B/param on the device, zero host RSS. `make_optimizer` builds it for `adamw8bit`; `OffloadAdamW` (`align/train/offload_adamw.py`) stays the `RunConfig` default but cannot launch 7.25B/4 ranks (4x over its pinned-staging cap, ~105GB host RSS) |
| `peft` | LoRA adapters for the PPO value head | `align/ppo/lora.py` (own LoRA on `nn.Linear`) |
| `deepspeed` | ZeRO in the old PRM-PPO | banned. FSDP2 via torch (`align/train/fsdp.py`). (Also: ZeRO-3 + Adafactor silently did not train, 2026-08.) |
| `trl` | -- | banned; never was a dependency of this lineage |
| `vllm` / `ray` | high-throughput rollout | optional, isolated in `extras/vllm_rollout/` with its own venv; talks to the trainer only via `align/rollout/protocol.py`. The core rollout is `align/rollout/native.py` (torch KV-cache sampler) |
| `matplotlib` | plots | dropped from the alignment package (plots live in the interp repo) |

## Rules

1. New runtime dependency => a dated entry in this file **and** a change to `ALLOWED` in the firewall test, in the same commit.
2. Optional extras live in exactly one file each, imported lazily, whitelisted by path.
3. Parity tests against the legacy stack (transformers tokenizer, math_verify) run under the
   legacy venv as an *oracle* (`pytest -m oracle`); the oracle is never imported at runtime.

## External binary: Lean 4 toolchain (2026-10-12)

Track L's strict proof verifier (`align/verify/lean/`, docs/LEAN_VERIFIER.md) needs a real Lean
kernel. This is **not a Python dependency**: nothing under `align/` or `scripts/` imports it,
the firewall's `ALLOWED` set is unchanged, and the Python side of the verifier is stdlib-only.
It is an external *toolchain* reached only through `subprocess`, the same category as `git`,
`nvidia-smi` or `du`.

| | |
|---|---|
| What | `elan` 4.2.4 -> Lean `v4.9.0-rc1` (`leanprover/lean4`), Mathlib `d1d1e4b72c6e` (2024-06-12), `leanprover-community/repl` tag `v4.9.0-rc1` (`7df77c0a1c91`) |
| Why this pin | It is the environment Goedel-Prover-V2 (primary prover base), DeepSeek-Prover-V1.5 and Goedel-Prover-V1 verify against: V2's `mathlib4` submodule is `xinhjBrant/mathlib4@2f65ba7f`, one commit on top of upstream `d1d1e4b7` that only adds DeepSeek's REPL fork to `lakefile.lean`. We pin the upstream parent because Mathlib's prebuilt `.olean` cache is keyed on a hash that includes `lakefile.lean`; the edit would orphan the cache and force a multi-hour Mathlib build |
| Where | `/apps/development/rlhf/lean/` (`ELAN_HOME`, `MATHLIB_CACHE_DIR`, `XDG_CACHE_HOME`, `TMPDIR` all redirected there; nothing is written to `$HOME`, `/` or `/tmp`). ~5.8 GB |
| Installed by | `scripts/lean_setup.sh` (idempotent, every network/build step under `timeout`, preflighted against the host rules); records `PINS.json` |
| How it is reached | `align.verify.lean.toolchain` finds the install; `repl_pool` spawns the REPL as a child with a **minimal environment** (`PATH`, `LEAN_PATH`, `LD_LIBRARY_PATH`, `LEAN_SYSROOT`; no tokens, no `$HOME`), its own session, `prlimit --as/--fsize/--core`, and a memory watchdog |
| Absent toolchain | `find_toolchain()` raises `ToolchainNotFound` naming the setup command. Tests that need it are `@pytest.mark.slow` and skip; everything else uses a fake REPL binary |
| Network | install time only (`github.com`, `releases.lean-lang.org`, Mathlib cache host); the verifier itself makes no network access |
| Licences | Lean 4, Mathlib, REPL, elan: Apache-2.0 |
| Rule it follows | Rule 1 above is about *runtime Python dependencies*. An external binary needs (a) a dated entry here, (b) a pinned version and an install script, (c) subprocess-only access, (d) a skip-when-absent test path, (e) a resource cap (the report counts a Lean pool as a heavy job: one REPL = 0.25 GiB private + a shared 3.4 GiB `.olean` mapping, measured) |

## Vendored assets (2026-10-12)

`assets/tokenizer/mistral-7b-v0.3/` holds the four Mistral-7B-v0.3 tokenizer files
(~2.6MB, Apache-2.0; md5s in its `PROVENANCE.md`). They are byte-identical to the HF
snapshot's, and exist because (a) the HF-cache snapshot the code names is not present on
every box and (b) the only other copies were inside rotating legacy SFT checkpoint
directories. `tests/conftest.py::resolve_v03_tokenizer_dir` and
`scripts/build_data.py::default_tokenizer_dir` prefer the real snapshot when it exists and
fall back to this copy. It is tokenizer-only: no `config.json`, no weights, so anything that
needs the base model's weights still requires the snapshot.

## Status (2026-10-12, second integration pass)

### Firewall and import health

- **Dependency firewall** (`tests/test_dependency_firewall.py`): 2 passed, 0 violators, no whitelist additions. `bitsandbytes` stays banned; `adamw8bit` is in-repo.
- **Import smoke** (15 modules incl. `train.adamw8bit`, `train.quant8`, `train.offload_adamw`, `utils.guard`; no GPU, no transformers): OK.
- **Py_compile scan** (every `scripts/*.py`): OK.

### Test coverage

- **Full CPU suite** (`pytest -m "not gpu and not oracle and not slow"`, align venv): **1089 passed, 13 skipped, 60 deselected** (61 s). Up from 997 passed at HEAD a4d4d5e.
  - Skips (unchanged): 5 (argilla_arm_a, HF cache missing), 3 (eval_set_build, MATH-500 missing), 1 (parquet ingest, no pyarrow), 3 (shard-tree data), 1 (manifest).
  - Deselected: GPU tests (51), slow tests (~9).
- **Oracle parity** (legacy venv, last measured in the first pass, not re-run): tokenizer 100% identical, verifier agreement 0.9975 on 400 pairs, tiny forward identical logits, checkpoint round-trip exact.

### Optimizer resolution

Every shipped full-FT config (`sft`, `pretrain`, `grpo_stageb`, `dpo` via `scripts/dpo.py`) names `optimizer.name = "adamw8bit"`. After bitsandbytes was removed that name raised `NotImplementedError`, and the only other optimizer (`offload_adamw`) cannot launch 7.25B on 4 ranks (4x over its pinned-staging cap, ~105GB host RSS). Resolved by `align/train/adamw8bit.py` (`BlockwiseAdamW8bit`): `make_optimizer(load_config("configs/sft.toml"), params)` now returns it and it steps (pinned by `tests/test_adamw8bit.py::test_every_shipped_adamw8bit_config_constructs_the_real_optimizer_and_steps`).

Projection for 7.25B / 4 ranks (`guard.projected_full_ft_rss_gb`, config-derived):

| Quantity | adamw8bit | offload_adamw |
|---|---|---|
| host optimizer state + master | 0 GB | 81 GB |
| on-GPU optimizer state, per rank | 3.38 GB | 0 |
| steady-state host RSS (`total_gb`) | 22.0 GB | 104.7 GB |
| checkpoint-save transient, rank-0 gather | 40.6 GB (3x state, estimate) | not modelled |
| checkpoint-save transient, `train.per_rank_optimizer_state` | 13.5 GB (1x) | not modelled |
| `peak_gb` preflight gates on | 62.6 GB (35.5 GB per-rank files) | 104.7 GB |

`preflight_rss` now requires `peak_gb + 45GB` of MemAvailable (the 45GB floor of `/apps/CLAUDE.md`). With the ~85GB the shipped configs were sized against, the default rank-0 gather is REFUSED (62.6 + 45 = 107.6) and per-rank files pass (35.5 + 45 = 80.5). `train.per_rank_optimizer_state = true` is now set in `configs/sft.toml` and `configs/pretrain.toml` (validated on 2 CPU/gloo ranks, bit-exact resume, `tests/test_trainer_dist.py`). The DPO and GRPO stage configs cannot carry it yet (the flag is not plumbed into the RunConfig their scripts build); see `.tmp/handoff/per_rank.md`.

### Resolved in this pass (was open in the first report)

| Area | Item | Resolution |
|---|---|---|
| Launch | No optimizer constructible for any shipped full-FT config | `adamw8bit` implemented, wired, tested (41 + 4 tests, 16 hand mutants killed) |
| Trainer | Logged lr was the post-step value | logs the last APPLIED lr; saved as `last_applied_lr` in `resume.pt` |
| Trainer | Resume lost `checkpoints_written` (stale-substrate rollout, keep-N blind to resumed checkpoint) | `resume_from` registers the resumed micro-step (idempotent); `scripts/ppo.py::_resume` is now a thin wrapper |
| Trainer | `extra_state` hook | `TrainModule.extra_state()/load_extra_state()` base stubs; Trainer saves/restores `resume.pt["extra_state"]` |
| Trainer | Rank-0 optimizer gather | opt-in `train.per_rank_optimizer_state` (field now exists, so struct-mode TOML can set it) |
| Checkpoint | Partial tokenizer copy | all-or-nothing, loud fallback to the vendored copy |
| GRPO | grad-norm tripwire measured post-clip | pre-clip norm captured from `clip_grad_norm_`'s return |
| GRPO | CollapseTripwire state lost on resume | `state_dict`/`load_state_dict`; rides `extra_state` |
| DPO | RefCacheKey could not tell checkpoints apart | key carries a streamed sha256 of the weights; old caches refused (re-precompute) |
| DPO | path-bias `forward_impl` mismatch; global_batch not validated | `dpo.py` reads `forward_impl` from the cache header; launch topology checked against `global_batch` |
| PPO | gamma != 1.0, lm_head without `label_token_ids` | refused at stage-config validation |
| PPO | optimizer/scheduler stepped on an all-skipped window | skipped (rank-consistent via one tiny all-reduce) |
| Config | `OptimizerConfig.max_grad_norm` | new optional field, forwarded to `adamw8bit` (None = no clipping, as before) |
| Guard | No device/save-transient projection for 8-bit | `device_optimizer_state_gb_per_rank`, `checkpoint_save_transient_gb`, `peak_gb`; preflight gates on peak |

### Still open

| Area | Issue | Why not closed |
|---|---|---|
| GRPO | `scripts/grpo.py` never calls `resume_from_latest()` | a correct resume needs the prompt cursor and dead-group-retry draws checkpointed (`GrpoBatchSource._cursor/_order` are not in any state); wiring only the Trainer resume would re-draw the earliest prompt blocks |
| GRPO | bf16 `log_softmax` under FSDP2 (rho misses it) | needs a GPU memory check at the 2-row micro shape |
| GRPO | restored latched tripwire does not re-populate `GrpoBatchSource.alarms` | `collapse_tripwire.fired` in `stage_result.json` still reports it |
| adamw8bit | no triton kernel; GPU step time at 1.8B local elements unmeasured; save peak is an estimate | no GPU jobs allowed |
| per-rank ckpt | barriers/resume now exercised on 2 CPU gloo ranks; NCCL/GPU 4-rank and DPO/GRPO plumbing still open | no GPU jobs allowed; DPO/GRPO need package edits (handoff) |
| PPO | `configs/ppo_stage.toml` `[TBD]` fields, PRM checkpoint absent, KL in both reward and loss | experiment-owner decisions |
| PPO | path-bias report is not bound to a checkpoint hash | pre-existing |
| Data | base model snapshot absent; data not migrated; corpus agreement figures stale | environment |

### Pre-flight checklist for next run

1. Firewall: `.venv/bin/python -m pytest tests/test_dependency_firewall.py -q` must pass.
2. CPU suite: `.venv/bin/python -m pytest tests -q -m "not gpu and not oracle and not slow" -p no:cacheprovider` (1365 passed, 13 skipped expected; `test_roundtrip_ceiling.py` x2 fail unless the PRM800K snapshot is cached).
3. Host: `free -g; uptime; awk '$3=="nvme2n1"{print $12}' /proc/diskstats`. Require >=45GB free, load < 2x cores, disk queue < 64 (`/apps/CLAUDE.md`); for adamw8bit the preflight additionally requires the projected save peak.
4. Re-run `scripts/precompute_ref_logprobs.py` for both DPO arms before `scripts/dpo.py` (old ref caches are refused by design).
5. `train.per_rank_optimizer_state = true` is set in sft/pretrain TOMLs; for DPO/GRPO 7B jobs it is not settable until plumbed (handoff) -- ensure MemAvailable >= ~108GB for those.
6. A pre-fix checkpoint resumes with the corrected LR schedule applied from the next step.
7. PPO: resolve the `[TBD]` fields in `configs/ppo_stage.toml` first.