AbstractPhil/mini-beatrix-1
12.8k
1"""Mission presets — the Mini-Beatrix ladder.2 3Naming convention (voyager style): numbered missions, each a fixed craft.4Small crafts are "mini-beatrix-N"; the BPE flagship is "beatrix-voyager".5Beatrix is the lineage collective name; missions are launched in order and6all upload to the one training repo (TRAINING_REPO), each craft under its7own path prefix (checkpoints + manifest + tensorboard).8 9 mini-beatrix-0 d512 L12 ctx1024 byte-trigram 37.6M gate craft:10 its first toggle evals ARE the anchored-bank-under-AR11 screen (P1) running live.12 mini-beatrix-1 d768 L16 ctx2048 byte-trigram 112.5M first Colab13 mission (default).14 mini-beatrix-2 d1024 L32 ctx8192 byte-trigram ~873.7M FULL SPLAT:15 a governed multi-constellation hub in EVERY block16 (2026-08-26 rescale; the v1 249M 3-hub shape retired17 untrained — plan 2026-08-26_mini_beatrix_v2_shape.md).18 mini-beatrix-2s d1024 L20 ctx4096 byte-trigram ~233M the lawful19 screen craft: every v2 gating cell runs here first.20 beatrix-voyager d1536 L24 ctx4096 BPE(gpt2 50k) 775.3M flagship;21 vocab-scale head + BPE screens (P2/P5) still open —22 launch only after mini-beatrix verdicts.23 24Every craft is inference-capable on consumer hardware in its shipped25form (fp8-e4m3 safetensors variants are exported alongside checkpoints).26"""27from __future__ import annotations28 29from dataclasses import dataclass, field, asdict30from typing import Optional31 32 33@dataclass34class AlephLMConfig:35 name: str = "mini-beatrix-0"36 d_model: int = 51237 n_layers: int = 1238 n_heads: int = 839 context: int = 102440 vocab_size: int = 256 # bytes; BPE presets override41 tokenizer: str = "byte-trigram" # "byte-trigram" | "hf:<repo or name>"42 hub_layers: tuple = (3, 7, 11) # CausalSplatHUB depths; () = pure sdpa control43 hub_K: int = 51244 hub_D: int = 3245 tau: float = 0.146 bank_experts: int = 3 # E1-validated fat-expert count47 bank_ff: Optional[int] = None # None -> d_model (E1 ratio)48 head_K: int = 51249 head_D: int = 3250 gate_init: float = -3.051 tie_embeddings: bool = False # BPE crafts tie; byte crafts cannot (trigram)52 hub_chunk: int = 128 # chunked-scan block for the hub prefix memories53 # v2 (2026-08-26): multi-constellation hubs — the product-code form at54 # lawful supply (K <= 2*hub_D per book; ROUND 5e). 1 = the v1 layout,55 # bit-identical state dict. Old manifests load via the default.56 hub_const: int = 157 # Activation checkpointing (training only; inference/decode untouched).58 # 0 = off (v1 verbatim). 1 = recompute the hub read in backward.59 # 2 = also recompute the bank branch. At v2 scale (16 books x ctx 819260 # x 32 layers) the retained scan tensors alone exceed a 95GB card —61 # measured OOM, Blackwell preflight 2026-08-26. ~2x hub recompute cost.62 hub_ckpt: int = 063 # v3 (2026-09-19): weak-token fusion at the input plane. None = the64 # byte-resolution trunk verbatim. A dict selects the hourglass form:65 # {"rule": "entropy" | "spacelike", "theta": bits, "witness_floor": n,66 # "table": "<npz path>", "k_lo": front blocks, "k_hi": back blocks} —67 # see model/fusion.py. Old manifests load via the default.68 fusion: Optional[dict] = None69 70 def to_dict(self):71 d = asdict(self)72 d["hub_layers"] = list(self.hub_layers)73 return d74 75 @staticmethod76 def from_dict(d):77 d = dict(d)78 d["hub_layers"] = tuple(d.get("hub_layers", ()))79 return AlephLMConfig(**d)80 81 82@dataclass83class TrainConfig:84 # Optimizer split (measured: momentum-geometric +.09 on the aleph;85 # the mechanism is ~20x more optimizer-sensitive than sdpa).86 muon_lr: float = 2e-287 muon_momentum: float = 0.9588 adam_lr: float = 3e-4 # pure Adam, wd=0 — never AdamW89 warmup_steps: int = 200 # scale insurance; flat after (flat-LR law)90 grad_clip: float = 1.091 micro_batch: int = 2492 grad_accum: int = 193 # Cadences (steps)94 log_every: int = 5095 health_every: int = 50096 eval_every: int = 200097 ckpt_every: int = 2000 # safetensors + resume .pt98 fp8_every_ckpts: int = 5 # every Nth checkpoint also exports fp899 tb_upload_every: int = 1000100 # Eval sizes101 val_tokens: int = 262144102 canary_episodes: int = 128103 seed: int = 1337104 compile: bool = False105 # The anchor governor (ROUND 5f, 2026-08-25): post-optimizer-step106 # min-separation projection over hub/head codebooks — preventive107 # anti-crowding, identity when slack, zero parameters, outside the108 # task gradient (the no-balance-machinery law is untouched).109 governor: str = "" # "" off (v1 verbatim) | "minsep"110 governor_theta: float = 45.0 # deg; scale ~ gamma*(D): 45 at D=256111 governor_every: int = 8 # steps between slack checks (~free)112 # Post-revival address freeze (0.8.2; RIDERS 11-12): after the113 # BOUNDARY-WRITE head revival, proj + head codebook freeze so the114 # self-burial channel (proj rotating to codebook-orthogonality,115 # measured 2/2 crafts) is structurally closed — only W_s trains.116 # requires_grad-only: optimizer param groups are UNCHANGED, so resume117 # state loads verbatim (Muon skips grad-less params).118 head_addr_frozen: bool = False119 # v3 (2026-09-19): per-phase LR multiplier keyed by phase-name PREFIX120 # ({"anneal": 0.5} scales both anneal phases). {} = the flat-LR form121 # verbatim — the v2 anneal ran at lr_scale 1.000 throughout (a diet122 # change, not an LR decay); the anneal as a LOWER-rate consolidation123 # stage is the v3 routine's term, its multiplier unmeasured (owed).124 phase_lr_scale: dict = field(default_factory=dict)125 # v3: open every phase's stream with a phase-specific seed offset so a126 # corpus that sits at the same recipe index in several stages does not127 # replay the identical shuffle head; False = the 2s form.128 phase_seed_offset: bool = False129 130 131# All missions upload to the one training repo, each under its own prefix132# (Phil's repo: checkpoints + manifests + tensorboard for every craft).133TRAINING_REPO = "AbstractPhil/alephllm-mini-beatrix-training"134 135 136@dataclass137class Preset:138 model: AlephLMConfig139 train: TrainConfig140 hf_repo: str = TRAINING_REPO # run repo (ckpts+manifest+tb)141 curriculum: list = field(default_factory=list) # [(phase, dataset, planned_tokens)]142 # v3: the curriculum-stage mixes are scaled (and rebalanced under the143 # epoch cap) by this factor when the trainer opens a stage — see144 # data/curriculum.py apply_curriculum_scale. 1.0 = the 2s schedule.145 data_scale: float = 1.0146 # v3: the two data-plane decisions a scale other than 1x needs (the147 # trainer refuses to open a scaled stage without them): the epoch cap148 # per finite corpus per stage (None = the audit threshold, flagged)149 # and the rebalance rule ('natural' | 'generators' | 'hold').150 epoch_cap: float | None = None151 rebalance_to: str | None = None152 153 @property154 def prefix(self) -> str: # path prefix inside hf_repo155 return self.model.name156 157 158def _curriculum(warm: int, main: int, ext: int):159 return [160 dict(name="warmup_wikitext", dataset="wikitext-103", planned_tokens=warm,161 status="planned"),162 dict(name="fineweb_main", dataset="fineweb-edu", planned_tokens=main,163 status="planned"),164 # Deliberately not prepped beyond a name — the full plan exists in the165 # manifest, the data work happens when the phase activates.166 dict(name="fineweb_extended", dataset="fineweb-edu", planned_tokens=ext,167 status="deferred"),168 # phase C: distribution shift toward chat format / simple register /169 # narrative (incl. moral texture) / binding demand — see streams.ANNEAL_MIX170 dict(name="anneal_mix", dataset="anneal-mix",171 planned_tokens=2_000_000_000, status="deferred"),172 ]173 174 175PRESETS: dict[str, Preset] = {176 "mini-beatrix-0": Preset(177 model=AlephLMConfig(name="mini-beatrix-0"),178 train=TrainConfig(micro_batch=96, grad_accum=1),179 curriculum=_curriculum(150_000_000, 1_000_000_000, 2_000_000_000),180 ),181 "mini-beatrix-1": Preset(182 model=AlephLMConfig(name="mini-beatrix-1", d_model=768, n_layers=16,183 n_heads=12, context=2048, hub_layers=(4, 9, 14)),184 train=TrainConfig(micro_batch=48, grad_accum=3),185 curriculum=_curriculum(300_000_000, 3_000_000_000, 6_000_000_000),186 ),187 # v2 (2026-08-26, Phil's draft off the Foundry console): FULL-SPLAT —188 # a hub in every block, multi-constellation product code at lawful189 # supply (16 books x 256 anchors in 256-dim spaces = 1.0x supply;190 # v1's single book ran 16x and crowded), governed from birth, ctx 8192191 # where the O(L) read is ~4.5x cheaper than the MHA equivalent.192 # ~873.7M params. Plan: history/plans/2026-08-26_mini_beatrix_v2_shape.md.193 "mini-beatrix-2": Preset(194 model=AlephLMConfig(name="mini-beatrix-2", d_model=1024, n_layers=32,195 n_heads=16, context=8192,196 hub_layers=tuple(range(32)),197 hub_K=256, hub_D=256, hub_const=16,198 bank_experts=6, bank_ff=1024,199 # chunk 1024 MEASURED on the mission card (C2e,200 # Blackwell 2026-08-26): 72.2 vs 83.2 ms/layer201 # fwd+bwd at chunk 256, peak 39.4 -> 26.6 GB.202 # S/P traffic ~ 1/C, att work ~ C; config-only,203 # checkpoint-compatible, exactness C-independent.204 head_K=256, head_D=256, hub_chunk=1024,205 hub_ckpt=2),206 train=TrainConfig(micro_batch=4, grad_accum=16,207 governor="minsep", governor_theta=45.0),208 curriculum=_curriculum(500_000_000, 8_000_000_000, 16_000_000_000),209 ),210 # THE ACTIVE MISSION (2026-08-26, Phil: "train the next stage up from211 # the beatrix v1; we can't train the large one currently"): the lawful212 # full-splat craft one rung above v1 — d1024 L20 ctx4096, governed213 # 4x64@128 books (4x supply headroom vs v1's crowded 16x). Also the214 # screen bed for every v2-era gating cell. hub_ckpt=0: at 237M the215 # retained scan fits the 96GB card, so the recompute tax is pure waste216 # (fallback: set hub_ckpt=2 if the preflight bench gate aborts >88GB).217 "mini-beatrix-2s": Preset(218 model=AlephLMConfig(name="mini-beatrix-2s", d_model=1024, n_layers=20,219 n_heads=16, context=4096,220 hub_layers=tuple(range(20)),221 hub_K=64, hub_D=128, hub_const=4,222 bank_experts=3, bank_ff=1024,223 head_K=256, head_D=256, hub_chunk=256,224 hub_ckpt=0),225 train=TrainConfig(micro_batch=16, grad_accum=4,226 governor="minsep", governor_theta=45.0,227 head_addr_frozen=True),228 curriculum=_curriculum(300_000_000, 5_000_000_000, 10_000_000_000),229 ),230 "beatrix-voyager": Preset(231 model=AlephLMConfig(name="beatrix-voyager", d_model=1536, n_layers=24,232 n_heads=16, context=4096, vocab_size=50257,233 tokenizer="hf:gpt2", tie_embeddings=True,234 hub_layers=(6, 13, 20)),235 train=TrainConfig(micro_batch=8, grad_accum=16),236 curriculum=_curriculum(500_000_000, 12_000_000_000, 24_000_000_000),237 ),238}239 240def make_v3_preset(n_layers: int = 24, d_model: int = 1024,241 data_scale: float = 4.0, epoch_cap: float | None = None,242 rebalance_to: str | None = None,243 name: str | None = None) -> Preset:244 """The v3 craft (plan of record 2026-09-15, S2/S14; sizing 09-15):245 the solidified all-splat form at d1024 — a governed hub in EVERY246 block, the certified hub geometry (4 books x 64 @ D128), banks247 3 x ff1024, head 256@256, ctx 4096 — at a depth the throughput bench248 priced (24 or 28 blocks; the choice is the program lead's, with the249 price beside it). Phases at `data_scale` x the 2s schedule (4x:250 warmup 0.3B, fineweb_main 20.9B, S0-S8 35.2B rebalanced under the251 epoch cap, anneal_nochat 4B, anneal_mix 4B = 64.4B bytes), listed252 CHRONOLOGICALLY and planned from birth (the two-phase anneal is part253 of the routine, not a post-hoc activation). Birth recipe: the head254 address trains (head_addr_frozen False — the 2s's True is a255 post-revival flag); no hub gain, no fusion (owed / the lead's).256 epoch_cap / rebalance_to: the data-plane decisions (the trainer257 refuses a scaled stage without a rebalance rule); under 'hold' the258 stages stay at 1x and the held budget goes to fineweb_main."""259 # VENDORED: the live package reads260 # from .data.curriculum import curriculum_phases, _BASE_STAGE_TOKENS261 # Written that way here it is a relative import of a SUBPACKAGE, and262 # transformers' remote-code loader resolves every relative import to a263 # flat file beside this one — it would demand "data.curriculum.py" and264 # refuse to load the model. This package ships the inference path only;265 # the curriculum lives in the installable geolip-alephllm, so the266 # import is deferred and resolved by name when that package is there.267 from importlib import import_module268 _curriculum = import_module("geolip.alephllm.data.curriculum")269 curriculum_phases = _curriculum.curriculum_phases270 _BASE_STAGE_TOKENS = _curriculum._BASE_STAGE_TOKENS271 if name is None:272 name = "mini-beatrix-3" if n_layers == 24 and d_model == 1024 \273 else f"mini-beatrix-3-d{d_model}-l{n_layers}"274 s = float(data_scale)275 model = AlephLMConfig(name=name, d_model=d_model, n_layers=n_layers,276 n_heads=max(1, d_model // 64), context=4096,277 hub_layers=tuple(range(n_layers)),278 hub_K=64, hub_D=128, hub_const=4,279 bank_experts=3, bank_ff=1024,280 head_K=256, head_D=256, hub_chunk=256, hub_ckpt=0)281 train = TrainConfig(micro_batch=16, grad_accum=4,282 governor="minsep", governor_theta=45.0,283 head_addr_frozen=False, phase_seed_offset=True)284 # the warmup phase stays at 300M (the LR warmup is 200 steps = 52M285 # tokens; wikitext-103 is a finite corpus the stage audit does not286 # cover) and its share of the scale moves to fineweb_main, so the287 # general-text total is (0.3 + 5.0) x scale exactly288 warm = 300_000_000289 main = int((300_000_000 + 5_000_000_000) * s) - warm290 if rebalance_to == "hold":291 # the stages stay at 1x bytes; the held (s-1) x 8.8B is general text292 main += int(round((s - 1.0) * sum(_BASE_STAGE_TOKENS.values())))293 phases = [294 dict(name="warmup_wikitext", dataset="wikitext-103",295 planned_tokens=warm, status="planned"),296 dict(name="fineweb_main", dataset="fineweb-edu",297 planned_tokens=main, status="planned"),298 *curriculum_phases(s, rebalance_to),299 dict(name="anneal_nochat", dataset="anneal-nochat",300 planned_tokens=int(1_000_000_000 * s), status="planned"),301 dict(name="anneal_mix", dataset="anneal-mix",302 planned_tokens=int(1_000_000_000 * s), status="planned"),303 ]304 return Preset(model=model, train=train, curriculum=phases, data_scale=s,305 epoch_cap=epoch_cap, rebalance_to=rebalance_to)306 307 308try:309 PRESETS["mini-beatrix-3"] = make_v3_preset(24)310 PRESETS["mini-beatrix-3-l28"] = make_v3_preset(28, name="mini-beatrix-3-l28")311except ImportError:312 # the vendored automodel copies (the mirror law) carry model/ +313 # presets.py without the data stack: the v3 presets need the314 # curriculum registry and are simply absent there315 pass316 317 318def _copy_train(t: TrainConfig) -> TrainConfig:319 """A field-wise copy with NO shared containers (the dict field would320 otherwise alias between a treatment and its twin)."""321 import copy as _copy322 return TrainConfig(**{k: _copy.deepcopy(getattr(t, k))323 for k in t.__dataclass_fields__})324 325 326# Pure-sdpa control crafts (hub layers removed) — the running architecture327# control for any mission: same params otherwise, suffix "-control".328for _name in list(PRESETS):329 _p = PRESETS[_name]330 _m = AlephLMConfig.from_dict(_p.model.to_dict())331 _m.name = _name + "-control"332 _m.hub_layers = ()333 # 0.8.7: twins get their OWN TrainConfig copy — the shared-instance334 # form let treatment-specific flags leak into controls (2s-control335 # inherited head_addr_frozen=True, a post-revival flag no control's336 # birth recipe may carry) and made cross-mutation possible.337 _t = _copy_train(_p.train)338 PRESETS[_name + "-control"] = Preset(339 model=_m, train=_t,340 curriculum=[dict(x) for x in _p.curriculum],341 data_scale=_p.data_scale, epoch_cap=_p.epoch_cap,342 rebalance_to=_p.rebalance_to)343 344# The 2s architecture control runs the BIRTH recipe verbatim: born-null345# unfrozen head (it buries, as the treatment's did for its first 24,860346# steps — measured 3/3; the +0.01 head term is immaterial at the ±3.4347# hub scale this control exists to judge).348PRESETS["mini-beatrix-2s-control"].train.head_addr_frozen = False349 350 351def get_preset(name: str) -> Preset:352 if name not in PRESETS:353 raise KeyError(f"unknown preset '{name}' — have: {sorted(PRESETS)}")354 return PRESETS[name]355