AbstractPhil's picture
mini-beatrix-3/code: re-cut at amoe-lora 0.2.11 (9bae609: the compiled adapter chain's recompile budget and shape gate) + the stop / install / relaunch / guard scripts of the 2026-09-28 relaunch at 122,000; the pod's mission file as running
564e4dd verified
|
Raw History Blame Contribute Delete
8.6 kB

mini-beatrix-3 — the training code, pinned to the commits that train it

This folder is a verbatim snapshot of the two libraries and the pod files behind the mini-beatrix-3 run, first taken on 2026-09-27 at step ~106,600 of 245,670 and re-cut on 2026-09-28 at step 122,000 when the arm library moved to 0.2.11 (the section on the compiled chain below). The libraries are copies of their public GitHub commits; the pod files are the launch, resume, stop, install and shipping scripts as they run on the training machine.

part what it is where it comes from
alephllm/ the model, the trainer, the data plane, the instruments, the arm program and the mission loop (geolip.alephllm, version 0.10.5) github.com/AbstractEyes/alephllm @ e8b7440
amoe-lora/ the arm library: the relay adapters, attach / detach / mask handles, recompute-in-backward, the compiled adapter chain, the anchor file format (amoe.__version__ 0.2.11; its pyproject.toml still says 0.10.0) github.com/AbstractEyes/amoe-lora @ 9bae609 (branch experimental; f3db5d4 = 0.2.10 trained steps 106,000 to 122,000)
pod/mission.json the mission file: the craft, the data plane, the stage arms, the knobs; the trainer re-reads the refreshable keys at every phase boundary the pod, as running
pod/launch_mission.sh the launch: torchrun data-parallel over every card on the machine, resume-first the pod
pod/relaunch_106000.sh the one-command relaunch used on 2026-09-27: sets the knobs, checks the installed versions, launches, restarts the sidecar, prints the proving log lines the pod
pod/stop_at_106000.sh a clean manual stop at a 2,000-step checkpoint (the resume state is exactly that checkpoint) the pod
pod/stop_at_checkpoint.sh the same stop generalized to any checkpoint step (used at 122,000 on 2026-09-28) the pod
pod/install_0211.sh installs the arm library at the pinned 0.2.11 commit and reads the versions back; refuses while a trainer is alive the pod
pod/relaunch_3arms.sh the relaunch used on 2026-09-28: the running form, the version gate, the proving log lines, the recompile-warning count the pod
pod/go_relaunch_0211.sh stop at a named checkpoint, install, relaunch, each stage gated on the previous one's report line the pod
pod/compile_guard.sh polls the log for torch's recompile-limit warning (the sign of the compiled chain falling back to eager) and arms the stop-and-relaunch at the next checkpoint the pod
pod/arm_anchor_sidecar.py ships every attached stage arm as a safetensors anchor at each checkpoint (the trainer's own boundary anchors ship at phase ends) the pod
pod/write_mission_json.py how the mission file was written from the pre-flight screens' verdicts the pod
pod/env.sh the shell environment (the token comes from the container's own environment; nothing secret is in the file) the pod
requirements.txt the pod's package versions (torch 2.8.0 + CUDA 12.8 on two RTX 5090) the pod

Reproducing the run

  1. Install torch 2.8.0 for your CUDA, then the two libraries at the pinned commits (requirements.txt has the lines), or pip install -e the two snapshots here.
  2. Copy pod/mission.json, set out_dir and local, and point ALEPHLLM_MISSION_CONFIG at it. A fresh run starts from step 0 (fresh: true); a resume pulls <craft>/manifest.json and resume/latest.pt from this repo and continues.
  3. Launch: python -m torch.distributed.run --standalone --nproc_per_node=<cards> -m geolip.alephllm.train.mission (pod/launch_mission.sh does exactly this and logs to logs/mission.log).
  4. Outputs land under out_dir/<craft>/: checkpoints/step_XXXXXXXX.safetensors (bf16) every 2,000 steps with an fp8 shipping copy, resume/latest.pt (fp32 model, Muon and Adam states, stream position, RNG), arms/<stage>_stepN.safetensors (the stage arms), reports/v3/<phase>_stepN.json (the boundary reads) and runs/ (TensorBoard).

The knobs that matter (all in mission.json)

key meaning value in this run
micro_batch, grad_accum rows per forward, chunks per step; with two cards that is 262,144 tokens per step at context 4096 2, 16
stage_arms one relay arm per curriculum stage: name, phase, the quiet-term weight (lambda), the stage index eight arms at lambda 2
arm_spec the arm geometry per block: 16 slots, K 16 atoms at D 8, hidden 256 LAWFUL_16x8
arm_spec_certified a gate: a phase that attaches an arm opens only when this is true; setting it false ends the session cleanly at the next boundary true
abstain_chunks abstention (quiet-term) chunks per step; null = one per task chunk (16 of 16); an integer = that many, each scaled so the term's weight per step is unchanged 4 since step 106,000
adapter_recompute the arms' intermediates are recomputed in the backward instead of kept (about 0.3 GB per arm instead of 3 GB at this shape) true since 106,000
adapter_compile the per-block adapter chain runs through torch.compile with inductor's precision-cast emulation (the plain compile produced a non-finite backward on the address stage); since 0.2.11 the library also raises torch's recompile budget and routes only the training shapes through the compiled chain (the section below) true since 106,000
minted_lexicon the curriculum's word pool for the rule-chain stages (216 minted words), its share of stage 3 and the pair rate see the file
head_birth how the aleph head is born (revival_birth fits it to the first batch at step 0) revival_birth
l6_ledger the in-run guard certification record (the rank-collapse guard halts the run on its own when certified) v3_preflight/l6_guards/...

The abstention term in one paragraph

Each stage arm is a 13.7M-parameter adapter applied in series after every block: a projection to 128, a closed-form signed address over 16 slots x 16 atoms, an MLP 128-256-1024, a sigmoid gate born at -3 and a zero-initialised head, so a fresh arm changes nothing at birth. Beside the task loss, every step runs abstention chunks on held-out web text: for each live arm, one forward with that arm masked and one armed forward, and the KL from the masked distribution to the armed one, summed over arms and scaled by lambda, pulls each arm toward silence off its stage. Masking an arm is an exact bypass, so the trunk with the arms off is bit-identical to the bare model, and each arm ships as a separate file.

Reading an arm file

from safetensors.torch import load_file
sd = load_file("arms/s1_perspective_step106000.safetensors")   # keys blocks.<block>.<param>; 320 tensors, 13,697,056 params
# or, with the arm library: from amoe.io import load_anchor; a = load_anchor(path) -> a.meta (phase, lambda, spec, base_model_id) + a.adapters

Anchors before 2026-09-25 were converted in place from a torch-serialised form to real safetensors; load_anchor reads either form by content.

The compiled chain and its recompile budget (0.2.11)

The adapter chain is compiled with static shapes, so every distinct (mask pattern, input shape) pair is one entry in torch's compile cache for that function. Torch's default budget is 8 recompiles per function, and when it is exceeded torch does not run that one call eagerly: it marks the function to bypass the compiler for the rest of the process, with a single warning line. The stage 2 boundary evaluation fed the armed model probe prompts of many different lengths, spent that budget in seconds, and from the stage 3 attach on every step ran the arms eager (10.5 s per step at three arms instead of about 7.5) until the run was stopped at 122,000 and relaunched on 0.2.11. The fix lives in the library, not in a launch flag: compiled_chain() raises the budget (64 per function, 4,096 accumulated) before the first compile, and chain_for(h) routes only the first two distinct (batch, length) shapes seen through the compiled chain, every other shape through the eager chain, which is the same function. The unit test tests/test_compiled_chain_gate.py in the GitHub repository stacks three arms, runs the training shape armed and with each member masked, then forty prompt lengths, with torch configured to raise on any budget hit. pod/compile_guard.sh watches the log for the warning and arms a stop-and-relaunch if it ever returns.