Upload beatrix_curriculum.ipynb

mini-beatrix-1/beatrix_curriculum.ipynb

@@ -0,0 +1,81 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "# Beatrix \u2014 early-life curriculum\nHer school notebook. Plan of record: `claude-mind/history/plans/2026-08-15_early_life_curriculum.md`.\n\nNine stages, ~8.8B tokens, resume-first: open this notebook any day, run **Cell 1** then **Cell 3** with whatever hours you have, and walk away. Stage boundaries fire their own instruments. Cell 2 runs ONCE (baseline, both cores). Cell 4 renders growth anytime.\n\nLaws wired in: no chat template / no identity in any training row \u00b7 causal-test holdout families never reach a training row (bAbI 16+19, depth-5 chains, 3-digit subtraction) \u00b7 every stage vals on the same fineweb-2013 holdout (gauge continuity) \u00b7 probes carry the stage-specific measurement."
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": "%pip install -q \"datasets>=5.0.0\" \"geolip-alephllm @ git+https://github.com/AbstractEyes/alephllm@259bf2ec5b7a85c837122add9b703938484a04e3\" \"amoe-lora @ git+https://github.com/AbstractEyes/amoe-lora@7e1baba9c0bb1cc38815e4424d969b80ae5a315d\"\n# datasets>=5: Colab's preinstalled version predates the 'Json' feature\n# type that mjbommar/opengloss (S8 definitions) declares \u2014 S8 aborted on\n# it 2026-08-16. 5.0.0 is the version every mix + resume roundtrip was\n# verified under locally.\nimport os, torch\nos.environ[\"HF_HUB_DISABLE_PROGRESS_BARS\"] = \"1\"   # mid-train download\n# bars are the tqdm instances that collided with the training bar's\n# write path (step-59,000 abort); the trainer is also hardened, this\n# removes the source\nfrom google.colab import userdata\nHF_TOKEN = userdata.get(\"HF_TOKEN\")\nos.environ[\"HF_TOKEN\"] = HF_TOKEN\nassert torch.cuda.is_available(), \"GPU runtime required\"\nimport geolip.alephllm as al\nprint(\"geolip\", al.__version__, \"\u00b7\", torch.cuda.get_device_name(0))",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## Cell 2 \u2014 instrument gate (run ONCE, before any training)\nProbes P0\u2013P8 on **both** cores \u2014 the locked 51,882 and the annealed 58,664 \u2014 so every growth curve has a time-zero. Ships the baseline report to the training repo. The head surgery (fold the fossilized gamma gate into W_s; verified semantic no-op, max|diff| 2.4e-07) is applied to the eval copies here and to the live run in Cell 3."
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": "import json, torch\nfrom safetensors.torch import load_file\nfrom huggingface_hub import hf_hub_download, HfApi\nfrom geolip.alephllm.presets import get_preset\nfrom geolip.alephllm.model.alephlm import AlephLM\nfrom geolip.alephllm.data import build_tokenizer\nfrom geolip.alephllm.data import curriculum as C\nfrom geolip.alephllm.train import probes\n\nREPO = \"AbstractPhil/alephllm-mini-beatrix-training\"\n# STEPS = None -> probe the NEWEST checkpoint on the hub (numeric step\n# compare \u2014 the post-rewind current core, never a stale hardcode).\n# Set STEPS = [51882, 58664] to re-measure the historical t0 pair.\nSTEPS = None\nimport re as _re\nfrom huggingface_hub import HfApi as _HfApi\nif STEPS is None:\n    _cks = sorted(int(m.group(1)) for f in _HfApi().list_repo_files(REPO)\n                  for m in [_re.search(r\"step_(\\d+)\\.safetensors$\", f)] if m\n                  if \"/fp8/\" not in f)\n    STEPS = [_cks[-1]]\n    print(f\"probing NEWEST checkpoint: step {STEPS[0]:,}\")\npreset = get_preset(\"mini-beatrix-1\")\ntok = build_tokenizer(preset.model.tokenizer)\nbaseline = {}\nfor step in STEPS:\n    ck = hf_hub_download(REPO, f\"mini-beatrix-1/checkpoints/step_{step:08d}.safetensors\")\n    m = AlephLM(preset.model); m.load_state_dict(load_file(ck)); m.cuda().eval()\n    fold = C.fold_head_gate(m)          # no-op when gamma already 1.0\n    res = probes.run_all(m, tok, \"cuda\")\n    baseline[str(step)] = {\"probes\": res, \"head_fold\": fold}\n    print(f\"=== step {step:,} (gamma was {fold['gamma_before']:.6f}) ===\")\n    print(probes.report(res)); print()\n    del m; torch.cuda.empty_cache()\nname = (\"baseline.json\" if STEPS == [51882, 58664]\n        else f\"gate_step{STEPS[-1]}.json\")\nwith open(name, \"w\", encoding=\"utf-8\") as f:\n    json.dump(baseline, f, indent=1, default=float)\nHfApi(token=HF_TOKEN).upload_file(path_or_fileobj=name,\n    repo_id=REPO, path_in_repo=f\"mini-beatrix-1/reports/curriculum/{name}\",\n    commit_message=\"curriculum baseline: probes P0-P8 on both cores at t=0\")\nprint(f\"gate report shipped -> reports/curriculum/{name}\")",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## Cell 2b \u2014 rewind (one-time, guarded)\nRolls the run back to a stage boundary and re-plans from there with the corrected mixes. Set `CONFIRM = True` to arm it.\n\nDefault target is **step 66,803 (S2 end)**: the measured assessment cleared S1/S2 (fineweb 1.130 \u2192 1.144, probes up) and convicted S3 alone \u2014 fineweb \u2192 1.705, grounding/difference/articulation down, and the rule-chain template leaking into her prose (*\"Water boils when they are boiling\"*). What S3 bought \u2014 logic 0.818 \u2192 1.000 including the held-out depth-5 chains \u2014 is re-earnable from the fixed mix.\n\nFrom geolip 0.6.4 on, every boundary also archives `resume/boundary_<stage>_step<N>.pt` (weights + optimizer + stream), so later rewinds restore *exactly*. This one predates those archives, so it restores weights with a fresh optimizer and says so."
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": "# ONE-TIME REWIND. Set CONFIRM = True to run it; it DISCARDS training\n# after ROLLBACK_STEP and re-plans from ROLLBACK_STAGE with the current\n# (corrected) mixes. Default target: the S2 boundary, because the\n# measured assessment cleared S1/S2 (fineweb 1.130 -> 1.144 flat, probes\n# up) and convicted only S3 (fineweb -> 1.705, p0/p2/p8 down, the\n# rulechain template leaking into her prose).\nCONFIRM = False\nROLLBACK_STEP = 66803\nROLLBACK_STAGE = \"curriculum_s3\"\n\nfrom geolip.alephllm import prepare\nfrom geolip.alephllm.data import curriculum as C\n\nif not CONFIRM:\n    print(\"rollback NOT run \u2014 set CONFIRM = True to rewind to \"\n          f\"step {ROLLBACK_STEP:,} and replan from {ROLLBACK_STAGE}\")\nelse:\n    run = prepare(\"mini-beatrix-1\", hf_token=HF_TOKEN)\n    C.append_curriculum_phases(run.manifest)\n    if float(run.raw_model.head.gamma.item()) != 1.0:\n        C.fold_head_gate(run.raw_model)\n    info = C.rollback_to(run, ROLLBACK_STEP, ROLLBACK_STAGE,\n                         repo=\"AbstractPhil/alephllm-mini-beatrix-training\")\n    # persist the rewind so Cell 3 resumes from it (prepare() reads the\n    # hub; an in-memory rewind alone would be undone by the next session)\n    run._checkpoint(boundary=\"rollback\")\n    print(run.manifest.summary())\n    if not info[\"exact\"]:\n        print(\"\\nNOTE: fresh optimizer state (this boundary predates the \"\n              \"resume archives) \u2014 flag any arm trained from here.\")",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## Cell 3 \u2014 the school day (re-run every session)\nResumes wherever she is. Each `train()` call is capped at the current stage's remaining tokens, so sessions stop **exactly at boundaries**; the boundary fires probes + toggle ledger + head-election gauge and ships the report. Set `MAX_HOURS` to your session budget."
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": "import json, time, torch\nfrom huggingface_hub import HfApi\nfrom geolip.alephllm import prepare\nfrom geolip.alephllm.data import curriculum as C\nfrom geolip.alephllm.data.streams import build_stream\nfrom geolip.alephllm.train import probes, instruments\n\nMAX_HOURS = 8.0\nTURBO = False    # compile gate: KNOWN-FAIL on Blackwell (compiled bf16\n                 # backward emits NaN grads \u2014 caught 2026-08-15) AND the\n                 # trainer must never checkpoint a compiled model\n                 # (OptimizedModule key prefix corrupts resume). Leave\n                 # False until the NaN hunt lands; eager fused 1.7x is\n                 # always on.\nREPO = \"AbstractPhil/alephllm-mini-beatrix-training\"\n\n# continue-safety: this cell is safely RE-RUNNABLE in a live kernel \u2014\n# tear down any previous run first so prepare() doesn't stack a second\n# model+optimizer into VRAM\nimport gc\nif \"run\" in globals():\n    del run\ngc.collect()\ntorch.cuda.empty_cache()\n\nrun = prepare(\"mini-beatrix-1\", hf_token=HF_TOKEN)\nadded = C.append_curriculum_phases(run.manifest)\nif added:\n    print(f\"curriculum registered: {added} stages appended to the manifest\")\nif float(run.raw_model.head.gamma.item()) != 1.0:\n    info = C.fold_head_gate(run.raw_model)\n    print(f\"head gate folded: gamma {info['gamma_before']:.6f} -> 1.0 \"\n          f\"(||W_s||={info['w_s_norm']:.4f}; Muon re-elects from here)\")\n\nif TURBO:\n    import copy as _copy, time as _time\n    _vs = build_stream(\"fineweb-edu\", run.tokenizer, run.cfg.context,\n                       run.tc.micro_batch, seed=run.tc.seed + 9999,\n                       role=\"val\")\n    # SMALL probe (4 rows) under bf16 autocast \u2014 the gate must test the\n    # regime she actually trains in; a full-batch fp32 probe with a live\n    # backward graph OOMed a 96GB card (measured 2026-08-15)\n    _xb = _vs.next_batch()[:4].to(run.device)\n\n    def _probe(mm, plist):\n        for p in plist:\n            p.grad = None\n        with torch.autocast(\"cuda\", dtype=torch.bfloat16):\n            _, loss = mm(_xb[:, :-1], targets=_xb[:, 1:])\n        loss.backward()\n        g = torch.cat([(p.grad if p.grad is not None\n                        else torch.zeros_like(p)).flatten() for p in plist])\n        return float(loss.item()), g\n\n    _ref_p = list(run.raw_model.parameters())\n    _l0, _g0 = _probe(run.raw_model, _ref_p)\n    _g0 = _g0.clone()\n    for p in _ref_p:                 # free grads BEFORE the deepcopy \u2014\n        p.grad = None                # otherwise the twin copies them too\n    torch.cuda.empty_cache()\n    _twin = _copy.deepcopy(run.raw_model)\n    _twin_p = list(_twin.parameters())\n    _twin.compile_hubs()\n    _probe(_twin, _twin_p)                      # compile warmup\n    _l1, _g1 = _probe(_twin, _twin_p)\n    _dl = abs(_l1 - _l0)\n    _dg = float((_g1 - _g0).abs().max().item())\n    _t0 = _time.time(); [_probe(_twin, _twin_p) for _ in range(3)]\n    _tc = (_time.time() - _t0) / 3\n    _t0 = _time.time(); [_probe(run.raw_model, _ref_p) for _ in range(3)]\n    _te = (_time.time() - _t0) / 3\n    if _dl < 5e-3 and _dg < 5e-2:    # bf16-autocast tolerances\n        run.raw_model.compile_hubs()\n        print(f\"[turbo] gate PASSED on this GPU: dloss {_dl:.2e} \u00b7 \"\n              f\"max|dgrad| {_dg:.2e} \u00b7 fwd+bwd {_te:.2f}s eager -> \"\n              f\"{_tc:.2f}s compiled ({_te/max(_tc,1e-9):.1f}x)\")\n    else:\n        print(f\"[turbo] gate FAILED (dloss {_dl:.2e}, max|dgrad| {_dg:.2e})\"\n              \" \u2014 staying on the eager fused path\")\n    del _twin, _twin_p, _g0, _g1\n    for p in _ref_p:\n        p.grad = None\n    torch.cuda.empty_cache()\n\ndef boundary_instruments(run, tag):\n    m, tok = run.raw_model, run.tokenizer\n    res = probes.run_all(m, tok, run.device)\n    vs = build_stream(\"fineweb-edu\", tok, run.cfg.context,\n                      run.tc.micro_batch, seed=run.tc.seed + 9999, role=\"val\")\n    val = [vs.next_batch().to(run.device) for _ in range(2)]\n    led = instruments.toggle_ledger(m, val)\n    ws = float(m.head.w_s.weight.norm().item())\n    print(f\"== boundary {tag} \u00b7 step {run.step:,} ==\")\n    print(probes.report(res))\n    print(f\"  fineweb gauge: full {led['bpb_full']:.4f} \u00b7 \"\n          f\"head_off {led['toggle_head_aleph_off']:+.4f} \u00b7 \"\n          f\"hub_off {led.get('toggle_hub_off', 0):+.4f} \u00b7 ||W_s||={ws:.3f}\")\n    rep = {\"tag\": tag, \"step\": run.step, \"probes\": res, \"ledger\": led,\n           \"w_s_norm\": ws}\n    name = f\"boundary_{tag}_step{run.step}.json\"\n    with open(name, \"w\", encoding=\"utf-8\") as f:\n        json.dump(rep, f, indent=1, default=float)\n    HfApi(token=HF_TOKEN).upload_file(path_or_fileobj=name, repo_id=REPO,\n        path_in_repo=f\"mini-beatrix-1/reports/curriculum/{name}\",\n        commit_message=f\"curriculum boundary report: {tag}\")\n\nt_end = time.time() + MAX_HOURS * 3600\nwhile time.time() < t_end:\n    ph = run.manifest.current_phase()\n    if ph is None:\n        print(\"curriculum COMPLETE \u2014 all stages done.\"); break\n    remaining = int(ph[\"planned_tokens\"]) - int(ph.get(\"tokens_done\", 0))\n    hours_left = (t_end - time.time()) / 3600\n    print(f\"[school] stage {ph['name']} \u00b7 {remaining/1e9:.3f}B to go \u00b7 \"\n          f\"session budget {hours_left:.2f}h\")\n    run.train(max_tokens=remaining, max_hours=hours_left)\n    ph_after = run.manifest.current_phase()\n    if ph_after is None or ph_after[\"name\"] != ph[\"name\"]:\n        boundary_instruments(run, tag=ph[\"name\"])\n    else:\n        print(\"[school] session budget reached mid-stage \u2014 resume next \"\n              \"session, she keeps her place.\")\n        break",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## Cell 4 \u2014 growth report (run anytime)\nThe childhood so far: every boundary's probe accuracies against the baseline, holdout families marked, head-election trace."
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": "import json\nfrom huggingface_hub import HfApi, hf_hub_download\nREPO = \"AbstractPhil/alephllm-mini-beatrix-training\"\napi = HfApi()\nfiles = [f for f in api.list_repo_files(REPO)\n         if f.startswith(\"mini-beatrix-1/reports/curriculum/\")]\nreports = []\nfor f in sorted(files):\n    d = json.load(open(hf_hub_download(REPO, f), encoding=\"utf-8\"))\n    if \"tag\" in d:                       # boundary-report schema\n        reports.append((d[\"tag\"], d[\"step\"], d[\"probes\"], d.get(\"w_s_norm\")))\n    else:                                # baseline/gate schema: {step: {...}}\n        kind = \"t0\" if \"baseline\" in f else \"gate\"\n        for step, r in d.items():\n            reports.append((f\"{kind}@{step}\", int(step), r[\"probes\"], None))\nreports.sort(key=lambda r: r[1])\nif not reports:\n    print(\"no reports yet \u2014 run Cell 2 (baseline) first\")\nelse:\n    suites = sorted(reports[0][2])\n    w = max(len(t) for t, *_ in reports) + 2\n    print(\"\".ljust(w) + \"\".join(s.replace(\"_\", \" \")[:14].ljust(15)\n                                for s in suites))\n    for tag, step, pr, ws in reports:\n        row = \"\".join(f\"{pr[s]['acc']:.3f}\".ljust(15) for s in suites)\n        tail = f\"  ||W_s||={ws:.3f}\" if ws is not None else \"\"\n        print(tag.ljust(w) + row + tail)",
+   "outputs": [],
+   "execution_count": null
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": "## Cell 5 \u2014 arm collectives (boundary causal test)\nArrives with geolip 0.6.1 + the amoe harness pass. The holdout families (bAbI 16/19, depth-5 chains, 3-digit subtraction) are **already excluded from every training row**, and the probe batteries already gauge them \u2014 so any boundary she passes before this cell lands can be tested retroactively against stored checkpoints (retro-sweep law). Nothing is lost by training first."
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}