| # 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. |
| |
