Add Mycel-LM 79M: architecture, tokenizer, training scripts, weights-only checkpoints

README.md

@@ -0,0 +1,317 @@
+---
+license: apache-2.0
+library_name: pytorch
+pipeline_tag: text-generation
+tags:
+  - custom-architecture
+  - kuramoto
+  - oscillator
+  - mycelium
+  - neighbour-sensing
+  - experimental
+  - research
+  - spikewhale-tokenizer
+---
+
+# Mycel-LM (79M)
+
+**Mycel-LM** is a **79.2M-parameter** research language model whose channel-mixing block is
+**not** an MLP. It is a differentiable **Neighbour-Sensing fungal-colony-growth** model:
+each token is expanded into a colony of *hyphal tips* that grow in a bounded latent
+region, sense a shared density field, and steer their own growth — the "MLP" is replaced
+by a few differentiable steps of colony growth, read back out into the hidden state.
+
+It is part of a family of models that ask a single question: **can the generalizing
+ability of a transformer be carried by an unusual, self-organizing dynamical system in
+place of the feed-forward block?** Mycel-LM keeps the family's tokenizer, traits, and
+data fixed and swaps *only* the mixer, so it is a controlled experiment against the
+sibling **Quazimoto** models (whose mixer is a bank of coupled Kuramoto oscillators).
+
+> ⚠️ **Research artifact, not a product.** At ~79M parameters it is fluent but small: it
+> models the *shape* of language well and generates coherent, grammatical text, but it is
+> **not factual** and will confidently hallucinate. See [Limitations](#limitations).
+
+---
+
+## Table of contents
+
+- [Highlights](#highlights)
+- [Architecture](#architecture)
+- [Repository layout](#repository-layout)
+- [Install](#install)
+- [Quickstart](#quickstart)
+- [Command-line usage](#command-line-usage)
+- [Live visualizer & Space](#live-visualizer--space)
+- [Training from scratch](#training-from-scratch)
+- [Fine-tuning (SFT)](#fine-tuning-sft)
+- [Checkpoints](#checkpoints)
+- [Limitations](#limitations)
+- [Citation / basis](#citation--basis)
+
+---
+
+## Highlights
+
+- **Novel mixer.** The per-layer feed-forward block is replaced by a **MycelBlock** — a
+  differentiable simulation of fungal colony growth (Neighbour-Sensing).
+- **Self-describing checkpoints.** Each `.pt` embeds a `family_config` recording the exact
+  geometry, so `generate.py` / `healthcheck.py` / `visualize.py` rebuild the model with no
+  external config.
+- **KV cache.** Incremental decoding is wired through the whole stack (attention presents
+  are threaded per layer); `generate()` prefills the prompt once and decodes one token per
+  forward.
+- **Self-speculative decoding.** Four MTP draft heads propose the next tokens and the main
+  head verifies them in one parallel forward — bit-identical to greedy, just fewer forwards.
+- **Live 3-D visualizer.** Watch the colony grow token-by-token as a Three.js filament web.
+
+---
+
+## Architecture
+
+Standard causal Transformer backbone (token-mixing = attention, tied LM head), with the
+per-layer feed-forward network replaced by a **MycelBlock**.
+
+### The MycelBlock (the novel part)
+
+Based on the Meškauskas / Fricker / Moore (2004) **Neighbour-Sensing** model of fungal
+colony growth:
+
+1. The hidden state projects to **N = 96 hyphal tips** per token, each with a **position**
+   in a bounded 3-D latent region and a **growth vector**.
+2. A few differentiable **growth steps** run: each tip senses the local **density field**,
+   steers *away* from the colony's own density (negative autotropism) with persistence,
+   moves, and is **re-clamped** into the bounded region (the colony can't grow unbounded).
+3. The final `[position, growth-vector, sensed-density]` of every tip is read out back into
+   the hidden state, behind a family gate.
+
+The density field is evaluated against **16 learnable field centres** (a low-rank sample of
+the field) so cost is **O(N·F)** per step, not O(N²) — the same mean-field trick that keeps
+the sibling oscillator block cheap. Health-checking a trained checkpoint shows the tropism
+parameter converges **strongly negative** across layers, i.e. the model genuinely learns the
+grow-away-from-density behaviour rather than leaving it at init.
+
+**Trait stations** (`MycelStations`): tiny memory specialists sit at fixed anchor positions
+in the colony. A tip interacts with a station by **proximity** — which is *emergent* from
+where the tip grew — so which tips use which trait "comes to be" during growth rather than
+being assigned to a fixed index. The stations hold test-time-writable input/output stores
+that act as an addressable context memory at inference.
+
+### Attention
+
+Family attention ported from the Quazimoto v2 stack:
+- **MLA** low-rank Q/O projections
+- **Partial RoPE** (nope + rope split), **QK-Norm**, **GQA** (4 KV heads)
+- optional DERF (erf attention) and XSA (value-subspace removal) — off in this checkpoint
+- **KV cache** for incremental decoding (per-layer `(k, v)` presents threaded through the stack)
+
+### Opt-in family traits (all live in this checkpoint)
+
+| Trait | Role |
+|---|---|
+| **HRM** | iterative gated hidden-state refinement (random init state, gates open) |
+| **MoE** | SwiGLU mixture (4 routed + 1 shared, top-2) refining the trunk |
+| **MTP** (×4) | multi-token-prediction draft heads → enables self-speculative decoding |
+| **JEPA** | representation-prediction aux loss (train-only; never runs at inference) |
+| **Ring Specialists** (7/ring) | the trait stations described above |
+| **Fractal Phase Seed** | seeds tip positions from each token's Mandelbrot orbit angles (gated) |
+
+### Config (this checkpoint)
+
+| | |
+|---|---|
+| params | **79.2M** |
+| layers | 10 |
+| d_model | 768 |
+| heads | 12 (4 KV) |
+| vocab | 16512 (SpikeWhale byte-merge) |
+| block size | 2048 |
+| tips / token | 96, in a 3-D bounded colony |
+| field centres | 16 · growth steps 3 · stations 16 |
+
+The checkpoint is **self-describing**: `family_config` inside the `.pt` records the exact
+geometry so the model rebuilds itself on load.
+
+---
+
+## Repository layout
+
+```
+model.py               QuazimotoLM + QuazimotoConfig — the transformer backbone, attention,
+                       KV cache, traits (HRM/MoE/MTP/JEPA), generate() and forward_drafts()
+mycel.py               MycelBlock (Neighbour-Sensing growth mixer) + MycelStations
+family.py              shared family layers (MoE, HRM, specialists, norms, ...)
+fractal.py             hierarchical Mandelbrot phase seeding (FractalSeed trait)
+instrument.py          zero-cost capture hooks the visualizer reads from
+special_tokens.py      ChatML / control-token definitions
+spike_tokenizer.py     SpikeWhale byte-merge tokenizer (subclasses PreTrainedTokenizer)
+tokenizer.json         the tokenizer vocab / merges (vocab 16,512)
+fractal_phase.pt       precomputed hierarchical Mandelbrot phase table (regenerable)
+
+generate.py            inference harness — KV cache + self-speculative decoding + sampling
+healthcheck.py         per-layer weight / gate / PPL diagnostics for a checkpoint
+visualize.py           builds the 3-D colony dashboard (viz.html) from a generation
+
+train.py               pretraining entry point (streamed multi-corpus blend)
+train_sft.py           supervised fine-tuning (ChatML, assistant-only loss masking)
+chat_sft.py            chat-format rendering / loss masking helpers used by SFT
+train_opd.py           OPD (on-policy distillation) training loop
+distill_uld.py         universal-logit-distillation utilities
+opd_teacher.py         teacher wrapper for distillation
+build_fractal_table.py regenerates fractal_phase.pt
+train.bat / train_sft.bat   Windows convenience launchers
+
+chkpt/quazimoto.pt       pretraining checkpoint (step 149,000)
+chkpt/quazimoto_sft.pt   SFT checkpoint (step 4,000, ChatML)
+```
+
+> **Note:** the Modal cloud launchers (`modal_train.py`, `modal_sft.py`) are intentionally
+> **not** part of this package. The scripts above run locally on CPU or a single GPU.
+
+---
+
+## Install
+
+```bash
+pip install -r requirements.txt
+```
+
+Requirements are minimal: `torch`, `numpy`, `transformers` (the tokenizer subclasses
+`PreTrainedTokenizer`). Training additionally uses `datasets` and `huggingface_hub`.
+Everything below runs on **CPU** (slow but functional) or a single GPU.
+
+---
+
+## Quickstart
+
+```python
+import torch
+from model import QuazimotoLM, QuazimotoConfig
+from spike_tokenizer import SpikeTokenizer
+
+ck  = torch.load("chkpt/quazimoto.pt", map_location="cpu", weights_only=False)
+cfg = QuazimotoConfig(**ck["family_config"])          # self-describing
+model = QuazimotoLM(cfg); model.load_state_dict(ck["model"], strict=False); model.eval()
+tok = SpikeTokenizer(vocab_file="tokenizer.json")
+
+ids = torch.tensor([tok.encode("The mycelium spreads through the soil", add_special_tokens=False)])
+out = model.generate(ids, n_new=80, temperature=0.8, top_k=40)   # KV cache on by default
+print(tok.decode(out[0].tolist(), skip_special_tokens=True))
+```
+
+For a **chat** turn, wrap the prompt in ChatML and stop on `<|im_end|>` (the SFT checkpoint
+was trained on this framing):
+
+```python
+prompt = "<|im_start|><|user|>\nWhat is mycelium?<|im_end|>\n<|im_start|><|assistant|>\n"
+ids = torch.tensor([tok.encode(prompt, add_special_tokens=False)])
+out = model.generate(ids, n_new=120, temperature=0.7, top_k=40)
+```
+
+---
+
+## Command-line usage
+
+```bash
+# plain completion (KV cache on by default)
+python generate.py --ckpt chkpt/quazimoto.pt --prompt "In the beginning" --max_new_tokens 80
+
+# chat turn (ChatML framing + stop on <|im_end|>)
+python generate.py --ckpt chkpt/quazimoto_sft.pt --chat --prompt "Hello, who are you?"
+
+# interactive REPL
+python generate.py --ckpt chkpt/quazimoto_sft.pt --interactive
+
+# self-speculative decoding (MTP heads draft, main head verifies; report acceptance)
+python generate.py --ckpt chkpt/quazimoto.pt --speculative --spec_stats
+
+# disable the KV cache (full recompute each step — for comparison)
+python generate.py --ckpt chkpt/quazimoto.pt --no_cache
+
+# per-layer diagnostics (weights / gates / PPL)
+python healthcheck.py --ckpt chkpt/quazimoto.pt
+```
+
+Sampling knobs: `--temperature`, `--top_k`, `--top_p`, `--repetition_penalty`, `--seed`.
+
+---
+
+## Live visualizer & Space
+
+`visualize.py` renders the colony growing in 3-D as the model generates, token by token —
+hyphal tips linked into a filament web, coloured by local density, with the trait stations
+shown as orange wire-spheres. It writes a self-contained `viz.html` (Three.js from a CDN):
+
+```bash
+python visualize.py --ckpt chkpt/quazimoto_sft.pt --prompt "the mycelium spreads" --tokens 50
+```
+
+A companion **Hugging Face Space (`Mycel-LM v1`)** wraps the same architecture in an
+interactive chat — KV-cache decoding drives the reply while the 3-D colony visualizer
+animates the growth for the generated tokens.
+
+---
+
+## Training from scratch
+
+```bash
+python train.py --device cuda --steps 160000 --batch 12 --block 2048 --amp \
+    --use-hrm --use-moe --use-mtp --use-jepa --use-ring-specialists --use-fractal-phase-seed \
+    --stream --math-frac 0.25 --out chkpt/quazimoto.pt --ckpt-every 500 --resume
+```
+
+- **Tokenizer:** SpikeWhale byte-merge, vocab 16,512. (Byte-merge perplexity is
+  tokenizer-inflated; **bits/byte** is the honest metric.)
+- **Pretraining blend:** 35% Ultra-FineWeb-L3 / 25% FineWeb-Edu / 25% FineMath /
+  15% Quazim0t0/PretrainNew, streamed. Streamed datasets are pulled with `datasets`;
+  gated corpora need `huggingface-cli login`.
+- `--resume` continues from the checkpoint at `--out`. The growth loop is activation-heavy,
+  so keep the batch modest; `--amp` gives a bf16 speedup on GPU.
+
+Pass `--help` to `train.py` for the full trait / optimiser / schedule surface.
+
+## Fine-tuning (SFT)
+
+```bash
+python train_sft.py --init chkpt/quazimoto.pt --out chkpt/quazimoto_sft.pt \
+    --steps 4000 --batch 8 --block 2048 --amp
+```
+
+- Renders a chat mix in **ChatML** with **assistant-only loss masking** (`chat_sft.py`).
+- **SFT blend:** ultrachat_200k_sft + ultrafeedback-sft + UltraData-SFT-2605/Knowledge +
+  OpenThoughts2-1M-ShortThink.
+- The bundled SFT checkpoint is only **4k steps** — the chat format transferred but the
+  model is still shallow.
+
+> The distributed checkpoints carry **weights only** (optimizer state stripped to keep the
+> download small). Fine-tuning starts a fresh optimizer from them, which is the normal path;
+> only exact *resumption* of the original pretraining run would need the optimizer state.
+
+---
+
+## Checkpoints
+
+- `chkpt/quazimoto.pt` — **pretraining** checkpoint, step **149,000**
+- `chkpt/quazimoto_sft.pt` — **SFT** checkpoint, step **4,000** (ChatML, early)
+
+Both embed `family_config` (self-describing) and load with `strict=False` so future trait
+additions stay backward-compatible.
+
+---
+
+## Limitations
+
+- **Not factual.** Small-model behaviour: fluent and grammatical, but it invents facts
+  ("the capital of France is the largest and most important part of the world").
+- **SFT is early** (4k steps) — answers follow the chat format but hallucinate.
+- **No safety tuning.** No RLHF/guardrails; do not deploy in user-facing settings.
+- **Custom architecture** — cannot be loaded with `AutoModel`; use the bundled `model.py`.
+- This is an **experiment in architecture**, released to study whether a self-organizing
+  growth process can carry a transformer's generalization. Treat outputs accordingly.
+
+## Citation / basis
+
+Neighbour-Sensing model of hyphal growth: Meškauskas, Fricker & Moore (2004),
+*Simulating colonial growth of fungi with the Neighbour-Sensing model of hyphal growth*,
+Mycological Research 108(11).
+
+*License: Apache-2.0.*

build_fractal_table.py

@@ -0,0 +1,106 @@
+"""
+build_fractal_table.py -- precompute the HIERARCHICAL (tokenizer-aware) Mandelbrot
+phase-seed table for Quazimoto-LM. CPU-only, run once; the model loads the result.
+
+Why offline: the hierarchical map needs the tokenizer (each token's byte string),
+which the model deliberately doesn't carry. So we build the [vocab, n_osc] table
+here and save it; model.py loads it when --use-fractal-phase-seed is set.
+
+Hierarchy via an IFS address (this is exactly how fractal self-similarity encodes
+prefixes):
+  * each of the 256 byte VALUES gets a fixed anchor point in the Mandelbrot region
+    (Halton spread, so anchors are evenly distributed and deterministic);
+  * a token's c starts at its FIRST byte's anchor and is contracted toward each
+    subsequent byte's anchor (p <- (1-w)*p + w*anchor[b]).
+Tokens sharing a prefix therefore land in nested, self-similar neighborhoods
+(prefix dominates; later bytes refine at shrinking scale). Specials (ids 0-3 and
+the appended <|...|> set) get their own deterministic anchors away from the bytes.
+Then phase_k = angle(z_k) of z <- z^2 + c, identical to the flat builder.
+
+Usage:
+  python build_fractal_table.py                 # default n_osc from QuazimotoConfig
+  python build_fractal_table.py --n-osc 508 --out fractal_phase.pt
+"""
+
+import argparse
+import os
+import sys
+
+import torch
+
+from fractal import _halton, phases_from_c
+from model import QuazimotoConfig
+
+PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+def load_tokenizer(tok_dir):
+    sys.path.insert(0, tok_dir)
+    from spike_tokenizer import SpikeTokenizer
+    return SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json"))
+
+
+def byte_anchors(region):
+    """256 fixed anchor points (one per byte value) spread by Halton over region."""
+    x0, x1, y0, y1 = region
+    ax = [x0 + _halton(b + 1, 2) * (x1 - x0) for b in range(256)]
+    ay = [y0 + _halton(b + 1, 3) * (y1 - y0) for b in range(256)]
+    return ax, ay
+
+
+def build(tok, n_osc, region=(-0.8, -0.2, 0.4, 0.9), contract=0.15):
+    # region: an off-origin band near the Mandelbrot boundary so the early orbit
+    #   angles vary SMOOTHLY with c (angle doesn't wrap wildly as it does across 0)
+    #   -> prefix locality survives into the low oscillators. contract: small, so the
+    #   FIRST byte dominates the address and shared-prefix tokens cluster tightly.
+    V = tok.vocab_size
+    ax, ay = byte_anchors(region)
+    specials = set(getattr(tok, "_extra_specials", []))
+    ids2tok = tok._ids_to_tokens
+    cr = torch.empty(V); ci = torch.empty(V)
+    n_byte = n_merge = n_special = 0
+    for tid in range(V):
+        s = ids2tok.get(tid, "")
+        is_special = tid < 4 or s in specials
+        if is_special or not s:
+            # own corner, away from the byte anchors (different Halton stream)
+            cr[tid] = region[0] + _halton(tid + 1, 5) * (region[1] - region[0])
+            ci[tid] = region[2] + _halton(tid + 1, 7) * (region[3] - region[2])
+            n_special += 1
+            continue
+        byts = [ord(ch) & 0xFF for ch in s]            # latin-1 byte values
+        px, py = ax[byts[0]], ay[byts[0]]              # first byte dominates
+        for b in byts[1:]:                             # contract toward each later byte
+            px = (1 - contract) * px + contract * ax[b]
+            py = (1 - contract) * py + contract * ay[b]
+        cr[tid], ci[tid] = px, py
+        if len(byts) == 1:
+            n_byte += 1
+        else:
+            n_merge += 1
+    print(f"mapped {V} tokens -> c  ({n_byte} single-byte, {n_merge} merges, {n_special} special)")
+    phases = phases_from_c(torch.complex(cr, ci), n_osc)
+    return phases
+
+
+def main():
+    p = argparse.ArgumentParser(description="Build hierarchical Mandelbrot phase table (CPU)")
+    p.add_argument("--tok_dir", default=PKG_DIR)
+    p.add_argument("--n-osc", type=int, default=QuazimotoConfig().n_osc,
+                   help="number of oscillators (sum of ring_sizes); default from config")
+    p.add_argument("--contract", type=float, default=0.15, help="IFS contraction toward later bytes (small = first byte dominates)")
+    p.add_argument("--out", default=os.path.join(PKG_DIR, "fractal_phase.pt"))
+    args = p.parse_args()
+
+    torch.set_default_device("cpu")
+    tok = load_tokenizer(args.tok_dir)
+    print(f"tokenizer vocab {tok.vocab_size} | building hierarchical table, n_osc={args.n_osc}")
+    phases = build(tok, args.n_osc, contract=args.contract).cpu()
+    torch.save({"phases": phases, "vocab_size": tok.vocab_size, "n_osc": args.n_osc,
+                "mode": "hierarchical", "contract": args.contract}, args.out)
+    print(f"saved {args.out}  shape {tuple(phases.shape)}  "
+          f"({os.path.getsize(args.out)//1024} KB) | range [{phases.min():.3f}, {phases.max():.3f}]")
+
+
+if __name__ == "__main__":
+    main()

chat_sft.py

@@ -0,0 +1,129 @@
+"""
+chat_sft.py -- interactive multi-turn chat REPL for an SFT'd Quazimoto-LM.
+
+Renders the running conversation in the student's ChatML and generates the
+assistant reply. Two decode paths:
+
+  * SPECULATIVE (default) -- DeepSpec/DSpark-style self-speculative decoding: the
+    MTP heads draft the next few tokens, the main head verifies them in one pass,
+    accepting the longest correct prefix (see model.forward_drafts / generate.py
+    generate_speculative). Greedy, so output is deterministic.
+  * SAMPLED (--temperature > 0) -- the KV-cache sampler (generate.py generate):
+    the conversation is prefilled once, then each new token is a single-position
+    cached forward. Supports temperature / top-k / top-p.
+
+Both reuse the model's KV cache; speculative additionally drafts multiple tokens
+per verify. Usage:
+    python chat_sft.py                       # auto-find newest chkpt, speculative
+    python chat_sft.py --temperature 0.7     # sampled instead
+    python chat_sft.py --ckpt chkpt/quazimoto_sft.pt --system "You are concise."
+"""
+import argparse, os, sys
+import torch
+
+from model import QuazimotoLM, QuazimotoConfig
+from generate import generate, generate_speculative, resolve_stop_ids
+
+PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+ROLE_TOK = {"system": "<|system|>", "user": "<|user|>", "assistant": "<|assistant|>"}
+
+
+def find_ckpt(path):
+    if path and os.path.isfile(path):
+        return path
+    folder = (os.path.dirname(path) if path else os.path.join(PKG_DIR, "chkpt")) or "."
+    if not os.path.isdir(folder):
+        return None
+    # prefer an SFT checkpoint, else newest *.pt
+    pts = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".pt")]
+    sft = [p for p in pts if "sft" in os.path.basename(p).lower()]
+    pool = sft or pts
+    return max(pool, key=os.path.getmtime) if pool else None
+
+
+def load_tokenizer(tok_dir):
+    sys.path.insert(0, tok_dir)
+    from spike_tokenizer import SpikeTokenizer
+    return SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json"))
+
+
+def render(history, tok):
+    """history: list of {role,content} -> token ids ending with the assistant header
+    so the model continues the assistant turn."""
+    parts = [f"<|im_start|>{ROLE_TOK.get(m['role'],'<|user|>')}\n{m['content']}<|im_end|>\n"
+             for m in history]
+    parts.append("<|im_start|><|assistant|>\n")
+    return tok.encode("".join(parts), add_special_tokens=False)
+
+
+def main():
+    p = argparse.ArgumentParser(description="Quazimoto-LM SFT chat (KV cache + speculative)")
+    p.add_argument("--ckpt", default="", help="SFT checkpoint (default: newest *sft*.pt in ./chkpt)")
+    p.add_argument("--tok_dir", default=PKG_DIR)
+    p.add_argument("--system", default="", help="optional system prompt")
+    p.add_argument("--max_new_tokens", type=int, default=256)
+    p.add_argument("--temperature", type=float, default=0.0, help="0 = speculative/greedy; >0 = sampled")
+    p.add_argument("--top_k", type=int, default=40)
+    p.add_argument("--top_p", type=float, default=0.95)
+    p.add_argument("--repetition_penalty", type=float, default=1.1)
+    p.add_argument("--no_speculative", action="store_true", help="force the cached sampler even at temp 0")
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    args = p.parse_args()
+
+    path = find_ckpt(args.ckpt)
+    if path is None:
+        print("No checkpoint found; pass --ckpt or train one."); return
+    ckpt = torch.load(path, map_location=args.device, weights_only=False)
+    cfg = QuazimotoConfig(**ckpt["family_config"])
+    model = QuazimotoLM(cfg); model.load_state_dict(ckpt["model"], strict=False)
+    model.to(args.device).eval()
+    tok = load_tokenizer(args.tok_dir)
+    stop = resolve_stop_ids(tok, chat=True)          # <|im_end|>, <eos>, <|endoftext|>
+    imend = tok.get_vocab().get("<|im_end|>")
+
+    spec = (args.temperature <= 1e-4 and not args.no_speculative and model.mtp_heads is not None)
+    mode = "speculative (DSpark draft+verify)" if spec else f"sampled (temp {args.temperature}, KV cache)"
+    print(f"loaded {os.path.basename(path)} (step {ckpt.get('step')}) | decode: {mode}")
+    print("type your message; /reset clears history, /quit exits.\n")
+
+    history = []
+    if args.system:
+        history.append({"role": "system", "content": args.system})
+
+    while True:
+        try:
+            user = input("you> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print(); break
+        if not user:
+            continue
+        if user in ("/quit", "/exit"):
+            break
+        if user == "/reset":
+            history = ([{"role": "system", "content": args.system}] if args.system else [])
+            print("(history cleared)\n"); continue
+
+        history.append({"role": "user", "content": user})
+        ids = render(history, tok)
+        x = torch.tensor([ids], device=args.device)
+        with torch.no_grad():
+            if spec:
+                out = generate_speculative(model, cfg, x, args.max_new_tokens, stop_ids=stop)
+            else:
+                out = generate(model, cfg, x, args.max_new_tokens, temperature=max(args.temperature, 1e-6),
+                               top_k=args.top_k or None, top_p=args.top_p,
+                               repetition_penalty=args.repetition_penalty, stop_ids=stop,
+                               device=args.device, use_cache=True)
+        gen = out[0, len(ids):].tolist()
+        if imend in gen:                              # trim at the assistant turn end
+            gen = gen[:gen.index(imend)]
+        reply = tok.decode(gen, skip_special_tokens=True).strip()
+        try:
+            print(f"bot> {reply}\n")
+        except UnicodeEncodeError:
+            print("bot> " + reply.encode("ascii", "replace").decode() + "\n")
+        history.append({"role": "assistant", "content": reply})
+
+
+if __name__ == "__main__":
+    main()

chkpt/quazimoto.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24688695d1d6540b31958fb143e2b1a34ab726a2bda4f53b79aa74f0c54b3d71
+size 317207627

chkpt/quazimoto_sft.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f5e88339e4d88140884d360e58952a9c6f867448d4a5608fbf98cd34cda9815
+size 317209099

distill_uld.py

@@ -0,0 +1,194 @@
+"""
+distill_uld.py -- OFF-policy cross-tokenizer distillation of MobileLLM-R1-140M-base
+(teacher) into Quazimoto-LM (student), using the ULD recipe from distill_mark2.py:
+
+  loss = CE(student, real tokens)  +  ULD_W * ULD(student_logits, teacher_logits)
+
+ULD (Universal Logit Distillation) is what makes this tokenizer-agnostic: at each
+position it softmaxes both models, SORTS each distribution descending, keeps the
+top-K, and takes an L1 between the sorted vectors -- comparing the *shape* of the
+distributions (how peaked vs. flat) while discarding token identity. So no vocab
+alignment is needed, and unlike the on-policy OPD script there is NO generation
+(teacher-forced), so it runs at normal-training speed.
+
+Differences from the HF student in distill_mark2.py, handled here:
+  * QuazimotoLM does NOT shift targets internally -> we build next-token targets.
+  * its CE uses ignore_index=-1 (not -100).
+  * right-padding is safe (causal attention; real tokens never attend the pad tail).
+  * we pass targets=None and compute CE manually, so the AR aux heads (MTP/JEPA)
+    don't run during distillation.
+
+Run on a rig with transformers + the teacher weights:
+  python distill_uld.py --student-ckpt chkpt/quazimoto.pt --device cuda --stream
+"""
+
+import argparse, math, os, time
+import torch
+import torch.nn.functional as F
+
+from model import QuazimotoLM, QuazimotoConfig
+import train as T
+
+PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+def uld_loss(s_logits, t_logits, uld_n, top_k):
+    """Sorted-top-K L1 between student and teacher next-token distributions,
+    averaged over valid positions. s_logits/t_logits: [B, S, V*] (each its own V)."""
+    total_L, total_n = 0.0, 0
+    for j, n in enumerate(uld_n):
+        if n <= 0:
+            continue
+        sp = F.softmax(s_logits[j, :n].float(), -1).sort(-1, descending=True).values[:, :top_k]
+        tp = F.softmax(t_logits[j, :n].float(), -1).sort(-1, descending=True).values[:, :top_k]
+        k = min(sp.size(-1), tp.size(-1))               # vocab could be < top_k
+        total_L = total_L + (sp[:, :k] - tp[:, :k]).abs().sum()
+        total_n += n
+    return total_L / max(total_n, 1)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--tok-dir", default=PKG_DIR)
+    p.add_argument("--student-ckpt", default=os.path.join(PKG_DIR, "chkpt", "quazimoto.pt"))
+    p.add_argument("--teacher", default="facebook/MobileLLM-R1-140M-base")
+    p.add_argument("--steps", type=int, default=6000)
+    p.add_argument("--seq", type=int, default=1024)
+    p.add_argument("--micro", type=int, default=2, help="micro-batch (sequences per fwd)")
+    p.add_argument("--accum", type=int, default=8, help="grad-accum micro-steps")
+    p.add_argument("--lr", type=float, default=1.5e-4)
+    p.add_argument("--warmup", type=int, default=100)
+    p.add_argument("--min-lr-frac", type=float, default=0.1)
+    p.add_argument("--uld-w", type=float, default=0.5, help="ULD loss weight")
+    p.add_argument("--top-k", type=int, default=128, help="ULD sorted top-K")
+    p.add_argument("--weight-decay", type=float, default=0.01)
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--teacher-device", default="")
+    p.add_argument("--amp", action="store_true", help="bf16 autocast on the student forward")
+    p.add_argument("--out", default=os.path.join(PKG_DIR, "chkpt", "quazimoto_uld.pt"))
+    p.add_argument("--ckpt-every", type=int, default=250)
+    p.add_argument("--resume", action="store_true")
+    # streaming blend (prose source) -- same flags/defaults as train.py
+    p.add_argument("--stream", action="store_true")
+    p.add_argument("--data", default="")
+    for name, d in [("ultra-frac", 0.35), ("edu-frac", 0.25),
+                    ("math-frac", 0.25), ("pretrain-frac", 0.15)]:
+        p.add_argument(f"--{name}", type=float, default=d)
+    p.add_argument("--ultra-dataset", default="openbmb/Ultra-FineWeb-L3")
+    p.add_argument("--ultra-config", default="Ultra-FineWeb-L3-en-Multi-Style-Synthetic")
+    p.add_argument("--fineweb-dataset", default="HuggingFaceFW/fineweb-edu")
+    p.add_argument("--fineweb-config", default="sample-10BT")
+    p.add_argument("--math-dataset", default="HuggingFaceTB/finemath")
+    p.add_argument("--math-config", default="finemath-4plus")
+    p.add_argument("--pretrain-dataset", default="Quazim0t0/PretrainNew")
+    p.add_argument("--pretrain-config", default=None)
+    p.add_argument("--split", default="train")
+    p.add_argument("--seed", type=int, default=0)
+    args = p.parse_args()
+
+    torch.manual_seed(args.seed)
+    dev, tdev = args.device, (args.teacher_device or args.device)
+    s_tok = T.load_tokenizer(args.tok_dir)
+    s_pad = s_tok.pad_token_id or 0
+
+    ck = torch.load(args.student_ckpt, map_location=dev, weights_only=False)
+    cfg = QuazimotoConfig(**ck["family_config"])
+    V = cfg.vocab_size
+    student = QuazimotoLM(cfg); student.load_state_dict(ck["model"], strict=False)
+    student.to(dev).train()
+    print(f"student: {args.student_ckpt} (step {ck.get('step')}), vocab {V}")
+
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    t_tok = AutoTokenizer.from_pretrained(args.teacher, trust_remote_code=True)
+    try:
+        teacher = AutoModelForCausalLM.from_pretrained(args.teacher, dtype=torch.bfloat16,
+                                                       trust_remote_code=True)
+    except TypeError:
+        teacher = AutoModelForCausalLM.from_pretrained(args.teacher, torch_dtype=torch.bfloat16,
+                                                       trust_remote_code=True)
+    teacher = teacher.to(tdev).eval()
+    for pp in teacher.parameters():
+        pp.requires_grad_(False)
+    t_pad = t_tok.pad_token_id if t_tok.pad_token_id is not None else t_tok.eos_token_id
+    print(f"teacher: {args.teacher} on {tdev} (frozen, bf16)")
+
+    # ---- data: same text tokenized two ways ----
+    if args.stream:
+        docs = T.doc_generator(args)
+        get_text = lambda: next(docs)
+    else:
+        ids_all, _ = T.load_ids(args.data, s_tok)
+        def get_text():
+            import numpy as np
+            i = np.random.randint(0, max(len(ids_all) - args.seq, 1))
+            return s_tok.decode(ids_all[i:i + args.seq].tolist(), skip_special_tokens=True)
+
+    def get_batch():
+        texts = [get_text() for _ in range(args.micro)]
+        s_rows = [s_tok.encode(t, add_special_tokens=False)[:args.seq] for t in texts]
+        t_rows = [t_tok(t, add_special_tokens=False).input_ids[:args.seq] for t in texts]
+        sL = max(len(r) for r in s_rows); tL = max(len(r) for r in t_rows)
+        s_ids = torch.full((args.micro, sL), s_pad, dtype=torch.long)
+        targets = torch.full((args.micro, sL), -1, dtype=torch.long)      # ignore_index=-1
+        t_ids = torch.full((args.micro, tL), t_pad, dtype=torch.long)
+        uld_n = []
+        for j, (sr, tr) in enumerate(zip(s_rows, t_rows)):
+            s_ids[j, :len(sr)] = torch.tensor(sr)
+            if len(sr) > 1:
+                targets[j, :len(sr) - 1] = torch.tensor(sr[1:])           # next-token targets
+            t_ids[j, :len(tr)] = torch.tensor(tr)
+            uld_n.append(min(len(sr), len(tr)) - 1)
+        return (s_ids.to(dev), targets.to(dev), t_ids.to(tdev),
+                (t_ids != t_pad).to(tdev), uld_n)
+
+    opt = torch.optim.AdamW(student.parameters(), lr=args.lr, betas=(0.9, 0.95),
+                            weight_decay=args.weight_decay)
+
+    start_step = 1
+    if args.resume and os.path.isfile(args.out):
+        r = torch.load(args.out, map_location=dev, weights_only=False)
+        student.load_state_dict(r["model"], strict=False)
+        if "optim" in r:
+            try: opt.load_state_dict(r["optim"])
+            except ValueError: pass
+        start_step = int(r.get("step", 0)) + 1
+        print(f"resumed ULD from {args.out} at step {r.get('step')} -> {start_step}")
+    if start_step > args.steps:
+        print(f"  nothing to do: already at {start_step-1} >= --steps {args.steps}."); return
+
+    t0 = time.time()
+    for step in range(start_step, args.steps + 1):
+        for g in opt.param_groups:
+            g["lr"] = T.lr_at(step, args.lr, args.warmup, args.steps, args.min_lr_frac)
+        opt.zero_grad(set_to_none=True)
+        tot = tul = tce = 0.0
+        for _ in range(args.accum):
+            s_ids, targets, t_ids, t_mask, uld_n = get_batch()
+            with torch.no_grad():
+                t_logits = teacher(t_ids, attention_mask=t_mask).logits.to(dev)
+            with torch.autocast("cuda", torch.bfloat16, enabled=args.amp and dev == "cuda"):
+                logits, _, _ = student(s_ids)                            # targets=None -> no aux
+            ce = F.cross_entropy(logits.reshape(-1, V), targets.reshape(-1), ignore_index=-1)
+            ul = uld_loss(logits, t_logits, uld_n, args.top_k)
+            loss = ce + args.uld_w * ul
+            (loss / args.accum).backward()
+            tot += loss.item() / args.accum; tul += ul.item() / args.accum; tce += ce.item() / args.accum
+            del t_logits, logits
+        gn = torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
+        if torch.isfinite(gn):
+            opt.step()
+        else:
+            print(f"step {step}: non-finite grad, skip")
+
+        if step % 25 == 0 or step == 1:
+            dt = (time.time() - t0) / (step - start_step + 1)
+            print(f"step {step}/{args.steps} | loss {tot:.3f} | ce {tce:.3f} | uld {tul:.3f} "
+                  f"| ppl {math.exp(min(tce,20)):.1f} | {dt:.1f}s/step", flush=True)
+        if args.ckpt_every and step % args.ckpt_every == 0:
+            T.save_ckpt(student, s_tok.vocab_size, step, args.out, opt)
+    T.save_ckpt(student, s_tok.vocab_size, args.steps, args.out, opt)
+    print("done ->", args.out)
+
+
+if __name__ == "__main__":
+    main()

family.py

@@ -0,0 +1,446 @@
+"""SpikeWhale / Byrne family traits, ported to Quazimoto-LM (attention backbone).
+
+These are the *transformer-native* family blocks (their origin is the transformer
+`modeling_byrne_embed.py`), so unlike the SNN port they operate directly on the
+sequence hidden state [B,T,d] -- no per-step adaptation needed. Each keeps the
+family's safe-at-init contract: a tanh/zero gate makes the block a no-op at start,
+while the content (`up`/`down`) weights are NON-zero so the gate still receives
+gradient (the double-zero saddle would deadlock it). DERF soft_clamp bounds any
+new instability surface, in line with the family's stability discipline.
+
+Included: HRMRefinementBlock (signature), MoESwiGLU, MTPHead, JEPAPredictorBlock.
+Engram / ProgSem are bio/SNN-specific and SpikingLinearAttention is the SNN's
+stand-in for the real attention this model already has, so they are omitted.
+"""
+from __future__ import annotations
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import instrument as _viz   # live-visualizer capture hooks (no-op unless a recorder is active)
+
+# sqrt(pi)/2: soft_clamp is the identity for small inputs and saturates smoothly to
+# +/-bound with a non-zero gradient everywhere (no dead-gradient zones).
+_ERF_K = math.sqrt(math.pi) / 2.0
+
+
+def soft_clamp(x, bound):
+    return bound * torch.erf(x * (_ERF_K / bound))
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+
+def sqrtsoftplus(x):
+    """Family expert-scoring function: sqrt(softplus(x))."""
+    return torch.sqrt(F.softplus(x) + 1e-8)
+
+
+class HRMRefinementBlock(nn.Module):
+    """Signature family block: iterative gated refinement that, in the canonical
+    HRM spirit, starts the reasoning from a RANDOM initial state z0 and iterates
+    toward a solution conditioned on the input (`anchor`).
+
+    Unlike the other family traits this block does NOT start as a no-op: the gate
+    is initialised OPEN (gate_init_open). With an input-anchored no-op start the
+    gate received ~zero gradient and never woke up; a random z0 forces the block
+    to actively reconcile the random state against the input, so the open gate
+    carries real signal from step 0. To keep the deep trunk intact we contribute
+    only the reasoning DELTA (h - z0) as a residual -- z0 itself is never dumped
+    into the trunk, and a closed gate (h == z0) degrades cleanly to a no-op."""
+
+    def __init__(self, hidden_size, refine_dim, steps, eps=1e-3, gate_init_open=0.1):
+        super().__init__()
+        self.steps = steps
+        self.norm = RMSNorm(hidden_size, eps)
+        self.down = nn.Linear(hidden_size * 2, refine_dim, bias=False)
+        self.up = nn.Linear(refine_dim, hidden_size, bias=False)
+        # random initial reasoning state (learnable), broadcast over batch/time
+        self.z0 = nn.Parameter(torch.empty(hidden_size))
+        nn.init.trunc_normal_(self.z0, std=1.0, a=-2.0, b=2.0)
+        # gates start OPEN so the random-state reasoning reaches the output at init
+        go = math.atanh(min(gate_init_open, 0.9)) if gate_init_open > 0 else 0.0
+        self.gate = nn.Parameter(torch.full((steps,), go))
+        nn.init.normal_(self.down.weight, std=0.02)
+        nn.init.normal_(self.up.weight, std=0.02)
+
+    def forward(self, x):                       # x: [B,T,d]
+        B, T, _ = x.shape
+        anchor = x
+        h = self.z0.expand(B, T, -1)            # random initial reasoning state
+        for t in range(self.steps):
+            inp = torch.cat([self.norm(h), anchor], dim=-1)
+            update = soft_clamp(self.up(F.silu(self.down(inp))), 10.0)
+            h = h + torch.tanh(self.gate[t]) * update
+        return x + (h - self.z0)                # add reasoning delta, keep trunk
+
+
+class ExpertFFN(nn.Module):
+    def __init__(self, hidden_size, intermediate_size):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class MoESwiGLU(nn.Module):
+    """Shared + top-k routed SwiGLU experts, sqrtsoftplus scoring, norm_topk_prob,
+    Switch-style load-balance aux (via `last_aux_loss`). down-projections zero-init
+    => no-op at start."""
+
+    def __init__(self, hidden_size, intermediate_size, n_routed=4, n_shared=1,
+                 top_k=2, aux_loss_coef=0.01):
+        super().__init__()
+        self.top_k = min(top_k, n_routed)
+        self.n_routed = n_routed
+        self.n_shared = n_shared
+        self.aux_loss_coef = aux_loss_coef
+        self.router = nn.Linear(hidden_size, n_routed, bias=False)
+        self.experts = nn.ModuleList([ExpertFFN(hidden_size, intermediate_size)
+                                      for _ in range(n_routed)])
+        self.shared = (ExpertFFN(hidden_size, intermediate_size * n_shared)
+                       if n_shared > 0 else None)
+        for e in self.experts:
+            nn.init.zeros_(e.down_proj.weight)
+        if self.shared is not None:
+            nn.init.zeros_(self.shared.down_proj.weight)
+        self.last_aux_loss = None
+
+    def forward(self, x):                       # x: [B,T,d]
+        flat = x.reshape(-1, x.shape[-1])
+        out = torch.zeros_like(flat)
+        if self.shared is not None:
+            s = self.shared(flat)
+            out = out + (s / self.n_shared if self.n_shared > 1 else s)
+
+        logits = self.router(flat)
+        scores = sqrtsoftplus(logits)
+        topv, topi = scores.topk(self.top_k, dim=-1)
+        topv = topv / (topv.sum(-1, keepdim=True) + 1e-8)
+        for slot in range(self.top_k):
+            idx = topi[:, slot]
+            w = topv[:, slot].unsqueeze(-1)
+            for e_id, expert in enumerate(self.experts):
+                mask = idx == e_id
+                if mask.any():
+                    out[mask] = out[mask] + w[mask] * expert(flat[mask])
+
+        probs = F.softmax(logits, dim=-1)
+        expert_mask = torch.zeros_like(probs)
+        expert_mask.scatter_(1, topi, 1.0)
+        self.last_aux_loss = (self.n_routed * (expert_mask.mean(0) * probs.mean(0)).sum()
+                              * self.aux_loss_coef)
+        return out.view_as(x)
+
+
+class MTPHead(nn.Module):
+    """Multi-token-prediction head: zero-init d->d residual, reuses the tied readout."""
+
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, hidden_size, bias=False)
+        nn.init.zeros_(self.proj.weight)
+
+    def forward(self, hidden):
+        return hidden + self.proj(hidden)
+
+
+class TokenCompressor(nn.Module):
+    """Frozen LSH-style projection (gradient never reaches it through the hash cast)."""
+    def __init__(self, hidden_size, compress_dim):
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, compress_dim, bias=False)
+        nn.init.normal_(self.proj.weight, std=0.02)
+        self.proj.weight.requires_grad_(False)
+
+    def forward(self, x):
+        return self.proj(x)
+
+
+class MultiHeadHashLookup(nn.Module):
+    """N-gram hash memory: for n=1..max_ngram, hash the n-token compressed window
+    into per-head tables and average. Ported from v2 EngramModule."""
+    def __init__(self, num_heads, table_size, compress_dim, out_dim, max_ngram=3):
+        super().__init__()
+        self.num_heads, self.table_size = num_heads, table_size
+        self.max_ngram, self.out_dim = max_ngram, out_dim
+        self.tables = nn.ModuleList([nn.Embedding(table_size, out_dim) for _ in range(num_heads)])
+        for t in self.tables:
+            nn.init.normal_(t.weight, std=0.01)
+        for n in range(1, max_ngram + 1):
+            for k in range(n):
+                proj = torch.randn(num_heads, compress_dim)
+                proj = proj / (proj.norm(dim=1, keepdim=True) + 1e-8)
+                self.register_buffer(f"hash_proj_n{n}_p{k}", proj, persistent=True)
+
+    def forward(self, compressed):
+        B, S, _ = compressed.shape
+        dev = compressed.device
+        out = torch.zeros(B, S, self.out_dim, device=dev, dtype=compressed.dtype)
+        norm = torch.zeros(S, device=dev)
+        for n in range(1, self.max_ngram + 1):
+            if S < n:
+                continue
+            valid, start = S - n + 1, n - 1
+            h = torch.zeros(B, valid, self.num_heads, device=dev)
+            for k in range(n):
+                proj = getattr(self, f"hash_proj_n{n}_p{k}")
+                h = h + torch.matmul(compressed[:, k:k + valid, :].float(), proj.t())
+            idx = h.abs().long() % self.table_size
+            for hi, table in enumerate(self.tables):
+                out[:, start:, :] = out[:, start:, :] + table(idx[:, :, hi])
+            norm[start:] += self.num_heads
+        return (out / norm.view(1, -1, 1).clamp(min=1)).to(compressed.dtype)
+
+
+class DERFContextGate(nn.Module):
+    def __init__(self, obs_size, init_bias=-4.0):
+        super().__init__()
+        self.proj = nn.Linear(obs_size * 2, obs_size)
+        self.alpha = nn.Parameter(torch.ones(obs_size))
+        self.bias = nn.Parameter(torch.full((obs_size,), init_bias))
+        self.gamma = nn.Parameter(torch.ones(obs_size))
+
+    def forward(self, retrieved, obs):
+        logits = self.proj(torch.cat([retrieved, obs], dim=-1))
+        gate = self.gamma * ((torch.erf(self.alpha * logits + self.bias) + 1.0) / 2.0)
+        return retrieved * gate
+
+
+class PhaseAttentionRing(nn.Module):
+    """Interstitial ATTENTION ring: attends causally over the sequence in
+    oscillator-PHASE space ([cos,sin] of the two neighbor oscillator rings) and
+    returns an injection current of width m = n_r + n_{r+1} for those neighbors.
+    Zero-init gate => no-op at start; soft_clamp bounds the injected drive."""
+
+    def __init__(self, m, n_heads=4, head_dim=16, bound=10.0):
+        super().__init__()
+        self.h, self.d, self.bound = n_heads, head_dim, bound
+        self.qkv = nn.Linear(2 * m, 3 * n_heads * head_dim, bias=False)
+        self.out = nn.Linear(n_heads * head_dim, m, bias=False)
+        self.gate = nn.Parameter(torch.zeros(1))
+        nn.init.normal_(self.qkv.weight, std=0.02)
+        nn.init.normal_(self.out.weight, std=0.02)
+
+    def forward(self, theta_slice):              # [B,T,m] phases of the neighbor rings
+        B, T, m = theta_slice.shape
+        feat = torch.cat([torch.cos(theta_slice), torch.sin(theta_slice)], dim=-1)
+        q, k, v = self.qkv(feat).split(self.h * self.d, dim=-1)
+        shp = lambda z: z.view(B, T, self.h, self.d).transpose(1, 2)
+        y = F.scaled_dot_product_attention(shp(q), shp(k), shp(v), is_causal=True)
+        y = y.transpose(1, 2).reshape(B, T, self.h * self.d)
+        return soft_clamp(self.out(y) * torch.tanh(self.gate), self.bound)
+
+
+class EngramRing(nn.Module):
+    """Interstitial ENGRAM ring: absorbs n-gram context from the hidden state via
+    hash memory + DERF gate, projected to an injection current of width m for the
+    two neighbor oscillator rings. Two no-op gates at init (DERF bias -4 + scale)."""
+
+    def __init__(self, hidden_size, m, compress_dim=32, num_heads=2,
+                 table_size=2048, max_ngram=3, bound=10.0):
+        super().__init__()
+        self.bound = bound
+        self.compressor = TokenCompressor(hidden_size, compress_dim)
+        self.lookup = MultiHeadHashLookup(num_heads, table_size, compress_dim, m, max_ngram)
+        self.to_obs = nn.Linear(hidden_size, m, bias=False)
+        self.gate = DERFContextGate(m, init_bias=-4.0)
+        self.scale = nn.Parameter(torch.zeros(1))   # extra no-op gate at init
+        nn.init.normal_(self.to_obs.weight, std=0.02)
+
+    def family_reinit(self):
+        """Re-apply the inits the model's global self.apply would clobber (table
+        std 0.01, frozen-random compressor, DERF bias -4)."""
+        for t in self.lookup.tables:
+            nn.init.normal_(t.weight, std=0.01)
+        nn.init.normal_(self.compressor.proj.weight, std=0.02)
+        self.compressor.proj.weight.requires_grad_(False)
+        nn.init.constant_(self.gate.bias, -4.0)
+
+    def forward(self, h):                        # h: [B,T,hidden]
+        retrieved = self.lookup(self.compressor(h.detach()))
+        gated = self.gate(retrieved, self.to_obs(h))
+        return soft_clamp(gated * torch.tanh(self.scale), self.bound)
+
+
+class RingController(nn.Module):
+    """Tiny per-ring manager that OPTIMIZES ITSELF by a predictive / free-energy rule.
+
+    Core = a fast-weight linear predictor `W` (a BUFFER, excluded from the global
+    optimizer) updated online by the delta rule  W += lr * (f - W@prev) outer prev,
+    which is exactly one gradient step on the squared prediction error -- the
+    controller learns to predict its ring's next state, minimizing surprise, with no
+    backprop. A small backprop-trained decoder maps the self-organized feature +
+    surprise into ring-control modulations; zero-init => exact no-op at start."""
+
+    def __init__(self, d_obs=4, feat=384, n_ctrl=4, local_lr=0.01):
+        super().__init__()
+        self.feat, self.local_lr = feat, local_lr
+        self.enc = nn.Linear(d_obs, feat)
+        self.dec = nn.Linear(feat + 1, n_ctrl)
+        nn.init.zeros_(self.dec.weight)
+        nn.init.zeros_(self.dec.bias)                 # control == 0 at init (no-op)
+        self.register_buffer("W", torch.zeros(feat, feat))   # self-organizing fast weights
+        self.register_buffer("prev_f", torch.zeros(feat))
+
+    def family_reinit(self):
+        nn.init.zeros_(self.dec.weight)
+        nn.init.zeros_(self.dec.bias)
+
+    def forward(self, obs):                            # obs: [d_obs] (detached ring stats)
+        f = torch.tanh(self.enc(obs))                  # [feat]
+        pred = self.W @ self.prev_f                    # predicted current feature
+        surprise = F.mse_loss(f.detach(), pred)
+        if self.training:
+            with torch.no_grad():                      # predictive self-organization (no global grad)
+                err = f.detach() - pred
+                self.W.add_(self.local_lr * torch.outer(err, self.prev_f)).clamp_(-3.0, 3.0)
+                self.prev_f.copy_(f.detach())
+        ctrl = self.dec(torch.cat([f, surprise.detach().reshape(1)]))   # [n_ctrl]
+        return ctrl, surprise.detach()
+
+
+class RingControllerBank(nn.Module):
+    """One RingController per oscillator ring (shared across all layers)."""
+
+    def __init__(self, n_rings, d_obs=4, feat=384, local_lr=0.01):
+        super().__init__()
+        self.controllers = nn.ModuleList(
+            [RingController(d_obs, feat, 4, local_lr) for _ in range(n_rings)])
+        self.last_surprise = None
+
+    def forward(self, obs):                            # obs: [R, d_obs] -> ctrl [R, 4]
+        ctrls, surps = [], []
+        for r, c in enumerate(self.controllers):
+            ct, sp = c(obs[r])
+            ctrls.append(ct)
+            surps.append(sp)
+        self.last_surprise = torch.stack(surps).mean()
+        return torch.stack(ctrls, dim=0)
+
+
+class RingSpecialists(nn.Module):
+    """A MoE-style bank of `n_spec` MINI MEMORY SPECIALISTS for ONE oscillator ring.
+
+    Each specialist owns two fast-weight stores (test-time-mutable BUFFERS, like
+    RingController.W -- excluded from the optimizer):
+      * store_in  -- a memory of the INPUT context that routes to it, and
+      * store_out -- the OUTPUT information it injects back into the ring.
+    Tokens are routed to the top-k specialists (a small MoE) by similarity to each
+    specialist's address = its learnable identity key + a read of its accumulated
+    input memory. The routed store_out is decoded into an injection current for the
+    ring, and BOTH stores are written online (train AND inference) by a gated EMA
+    rule -- so a generation accumulates an addressable context memory as it runs.
+
+    Slow (backprop) weights -- q_proj/in_enc/val_enc/out_dec/in_read/key/active --
+    learn to route, encode, retrieve and decode; the stores are the fast memory.
+    Family contract: zero-init `scale` => exact no-op at start, and an empty
+    store_out is zero anyway, so the block is doubly safe until it learns to write
+    and open the gate. `active` is a per-specialist usage gate biasing the router."""
+
+    def __init__(self, ring_size, hidden_size, n_spec=7, key_dim=32, slot_dim=64,
+                 top_k=2, write_lr=0.1, bound=10.0):
+        super().__init__()
+        self.n_spec = n_spec
+        self.top_k = min(top_k, n_spec)
+        self.write_lr, self.bound = write_lr, bound
+        self.write_enabled = True
+        self.q_proj  = nn.Linear(hidden_size, key_dim, bias=False)   # router query   (learned)
+        self.out_dec = nn.Linear(slot_dim, ring_size, bias=False)    # retrieved -> injection (learned)
+        self.in_read = nn.Linear(slot_dim, key_dim, bias=False)      # input-store -> addr  (learned)
+        # write-side encoders are FROZEN RANDOM projections (cf. EngramRing's frozen
+        # compressor): they only ever run inside the no-grad write, so backprop can't
+        # train them -- as fixed random features the stores hold a stable encoding the
+        # learned read/route/decode path can address.
+        self.in_enc  = nn.Linear(hidden_size, slot_dim, bias=False)  # input  -> input-store  (frozen)
+        self.val_enc = nn.Linear(hidden_size, slot_dim, bias=False)  # input  -> output-store (frozen)
+        self.key     = nn.Parameter(torch.randn(n_spec, key_dim) * 0.02)  # specialist identity
+        self.active  = nn.Parameter(torch.zeros(n_spec))             # per-specialist usage gate
+        self.scale   = nn.Parameter(torch.zeros(1))                  # no-op output gate at init
+        self.register_buffer("store_in",  torch.zeros(n_spec, slot_dim))
+        self.register_buffer("store_out", torch.zeros(n_spec, slot_dim))
+        for m in (self.q_proj, self.out_dec, self.in_read, self.in_enc, self.val_enc):
+            nn.init.normal_(m.weight, std=0.02)
+        self.in_enc.weight.requires_grad_(False)
+        self.val_enc.weight.requires_grad_(False)
+
+    def family_reinit(self):
+        """Re-apply inits the model's global self.apply clobbers (key, gates, and
+        the frozen write encoders)."""
+        nn.init.normal_(self.key, std=0.02)
+        nn.init.zeros_(self.active)
+        nn.init.zeros_(self.scale)
+        nn.init.normal_(self.in_enc.weight, std=0.02)
+        nn.init.normal_(self.val_enc.weight, std=0.02)
+        self.in_enc.weight.requires_grad_(False)
+        self.val_enc.weight.requires_grad_(False)
+
+    def reset_memory(self):
+        """Clear both stores -- call between independent prompts/sequences so
+        context memory does not bleed across them."""
+        self.store_in.zero_()
+        self.store_out.zero_()
+
+    def forward(self, h):                        # h: [B,T,hidden]
+        B, T, _ = h.shape
+        # snapshot the fast-weight stores: the graph must hold an immutable copy
+        # because we mutate the buffers in-place for the online write below.
+        store_in, store_out = self.store_in.clone(), self.store_out.clone()
+        q = self.q_proj(h)                                           # [B,T,key_dim]
+        addr = self.key + self.in_read(store_in)                     # [n_spec,key_dim]
+        logits = q @ addr.t()                                        # [B,T,n_spec]
+        logits = logits + F.logsigmoid(self.active)                  # usage gate biases routing
+        if self.top_k < self.n_spec:                                 # top-k MoE sparsity
+            tv = torch.topk(logits, self.top_k, dim=-1).values
+            logits = logits.masked_fill(logits < tv[..., [-1]], float("-inf"))
+        route = torch.softmax(logits, dim=-1)                        # [B,T,n_spec]
+
+        retrieved = route @ store_out                                # [B,T,slot_dim]
+        inject = self.out_dec(retrieved) * torch.tanh(self.scale)    # [B,T,ring_size]
+
+        rec = _viz.get_rec()
+        if rec is not None and rec.enabled:                          # last-token routing
+            rec.push_spec(route[0, -1].tolist())
+
+        # online write: blend this step's input/value into the routed specialists
+        if self.write_enabled and self.write_lr > 0:
+            with torch.no_grad():
+                w = route.reshape(-1, self.n_spec)                   # [BT,n_spec]
+                denom = w.sum(0).clamp(min=1e-3).unsqueeze(1)        # [n_spec,1]
+                in_info  = (w.t() @ self.in_enc(h).reshape(-1, self.in_enc.out_features)) / denom
+                val_info = (w.t() @ self.val_enc(h).reshape(-1, self.val_enc.out_features)) / denom
+                a = self.write_lr
+                self.store_in.mul_(1 - a).add_(a * in_info).clamp_(-self.bound, self.bound)
+                self.store_out.mul_(1 - a).add_(a * val_info).clamp_(-self.bound, self.bound)
+        return soft_clamp(inject, self.bound)
+
+
+class JEPAPredictorBlock(nn.Module):
+    """Representation-space k-ahead prediction with stop-grad target (JEPA asymmetry).
+    Zero-init gate => identity at init; `up` normal so the gate gets gradient."""
+
+    def __init__(self, dim, pred_dim, horizon, eps=1e-3):
+        super().__init__()
+        self.horizon = horizon
+        self.norm = RMSNorm(dim, eps)
+        self.down = nn.Linear(dim, pred_dim, bias=False)
+        self.up = nn.Linear(pred_dim, dim, bias=False)
+        self.gate = nn.Parameter(torch.zeros(horizon))
+        nn.init.normal_(self.down.weight, std=0.02)
+        nn.init.normal_(self.up.weight, std=0.02)
+
+    def forward(self, h, k):                    # h: [B,T,dim]
+        update = self.up(F.silu(self.down(self.norm(h))))
+        return h + torch.tanh(self.gate[k - 1]) * update

fractal.py

@@ -0,0 +1,79 @@
+"""
+fractal.py -- Mandelbrot phase seeding for Quazimoto-LM's oscillator bank.
+
+The recommended (and only) fractal integration: instead of a generic learned
+phase initializer, give each TOKEN a characteristic dynamical signature drawn
+from the Mandelbrot iteration z <- z^2 + c, and use the ANGLES of that orbit to
+seed the N oscillator phases. The Mandelbrot map is itself an iterated dynamical
+system, so this hands the Kuramoto ring bank a token-specific, deterministic,
+parameter-free phase prior congruent with what the block already does.
+
+We build a frozen [vocab_size, n_osc] table once: token id -> a complex point c
+(spread over the Mandelbrot region by a 2D Halton low-discrepancy sequence so
+coverage is even and deterministic) -> phase_k = angle(z_k) for the first n_osc
+orbit points. The model adds this, through a zero-init gate, to to_theta(h) inside
+each QuazimotoBlock -- a no-op at init that the optimizer can choose to open.
+
+Smooth-by-construction: we read the orbit ANGLE (always defined, bounded to
+(-pi, pi]) rather than escape-time, avoiding the chaotic boundary discontinuities
+that raw escape counts would inject.
+"""
+
+import torch
+
+
+def _halton(i, base):
+    """Radical-inverse (van der Corput) value of i in the given base, in [0,1)."""
+    f, r = 1.0, 0.0
+    while i > 0:
+        f /= base
+        r += f * (i % base)
+        i //= base
+    return r
+
+
+@torch.no_grad()
+def mandelbrot_phase_table(vocab_size, n_osc, region=(-2.5, 1.0, -1.25, 1.25),
+                           clamp_mag=1e3):
+    """Return a frozen [vocab_size, n_osc] tensor of orbit-angle phase seeds.
+
+    token id -> c via 2D Halton(base 2,3) over `region`; phase_k = angle(z_k) for
+    k = 0..n_osc-1 of the iteration z <- z^2 + c (z0 = 0). This is the FLAT
+    (tokenizer-agnostic) map; for the hierarchical byte-merge map build the table
+    offline with build_fractal_table.py and load it via load_phase_table()."""
+    x0, x1, y0, y1 = region
+    ids = torch.arange(1, vocab_size + 1)
+    hx = torch.tensor([_halton(int(i), 2) for i in ids])
+    hy = torch.tensor([_halton(int(i), 3) for i in ids])
+    cr = x0 + hx * (x1 - x0)
+    ci = y0 + hy * (y1 - y0)
+    return phases_from_c(torch.complex(cr, ci), n_osc, clamp_mag)
+
+
+@torch.no_grad()
+def phases_from_c(c, n_osc, clamp_mag=1e3):
+    """Orbit-angle phases for a batch of complex seeds c [V] -> [V, n_osc].
+    phase_k = angle(z_k) of z <- z^2 + c (z0=0); magnitude clamped (angle kept)."""
+    z = torch.zeros_like(c)
+    phases = torch.empty(c.shape[0], n_osc)
+    for k in range(n_osc):
+        z = z * z + c
+        mag = z.abs()
+        over = mag > clamp_mag                        # rescale escaped orbits, keep direction
+        if over.any():
+            z = torch.where(over, z / mag * clamp_mag, z)
+        phases[:, k] = torch.angle(z)                 # in (-pi, pi]
+    return phases
+
+
+def load_phase_table(vocab_size, n_osc, path):
+    """Load a precomputed phase table if it matches (vocab_size, n_osc); else None.
+    Returns (phases, mode) or (None, None)."""
+    import os
+    if not os.path.exists(path):
+        return None, None
+    d = torch.load(path, map_location="cpu", weights_only=False)
+    ph = d.get("phases")
+    if ph is not None and tuple(ph.shape) == (vocab_size, n_osc):
+        return ph.float(), d.get("mode", "precomputed")
+    return None, None

fractal_phase.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee042ce925f1a47e0a7b129fb78f95ef2f5fe979cfd2bb9e3b410185e8dea952
+size 19023162

generate.py

@@ -0,0 +1,326 @@
+"""
+generate.py -- inference / generation harness for Quazimoto-LM.
+
+Loads a checkpoint saved by train.py (save_ckpt writes
+    {"model", "family_config", "vocab_size", "step"})
+rebuilds the exact model from the embedded family_config, restores weights, and
+samples completions with the bundled SpikeWhale tokenizer.
+
+This is tuned to the Quazimoto architecture's quirks:
+  * The model has NO KV cache (attention is training-only / full recompute), so
+    each new token re-runs the whole context. We respect cfg.block_size by
+    cropping the context window, exactly like model.generate().
+  * We reuse the model's NaN/inf-sanitised sampling path (the oscillator stack
+    can push logits to +/-inf over long rollouts), and add nucleus (top-p),
+    top-k, and repetition-penalty controls on top.
+  * Generation stops early on <eos> / <|im_end|> / <|endoftext|> when present.
+
+Examples
+--------
+  # plain completion
+  python generate.py --ckpt chkpt/quazimoto.pt --prompt "the quazimoto oscillator"
+
+  # chat-style turn (wraps prompt in ChatML and stops on <|im_end|>)
+  python generate.py --ckpt chkpt/quazimoto.pt --chat --prompt "Hello, who are you?"
+
+  # interactive REPL
+  python generate.py --ckpt chkpt/quazimoto.pt --interactive
+"""
+
+import argparse
+import os
+import sys
+
+import torch
+import torch.nn.functional as F
+
+from model import QuazimotoLM, QuazimotoConfig
+
+PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+# --------------------------------------------------------------------------- #
+# loading
+# --------------------------------------------------------------------------- #
+def load_tokenizer(tok_dir):
+    sys.path.insert(0, tok_dir)
+    from spike_tokenizer import SpikeTokenizer
+    return SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json"))
+
+
+def load_model(ckpt_path, device):
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+    fam = ckpt.get("family_config")
+    if fam is None:
+        raise ValueError(f"{ckpt_path} has no family_config; can't rebuild the model.")
+    cfg = QuazimotoConfig(**fam)
+    model = QuazimotoLM(cfg)
+    state = ckpt["model"]
+    missing, unexpected = model.load_state_dict(state, strict=False)
+    if missing:
+        print(f"  [warn] missing keys: {len(missing)} (e.g. {missing[:3]})")
+    if unexpected:
+        print(f"  [warn] unexpected keys: {len(unexpected)} (e.g. {unexpected[:3]})")
+    model.to(device).eval()
+    step = ckpt.get("step", "?")
+    print(f"loaded {ckpt_path} | step {step} | vocab {cfg.vocab_size} | "
+          f"block_size {cfg.block_size}")
+    return model, cfg
+
+
+# --------------------------------------------------------------------------- #
+# sampling -- tuned for the oscillator stack (no KV cache, sanitised logits)
+# --------------------------------------------------------------------------- #
+def _filter_logits(logits, top_k, top_p):
+    """Apply top-k then nucleus (top-p) filtering in place; returns logits."""
+    if top_k:
+        k = min(top_k, logits.size(-1))
+        kth = torch.topk(logits, k).values[..., -1, None]
+        logits = logits.masked_fill(logits < kth, float("-inf"))
+    if top_p and 0.0 < top_p < 1.0:
+        sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
+        cum = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        remove = cum > top_p
+        remove[..., 1:] = remove[..., :-1].clone()  # keep the first over-threshold token
+        remove[..., 0] = False
+        scatter_remove = remove.scatter(-1, sorted_idx, remove)
+        logits = logits.masked_fill(scatter_remove, float("-inf"))
+    return logits
+
+
+@torch.no_grad()
+def generate(model, cfg, idx, n_new, temperature=0.9, top_k=40, top_p=0.95,
+             repetition_penalty=1.1, stop_ids=None, device="cpu", use_cache=True):
+    """Autoregressive sampling. Mirrors model.generate()'s NaN/inf sanitisation
+    (long oscillator rollouts can emit +/-inf logits whose softmax -> NaN ->
+    multinomial device-side assert) and adds top-p + repetition penalty.
+
+    Uses the attention KV cache by default: the prompt is prefilled once, then
+    each new token is a single-position forward. Once the cache reaches
+    block_size we re-prime with a windowed recompute (absolute RoPE positions
+    must stay within max_position_embeddings)."""
+    model.eval()
+    stop_ids = set(stop_ids or [])
+    past = None
+    for _ in range(n_new):
+        if use_cache and past is not None and past[0][0].size(2) < cfg.block_size:
+            logits, _, _, past = model(idx[:, -1:], past_key_values=past, use_cache=True)
+        elif use_cache:
+            logits, _, _, past = model(idx[:, -cfg.block_size:], use_cache=True)
+        else:
+            logits, _, _ = model(idx[:, -cfg.block_size:])
+        lg = torch.nan_to_num(logits[:, -1, :].float(),
+                              nan=0.0, posinf=1e4, neginf=-1e4)
+
+        # repetition penalty over tokens already in the context
+        if repetition_penalty and repetition_penalty != 1.0:
+            for b in range(lg.size(0)):
+                seen = torch.unique(idx[b])
+                vals = lg[b, seen]
+                lg[b, seen] = torch.where(vals > 0, vals / repetition_penalty,
+                                          vals * repetition_penalty)
+
+        lg = lg / max(temperature, 1e-6)
+        lg = _filter_logits(lg, top_k, top_p)
+        probs = torch.softmax(lg, dim=-1)
+
+        if not torch.isfinite(probs).all() or float(probs.sum()) <= 0.0:
+            nxt = torch.argmax(torch.nan_to_num(lg, neginf=-1e4), dim=-1, keepdim=True)
+        else:
+            nxt = torch.multinomial(probs, 1)
+
+        idx = torch.cat([idx, nxt], dim=1)
+        if stop_ids and int(nxt[0]) in stop_ids and idx.size(0) == 1:
+            break
+    return idx
+
+
+# --------------------------------------------------------------------------- #
+# self-speculative decoding (DeepSpec draft+verify, drafted by the MTP heads)
+# --------------------------------------------------------------------------- #
+@torch.no_grad()
+def generate_speculative(model, cfg, idx, n_new, stop_ids=None, verbose=False, trace=None):
+    """Greedy self-speculative decoding (batch size 1).
+
+    DeepSpec-style draft+verify, but the model drafts on ITSELF: the `mtp_layers`
+    MTP heads propose the next `mtp_layers` future tokens from one hidden state,
+    and the main head verifies them in a single parallel forward, accepting the
+    longest correct prefix. Output is BIT-IDENTICAL to greedy generate() (top_k=1)
+    -- speculation only changes the number of forwards, never the tokens.
+
+    Per cycle (one forward over committed seq + the pending drafts):
+      base = index of the last committed token; main_logits[base+j] is the
+      verifier's genuine token AFTER draft j. We accept draft j while it matches,
+      commit one correction token at the first mismatch, then read the next round's
+      drafts from the MTP heads at the deepest still-valid position (base+a)."""
+    assert idx.size(0) == 1, "speculative decoder is batch-1"
+    model.eval()
+    if model.mtp_heads is None:
+        raise ValueError("speculative decoding needs MTP heads (use_mtp=True).")
+    stop_ids = set(stop_ids or [])
+    K = len(model.mtp_heads)
+    start_len = idx.size(1)
+    forwards = 0
+
+    def crop(seq):                                   # honor block_size context window
+        return seq[:, -cfg.block_size:]
+
+    # bootstrap: one forward over the prompt -> genuine next token + K drafts
+    main_logits, mtp = model.forward_drafts(crop(idx)); forwards += 1
+    g = main_logits[:, -1].argmax(-1, keepdim=True)
+    idx = torch.cat([idx, g], dim=1)
+    drafts = [mtp[k][:, -1].argmax(-1, keepdim=True) for k in range(K)]  # guesses after g
+
+    while idx.size(1) - start_len < n_new:
+        if stop_ids and int(idx[0, -1]) in stop_ids:
+            break
+        seq = torch.cat([idx] + drafts, dim=1)        # committed + K speculative tokens
+        base = idx.size(1) - 1                         # position of the last committed token
+        cseq = crop(seq)
+        off = cseq.size(1) - seq.size(1)              # crop shift (<=0) to realign `base`
+        main_logits, mtp = model.forward_drafts(cseq); forwards += 1
+        b = base + off
+
+        # verify: accept draft[j] while it matches the verifier's token after it
+        a = 0
+        for j in range(K):
+            v = main_logits[:, b + j].argmax(-1, keepdim=True)   # genuine token after draft j-1 / g
+            if int(v) == int(drafts[j]):
+                a += 1
+            else:
+                break
+        accepted = drafts[:a]
+        correction = main_logits[:, b + a].argmax(-1, keepdim=True)  # genuine token at first miss
+        new_toks = accepted + [correction]
+        idx = torch.cat([idx] + new_toks, dim=1)
+        if trace is not None:                          # per-cycle record for the visualizer
+            trace.append({"drafts": [int(d) for d in drafts], "accepted": a,
+                          "correction": int(correction)})
+
+        # next drafts from the deepest still-valid hidden (position b+a holds the last
+        # accepted token / g, whose MTP heads predict the tokens AFTER `correction`).
+        drafts = [mtp[k][:, b + a].argmax(-1, keepdim=True) for k in range(K)]
+
+        if stop_ids and any(int(t) in stop_ids for t in new_toks):
+            break
+
+    idx = idx[:, :start_len + n_new]
+    if verbose:
+        produced = idx.size(1) - start_len
+        print(f"  [spec] {produced} tokens in {forwards} forwards "
+              f"({produced / max(forwards,1):.2f} tok/forward; "
+              f"vs 1.00 for plain greedy)")
+    return idx
+
+
+# --------------------------------------------------------------------------- #
+# prompt helpers
+# --------------------------------------------------------------------------- #
+def build_prompt(tok, text, chat, system):
+    """Return token ids for the prompt. In chat mode, wrap in ChatML so the
+    model is cued to produce an assistant turn (only useful if it was trained
+    on that framing -- harmless otherwise)."""
+    if chat:
+        parts = []
+        if system:
+            parts.append(f"<|im_start|><|system|>\n{system}<|im_end|>\n")
+        parts.append(f"<|im_start|><|user|>\n{text}<|im_end|>\n")
+        parts.append("<|im_start|><|assistant|>\n")
+        text = "".join(parts)
+    ids = tok.encode(text, add_special_tokens=False)
+    return ids, text
+
+
+def resolve_stop_ids(tok, chat):
+    """Collect ids of end-of-turn / end-of-text markers that exist in the vocab."""
+    names = ["<eos>", "<|endoftext|>"]
+    if chat:
+        names.insert(0, "<|im_end|>")
+    vocab = tok.get_vocab()
+    return [vocab[n] for n in names if n in vocab]
+
+
+# --------------------------------------------------------------------------- #
+# main
+# --------------------------------------------------------------------------- #
+def run_once(model, cfg, tok, args, device, prompt_text):
+    stop_ids = resolve_stop_ids(tok, args.chat)
+    ids, shown = build_prompt(tok, prompt_text, args.chat, args.system)
+    x = torch.tensor([ids], dtype=torch.long, device=device)
+    if args.speculative:
+        # greedy self-speculative decoding (MTP heads draft, main head verifies);
+        # output equals greedy decoding, only faster. Ignores sampling knobs.
+        out = generate_speculative(model, cfg, x, args.max_new_tokens,
+                                   stop_ids=stop_ids, verbose=args.spec_stats)
+    else:
+        out = generate(model, cfg, x, args.max_new_tokens,
+                       temperature=args.temperature, top_k=args.top_k,
+                       top_p=args.top_p, repetition_penalty=args.repetition_penalty,
+                       stop_ids=stop_ids, device=device, use_cache=not args.no_cache)
+    gen_ids = out[0, len(ids):].tolist()
+    completion = tok.decode(gen_ids, skip_special_tokens=not args.show_special)
+    if args.echo_prompt:
+        print(shown, end="")
+    print(completion)
+    return completion
+
+
+def main():
+    p = argparse.ArgumentParser(description="Quazimoto-LM generation harness")
+    p.add_argument("--ckpt", default=os.path.join(PKG_DIR, "chkpt", "quazimoto.pt"),
+                   help="checkpoint .pt saved by train.py")
+    p.add_argument("--tok_dir", default=PKG_DIR, help="dir holding tokenizer.json")
+    p.add_argument("--prompt", default="the quazimoto oscillator")
+    p.add_argument("--max_new_tokens", type=int, default=200)
+    p.add_argument("--temperature", type=float, default=0.9)
+    p.add_argument("--top_k", type=int, default=40, help="0 to disable")
+    p.add_argument("--top_p", type=float, default=0.95, help="1.0 to disable")
+    p.add_argument("--repetition_penalty", type=float, default=1.1, help="1.0 to disable")
+    p.add_argument("--seed", type=int, default=None)
+    p.add_argument("--chat", action="store_true", help="wrap prompt in ChatML framing")
+    p.add_argument("--system", default="", help="system prompt (chat mode only)")
+    p.add_argument("--echo_prompt", action="store_true", help="print the prompt before output")
+    p.add_argument("--show_special", action="store_true", help="don't strip special tokens")
+    p.add_argument("--interactive", action="store_true", help="REPL: read prompts from stdin")
+    p.add_argument("--no_cache", action="store_true",
+                   help="disable the KV cache (full recompute each step)")
+    p.add_argument("--speculative", action="store_true",
+                   help="greedy self-speculative decoding via MTP draft heads (DeepSpec-style)")
+    p.add_argument("--spec_stats", action="store_true",
+                   help="print tokens/forward acceptance stats for speculative decoding")
+    p.add_argument("--device", default="cpu" if torch.cuda.is_available() else "cpu")
+    args = p.parse_args()
+
+    if args.top_k == 0:
+        args.top_k = None
+    if args.seed is not None:
+        torch.manual_seed(args.seed)
+
+    device = args.device
+    tok = load_tokenizer(args.tok_dir)
+    model, cfg = load_model(args.ckpt, device)
+
+    if cfg.vocab_size != tok.vocab_size:
+        print(f"  [warn] model vocab {cfg.vocab_size} != tokenizer vocab "
+              f"{tok.vocab_size}; decode may be misaligned.")
+
+    if args.interactive:
+        print("interactive mode -- type a prompt, Ctrl-C / empty line + EOF to quit.\n")
+        try:
+            while True:
+                try:
+                    line = input(">>> ").strip()
+                except EOFError:
+                    break
+                if not line:
+                    continue
+                run_once(model, cfg, tok, args, device, line)
+                print()
+        except KeyboardInterrupt:
+            print("\nbye.")
+    else:
+        run_once(model, cfg, tok, args, device, args.prompt)
+
+
+if __name__ == "__main__":
+    main()

healthcheck.py

@@ -0,0 +1,258 @@
+"""
+healthcheck.py -- checkpoint health report for Quazimoto-LM.
+
+Auto-finds the checkpoint (same rule as train.py --resume: the --out/--ckpt file
+if present, else the newest *.pt in the folder), then reports:
+  1. load + global tensor health (NaN/Inf, dtypes, param count, step)
+  2. per-(grouped)-layer weight stats with dead/exploded flags
+  3. family-trait layer status -- each trait's gates (tanh) so you can see which
+     blocks have actually woken up, plus their key weight stats
+  4. current perplexity on a sample of text (the local --data file, or a built-in
+     fallback string), with the random-vocab baseline for reference.
+
+Usage:
+  python healthcheck.py                       # auto-find in ./chkpt
+  python healthcheck.py --ckpt path/to.pt
+  python healthcheck.py --data corpus.txt     # measure PPL on your own text
+"""
+
+import argparse
+import math
+import os
+import re
+import sys
+from collections import Counter
+
+import torch
+
+from model import QuazimotoLM, QuazimotoConfig
+from family import (HRMRefinementBlock, JEPAPredictorBlock, MoESwiGLU, MTPHead,
+                    EngramRing, PhaseAttentionRing, RingControllerBank, RingSpecialists)
+
+PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+# --------------------------------------------------------------------------- #
+# discovery / loading
+# --------------------------------------------------------------------------- #
+def find_ckpt(path):
+    """Prefer an exact file; else newest *.pt in its folder (or ./chkpt)."""
+    if path and os.path.isfile(path):
+        return path
+    folder = os.path.dirname(path) if path else os.path.join(PKG_DIR, "chkpt")
+    folder = folder or "."
+    if not os.path.isdir(folder):
+        return None
+    pts = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".pt")]
+    return max(pts, key=os.path.getmtime) if pts else None
+
+
+def load_tokenizer(tok_dir):
+    sys.path.insert(0, tok_dir)
+    from spike_tokenizer import SpikeTokenizer
+    return SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json"))
+
+
+def section(title):
+    print("\n" + "=" * 72 + f"\n{title}\n" + "=" * 72)
+
+
+def fmt_gate(t):
+    return "[" + ", ".join(f"{x:+.3f}" for x in t) + "]"
+
+
+# --------------------------------------------------------------------------- #
+# 1. global tensor health
+# --------------------------------------------------------------------------- #
+def report_load(ckpt, path):
+    section("1. LOAD + GLOBAL TENSOR HEALTH")
+    print(f"checkpoint : {path}")
+    print(f"keys       : {list(ckpt.keys())}")
+    print(f"step       : {ckpt.get('step')} | saved vocab_size: {ckpt.get('vocab_size')}")
+    print(f"optimizer  : {'present (resumable)' if 'optim' in ckpt else 'absent'}")
+    sd = ckpt["model"]
+    nan = inf = total = 0
+    dtypes = Counter()
+    for v in sd.values():
+        if not torch.is_tensor(v):
+            continue
+        total += v.numel(); dtypes[str(v.dtype)] += 1
+        nan += torch.isnan(v).sum().item(); inf += torch.isinf(v).sum().item()
+    print(f"tensors    : {len(sd)} | params: {total/1e6:.1f}M | dtypes: {dict(dtypes)}")
+    flag = "OK" if (nan == 0 and inf == 0) else "*** PROBLEM ***"
+    print(f"NaNs       : {nan} | Infs: {inf}   -> {flag}")
+    return sd
+
+
+# --------------------------------------------------------------------------- #
+# 2. per-layer weight stats
+# --------------------------------------------------------------------------- #
+def report_layers(sd):
+    section("2. PER-LAYER WEIGHT HEALTH (grouped over repeated layers)")
+    buckets = {}
+    for k, v in sd.items():
+        if not torch.is_tensor(v) or v.dtype != torch.float32:
+            continue
+        base = re.sub(r"\.\d+\.", ".N.", k)
+        buckets.setdefault(base, []).append(v)
+    print(f"{'param group':46s} {'std':>8} {'absmax':>9} {'mean':>9}  flag")
+    for base, vs in buckets.items():
+        if "ring_bank" in base or "specialists" in base:
+            continue                                 # shown in the traits section
+        s = torch.cat([x.flatten() for x in vs]).float()
+        if s.numel() < 2:
+            continue                                 # scalars -> traits section
+        flag = ""
+        if s.std() < 1e-4 and s.abs().max() < 1e-4: flag = "DEAD?"   # all ~0 (not just constant)
+        if s.abs().max() > 50: flag = "EXPLODED?"
+        print(f"{base:46s} {s.std():8.4f} {s.abs().max():9.3f} {s.mean():+9.4f}  {flag}")
+
+
+# --------------------------------------------------------------------------- #
+# 3. family-trait layer status
+# --------------------------------------------------------------------------- #
+def _wstat(t):
+    t = t.float()
+    return f"std {t.std():.4f} absmax {t.abs().max():.3f}" if t.numel() > 1 else f"val {t.item():+.4f}"
+
+
+def report_traits(model, cfg):
+    section("3. FAMILY-TRAIT LAYER STATUS  (tanh-gate ~0 => still a no-op)")
+    any_trait = False
+
+    # --- per-layer Quazimoto oscillator block gate + interstitial/specialist rings ---
+    for li, layer in enumerate(model.layers):
+        q = layer.quaz
+        tag = f"layer{li}.quaz"
+        print(f"\n[{tag}] oscillator block gate tanh = {torch.tanh(q.gate).item():+.4f} "
+              f"| block_gain {_wstat(q.block_gain)} | alpha {_wstat(q.alpha)}")
+        if hasattr(q, "phase_seed_gate"):
+            print(f"   fractal phase-seed gate : {torch.tanh(q.phase_seed_gate).item():+.4f}  "
+                  f"(0 => seed unused)")
+        if getattr(q, "use_rings", False):
+            ar = [torch.tanh(r.gate).item() for r in q.attn_rings]
+            er = [torch.tanh(r.scale).item() for r in q.engram_rings]
+            print(f"   phase-attn ring gates  : {fmt_gate(ar)}")
+            print(f"   engram ring scales     : {fmt_gate(er)}")
+        if getattr(q, "use_ring_specialists", False):
+            sc = [torch.tanh(b.scale).item() for b in q.specialists]
+            act = [b.active.mean().item() for b in q.specialists]
+            filled = [float(b.store_out.abs().sum() > 0) for b in q.specialists]
+            print(f"   specialist out-gates    : {fmt_gate(sc)}")
+            print(f"   specialist mean 'active': {fmt_gate(act)}")
+            print(f"   specialist store filled : {fmt_gate(filled)}  (0 until written at inference)")
+        if li == 0 and (getattr(q, 'use_rings', False) or getattr(q, 'use_ring_specialists', False)):
+            any_trait = True
+
+    # --- trunk-level traits ---
+    if model.hrm is not None:
+        any_trait = True
+        g = torch.tanh(model.hrm.gate)
+        print(f"\n[HRM]  gate tanh = {fmt_gate(g.tolist())}  "
+              f"({'OPEN' if g.abs().mean() > 0.01 else 'no-op'}) | "
+              f"z0 {_wstat(model.hrm.z0)} | down {_wstat(model.hrm.down.weight)}")
+    if model.moe is not None:
+        any_trait = True
+        print(f"\n[MoE]  router {_wstat(model.moe.router.weight)} | "
+              f"experts={model.moe.n_routed} shared={model.moe.n_shared} top_k={model.moe.top_k}")
+    if model.mtp_heads is not None:
+        any_trait = True
+        norms = [h.proj.weight.norm().item() for h in model.mtp_heads]
+        print(f"\n[MTP]  {len(model.mtp_heads)} draft heads | proj-weight norms "
+              f"{fmt_gate(norms)}  (0 => head still identity / untrained)")
+    if model.jepa is not None:
+        any_trait = True
+        g = torch.tanh(model.jepa.gate)
+        print(f"\n[JEPA] gate tanh = {fmt_gate(g.tolist())} (train-only aux; never runs in generate)")
+    if model.ring_bank is not None:
+        any_trait = True
+        decs = [c.dec.weight.norm().item() for c in model.ring_bank.controllers]
+        print(f"\n[Controllers] {len(model.ring_bank.controllers)} ring controllers | "
+              f"decoder-weight norms {fmt_gate(decs)}  (0 => no-op self-organizer)")
+    if getattr(cfg, "use_fractal_phase_seed", False):
+        any_trait = True
+        gates = [torch.tanh(layer.quaz.phase_seed_gate).item() for layer in model.layers]
+        mag = sum(abs(g) for g in gates) / max(len(gates), 1)
+        tbl = os.path.join(PKG_DIR, "fractal_phase.pt")
+        if os.path.exists(tbl):
+            d = torch.load(tbl, map_location="cpu", weights_only=False)
+            ok = tuple(d.get("phases").shape) == (cfg.vocab_size, cfg.n_osc)
+            tinfo = f"{d.get('mode','?')} table {'MATCHES' if ok else '** SHAPE MISMATCH **'} " \
+                    f"{tuple(d.get('phases').shape)}"
+        else:
+            tinfo = "no fractal_phase.pt found -> model uses FLAT Halton fallback"
+        print(f"\n[FractalSeed] per-layer phase-seed gate tanh {fmt_gate(gates)}")
+        print(f"   status: {'OPEN' if mag > 0.01 else 'no-op (gates ~0)'} | {tinfo}")
+
+    if not any_trait:
+        print("\n(no opt-in family traits enabled in this checkpoint)")
+
+
+# --------------------------------------------------------------------------- #
+# 4. perplexity
+# --------------------------------------------------------------------------- #
+@torch.no_grad()
+def report_ppl(model, cfg, tok, data_path, device):
+    section("4. PERPLEXITY")
+    if data_path and os.path.exists(data_path):
+        with open(data_path, "r", encoding="utf-8", errors="replace") as f:
+            text = f.read()[:200000]
+        src = data_path
+    else:
+        text = ("the quazimoto oscillator learns to synchronize phases across rings. "
+                "coupled clocks fall into step when coupling is strong enough. ") * 50
+        src = "built-in fallback string"
+    ids = tok.encode(text, add_special_tokens=False)
+    ids = ids[: cfg.block_size * 8 + 1]              # cap cost; a few windows is enough
+    if len(ids) < 2:
+        print("not enough tokens to score."); return
+    x = torch.tensor([ids], device=device)
+    losses = []
+    B = cfg.block_size
+    for s in range(0, x.size(1) - 1, B):             # non-overlapping windows
+        xb = x[:, s:s + B]; yb = x[:, s + 1:s + 1 + B]
+        n = min(xb.size(1), yb.size(1))
+        if n < 1: break
+        _, loss, _ = model(xb[:, :n], yb[:, :n])
+        losses.append(loss.item())
+    ce = sum(losses) / len(losses)
+    print(f"text source : {src}  ({len(ids)} tokens, {len(losses)} windows)")
+    print(f"CE loss     : {ce:.4f}")
+    print(f"bits/token  : {ce/math.log(2):.4f}")
+    print(f"perplexity  : {math.exp(ce):.1f}   (random-vocab baseline: {cfg.vocab_size})")
+
+
+# --------------------------------------------------------------------------- #
+def main():
+    p = argparse.ArgumentParser(description="Quazimoto-LM checkpoint health check")
+    p.add_argument("--ckpt", default="", help="checkpoint path (default: auto-find newest in ./chkpt)")
+    p.add_argument("--data", default="", help="text file to measure PPL on (default: built-in string)")
+    p.add_argument("--tok_dir", default=PKG_DIR)
+    p.add_argument("--device", default="cpu")
+    args = p.parse_args()
+
+    path = find_ckpt(args.ckpt)
+    if path is None:
+        print("No checkpoint found. Pass --ckpt or train one first.")
+        return
+
+    ckpt = torch.load(path, map_location=args.device, weights_only=False)
+    sd = report_load(ckpt, path)
+    report_layers(sd)
+
+    cfg = QuazimotoConfig(**ckpt["family_config"])
+    model = QuazimotoLM(cfg)
+    miss, unexp = model.load_state_dict(sd, strict=False)
+    model.to(args.device).eval()
+    if miss or unexp:
+        print(f"\n[load] missing {len(miss)} / unexpected {len(unexp)} keys "
+              f"(arch drift vs this code; e.g. {(miss or unexp)[:2]})")
+
+    tok = load_tokenizer(args.tok_dir)
+    report_traits(model, cfg)
+    report_ppl(model, cfg, tok, args.data, args.device)
+    print()
+
+
+if __name__ == "__main__":
+    main()

instrument.py

@@ -0,0 +1,84 @@
+"""
+instrument.py -- zero-overhead-when-off capture hooks for the live visualizer.
+
+A single global Recorder is consulted by the model's hot paths. When no recorder
+is active (the default), every call is a single `is None` check, so training and
+normal generation are unaffected. visualize.py activates a recorder, runs a
+forward per generated token, and reads back per-token state.
+"""
+
+_REC = None
+
+
+def get_rec():
+    return _REC
+
+
+def set_rec(r):
+    global _REC
+    _REC = r
+
+
+def _r(x, nd=4):
+    return round(float(x), nd)
+
+
+class Recorder:
+    """Collects one frame per generated token. Modules append in execution order,
+    so list position == layer index (attention then quaz, per layer)."""
+
+    def __init__(self, phase_layer=None, attn_layer=None):
+        self.frames = []
+        self.cur = None
+        self.enabled = False
+        self.phase_layer = phase_layer      # which layer's individual phases to keep
+        self.attn_layer = attn_layer        # which layer's attention map to keep
+        self._spec_tmp = []
+
+    # ---- frame lifecycle (driven by visualize.py) ----
+    def begin(self):
+        self.cur = {"rings": [], "phases": None, "attn": None,
+                    "spec": [], "quaz_norm": [], "traits": {}}
+        self._spec_tmp = []
+        self.enabled = True
+
+    def end(self, **meta):
+        self.enabled = False
+        if self.cur is not None:
+            self.cur.update(meta)
+            self.frames.append(self.cur)
+            self.cur = None
+
+    # ---- module callbacks (guarded by `enabled` at the call site) ----
+    def log_ring(self, R, psi, phases):
+        li = len(self.cur["rings"])
+        self.cur["rings"].append({"R": [_r(x, 4) for x in R],
+                                  "psi": [_r(x, 4) for x in psi]})
+        if phases is not None and (self.phase_layer is None or li == self.phase_layer):
+            self.cur["phases"] = {"layer": li, "theta": [_r(x, 3) for x in phases]}
+
+    def log_traj(self, traj):
+        """Attach the tip growth trajectory (list over growth steps, each a flattened
+        [N*pos_dim] snapshot) to THIS layer's colony frame (phase_layer only). Lets the
+        3D view draw each tip's hypha extending across the growth steps."""
+        ph = self.cur.get("phases")
+        if ph is not None and ph.get("layer") == len(self.cur["rings"]) - 1:
+            ph["traj"] = [[_r(x, 3) for x in step] for step in traj]
+
+    def log_attn(self, layer_idx, w):
+        if self.attn_layer is None or layer_idx == self.attn_layer:
+            self.cur["attn"] = {"layer": layer_idx, "w": [_r(x, 4) for x in w]}
+
+    def log_quaz_norm(self, n):
+        self.cur["quaz_norm"].append(_r(n, 4))
+
+    def push_spec(self, route):                 # one ring's [n_spec] route weights
+        self._spec_tmp.append([_r(x, 4) for x in route])
+
+    def flush_spec(self):                        # called once per quaz block (a layer)
+        if self._spec_tmp:
+            self.cur["spec"].append(self._spec_tmp)
+            self._spec_tmp = []
+
+    def log_trait(self, name, val):
+        self.cur["traits"][name] = _r(val, 4)

model.py

@@ -0,0 +1,773 @@
+"""
+Quazimoto-LM — a language model whose channel-mixing block is a bank of
+coupled phase oscillators (Kuramoto dynamics) arranged in concentric rings.
+
+Backbone: standard causal transformer (token mixing = attention, tied LM head).
+Novelty:  the per-layer MLP is replaced by a QuazimotoBlock that maps the hidden
+          state to N oscillator phases on rings [4, 8, 16, 32, ...], runs a few
+          differentiable Euler steps of structured Kuramoto coupling with a
+          learnable frustration alpha, then reads out [cos, sin] of the phases.
+
+Set ring coupling to dense + alpha=0 and it degenerates toward AKOrN-style
+oscillatory neurons; the ring structure + hierarchy-to-center is the Quazimoto bit.
+"""
+
+from dataclasses import dataclass, field
+import math
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+_PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+
+# family traits (DERF soft_clamp, RMSNorm, and the opt-in refinement/aux blocks)
+from family import (soft_clamp, RMSNorm, HRMRefinementBlock, MoESwiGLU,
+                    MTPHead, JEPAPredictorBlock, PhaseAttentionRing, EngramRing,
+                    RingControllerBank, RingSpecialists)
+import instrument as _viz   # live-visualizer capture hooks (no-op unless a recorder is active)
+from mycel import MycelBlock, MycelStations   # Neighbour-Sensing growth mixer
+
+
+@dataclass
+class QuazimotoConfig:
+    vocab_size: int = 32000
+    n_layer: int = 10
+    n_head: int = 12
+    d_model: int = 768
+    block_size: int = 1024
+    dropout: float = 0.0
+    # Quazimoto oscillator bank
+    ring_sizes: tuple = (4, 8, 16, 32, 64, 128, 256)  # (unused by the mycelium mixer)
+    osc_steps: int = 4          # differentiable Kuramoto Euler steps
+    osc_dt: float = 0.5
+    readout_mult: int = 3       # MLP expansion on the readout
+    # Neighbour-Sensing (mycelium) mixer: N growing tips in a bounded latent region.
+    n_tips: int = 96            # hyphal tips per token (analog of n_osc)
+    mycel_pos_dim: int = 3      # latent dimensionality tips grow in
+    mycel_field_centers: int = 16  # low-rank density-field sample sites (O(N*F) not O(N^2))
+    mycel_n_stations: int = 16  # trait stations spread through the colony
+    growth_steps: int = 3       # differentiable growth iterations (each retains ~5 [B,T,N,F])
+    growth_dt: float = 0.3
+    # Mandelbrot phase seeding: seed the N oscillator phases from each token's
+    # fractal orbit angles (frozen, parameter-free), added through a zero-init gate.
+    use_fractal_phase_seed: bool = False
+    # family-trait discipline
+    osc_bound: float = 10.0     # DERF soft_clamp bound on block output (BPTT-stability)
+    gate_init_open: float = 0.9 # tanh-gate opening at init (Quazimoto is the main mixer,
+                                #   so it starts OPEN; set 0.0 for a pure no-op refinement)
+    # interstitial rings: between each pair of oscillator rings, insert a phase-attention
+    # ring + an engram ring that ABSORB info and INJECT it into the two neighbor rings.
+    use_rings: bool = False
+    ring_attn_heads: int = 4
+    ring_attn_head_dim: int = 16
+    ring_engram_compress: int = 32
+    ring_engram_heads: int = 2
+    ring_engram_table: int = 2048
+    ring_engram_ngram: int = 3
+    # per-ring memory specialists: a small MoE of `ring_n_specialists` mini experts
+    # PER oscillator ring, each holding test-time-writable input/output stores that
+    # act as an addressable context memory at inference (see family.RingSpecialists).
+    use_ring_specialists: bool = False
+    ring_n_specialists: int = 7
+    ring_spec_key_dim: int = 32
+    ring_spec_slot_dim: int = 64
+    ring_spec_top_k: int = 2
+    ring_spec_write_lr: float = 0.1
+    # per-ring self-organizing controllers (~1M total: one tiny net per oscillator
+    # ring, shared across layers, weights self-optimize by a predictive/free-energy rule)
+    use_ring_controllers: bool = False
+    ring_ctrl_feat: int = 384        # fast-weight predictor width (dominates the ~150k/ring)
+    ring_ctrl_local_lr: float = 0.01 # delta-rule (surprise-minimization) step size
+    # opt-in family trait modules (all safe no-ops at init)
+    use_hrm: bool = False
+    hrm_steps: int = 3
+    hrm_dim: int = 256
+    hrm_gate_init: float = 0.1   # HRM gates start OPEN (random initial state needs a path out)
+    use_moe: bool = False
+    moe_intermediate: int = 768
+    moe_n_routed: int = 4
+    moe_n_shared: int = 1
+    moe_top_k: int = 2
+    use_mtp: bool = False
+    mtp_layers: int = 4   # depth of MTP draft heads (predict +1..+mtp_layers); used for spec-decode
+    mtp_loss_weight: float = 0.3
+    use_jepa: bool = False
+    jepa_pred_dim: int = 256
+    jepa_horizon: int = 1
+    jepa_loss_weight: float = 0.1
+    # ---- attention: ported from model_v2.MLADerfXSAAttention ----
+    head_dim: int = 64
+    qk_rope_head_dim: int = 32       # partial RoPE: rope part of each head
+    nope_head_dim: int = 32          #               no-pos part (sum == head_dim)
+    num_key_value_heads: int = 4     # GQA (== n_head for full MHA)
+    q_lora_rank: int = 384           # MLA low-rank q projection
+    o_lora_rank: int = 384           # MLA low-rank output projection
+    use_qk_norm: bool = True         # per-head RMSNorm on Q/K before RoPE
+    use_derf: bool = False           # erf attention instead of softmax (ablate)
+    use_xsa: bool = False            # value-subspace removal (ablate)
+    rope_theta: float = 10000.0
+    max_position_embeddings: int = 4096
+    rms_norm_eps: float = 1e-6
+    initializer_range: float = 0.02
+    # ---- model-level v2 features ----
+    zloss_coef: float = 1e-4         # z-loss on lm logits
+    use_value_embed: bool = False    # per-block value-embedding residual (zero-init gate)
+    use_hyper_connections: bool = False
+    hc_mult: int = 2
+
+    @property
+    def n_osc(self):
+        return self.n_tips * self.mycel_pos_dim   # fractal seeds tip POSITIONS here
+
+
+def _ring_index(ring_sizes):
+    """Return a LongTensor of length N mapping each oscillator -> its ring id."""
+    idx = []
+    for r, n in enumerate(ring_sizes):
+        idx += [r] * n
+    return torch.tensor(idx, dtype=torch.long)
+
+
+class QuazimotoBlock(nn.Module):
+    """Oscillatory channel-mixing block: hidden state -> ring phases -> Kuramoto -> readout."""
+
+    def __init__(self, cfg: QuazimotoConfig):
+        super().__init__()
+        self.cfg = cfg
+        N = cfg.n_osc
+        R = len(cfg.ring_sizes)
+        self.N, self.R = N, R
+
+        self.norm = RMSNorm(cfg.d_model)
+        # project hidden -> initial phases and natural frequencies
+        self.to_theta = nn.Linear(cfg.d_model, N)
+        self.to_omega = nn.Linear(cfg.d_model, N)
+
+        # structured, LEARNABLE coupling: an RxR gain between ring blocks.
+        # Block-constant coupling lets us use per-ring mean fields (order
+        # parameters) -> O(N*R) instead of the O(N^2) pairwise sum.
+        rid = _ring_index(cfg.ring_sizes)
+        self.register_buffer("ring_id", rid)
+        onehot = F.one_hot(rid, R).float()       # [N,R] ring membership
+        self.register_buffer("onehot", onehot)
+        self.block_gain = nn.Parameter(self._init_gain(R))
+        self.center_phase = nn.Parameter(torch.zeros(()))  # the rotating center ball
+        self.center_omega = nn.Parameter(torch.tensor(0.3))
+        self.center_gain = nn.Parameter(torch.full((R,), 0.5))  # ring -> center pull
+        self.alpha = nn.Parameter(torch.zeros(R))              # per-ring frustration
+
+        # readout MLP on [cos theta, sin theta]
+        hidden = cfg.readout_mult * cfg.d_model
+        self.readout = nn.Sequential(
+            nn.Linear(2 * N, hidden),
+            nn.GELU(),
+            nn.Linear(hidden, cfg.d_model),
+        )
+        self.drop = nn.Dropout(cfg.dropout)
+
+        # family-trait gate: out = readout * tanh(gate). gate_init_open seeds it so the
+        # block is a warm-startable signal whose magnitude the optimizer controls; the
+        # content weights are nonzero so the gate always receives gradient.
+        go = math.atanh(min(cfg.gate_init_open, 0.9)) if cfg.gate_init_open > 0 else 0.0
+        self.gate = nn.Parameter(torch.full((1,), go))
+
+        # Mandelbrot phase-seed gate: theta <- to_theta(h) + tanh(seed_gate) * orbit.
+        # Zero-init -> no-op at start; opens only if the fractal prior helps.
+        if cfg.use_fractal_phase_seed:
+            self.phase_seed_gate = nn.Parameter(torch.zeros(1))
+
+        # std=0.02 init discipline (matches the family's linear-layer init)
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+        # interstitial rings (instantiated AFTER the std-init loop). Oscillators are
+        # ordered ring-by-ring, so rings r and r+1 are the contiguous slice
+        # [off_r, off_{r+2}); the slot between them injects into exactly that block.
+        self.use_rings = cfg.use_rings
+        self.slot_bounds = []
+        if self.use_rings:
+            off = [0]
+            for n in cfg.ring_sizes:
+                off.append(off[-1] + n)
+            self.slot_bounds = [(off[r], off[r + 2]) for r in range(R - 1)]
+            self.attn_rings = nn.ModuleList([
+                PhaseAttentionRing(e - s, cfg.ring_attn_heads, cfg.ring_attn_head_dim, cfg.osc_bound)
+                for (s, e) in self.slot_bounds])
+            self.engram_rings = nn.ModuleList([
+                EngramRing(cfg.d_model, e - s, cfg.ring_engram_compress, cfg.ring_engram_heads,
+                           cfg.ring_engram_table, cfg.ring_engram_ngram, cfg.osc_bound)
+                for (s, e) in self.slot_bounds])
+
+        # one memory-specialist MoE bank PER oscillator ring (injects into that ring)
+        self.use_ring_specialists = cfg.use_ring_specialists
+        if self.use_ring_specialists:
+            self.specialists = nn.ModuleList([
+                RingSpecialists(n, cfg.d_model, cfg.ring_n_specialists, cfg.ring_spec_key_dim,
+                                cfg.ring_spec_slot_dim, cfg.ring_spec_top_k,
+                                cfg.ring_spec_write_lr, cfg.osc_bound)
+                for n in cfg.ring_sizes])
+
+    @staticmethod
+    def _init_gain(R):
+        # within-ring coupling strong on the diagonal, weak neighbor coupling
+        g = torch.zeros(R, R)
+        for i in range(R):
+            g[i, i] = 1.0
+            if i + 1 < R:
+                g[i, i + 1] = g[i + 1, i] = 0.3
+        return g
+
+    def forward(self, x, ring_ctl=None, phase_seed=None):
+        cfg = self.cfg
+        B, T, _ = x.shape
+        h = self.norm(x)
+
+        theta = self.to_theta(h)                 # [B,T,N]  initial phases
+        if phase_seed is not None:               # Mandelbrot orbit phase prior (gated, no-op at init)
+            theta = theta + torch.tanh(self.phase_seed_gate) * phase_seed
+        omega = torch.tanh(self.to_omega(h))     # [B,T,N]  natural frequencies (bounded)
+
+        # per-ring self-organizing controller: observe each ring (order-parameter
+        # magnitude/phase + mean freq), get no-op-at-init modulations of that ring's
+        # coupling / frustration / center-pull / injection.
+        d_coup = d_alpha = d_center = d_inj = None
+        if ring_ctl is not None:
+            with torch.no_grad():
+                c0, s0 = torch.cos(theta), torch.sin(theta)
+                cnt = self.onehot.sum(0).clamp(min=1.0)       # [R] ring sizes
+                zc = (c0 @ self.onehot) / cnt                 # [B,T,R]
+                zs = (s0 @ self.onehot) / cnt
+                mag = torch.sqrt(zc ** 2 + zs ** 2 + 1e-8)
+                om = (omega @ self.onehot) / cnt
+                obs = torch.stack([mag.mean((0, 1)), (zc / mag).mean((0, 1)),
+                                   (zs / mag).mean((0, 1)), om.mean((0, 1))], dim=-1)  # [R,4]
+            ctrl = ring_ctl(obs)                              # [R,4] (grad -> enc/dec only)
+            d_coup = torch.tanh(ctrl[:, 0])                   # ring coupling scale (1 + .)
+            d_alpha = 0.5 * torch.tanh(ctrl[:, 1])            # frustration shift
+            d_center = torch.tanh(ctrl[:, 2])                 # center-pull scale (1 + .)
+            d_inj = torch.tanh(ctrl[:, 3])                    # injection scale (1 + .)
+
+        # interstitial rings absorb info ONCE (from the initial phases / hidden) and
+        # inject a constant drive into their two neighbor oscillator rings' phase update.
+        inject = 0.0
+        if self.use_rings:
+            inject = torch.zeros_like(theta)
+            for (s, e), ar, er in zip(self.slot_bounds, self.attn_rings, self.engram_rings):
+                contrib = ar(theta[..., s:e]) + er(h)         # [B,T,e-s]
+                inject = inject + F.pad(contrib, (s, self.N - e))  # place at [s:e], sum overlaps
+            if d_inj is not None:
+                inject = inject * (1.0 + d_inj)[self.ring_id]      # controller scales the drive
+
+        # memory specialists: each per-ring bank injects retrieved store_out into its
+        # ring's oscillators. Rings are contiguous and tile [0,N), so cat == full drive.
+        if self.use_ring_specialists:
+            spec_inj = torch.cat([bank(h) for bank in self.specialists], dim=-1)  # [B,T,N]
+            inject = inject + spec_inj
+
+        rid = self.ring_id                        # [N]
+        gain = self.block_gain if d_coup is None else self.block_gain * (1.0 + d_coup).unsqueeze(1)
+        G = gain[rid]                             # [N,R]  per-oscillator ring gains
+        alpha_vec = self.alpha if d_alpha is None else self.alpha + d_alpha
+        alpha_i = alpha_vec[rid]                  # [N]
+        center = self.center_phase                # scalar phase, advances in time
+        cpull = self.center_gain if d_center is None else self.center_gain * (1.0 + d_center)
+        center_pull = cpull[rid]                  # [N]
+        invN = 1.0 / self.N
+
+        for s in range(cfg.osc_steps):
+            ctr = center + self.center_omega * (s * cfg.osc_dt)
+            c, sn = torch.cos(theta), torch.sin(theta)            # [B,T,N]
+            # per-ring summed mean fields, then routed back through gains G
+            zc = (c @ self.onehot) @ G.t()                        # [B,T,N]
+            zs = (sn @ self.onehot) @ G.t()                       # [B,T,N]
+            A = alpha_i - theta                                   # [B,T,N]
+            # (1/N) * sum_r G_ir * Im( e^{i(alpha_i - theta_i)} * Zsum_r )
+            coupling = invN * (torch.sin(A) * zc + torch.cos(A) * zs)
+            to_center = center_pull * torch.sin(ctr - theta + alpha_i)            # [B,T,N]
+            theta = theta + cfg.osc_dt * (omega + coupling + to_center + inject)  # ring drive
+
+        feat = torch.cat([torch.cos(theta), torch.sin(theta)], dim=-1)  # [B,T,2N]
+        out = self.readout(feat) * torch.tanh(self.gate)
+        out = self.drop(soft_clamp(out, self.cfg.osc_bound))           # bounded, live grad
+
+        rec = _viz.get_rec()
+        if rec is not None and rec.enabled:                            # live-viz capture
+            th = theta[0, -1]                                          # last-token phases [N]
+            cnt = self.onehot.sum(0).clamp(min=1.0)
+            zc = (torch.cos(th) @ self.onehot) / cnt                   # per-ring order param
+            zs = (torch.sin(th) @ self.onehot) / cnt
+            R = torch.sqrt(zc ** 2 + zs ** 2 + 1e-9)
+            psi = torch.atan2(zs, zc)
+            rec.log_ring(R.tolist(), psi.tolist(), th.tolist())
+            rec.log_quaz_norm(out[0, -1].norm().item())
+            rec.flush_spec()                                           # group this layer's routes
+        return out
+
+
+class RotaryEmbedding(nn.Module):
+    """RoPE for the rope partition of Q/K (qk_rope_head_dim dims). Ported from v2."""
+
+    def __init__(self, dim, max_positions=4096, theta=10000.0):
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        t = torch.arange(max_positions).float()
+        freqs = torch.outer(t, inv_freq)
+        self.register_buffer("cos_cache", freqs.cos(), persistent=False)
+        self.register_buffer("sin_cache", freqs.sin(), persistent=False)
+
+    def forward(self, x, position_ids):
+        cos = self.cos_cache[position_ids].unsqueeze(1)
+        sin = self.sin_cache[position_ids].unsqueeze(1)
+        d = cos.shape[-1]
+        x1, x2 = x[..., :d], x[..., d:]
+        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+
+class MLADerfXSAAttention(nn.Module):
+    """Family attention, ported from model_v2.MLADerfXSAAttention: MLA low-rank
+    q/o projections, partial RoPE (nope+rope split), QK-Norm, optional DERF erf
+    attention, optional XSA value-subspace removal, GQA. Supports an optional KV
+    cache for incremental decoding (cache holds post-RoPE, pre-GQA-expansion k/v
+    of shape [B, num_kv_heads, S, head_dim])."""
+
+    def __init__(self, cfg: QuazimotoConfig):
+        super().__init__()
+        self.num_heads = cfg.n_head
+        self.num_kv_heads = cfg.num_key_value_heads
+        self.head_dim = cfg.head_dim
+        self.nope_head_dim = cfg.nope_head_dim
+        self.use_qk_norm = cfg.use_qk_norm
+        self.use_derf = cfg.use_derf
+        self.use_xsa = cfg.use_xsa
+        self.dropout_p = cfg.dropout
+        self.kv_groups = self.num_heads // self.num_kv_heads
+        assert self.nope_head_dim + cfg.qk_rope_head_dim == self.head_dim, \
+            "nope_head_dim + qk_rope_head_dim must equal head_dim"
+
+        self.q_a_proj = nn.Linear(cfg.d_model, cfg.q_lora_rank, bias=False)
+        self.q_a_norm = RMSNorm(cfg.q_lora_rank, cfg.rms_norm_eps)
+        self.q_b_proj = nn.Linear(cfg.q_lora_rank, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(cfg.d_model, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(cfg.d_model, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_a_proj = nn.Linear(self.num_heads * self.head_dim, cfg.o_lora_rank, bias=False)
+        self.o_b_proj = nn.Linear(cfg.o_lora_rank, cfg.d_model, bias=False)
+
+        if self.use_qk_norm:
+            self.q_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
+            self.k_norm = RMSNorm(self.head_dim, cfg.rms_norm_eps)
+        self.rope = RotaryEmbedding(cfg.qk_rope_head_dim, cfg.max_position_embeddings,
+                                    cfg.rope_theta)
+        if self.use_derf:
+            self.derf_alpha = nn.Parameter(torch.ones(self.num_heads))
+            self.derf_bias = nn.Parameter(torch.zeros(self.num_heads))
+            self.derf_gamma = nn.Parameter(torch.ones(self.num_heads))
+        for m in (self.q_a_proj, self.q_b_proj, self.k_proj, self.v_proj,
+                  self.o_a_proj, self.o_b_proj):
+            nn.init.normal_(m.weight, std=cfg.initializer_range)
+
+    def forward(self, x, position_ids, past_kv=None, use_cache=False):
+        B, S, _ = x.shape
+        q = self.q_b_proj(self.q_a_norm(self.q_a_proj(x)))
+        q = q.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(B, S, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        if self.use_qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+
+        d = self.nope_head_dim
+        q = torch.cat([q[..., :d], self.rope(q[..., d:], position_ids)], dim=-1)
+        k = torch.cat([k[..., :d], self.rope(k[..., d:], position_ids)], dim=-1)
+
+        # KV cache: prepend previously-seen post-RoPE k/v (pre-GQA-expansion so
+        # the cache stays GQA-compact), then this step's tokens become the tail.
+        if past_kv is not None:
+            past_k, past_v = past_kv
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+        present = (k, v) if use_cache else None
+        L = k.size(2)                                  # total keys (cache + current)
+
+        if self.kv_groups > 1:
+            k = k.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape(
+                B, self.num_heads, L, self.head_dim)
+            v = v.unsqueeze(2).expand(-1, -1, self.kv_groups, -1, -1).reshape(
+                B, self.num_heads, L, self.head_dim)
+
+        # Causal mask over the [S queries x L keys] block. When there is no cache
+        # and S == L this is the plain lower triangle; with a cache the S new
+        # queries sit at absolute positions [L-S, L) and may attend all keys <=
+        # their own position. allowed[i,j] = j <= (L - S + i).
+        offset = L - S
+        if self.use_derf:
+            qpos = torch.arange(S, device=x.device).view(S, 1) + offset
+            kpos = torch.arange(L, device=x.device).view(1, L)
+            is_masked = (kpos > qpos).unsqueeze(0).unsqueeze(0)   # [1,1,S,L]
+            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            safe = scores.masked_fill(is_masked, -10000.0)
+            a, b, g = (self.derf_alpha.view(1, -1, 1, 1), self.derf_bias.view(1, -1, 1, 1),
+                       self.derf_gamma.view(1, -1, 1, 1))
+            w = g * torch.erf(a * safe + b)
+            w = (w + g) / 2.0
+            w = w.masked_fill(is_masked, 0.0)
+            w = w / (w.sum(-1, keepdim=True) + 1e-8)
+            if self.dropout_p > 0 and self.training:
+                w = F.dropout(w, p=self.dropout_p)
+            y = torch.matmul(w, v)
+        else:
+            if offset == 0:
+                attn_mask, causal = None, True
+            else:
+                qpos = torch.arange(S, device=x.device).view(S, 1) + offset
+                kpos = torch.arange(L, device=x.device).view(1, L)
+                attn_mask = (kpos <= qpos).unsqueeze(0).unsqueeze(0)  # bool: True=keep
+                causal = False
+            y = F.scaled_dot_product_attention(
+                q.contiguous(), k.contiguous(), v.contiguous(), attn_mask=attn_mask,
+                is_causal=causal, dropout_p=self.dropout_p if self.training else 0.0)
+
+        rec = _viz.get_rec()
+        if rec is not None and rec.enabled:                  # last-query attention, mean over heads
+            sc = (q[:, :, -1:] @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)  # [B,H,1,L]
+            w = torch.softmax(sc, dim=-1).mean(1)[0, 0]      # [L]
+            rec.log_attn(getattr(self, "layer_idx", 0), w.tolist())
+
+        if self.use_xsa:
+            # remove each query's component along its OWN value direction; with a
+            # cache the queries are the last S of the L cached value positions.
+            vq = v[:, :, -S:, :]
+            vn = vq / (vq.norm(dim=-1, keepdim=True) + 1e-8)
+            y = y - (y * vn).sum(-1, keepdim=True) * vn
+
+        y = y.transpose(1, 2).contiguous().view(B, S, self.num_heads * self.head_dim)
+        out = self.o_b_proj(self.o_a_proj(y))
+        return (out, present) if use_cache else out
+
+
+class HyperConnectionLayer(nn.Module):
+    """Softmax pre-mix / post-distribute over hc_mult residual streams (v2)."""
+
+    def __init__(self, hc_mult):
+        super().__init__()
+        self.pre_weight = nn.Parameter(torch.linspace(0.5, -0.5, hc_mult) / max(hc_mult, 1))
+        self.post_weight = nn.Parameter(torch.linspace(-0.5, 0.5, hc_mult) / max(hc_mult, 1))
+
+    def pre_op(self, copies):
+        w = F.softmax(self.pre_weight, dim=0)
+        return (copies * w.view(1, -1, 1, 1)).sum(dim=1)
+
+    def post_op(self, copies, delta):
+        w = F.softmax(self.post_weight, dim=0)
+        return copies + delta.unsqueeze(1) * w.view(1, -1, 1, 1)
+
+
+class HCOutputMix(nn.Module):
+    """Learned softmax mix over hc_mult streams at the output (init == mean)."""
+
+    def __init__(self, hc_mult):
+        super().__init__()
+        self.weight = nn.Parameter(torch.zeros(hc_mult))
+
+    def forward(self, copies):
+        w = F.softmax(self.weight, dim=0)
+        return (copies * w.view(1, -1, 1, 1)).sum(dim=1)
+
+
+class Layer(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.use_hc = cfg.use_hyper_connections
+        self.use_value_embed = cfg.use_value_embed
+        self.attn_norm = RMSNorm(cfg.d_model, cfg.rms_norm_eps)
+        self.attn = MLADerfXSAAttention(cfg)
+        self.quaz = MycelBlock(cfg)                  # the Neighbour-Sensing growth sub-layer
+        if self.use_hc:
+            self.hc_attn = HyperConnectionLayer(cfg.hc_mult)
+            self.hc_ffn = HyperConnectionLayer(cfg.hc_mult)
+        if self.use_value_embed:
+            self.ve_gate = nn.Parameter(torch.zeros(1))   # zero-init -> no-op
+
+    def forward(self, x, position_ids, token_embed=None, ring_ctl=None,
+                past_kv=None, use_cache=False, phase_seed=None):
+        h = self.hc_attn.pre_op(x) if self.use_hc else x
+        if self.use_value_embed and token_embed is not None:
+            h = h + torch.tanh(self.ve_gate) * token_embed
+        attn_out = self.attn(self.attn_norm(h), position_ids,
+                             past_kv=past_kv, use_cache=use_cache)
+        present = None
+        if use_cache:
+            attn_out, present = attn_out
+        if self.use_hc:
+            x = self.hc_attn.post_op(x, attn_out)
+            h = self.hc_ffn.pre_op(x)
+        else:
+            h = h + attn_out
+        ffn_out = self.quaz(h, ring_ctl, phase_seed)  # QuazimotoBlock norms internally
+        if self.use_hc:
+            x = self.hc_ffn.post_op(x, ffn_out)
+        else:
+            x = h + ffn_out
+        return (x, present) if use_cache else x
+
+
+class QuazimotoLM(nn.Module):
+    def __init__(self, cfg: QuazimotoConfig):
+        super().__init__()
+        self.cfg = cfg
+        self.tok = nn.Embedding(cfg.vocab_size, cfg.d_model)
+        self.drop = nn.Dropout(cfg.dropout)        # positions come from RoPE (no abs embed)
+        # frozen Mandelbrot orbit-angle phase table [vocab, n_osc]; non-persistent
+        # (deterministic, regenerated on load -> not stored in checkpoints).
+        if cfg.use_fractal_phase_seed:
+            from fractal import mandelbrot_phase_table, load_phase_table
+            pt, mode = load_phase_table(cfg.vocab_size, cfg.n_osc,
+                                        os.path.join(_PKG_DIR, "fractal_phase.pt"))
+            if pt is None:
+                pt = mandelbrot_phase_table(cfg.vocab_size, cfg.n_osc)
+                print("  [fractal] flat Halton seed (run build_fractal_table.py for hierarchical)")
+            else:
+                print(f"  [fractal] loaded {mode} phase table")
+            self.register_buffer("fractal_phase", pt, persistent=False)
+        self.layers = nn.ModuleList([Layer(cfg) for _ in range(cfg.n_layer)])
+        for i, layer in enumerate(self.layers):
+            layer.attn.layer_idx = i                 # for live-viz attention capture
+        self.norm_f = RMSNorm(cfg.d_model, cfg.rms_norm_eps)
+        self.head = nn.Linear(cfg.d_model, cfg.vocab_size, bias=False)
+        self.head.weight = self.tok.weight  # weight tying
+        self.logit_scale = nn.Parameter(torch.tensor(1.0))  # family trait: learned logit temp
+        self.use_value_embed = cfg.use_value_embed
+        self.hc_out_mix = HCOutputMix(cfg.hc_mult) if cfg.use_hyper_connections else None
+        self.apply(self._init)
+        # engram rings manage their own inits (frozen compressor, std-0.01 tables,
+        # DERF bias -4); re-apply them after the global init clobbers those.
+        for m in self.modules():
+            if isinstance(m, (EngramRing, RingSpecialists, MycelStations)):
+                m.family_reinit()
+
+        # opt-in trait modules -- instantiated AFTER self.apply so their internal
+        # zero-inits (MoE down-proj, MTP proj) and gate inits survive the global init.
+        # HRM/MoE refine the trunk after the stack (the family's "top membrane
+        # refinement" point); MTP/JEPA are train-time aux heads on the trunk rep.
+        self.hrm = (HRMRefinementBlock(cfg.d_model, cfg.hrm_dim, cfg.hrm_steps,
+                                       gate_init_open=cfg.hrm_gate_init)
+                    if cfg.use_hrm else None)
+        self.moe = MoESwiGLU(cfg.d_model, cfg.moe_intermediate, cfg.moe_n_routed,
+                             cfg.moe_n_shared, cfg.moe_top_k) if cfg.use_moe else None
+        self.mtp_heads = (nn.ModuleList([MTPHead(cfg.d_model) for _ in range(cfg.mtp_layers)])
+                          if cfg.use_mtp else None)
+        self.jepa = (JEPAPredictorBlock(cfg.d_model, cfg.jepa_pred_dim, cfg.jepa_horizon)
+                     if cfg.use_jepa else None)
+        self.ring_bank = (RingControllerBank(len(cfg.ring_sizes), 4, cfg.ring_ctrl_feat,
+                                             cfg.ring_ctrl_local_lr)
+                          if cfg.use_ring_controllers else None)
+
+        # exact kwargs to rebuild on resume (family `family_config` convention)
+        self.family_config = {k: getattr(cfg, k) for k in (
+            "vocab_size", "n_layer", "n_head", "d_model", "block_size",
+            "ring_sizes", "osc_steps", "osc_dt", "readout_mult", "osc_bound", "gate_init_open",
+            "use_hrm", "hrm_steps", "hrm_dim", "hrm_gate_init",
+            "use_moe", "moe_intermediate", "moe_n_routed",
+            "moe_n_shared", "moe_top_k", "use_mtp", "mtp_layers", "mtp_loss_weight",
+            "use_jepa", "jepa_pred_dim", "jepa_horizon", "jepa_loss_weight",
+            "head_dim", "qk_rope_head_dim", "nope_head_dim", "num_key_value_heads",
+            "q_lora_rank", "o_lora_rank", "use_qk_norm", "use_derf", "use_xsa",
+            "rope_theta", "max_position_embeddings", "rms_norm_eps", "initializer_range",
+            "zloss_coef", "use_value_embed", "use_hyper_connections", "hc_mult",
+            "use_rings", "ring_attn_heads", "ring_attn_head_dim", "ring_engram_compress",
+            "ring_engram_heads", "ring_engram_table", "ring_engram_ngram",
+            "use_ring_controllers", "ring_ctrl_feat", "ring_ctrl_local_lr",
+            "use_ring_specialists", "ring_n_specialists", "ring_spec_key_dim",
+            "ring_spec_slot_dim", "ring_spec_top_k", "ring_spec_write_lr",
+            "use_fractal_phase_seed")}
+        n = sum(p.numel() for p in self.parameters())
+        print(f"Mycel-LM: {n/1e6:.1f}M params | {cfg.n_tips} tips/token growing in a "
+              f"{cfg.mycel_pos_dim}D bounded colony ({cfg.growth_steps} growth steps) | "
+              f"traits: {self.active_traits()}")
+
+    def active_traits(self):
+        t = []
+        if self.cfg.use_hrm: t.append("HRM")
+        if self.cfg.use_moe: t.append("MoE")
+        if self.cfg.use_mtp: t.append("MTP")
+        if self.cfg.use_jepa: t.append("JEPA")
+        if self.cfg.use_rings: t.append("Rings")
+        if self.cfg.use_ring_controllers: t.append("Controllers")
+        if self.cfg.use_ring_specialists:
+            t.append(f"Specialists({self.cfg.ring_n_specialists}/ring)")
+        if self.cfg.use_fractal_phase_seed: t.append("FractalSeed")
+        return ",".join(t) or "none"
+
+    def reset_ring_memory(self):
+        """Zero every RingSpecialists store. Call between independent prompts so
+        the test-time context memory does not carry over from a previous sequence."""
+        for m in self.modules():
+            if isinstance(m, RingSpecialists):
+                m.reset_memory()
+
+    def set_ring_memory_writing(self, enabled: bool):
+        """Toggle online writes to the specialist stores (e.g. freeze the memory
+        during evaluation, or disable it to ablate the test-time-write behavior)."""
+        for m in self.modules():
+            if isinstance(m, RingSpecialists):
+                m.write_enabled = enabled
+
+    def _init(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Embedding):
+            nn.init.normal_(m.weight, std=0.02)
+
+    def _compute_trunk(self, idx, past_key_values=None, use_cache=False):
+        """Shared backbone: tokens -> layers -> trunk refinements -> RMSNorm trunk.
+        Returns (trunk, presents). Used by both forward() and forward_drafts()."""
+        B, T = idx.shape
+        # with a cache, the new tokens sit at absolute positions [past_len, past_len+T)
+        past_len = past_key_values[0][0].size(2) if past_key_values else 0
+        position_ids = (torch.arange(past_len, past_len + T, device=idx.device)
+                        .unsqueeze(0).expand(B, -1))
+        x = self.drop(self.tok(idx))
+        token_embed = x if self.use_value_embed else None
+        phase_seed = self.fractal_phase[idx] if self.cfg.use_fractal_phase_seed else None
+        if self.cfg.use_hyper_connections:
+            x = x.unsqueeze(1).expand(-1, self.cfg.hc_mult, -1, -1).clone()  # [B,M,T,H]
+        presents = [] if use_cache else None
+        for i, layer in enumerate(self.layers):
+            past = past_key_values[i] if past_key_values is not None else None
+            x = layer(x, position_ids, token_embed, self.ring_bank,
+                      past_kv=past, use_cache=use_cache, phase_seed=phase_seed)
+            if use_cache:
+                x, present = x
+                presents.append(present)
+        if self.hc_out_mix is not None:
+            x = self.hc_out_mix(x)               # learned mix over streams (init == mean)
+        # trunk refinements (no-op at init): HRM iterative gated refine, MoE expert mix
+        rec = _viz.get_rec()
+        if self.hrm is not None:
+            hx = self.hrm(x)
+            if rec is not None and rec.enabled:
+                rec.log_trait("hrm", (hx - x)[0, -1].norm().item())
+            x = hx
+        if self.moe is not None:
+            mx = self.moe(x)
+            if rec is not None and rec.enabled:
+                rec.log_trait("moe", mx[0, -1].norm().item())
+            x = x + mx
+        return self.norm_f(x), presents          # v2: RMSNorm trunk (no tanh)
+
+    @torch.no_grad()
+    def forward_drafts(self, idx, past_key_values=None, use_cache=False):
+        """Speculative-decoding support (DeepSpec-style draft+verify, self-drafted):
+        return (main_logits, mtp_logits_list[, presents]) over ALL T positions.
+          * main_logits[:, j]        predicts the token at position j+1   (verifier)
+          * mtp_logits_list[k][:, j] predicts the token at position j+2+k  (k=0..mtp-1)
+        So one forward yields, at the last position, the genuine next token (main)
+        plus `mtp_layers` drafted future tokens to be verified next round."""
+        trunk, presents = self._compute_trunk(idx, past_key_values, use_cache)
+        main_logits = self.head(trunk) * self.logit_scale
+        mtp_logits = ([self.head(h(trunk)) * self.logit_scale for h in self.mtp_heads]
+                      if self.mtp_heads is not None else [])
+        return (main_logits, mtp_logits, presents) if use_cache else (main_logits, mtp_logits)
+
+    def forward(self, idx, targets=None, past_key_values=None, use_cache=False):
+        B, T = idx.shape
+        trunk, presents = self._compute_trunk(idx, past_key_values, use_cache)
+        logits = self.head(trunk) * self.logit_scale
+
+        loss, aux = None, {}
+        if targets is not None:
+            flat = logits.view(-1, logits.size(-1))
+            tflat = targets.reshape(-1)
+            loss = F.cross_entropy(flat, tflat, ignore_index=-1)
+            # z-loss (v2): penalise log^2 of the partition function -> no logit drift
+            if self.cfg.zloss_coef > 0:
+                valid = tflat != -1
+                if valid.any():
+                    log_z = torch.logsumexp(flat[valid].float(), dim=-1)
+                    aux["zloss"] = self.cfg.zloss_coef * (log_z ** 2).mean()
+            if self.moe is not None and self.moe.last_aux_loss is not None:
+                aux["moe_aux"] = self.moe.last_aux_loss
+            if self.mtp_heads is not None:
+                mtp_total, n_active = 0.0, 0
+                for k, head in enumerate(self.mtp_heads, start=1):
+                    if T - k <= 0:
+                        break
+                    mlog = self.head(head(trunk[:, :T - k])) * self.logit_scale
+                    mtp_total = mtp_total + F.cross_entropy(
+                        mlog.reshape(-1, self.cfg.vocab_size),
+                        targets[:, k:].reshape(-1), ignore_index=-1)
+                    n_active += 1
+                if n_active:
+                    aux["mtp_loss"] = self.cfg.mtp_loss_weight * mtp_total / n_active
+            if self.jepa is not None and T > 1:
+                jepa_total, n_j = 0.0, 0
+                for k in range(1, self.cfg.jepa_horizon + 1):
+                    if T - k <= 0:
+                        break
+                    pred = self.jepa(trunk[:, :T - k], k)
+                    tgt = trunk[:, k:].detach()      # JEPA stop-grad target
+                    cos = F.cosine_similarity(pred.float(), tgt.float(), dim=-1)
+                    jepa_total = jepa_total + (1.0 - cos).mean()
+                    n_j += 1
+                if n_j:
+                    aux["jepa_loss"] = self.cfg.jepa_loss_weight * jepa_total / n_j
+        # use_cache returns the per-layer (k, v) presents as a 4th item; without
+        # it the signature is the original 3-tuple so train.py is unaffected.
+        if use_cache:
+            return logits, loss, aux, presents
+        return logits, loss, aux
+
+    @torch.no_grad()
+    def generate(self, idx, n_new, temperature=1.0, top_k=None, use_cache=True):
+        """Autoregressive sampling with the family's NaN/inf sanitisation: long
+        accumulation can push a logit to +/-inf or NaN, whose softmax -> NaN ->
+        multinomial fires a CUDA device-side assert. Clamp first, greedy-fallback.
+
+        With use_cache=True (default) the attention KV cache is kept across steps
+        so each new token costs one forward over a single position instead of a
+        full recompute. The cache holds absolute positions, so once it fills past
+        block_size we drop the cache and fall back to a windowed recompute (RoPE
+        positions would otherwise exceed max_position_embeddings)."""
+        self.eval()
+        past = None
+        for _ in range(n_new):
+            if use_cache and past is not None and past[0][0].size(2) < self.cfg.block_size:
+                step_in = idx[:, -1:]                       # only the newest token
+                logits, _, _, past = self(step_in, past_key_values=past, use_cache=True)
+            elif use_cache:
+                cond = idx[:, -self.cfg.block_size:]        # prefill / re-prime cache
+                logits, _, _, past = self(cond, use_cache=True)
+            else:
+                cond = idx[:, -self.cfg.block_size:]
+                logits, _, _ = self(cond)
+            lg = torch.nan_to_num(logits[:, -1, :].float(), nan=0.0,
+                                  posinf=1e4, neginf=-1e4) / max(temperature, 1e-6)
+            if top_k:
+                k = min(top_k, lg.size(-1))
+                v, _ = torch.topk(lg, k)
+                lg[lg < v[:, [-1]]] = -float("inf")
+            probs = torch.softmax(lg, dim=-1)
+            if not torch.isfinite(probs).all() or float(probs.sum()) <= 0.0:
+                nxt = torch.argmax(lg, dim=-1, keepdim=True)
+            else:
+                nxt = torch.multinomial(probs, 1)
+            idx = torch.cat([idx, nxt], dim=1)
+        return idx
+
+
+if __name__ == "__main__":
+    # exercise the base model and all trait modules at once
+    cfg = QuazimotoConfig(use_hrm=True, use_moe=True, use_mtp=True, use_jepa=True)
+    model = QuazimotoLM(cfg)
+    x = torch.randint(0, cfg.vocab_size, (2, 64))
+    logits, loss, aux = model(x, x)
+    total = loss + sum(aux.values())
+    print("forward ok:", tuple(logits.shape), "loss", round(loss.item(), 3),
+          "| aux:", {k: round(float(v), 4) for k, v in aux.items()})
+    total.backward()
+    print("backward ok")

mycel.py

@@ -0,0 +1,178 @@
+"""
+mycel.py -- the Neighbour-Sensing (fungal colony growth) channel-mixing block for
+Mycel-LM, the structural analog of Quazimoto's Kuramoto oscillator block.
+
+Meskauskas/Fricker/Moore (2004) Neighbour-Sensing model: each hyphal TIP is an agent
+with a position and a growth vector; tips generate a density FIELD; each tip SENSES
+the local field and steers its growth (negative autotropism -- grow away from the
+colony's own density) plus persistence. We turn that into a differentiable mixer:
+
+  hidden -> N tips (position p in a BOUNDED latent region + growth direction v)
+         -> a few differentiable growth steps: sense the density field, steer, move,
+            RE-CLAMP into the bounded region (the colony can't grow unbounded)
+         -> readout of [p, v, sensed-density] -> hidden, behind a family gate.
+
+Traits are "spread through the mycelium" (MycelStations): tiny memory specialists sit
+at fixed anchor positions in the bounded region; a tip interacts with a station by
+PROXIMITY, which is emergent from where it grew -- so which tips use which trait
+"comes to be" during growth rather than being assigned to a fixed index.
+
+The density field is computed against F learnable field centres (a low-rank sample of
+the field) so cost is O(N*F) per step, not O(N^2) -- the same trick that keeps the
+oscillator block's mean-field cheap.
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from family import soft_clamp, RMSNorm
+import instrument as _viz
+
+
+class MycelStations(nn.Module):
+    """Memory specialists placed at anchor positions in the colony's bounded region.
+    A tip reads/writes a station by proximity (softmax over -distance). Slow weights
+    (anchors, encoders, decoder) are learned; the stores are test-time fast memory
+    (buffers). Zero-init scale => no-op at start (family contract)."""
+
+    def __init__(self, n_stations, pos_dim, d_model, slot_dim=64, top_k=2,
+                 write_lr=0.1, bound=10.0):
+        super().__init__()
+        self.n = n_stations
+        self.top_k = min(top_k, n_stations)
+        self.write_lr, self.bound, self.write_enabled = write_lr, bound, True
+        self.anchors = nn.Parameter(torch.randn(n_stations, pos_dim) * 0.5)  # station sites
+        self.log_bw = nn.Parameter(torch.zeros(1))              # proximity bandwidth
+        self.in_enc = nn.Linear(d_model, slot_dim, bias=False)  # frozen write encoders
+        self.val_enc = nn.Linear(d_model, slot_dim, bias=False)
+        self.out_dec = nn.Linear(slot_dim, pos_dim, bias=False)  # station -> steering drive
+        self.scale = nn.Parameter(torch.zeros(1))
+        self.register_buffer("store_in", torch.zeros(n_stations, slot_dim))
+        self.register_buffer("store_out", torch.zeros(n_stations, slot_dim))
+        for m in (self.in_enc, self.val_enc, self.out_dec):
+            nn.init.normal_(m.weight, std=0.02)
+        self.in_enc.weight.requires_grad_(False)
+        self.val_enc.weight.requires_grad_(False)
+
+    def family_reinit(self):
+        nn.init.normal_(self.anchors, std=0.5)
+        nn.init.zeros_(self.scale)
+        for m in (self.in_enc, self.val_enc):
+            nn.init.normal_(m.weight, std=0.02); m.weight.requires_grad_(False)
+
+    def reset_memory(self):
+        self.store_in.zero_(); self.store_out.zero_()
+
+    def forward(self, p, h):
+        # p: [B,T,N,pos_dim] tip positions ; h: [B,T,d] token hidden
+        store_in, store_out = self.store_in.clone(), self.store_out.clone()  # noqa: F841
+        bw = F.softplus(self.log_bw).clamp(min=1e-3)
+        # tip->station squared distance via the expansion (no [..,N,n,pos_dim] intermediate)
+        p2 = (p * p).sum(-1, keepdim=True)                       # [B,T,N,1]
+        a2 = (self.anchors * self.anchors).sum(-1)               # [n]
+        d2 = (p2 + a2 - 2.0 * (p @ self.anchors.t())).clamp(min=0.0)  # [B,T,N,n]
+        route = torch.softmax(-d2 / bw, dim=-1)
+        if self.top_k < self.n:
+            tv = torch.topk(route, self.top_k, dim=-1).values
+            route = route.masked_fill(route < tv[..., [-1]], 0.0)
+            route = route / route.sum(-1, keepdim=True).clamp(min=1e-6)
+        drive = self.out_dec(route @ store_out) * torch.tanh(self.scale)  # [B,T,N,pos_dim]
+
+        if self.write_enabled and self.write_lr > 0:                # stations absorb nearby tips
+            with torch.no_grad():
+                occ = route.reshape(-1, self.n)                     # [BTN, n] proximity mass
+                info_in = self.in_enc(h).unsqueeze(2).expand(-1, -1, p.size(2), -1).reshape(-1, self.in_enc.out_features)
+                info_v = self.val_enc(h).unsqueeze(2).expand(-1, -1, p.size(2), -1).reshape(-1, self.val_enc.out_features)
+                denom = occ.sum(0).clamp(min=1e-3).unsqueeze(1)
+                a = self.write_lr
+                self.store_in.mul_(1 - a).add_(a * (occ.t() @ info_in) / denom).clamp_(-self.bound, self.bound)
+                self.store_out.mul_(1 - a).add_(a * (occ.t() @ info_v) / denom).clamp_(-self.bound, self.bound)
+        return drive
+
+
+class MycelBlock(nn.Module):
+    """Neighbour-Sensing growth mixer: hidden -> growing tips in a bounded region ->
+    readout. Structural analog of QuazimotoBlock.forward(x, ring_ctl, phase_seed)."""
+
+    def __init__(self, cfg):
+        super().__init__()
+        self.cfg = cfg
+        self.N, self.pd, self.F = cfg.n_tips, cfg.mycel_pos_dim, cfg.mycel_field_centers
+        self.norm = RMSNorm(cfg.d_model)
+        self.to_pos = nn.Linear(cfg.d_model, self.N * self.pd)   # initial tip positions
+        self.to_dir = nn.Linear(cfg.d_model, self.N * self.pd)   # initial growth vectors
+        self.centers = nn.Parameter(torch.randn(self.F, self.pd) * 0.5)  # field sample sites
+        self.log_persist = nn.Parameter(torch.tensor(1.4))       # sigmoid -> ~0.8 persistence
+        self.tropism = nn.Parameter(torch.zeros(1))              # negative-autotropism strength
+        self.log_field_bw = nn.Parameter(torch.zeros(1))
+        hidden = cfg.readout_mult * cfg.d_model
+        self.readout = nn.Sequential(
+            nn.Linear(self.N * (2 * self.pd + 1), hidden), nn.GELU(),
+            nn.Linear(hidden, cfg.d_model))
+        self.drop = nn.Dropout(cfg.dropout)
+        go = math.atanh(min(cfg.gate_init_open, 0.9)) if cfg.gate_init_open > 0 else 0.0
+        self.gate = nn.Parameter(torch.full((1,), go))
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=0.02)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+        if cfg.use_fractal_phase_seed:
+            self.pos_seed_gate = nn.Parameter(torch.zeros(1))   # zero-init -> no-op
+        self.use_stations = cfg.use_ring_specialists
+        if self.use_stations:
+            self.stations = MycelStations(cfg.mycel_n_stations, self.pd, cfg.d_model,
+                                          cfg.ring_spec_slot_dim, cfg.ring_spec_top_k,
+                                          cfg.ring_spec_write_lr, cfg.osc_bound)
+
+    def _clamp(self, p):
+        b = self.cfg.osc_bound
+        return b * torch.tanh(p / b)                              # bound the colony region
+
+    @staticmethod
+    def _cdist2(p, c):
+        # squared distances [..,N,F] via ||p||^2 + ||c||^2 - 2 p.c -- NO [..,N,F,pd]
+        # intermediate (that broadcast is the memory hog; this keeps it O(N*F)).
+        p2 = (p * p).sum(-1, keepdim=True)                       # [..,N,1]
+        c2 = (c * c).sum(-1)                                     # [F]
+        return (p2 + c2 - 2.0 * (p @ c.t())).clamp(min=0.0)      # [..,N,F]
+
+    def _grow(self, p, v, h):
+        cfg = self.cfg
+        pers = torch.sigmoid(self.log_persist)
+        trop = torch.tanh(self.tropism)
+        bw = F.softplus(self.log_field_bw).clamp(min=1e-3)
+        for _ in range(cfg.growth_steps):
+            K = torch.exp(-self._cdist2(p, self.centers) / bw)   # [B,T,N,F]
+            w = K.mean(-2, keepdim=True) * K                     # rho_f * K_if  [B,T,N,F]
+            # steer away from dense centres, factored to avoid the [..,N,F,pd] tensor:
+            #   sum_f w_if (p_i - c_f) = p_i * sum_f w_if  -  (w @ centres)
+            away = p * w.sum(-1, keepdim=True) - torch.matmul(w, self.centers)   # [B,T,N,pd]
+            v = pers * v + trop * away
+            if self.use_stations:
+                v = v + self.stations(p, h)                       # traits steer nearby tips
+            p = self._clamp(p + cfg.growth_dt * v)
+        dens = torch.exp(-self._cdist2(p, self.centers) / bw).mean(-1, keepdim=True)
+        return p, v, dens
+
+    def forward(self, x, ring_ctl=None, phase_seed=None):
+        cfg = self.cfg
+        B, T, _ = x.shape
+        h = self.norm(x)
+        p = self._clamp(self.to_pos(h).view(B, T, self.N, self.pd))
+        if phase_seed is not None:                               # optional fractal position seed
+            seed = phase_seed[..., :self.N * self.pd].view(B, T, self.N, self.pd)
+            p = self._clamp(p + torch.tanh(self.pos_seed_gate) * seed)
+        v = self.to_dir(h).view(B, T, self.N, self.pd)
+        p, v, dens = self._grow(p, v, h)
+        feat = torch.cat([p, v, dens], dim=-1).flatten(2)        # [B,T,N*(2pd+1)]
+        out = self.drop(soft_clamp(self.readout(feat) * torch.tanh(self.gate), cfg.osc_bound))
+
+        rec = _viz.get_rec()
+        if rec is not None and rec.enabled:                      # live-viz: colony spread
+            spread = dens[0, -1, :, 0]
+            rec.log_ring([float(spread.mean())], [0.0], p[0, -1].flatten().tolist())
+            rec.log_quaz_norm(out[0, -1].norm().item())
+            rec.flush_spec()
+        return out

opd_teacher.py

@@ -0,0 +1,79 @@
+"""
+opd_teacher.py -- frozen teacher (facebook/MobileLLM-R1-140M-base) for on-policy
+distillation into Quazimoto-LM.
+
+The teacher and student use DIFFERENT tokenizers, so a token-level KL is undefined.
+This scorer instead returns the teacher's per-token log-probabilities of a piece of
+text together with each token's BYTE span, so train_opd.py can re-aggregate that
+dense signal onto the student's own tokens (byte-space alignment). The teacher is
+never trained; it only scores the student's rollouts.
+
+Requires: transformers, torch. The teacher is ~140M params -> fits on a small GPU
+(or CPU, slowly). Run on the rig that has the weights / network access.
+"""
+
+import torch
+import torch.nn.functional as F
+
+
+class Teacher:
+    def __init__(self, name="facebook/MobileLLM-R1-140M-base", device="cuda",
+                 dtype=torch.float32):
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        self.device = device
+        self.tok = AutoTokenizer.from_pretrained(name, trust_remote_code=True)
+        try:                                              # newer transformers: dtype=
+            self.model = AutoModelForCausalLM.from_pretrained(
+                name, dtype=dtype, trust_remote_code=True)
+        except TypeError:                                 # older transformers: torch_dtype=
+            self.model = AutoModelForCausalLM.from_pretrained(
+                name, torch_dtype=dtype, trust_remote_code=True)
+        self.model = self.model.to(device).eval()
+        self.has_offsets = bool(getattr(self.tok, "is_fast", False))
+        if not self.has_offsets:
+            print("  [teacher] WARNING: slow tokenizer (no offset mapping) -> "
+                  "byte alignment unavailable; train_opd will fall back to seq reward.")
+
+    @torch.no_grad()
+    def score(self, text):
+        """Return (byte_logprob, n_bytes): a per-BYTE teacher log-prob vector for
+        `text` (length = len(text.encode('utf-8'))). Each teacher token's log-prob
+        (of being generated given its prefix) is spread uniformly over the bytes it
+        covers. byte_logprob[k] is the teacher's per-byte score for byte k -- the
+        dense signal the student is distilled toward.
+
+        If offsets are unavailable, returns (mean_logprob_scalar, n_bytes) with the
+        scalar broadcast (sequence-level fallback)."""
+        raw = text.encode("utf-8")
+        nb = len(raw)
+        if nb == 0:
+            return torch.zeros(0), 0
+        enc = self.tok(text, return_offsets_mapping=self.has_offsets,
+                       return_tensors="pt", add_special_tokens=False)
+        ids = enc["input_ids"].to(self.device)
+        T = ids.size(1)
+        if T < 2:
+            return torch.zeros(nb), nb
+        logits = self.model(ids).logits[0].float()           # [T, V]
+        logp = F.log_softmax(logits, dim=-1)
+        # logprob of token t (t>=1) given prefix = logp[t-1, ids[t]]
+        tok_lp = torch.zeros(T)
+        tok_lp[1:] = logp[torch.arange(T - 1), ids[0, 1:]].cpu()
+
+        if not self.has_offsets:
+            return torch.full((nb,), float(tok_lp[1:].mean())), nb
+
+        # char offsets -> byte offsets (UTF-8 widths of the prefix)
+        char_to_byte = [0]
+        for ch in text:
+            char_to_byte.append(char_to_byte[-1] + len(ch.encode("utf-8")))
+        offs = enc["offset_mapping"][0].tolist()             # [(c0,c1), ...]
+        byte_lp = torch.zeros(nb)
+        for t in range(T):
+            c0, c1 = offs[t]
+            if c1 <= c0:
+                continue
+            b0, b1 = char_to_byte[c0], char_to_byte[c1]
+            span = max(b1 - b0, 1)
+            byte_lp[b0:b1] = tok_lp[t] / span                # spread token logprob over its bytes
+        return byte_lp, nb

requirements.txt

@@ -0,0 +1,11 @@
+# Core (inference, generation, visualizer, health-check)
+torch>=2.1
+numpy
+transformers>=4.40          # SpikeTokenizer subclasses PreTrainedTokenizer
+
+# Training / fine-tuning (train.py, train_sft.py, distillation)
+datasets>=2.18              # streamed pretrain / SFT blends
+huggingface_hub>=0.23       # gated dataset + repo access
+
+# Optional
+# safetensors               # not required; checkpoints are torch .pt

special_tokens.py

@@ -0,0 +1,85 @@
+"""
+special_tokens.py -- central, *append-only* registry of special tokens for the
+SpikeWhale length-max tokenizer.
+
+WHY THIS FILE EXISTS
+--------------------
+The base vocab (tokenizer.json) is 16384 contiguous ids:
+    0..3   -> <pad> <unk> <bos> <eos>
+    4..259 -> the 256 raw bytes
+    260..  -> learned byte-merges
+Adding tokens "without breaking the model" has exactly one rule:
+
+    ***APPEND ONLY. NEVER REORDER OR REMOVE AN EXISTING ID.***
+
+Every existing id keeps pointing at the same embedding row and the same logit
+column, so the model's behaviour on already-seen tokens is bit-for-bit
+unchanged. New tokens are appended at ids >= 16384 and their embedding / lm_head
+/ mtp rows are freshly initialised (near-zero contribution) so they are
+no-ops until you train them.
+
+To stay tensor-core friendly the final vocab is padded up to a multiple of
+`VOCAB_MULTIPLE` (128) with `<|reserved_N|>` slots. Those reserves let you name
+*future* tokens later by editing the registry WITHOUT another model resize, as
+long as the total stays <= the padded size.
+
+HOW TO ADD MORE LATER
+---------------------
+Append new names to NAMED_SPECIAL_TOKENS (at the END), then either:
+  * if you still have <|reserved_*|> slots free, just rename a reserved id in
+    tokenizer.json (no model change needed), or
+  * re-run add_special_tokens.py to grow + re-pad the vocab (model resized).
+"""
+
+# Tensor-core / matmul friendly vocab alignment. 16384 is already 128*128.
+VOCAB_MULTIPLE = 128
+
+# ---------------------------------------------------------------------------
+# The universal named set. ORDER IS PERMANENT -- append only, never reorder.
+# Mixing the common conventions so the same model can do chat, reasoning,
+# agentic tool use, and code infilling.
+# ---------------------------------------------------------------------------
+NAMED_SPECIAL_TOKENS = [
+    # ChatML turn framing
+    "<|im_start|>",
+    "<|im_end|>",
+    # Reasoning / scratchpad
+    "<think>",
+    "</think>",
+    # Explicit solution block
+    "<begin_solution>",
+    "<end_solution>",
+    # Agentic tool calling
+    "<tool_call>",
+    "</tool_call>",
+    "<tool_response>",
+    "</tool_response>",
+    # Role markers (usable standalone or inside an im_start header)
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    # Fill-in-the-middle (code)
+    "<|fim_prefix|>",
+    "<|fim_middle|>",
+    "<|fim_suffix|>",
+    # Generic document separator
+    "<|endoftext|>",
+]
+
+
+def build_special_token_list(base_vocab_size: int,
+                             multiple: int = VOCAB_MULTIPLE):
+    """
+    Return the ordered list of tokens to APPEND after `base_vocab_size`:
+    the named set followed by enough <|reserved_N|> slots to pad the final
+    vocab size up to the next multiple of `multiple`.
+
+    The returned list's element i gets id (base_vocab_size + i).
+    """
+    tokens = list(NAMED_SPECIAL_TOKENS)
+    target = base_vocab_size + len(tokens)
+    # round up to the next multiple (or stay put if already aligned)
+    padded = ((target + multiple - 1) // multiple) * multiple
+    n_reserved = padded - target
+    tokens += [f"<|reserved_{i}|>" for i in range(n_reserved)]
+    return tokens

spike_tokenizer.py

@@ -0,0 +1,124 @@
+"""
+spike_tokenizer.py -- HuggingFace-compatible wrapper for the custom
+byte-level "length-max" (greedy longest-match) tokenizer in tokenizer.json.
+
+The raw tokenizer.json is NOT a HuggingFace `tokenizers` file; it is a plain
+dict {vocab, vocab_size, max_token_len, algorithm:"length-max"}. This wrapper
+makes it loadable by AutoTokenizer.from_pretrained / save_pretrained and
+exposes encode/decode + the bos/eos/pad/unk ids the training scripts expect.
+
+Encoding scheme (verified): byte-level. Text is UTF-8 encoded, each byte mapped
+to its latin-1 character, then greedily matched against the vocab using the
+longest key that matches at each position (max key length = max_token_len).
+"""
+import json, os
+from typing import List, Optional
+
+# We only use transformers' tokenizer base class. Stop it from probing
+# TensorFlow/Flax (whose builds are often incompatible with numpy 2 / protobuf and
+# would crash the import). setdefault so an explicit user override still wins.
+os.environ.setdefault("USE_TF", "0")
+os.environ.setdefault("USE_FLAX", "0")
+
+from transformers import PreTrainedTokenizer
+
+
+class SpikeTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "tokenizer.json"}
+    model_input_names = ["input_ids"]
+
+    def __init__(self, vocab_file=None, **kwargs):
+        with open(vocab_file, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        self._vocab = data["vocab"]                       # str -> id
+        self._ids_to_tokens = {i: t for t, i in self._vocab.items()}
+        self.max_token_len = int(data.get("max_token_len", 24))
+        # length-bucketed keys for fast greedy match (longest length first)
+        self._lengths = sorted({len(k) for k in self._vocab}, reverse=True)
+
+        # Appended special tokens (im_start / <think> / <begin_solution> / ...).
+        # They already live in self._vocab at their real ids; we hand them to the
+        # HF base class as `additional_special_tokens` so its AddedToken trie:
+        #   (1) splits them out ATOMICALLY before our byte-level greedy match
+        #       (verified: each maps back to its existing vocab id, no phantom id), and
+        #   (2) drops them on decode(skip_special_tokens=True).
+        # The set is stored in tokenizer.json under "special_tokens" so it
+        # survives save_pretrained/from_pretrained round-trips.
+        self._extra_specials = [
+            t for t in data.get("special_tokens", []) if t in self._vocab
+        ]
+        if self._extra_specials:
+            existing = list(kwargs.get("additional_special_tokens", []) or [])
+            merged = existing + [t for t in self._extra_specials if t not in existing]
+            kwargs["additional_special_tokens"] = merged
+
+        kwargs.setdefault("bos_token", "<bos>")
+        kwargs.setdefault("eos_token", "<eos>")
+        kwargs.setdefault("unk_token", "<unk>")
+        kwargs.setdefault("pad_token", "<pad>")
+        super().__init__(**kwargs)
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self._vocab)
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    # --- core byte-level greedy tokenization ---
+    def _tokenize(self, text: str) -> List[str]:
+        s = text.encode("utf-8").decode("latin-1")  # one char per byte
+        out, i, n = [], 0, len(s)
+        while i < n:
+            matched = None
+            hi = min(self.max_token_len, n - i)
+            for L in range(hi, 0, -1):
+                sub = s[i:i + L]
+                if sub in self._vocab:
+                    matched = sub
+                    break
+            if matched is None:           # single byte always exists in vocab
+                matched = s[i]
+            out.append(matched)
+            i += len(matched)
+        return out
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self._vocab.get(token, self._vocab["<unk>"])
+
+    def _convert_id_to_token(self, index: int) -> str:
+        return self._ids_to_tokens.get(index, "<unk>")
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        # transformers 5.x hands the FULL token list here (special tokens
+        # included; skip_special_tokens is already applied upstream via
+        # convert_ids_to_tokens). So we can't just byte-decode everything: a
+        # special token like "<|im_start|>" is a literal marker, not latin-1
+        # bytes. Decode runs of ordinary byte-tokens together (needed so
+        # multi-byte UTF-8 sequences reassemble) and emit any special token
+        # inline as its literal string.
+        specials = {"<pad>", "<unk>", "<bos>", "<eos>", *self._extra_specials}
+        out, buf = [], []
+        for tok in tokens:
+            if tok in specials:
+                if buf:
+                    out.append("".join(buf).encode("latin-1").decode("utf-8", errors="replace"))
+                    buf = []
+                out.append(tok)
+            else:
+                buf.append(tok)
+        if buf:
+            out.append("".join(buf).encode("latin-1").decode("utf-8", errors="replace"))
+        return "".join(out)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None):
+        os.makedirs(save_directory, exist_ok=True)
+        fn = (filename_prefix + "-" if filename_prefix else "") + "tokenizer.json"
+        path = os.path.join(save_directory, fn)
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump({"vocab": self._vocab, "vocab_size": self.vocab_size,
+                       "max_token_len": self.max_token_len,
+                       "algorithm": "length-max",
+                       "special_tokens": list(self._extra_specials)},
+                      f, ensure_ascii=False)
+        return (path,)

train.bat

@@ -0,0 +1,24 @@
+@echo off
+REM Mycel-LM pretraining -- the Neighbour-Sensing (fungal colony growth) mixer, same
+REM harness/traits/data as the Quazimoto lm3 run, so the two are a controlled swap
+REM (mixer is the only architectural difference). Ring-specific traits (--use-rings,
+REM --use-ring-controllers) are OMITTED: the mycelium has no rings -- its analog is
+REM --use-ring-specialists, which places the trait STATIONS through the colony.
+cd /d "%~dp0"
+REM the growth loop is activation-heavy; expandable_segments cuts fragmentation.
+set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+python train.py ^
+  --steps 130000 ^
+  --batch 4 ^
+  --block 512 ^
+  --math-frac 0.25 ^
+  --use-hrm ^
+  --use-moe ^
+  --use-mtp ^
+  --use-jepa ^
+  --use-ring-specialists ^
+  --use-fractal-phase-seed ^
+  --amp ^
+  --stream ^
+  --resume
+pause

train.py

@@ -0,0 +1,297 @@
+"""
+Train Quazimoto-LM with the SpikeWhale tokenizer and the family's AdamW regime.
+
+Tokenizer: the custom SpikeWhale length-max byte tokenizer (tokenizer.json +
+spike_tokenizer.py) bundled in this package -- self-contained, no external folder.
+Data is tokenized fresh in Python from a local UTF-8 .txt (or a synthetic fallback).
+
+Optimizer matches the family's logged-stable choice (LEVEL6_PLAN): AdamW, peak lr
+3e-4, linear warmup, cosine decay to min_lr_frac*peak (0.3 floor), grad-clip 1.0.
+MuonEq-R is deliberately NOT used here -- it caused weight runaway in this family.
+"""
+
+import argparse, time, math, os, sys
+import numpy as np
+import torch
+from model import QuazimotoLM, QuazimotoConfig
+
+# SpikeWhale tokenizer is bundled in this package (self-contained, no level6 needed).
+PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+def load_tokenizer(tok_dir):
+    sys.path.insert(0, tok_dir)
+    from spike_tokenizer import SpikeTokenizer
+    tok = SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json"))
+    return tok
+
+
+def load_ids(path, tok):
+    if path and os.path.exists(path):
+        with open(path, "r", encoding="utf-8", errors="replace") as f:
+            text = f.read()
+    else:
+        text = ("the quazimoto oscillator learns to synchronize phases across rings. "
+                "coupled clocks fall into step when coupling is strong enough. ") * 4000
+    ids = np.array(tok.encode(text, add_special_tokens=False), dtype=np.int64)
+    n = int(0.9 * len(ids))
+    return ids[:n], ids[n:]
+
+
+def get_batch(ids, bs, block, device):
+    ix = np.random.randint(0, len(ids) - block - 1, size=bs)
+    x = np.stack([ids[i:i + block] for i in ix])
+    y = np.stack([ids[i + 1:i + 1 + block] for i in ix])
+    return (torch.from_numpy(x).to(device), torch.from_numpy(y).to(device))
+
+
+def _doc_text(ex):
+    """Pull the text field from a dataset example, tolerating naming differences
+    across datasets (text / content / raw_content / document), else the first
+    non-trivial string value."""
+    for k in ("text", "content", "raw_content", "document"):
+        v = ex.get(k)
+        if isinstance(v, str) and v.strip():
+            return v
+    for v in ex.values():
+        if isinstance(v, str) and len(v.strip()) > 1:
+            return v
+    return None
+
+
+def _stream_docs(path, config, split):
+    """Infinite generator of non-empty doc strings from a streaming HF dataset.
+    Re-opens the stream when a shard pass is exhausted so training never runs dry.
+    (Mirrors the level6 train_snn_stream._stream_docs convention.)"""
+    from datasets import load_dataset
+    while True:
+        ds = load_dataset(path, name=config, split=split, streaming=True)
+        any_yielded = False
+        for ex in ds:
+            text = _doc_text(ex)
+            if text:
+                any_yielded = True
+                yield text
+        if not any_yielded:
+            raise RuntimeError(f"stream {path}:{config} yielded no usable text")
+
+
+def _blend_sources(args):
+    """Return the active (label, path, config, weight) sources for the blend.
+    Default mix: 35% Ultra-FineWeb-L3 / 25% FineWeb-Edu / 25% FineMath / 15% PretrainNew."""
+    srcs = [
+        ("ultra",    args.ultra_dataset,    args.ultra_config,    args.ultra_frac),
+        ("edu",      args.fineweb_dataset,  args.fineweb_config,  args.edu_frac),
+        ("math",     args.math_dataset,     args.math_config,     args.math_frac),
+        ("pretrain", args.pretrain_dataset, args.pretrain_config, args.pretrain_frac),
+    ]
+    return [s for s in srcs if s[3] > 0]
+
+
+def doc_generator(args):
+    """Blend multiple HF datasets by per-document probability (weights normalized)."""
+    import random
+    rng = random.Random(args.seed)
+    srcs = _blend_sources(args)
+    gens = [_stream_docs(p, c, args.split) for (_, p, c, _) in srcs]
+    weights = [s[3] for s in srcs]
+    tot = sum(weights)
+    cum, acc = [], 0.0
+    for w in weights:                                  # cumulative thresholds for sampling
+        acc += w / tot
+        cum.append(acc)
+    while True:
+        r = rng.random()
+        gi = next(k for k, c in enumerate(cum) if r <= c)
+        yield next(gens[gi]) + "\n"
+
+
+def stream_batches(tok, args, device):
+    """Infinite [B,T] next-token batches from the streamed blend via a token buffer.
+    Docs are EOS-separated and packed; each batch is B contiguous T+1 windows."""
+    eos = tok.eos_token_id
+    sep = [eos] if eos is not None else []
+    need = args.batch * (args.block + 1)
+    buf, docs = [], doc_generator(args)
+    srcs = _blend_sources(args)
+    tot = sum(s[3] for s in srcs)
+    blend = " / ".join(f"{int(round(100*s[3]/tot))}% {s[1]}" for s in srcs)
+    print(f"streaming blend: {blend}")
+    while True:
+        while len(buf) < need:
+            buf.extend(tok.encode(next(docs), add_special_tokens=False))
+            buf.extend(sep)
+        chunk = np.array(buf[:need], dtype=np.int64).reshape(args.batch, args.block + 1)
+        del buf[:need]
+        x = torch.from_numpy(chunk[:, :-1]).to(device)
+        y = torch.from_numpy(chunk[:, 1:]).to(device)
+        yield x, y
+
+
+def save_ckpt(model, tok_vocab, step, path, opt=None):
+    if not path:
+        return
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    ckpt = {"model": model.state_dict(), "family_config": model.family_config,
+            "vocab_size": tok_vocab, "step": step}
+    if opt is not None:
+        ckpt["optim"] = opt.state_dict()         # so --resume continues the optimizer too
+    torch.save(ckpt, path)
+    print(f"  saved checkpoint -> {path} (step {step})")
+
+
+def find_latest_ckpt(out_path):
+    """Auto-locate a checkpoint to resume from: prefer the exact --out file, else
+    the most recently modified *.pt in the same folder. Returns None if none."""
+    if out_path and os.path.isfile(out_path):
+        return out_path
+    folder = os.path.dirname(out_path) or "."
+    if not os.path.isdir(folder):
+        return None
+    pts = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".pt")]
+    return max(pts, key=os.path.getmtime) if pts else None
+
+
+def lr_at(step, peak, warmup, total, min_frac):
+    if step < warmup:
+        return peak * step / max(warmup, 1)
+    if step >= total:
+        return peak * min_frac
+    prog = (step - warmup) / max(total - warmup, 1)
+    return peak * (min_frac + (1 - min_frac) * 0.5 * (1 + math.cos(math.pi * prog)))
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data", default="")
+    p.add_argument("--tok-dir", default=PKG_DIR, help="dir with bundled tokenizer.json + spike_tokenizer.py")
+    p.add_argument("--steps", type=int, default=200)
+    p.add_argument("--batch", type=int, default=8)
+    p.add_argument("--block", type=int, default=256)
+    p.add_argument("--lr", type=float, default=3e-4)          # family peak
+    p.add_argument("--warmup", type=int, default=1000)
+    p.add_argument("--min-lr-frac", type=float, default=0.3)  # family floor (~9e-5)
+    p.add_argument("--weight-decay", type=float, default=0.01)
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--use-hrm", action="store_true")
+    p.add_argument("--use-moe", action="store_true")
+    p.add_argument("--use-mtp", action="store_true")
+    p.add_argument("--mtp-layers", type=int, default=4, help="MTP draft-head depth (spec-decode)")
+    p.add_argument("--use-jepa", action="store_true")
+    p.add_argument("--use-rings", action="store_true")
+    p.add_argument("--use-ring-controllers", action="store_true")
+    p.add_argument("--use-ring-specialists", action="store_true",
+                   help="per-ring MoE memory specialists with test-time input/output stores")
+    p.add_argument("--use-fractal-phase-seed", action="store_true",
+                   help="seed oscillator phases from each token's Mandelbrot orbit (gated)")
+    p.add_argument("--out", default=os.path.join(PKG_DIR, "chkpt", "quazimoto.pt"))  # checkpoint file path
+    p.add_argument("--ckpt-every", type=int, default=250)
+    p.add_argument("--resume", action="store_true",
+                   help="auto-find the latest checkpoint in the --out folder and continue training")
+    # streaming blend (HF datasets); off by default so local runs stay dependency-free.
+    # Default 4-way mix: 35% Ultra-FineWeb-L3 / 25% FineWeb-Edu / 25% FineMath / 15% PretrainNew.
+    p.add_argument("--stream", action="store_true", help="stream the weighted dataset blend")
+    p.add_argument("--ultra-frac", type=float, default=0.35)
+    p.add_argument("--ultra-dataset", default="openbmb/Ultra-FineWeb-L3")
+    p.add_argument("--ultra-config", default="Ultra-FineWeb-L3-en-Multi-Style-Synthetic")
+    p.add_argument("--edu-frac", type=float, default=0.25)
+    p.add_argument("--fineweb-dataset", default="HuggingFaceFW/fineweb-edu")
+    p.add_argument("--fineweb-config", default="sample-10BT")
+    p.add_argument("--math-frac", type=float, default=0.25)
+    p.add_argument("--math-dataset", default="HuggingFaceTB/finemath")
+    p.add_argument("--math-config", default="finemath-4plus")
+    p.add_argument("--pretrain-frac", type=float, default=0.15)
+    p.add_argument("--pretrain-dataset", default="nvidia/Nemotron-Pretraining-Specialized-v1.1")
+    p.add_argument("--pretrain-config", default="Nemotron-Pretraining-Formal-Logic")
+    p.add_argument("--split", default="train")
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--amp", action="store_true",
+                   help="bf16 autocast (biggest single-GPU speedup; also frees memory)")
+    args = p.parse_args()
+
+    # single-GPU speed: TF32 matmuls (free, ~fp32 accuracy on Ampere+/Blackwell).
+    if args.device == "cuda":
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+    tok = load_tokenizer(args.tok_dir)
+    cfg = QuazimotoConfig(vocab_size=tok.vocab_size, block_size=args.block,
+                          use_hrm=args.use_hrm, use_moe=args.use_moe,
+                          use_mtp=args.use_mtp, mtp_layers=args.mtp_layers,
+                          use_jepa=args.use_jepa,
+                          use_rings=args.use_rings,
+                          use_ring_controllers=args.use_ring_controllers,
+                          use_ring_specialists=args.use_ring_specialists,
+                          use_fractal_phase_seed=args.use_fractal_phase_seed)
+    model = QuazimotoLM(cfg).to(args.device)
+    print(f"tokenizer: SpikeWhale (vocab {tok.vocab_size})")
+
+    if args.stream:
+        batches = stream_batches(tok, args, args.device)     # infinite blended stream
+        train = None
+    else:
+        train, _ = load_ids(args.data, tok)
+        print(f"data: {len(train)} local tokens | device {args.device}")
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95),
+                            weight_decay=args.weight_decay,
+                            fused=(args.device == "cuda"))   # fused CUDA optimizer step
+
+    start_step = 1
+    if args.resume:
+        ckpt_path = find_latest_ckpt(args.out)
+        if ckpt_path is None:
+            print(f"--resume: no checkpoint found near {args.out}; starting fresh.")
+        else:
+            ck = torch.load(ckpt_path, map_location=args.device, weights_only=False)
+            miss, unexp = model.load_state_dict(ck["model"], strict=False)
+            if miss:   print(f"  [resume warn] missing keys: {len(miss)} (e.g. {miss[:2]})")
+            if unexp:  print(f"  [resume warn] unexpected keys: {len(unexp)} (e.g. {unexp[:2]})")
+            if "optim" in ck:
+                try:
+                    opt.load_state_dict(ck["optim"])
+                except ValueError as e:
+                    print(f"  [resume warn] optimizer state not restored ({e}); using fresh optimizer.")
+            start_step = int(ck.get("step", 0)) + 1
+            print(f"resumed from {ckpt_path} at step {ck.get('step')} -> continuing at {start_step}")
+            if start_step > args.steps:
+                print(f"  already at/past --steps ({args.steps}); nothing to do.")
+
+    model.train()
+    t0 = time.time()
+    for step in range(start_step, args.steps + 1):
+        lr = lr_at(step, args.lr, args.warmup, args.steps, args.min_lr_frac)
+        for g in opt.param_groups:
+            g["lr"] = lr
+        x, y = (next(batches) if args.stream
+                else get_batch(train, args.batch, args.block, args.device))
+        # bf16 autocast: ~1.5-2x on Blackwell tensor cores. bf16 keeps fp32 range so
+        # no grad scaler is needed, and it halves activation memory (room for bigger
+        # --batch). Params stay fp32; the oscillator trig/erf math is safe in bf16.
+        with torch.autocast("cuda", dtype=torch.bfloat16,
+                            enabled=args.amp and args.device == "cuda"):
+            _, loss, aux = model(x, y)
+            total = loss + sum(aux.values())     # add active trait aux losses
+        opt.zero_grad(set_to_none=True)
+        total.backward()
+        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        if not torch.isfinite(gn):                # family grad-finiteness guard
+            print(f"step {step}: non-finite grad norm, skipping")
+            opt.zero_grad(set_to_none=True)
+            continue
+        opt.step()
+        if step % 10 == 0 or step == 1:
+            dt = time.time() - t0
+            extra = " ".join(f"{k} {float(v):.3f}" for k, v in aux.items())
+            bank = getattr(model, "ring_bank", None)
+            surp = f" surprise {float(bank.last_surprise):.3f}" if bank is not None else ""
+            print(f"step {step:4d} | loss {loss.item():.3f} | "
+                  f"bpt {loss.item()/math.log(2):.3f} | lr {lr:.2e} | {extra}{surp} | {dt:.1f}s")
+        if args.ckpt_every and step % args.ckpt_every == 0:
+            save_ckpt(model, tok.vocab_size, step, args.out, opt)
+    save_ckpt(model, tok.vocab_size, args.steps, args.out, opt)
+    print("done.")
+
+
+if __name__ == "__main__":
+    main()

train_opd.py

@@ -0,0 +1,290 @@
+"""
+train_opd.py -- On-Policy Distillation of MobileLLM-R1-140M-base (teacher) into
+Quazimoto-LM (student), across different tokenizers.
+
+Loop (per step):
+  1. draw P text prompts from the streaming blend; take each prompt's first
+     --prompt-len student tokens.
+  2. the STUDENT samples G on-policy rollouts per prompt (its own tokens).
+  3. decode each rollout to text; the frozen TEACHER scores that text and returns a
+     per-BYTE log-prob; we re-aggregate it onto the student's tokens (byte-space
+     alignment in opd_teacher.Teacher.score) -> a DENSE per-student-token reward.
+  4. GRPO-normalize rewards within each prompt's group of G rollouts -> advantages.
+  5. policy-gradient update WITH a KL-to-reference anchor:
+        loss = -mean(advantage_t * logprob_t) + kl_beta * KL(student || frozen_ref)
+     The KL term is mandatory, not optional: unanchored on-policy RL against a reward
+     model reliably collapses into degenerate high-reward text (reward hacking) --
+     reward shoots past the teacher's own self-entropy while generations turn to mush.
+
+WATCH: teacher-reward/byte should rise then PLATEAU while kl stays bounded (small,
+non-exploding). If reward keeps climbing past the teacher's self-score AND kl blows
+up, lower --lr or raise --kl-beta. Do a ~300-step canary and eyeball a generation
+before committing to a long run.
+
+This keeps OPD's dense per-token teacher signal on the student's own samples while
+remaining tokenizer-agnostic (a token-level KL is impossible across the two vocabs).
+Resumable; checkpoints are interchangeable with train.py's format.
+
+Run on a rig with the teacher weights + transformers installed:
+  python train_opd.py --steps 5000 --prompts 4 --group 4 --gen-len 64 \
+      --student-ckpt chkpt/quazimoto.pt --device cuda --stream
+"""
+
+import argparse, math, os, time
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from model import QuazimotoLM, QuazimotoConfig
+import train as T
+
+PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+# --------------------------------------------------------------------------- #
+# student helpers
+# --------------------------------------------------------------------------- #
+def token_byte_lens(tok):
+    """Map student id -> number of bytes its decoding contributes (for alignment)."""
+    specials = set(getattr(tok, "_extra_specials", []))
+    ids2 = tok._ids_to_tokens
+    cache = {}
+    for tid, s in ids2.items():
+        if tid < 4 or s in specials:
+            cache[tid] = len(s.encode("utf-8"))
+        else:
+            cache[tid] = len(s.encode("latin-1"))
+    return cache
+
+
+@torch.no_grad()
+def sample_rollouts(model, cfg, idx, gen_len, temperature, top_k):
+    """Sample gen_len continuation tokens for a BATCH of equal-length starts.
+    idx: [B, Pl] on device. Returns [B, gen_len]. Uses the batched KV cache: the
+    starts are prefilled once, then each new token is a single-position forward.
+    Batching all prompts x rollouts here collapses what were P separate generation
+    loops into ONE, which is the main OPD throughput lever."""
+    past, gen = None, []
+    for _ in range(gen_len):
+        if past is not None and past[0][0].size(2) < cfg.block_size:
+            logits, _, _, past = model(idx[:, -1:], past_key_values=past, use_cache=True)
+        else:                                            # prefill / re-prime
+            logits, _, _, past = model(idx[:, -cfg.block_size:], use_cache=True)
+        lg = torch.nan_to_num(logits[:, -1, :].float(), nan=0.0, posinf=1e4, neginf=-1e4)
+        lg = lg / max(temperature, 1e-6)
+        if top_k:
+            v = torch.topk(lg, min(top_k, lg.size(-1)), dim=-1).values
+            lg = lg.masked_fill(lg < v[:, [-1]], float("-inf"))
+        nxt = torch.multinomial(torch.softmax(lg, dim=-1), 1)
+        idx = torch.cat([idx, nxt], dim=1)
+        gen.append(nxt)
+    return torch.cat(gen, dim=1)                          # [B, gen_len]
+
+
+def student_logprobs(model, cfg, full, prompt_len):
+    """Per-token log-probs (WITH grad) the student assigns to the rollout tokens,
+    batched. full: [B, Pl+gen_len] (prompt ++ rollout). Returns [B, gen_len].
+    Assumes Pl+gen_len <= block_size (the normal OPD regime)."""
+    B, L = full.shape
+    logits, _, _ = model(full[:, -cfg.block_size:])
+    logp = F.log_softmax(logits.float(), dim=-1)         # [B, L, V]
+    gl = L - prompt_len
+    tokens = full[:, prompt_len:]                        # [B, gen_len]
+    pred = logp[:, prompt_len - 1:L - 1, :]              # dist predicting each rollout token
+    return pred.gather(-1, tokens.unsqueeze(-1)).squeeze(-1)  # [B, gen_len]
+
+
+def rollout_token_byte_spans(blen, rollout):
+    """Byte span [a,b) each rollout token occupies in the decoded text -- O(T) via
+    the precomputed per-id byte lengths (byte tokens concatenate, so cumulative
+    lengths match the teacher's text bytes; the caller guards the tail with b<=nb
+    in case a trailing partial multi-byte char shifts the decoded length)."""
+    spans, prev = [], 0
+    for tid in rollout.tolist():
+        nb = prev + blen.get(int(tid), 0)
+        spans.append((prev, nb))
+        prev = nb
+    return spans
+
+
+# --------------------------------------------------------------------------- #
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--tok-dir", default=PKG_DIR)
+    p.add_argument("--student-ckpt", default=os.path.join(PKG_DIR, "chkpt", "quazimoto.pt"))
+    p.add_argument("--teacher", default="facebook/MobileLLM-R1-140M-base")
+    p.add_argument("--steps", type=int, default=2000)
+    p.add_argument("--prompts", type=int, default=4, help="distinct prompts per step (P)")
+    p.add_argument("--group", type=int, default=4, help="rollouts per prompt (G, GRPO group)")
+    p.add_argument("--prompt-len", type=int, default=32)
+    p.add_argument("--gen-len", type=int, default=64)
+    p.add_argument("--temperature", type=float, default=1.0)
+    p.add_argument("--top-k", type=int, default=50)
+    p.add_argument("--lr", type=float, default=1e-5)
+    p.add_argument("--warmup", type=int, default=200)
+    p.add_argument("--min-lr-frac", type=float, default=0.3)
+    p.add_argument("--weight-decay", type=float, default=0.0)
+    p.add_argument("--kl-beta", type=float, default=0.1,
+                   help="KL-to-reference anchor weight (prevents reward hacking; keep > 0)")
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--teacher-device", default="", help="default: same as --device")
+    p.add_argument("--out", default=os.path.join(PKG_DIR, "chkpt", "quazimoto_opd.pt"))
+    p.add_argument("--ckpt-every", type=int, default=200)
+    p.add_argument("--resume", action="store_true")
+    # streaming blend (prompts source) -- same flags/defaults as train.py
+    p.add_argument("--stream", action="store_true")
+    p.add_argument("--data", default="")
+    for name, default in [("ultra-frac", 0.35), ("edu-frac", 0.25), ("math-frac", 0.25),
+                          ("pretrain-frac", 0.15)]:
+        p.add_argument(f"--{name}", type=float, default=default)
+    p.add_argument("--ultra-dataset", default="openbmb/Ultra-FineWeb-L3")
+    p.add_argument("--ultra-config", default="Ultra-FineWeb-L3-en-Multi-Style-Synthetic")
+    p.add_argument("--fineweb-dataset", default="HuggingFaceFW/fineweb-edu")
+    p.add_argument("--fineweb-config", default="sample-10BT")
+    p.add_argument("--math-dataset", default="HuggingFaceTB/finemath")
+    p.add_argument("--math-config", default="finemath-4plus")
+    p.add_argument("--pretrain-dataset", default="Quazim0t0/PretrainNew")
+    p.add_argument("--pretrain-config", default=None)
+    p.add_argument("--split", default="train")
+    p.add_argument("--seed", type=int, default=0)
+    args = p.parse_args()
+
+    dev = args.device
+    tdev = args.teacher_device or dev
+    tok = T.load_tokenizer(args.tok_dir)
+
+    # student: rebuild from its checkpoint and continue training
+    ck = torch.load(args.student_ckpt, map_location=dev, weights_only=False)
+    cfg = QuazimotoConfig(**ck["family_config"])
+    student = QuazimotoLM(cfg)
+    student.load_state_dict(ck["model"], strict=False)
+    student.to(dev)
+    print(f"student: step {ck.get('step')} ckpt, vocab {cfg.vocab_size}")
+
+    # frozen REFERENCE policy = the pristine pre-OPD student. The KL anchor keeps the
+    # trained policy from drifting into degenerate high-reward text (reward hacking).
+    # Built from ck["model"] BEFORE any resume mutates `student`, so it stays stable.
+    reference = QuazimotoLM(cfg)
+    reference.load_state_dict(ck["model"], strict=False)
+    reference.to(dev).eval()
+    for pr in reference.parameters():
+        pr.requires_grad_(False)
+    # BUGFIX: the specialist memory stores are written on every forward regardless of
+    # train/eval mode, so without this the "frozen" reference silently drifts as it
+    # scores rollouts -> a moving anchor, weakening the KL. Freeze + reset its memory.
+    if hasattr(reference, "set_ring_memory_writing"):
+        reference.set_ring_memory_writing(False)
+        reference.reset_ring_memory()
+    print(f"reference: frozen pre-OPD student (KL anchor, beta={args.kl_beta})")
+
+    from opd_teacher import Teacher
+    teacher = Teacher(args.teacher, device=tdev)
+    print(f"teacher: {args.teacher} on {tdev} (frozen)")
+
+    blen = token_byte_lens(tok)  # id -> byte length, for O(T) rollout byte spans
+    opt = torch.optim.AdamW(student.parameters(), lr=args.lr, betas=(0.9, 0.95),
+                            weight_decay=args.weight_decay)
+
+    # OPD has its OWN step counter (independent of the student's pretraining step).
+    # Resume ONLY from this run's own output file -- never from the base checkpoint
+    # that may sit in the same folder (that's what --student-ckpt is for).
+    start_step = 1
+    if args.resume:
+        rp = args.out if os.path.isfile(args.out) else None
+        if rp:
+            r = torch.load(rp, map_location=dev, weights_only=False)
+            student.load_state_dict(r["model"], strict=False)
+            if "optim" in r:
+                try: opt.load_state_dict(r["optim"])
+                except ValueError: pass
+            start_step = int(r.get("step", 0)) + 1
+            print(f"resumed OPD from {rp} at step {r.get('step')} -> {start_step}")
+        else:
+            print(f"--resume: no prior OPD checkpoint at {args.out}; starting OPD at step 1")
+    if start_step > args.steps:
+        print(f"  nothing to do: OPD already at step {start_step-1} >= --steps {args.steps}. "
+              f"Raise --steps to train more."); return
+
+    # prompt source: first --prompt-len student tokens of streamed docs
+    if args.stream:
+        docs = T.doc_generator(args)
+    else:
+        ids_all, _ = T.load_ids(args.data, tok)
+        docs = None
+
+    def next_prompt():
+        if docs is not None:
+            while True:
+                t = tok.encode(next(docs), add_special_tokens=False)
+                if len(t) >= args.prompt_len + 1:
+                    return t[:args.prompt_len]
+        i = np.random.randint(0, len(ids_all) - args.prompt_len - 1)
+        return ids_all[i:i + args.prompt_len].tolist()
+
+    student.train()
+    t0 = time.time()
+    for step in range(start_step, args.steps + 1):
+        lr = T.lr_at(step, args.lr, args.warmup, args.steps, args.min_lr_frac)
+        for g in opt.param_groups:
+            g["lr"] = lr
+
+        opt.zero_grad(set_to_none=True)
+        P, G, GL = args.prompts, args.group, args.gen_len
+
+        # ONE batched generation over all P*G rollouts (prompts repeated G times)
+        prompts = torch.tensor([next_prompt() for _ in range(P)], device=dev)  # [P, Pl]
+        idx0 = prompts.repeat_interleave(G, dim=0)                              # [P*G, Pl]
+        rolls = sample_rollouts(student, cfg, idx0, GL, args.temperature, args.top_k)
+
+        # dense per-token teacher reward via byte alignment (one teacher fwd per rollout)
+        rewards = torch.zeros(P * G, GL, device=dev)
+        for b in range(P * G):
+            text = tok.decode(rolls[b].tolist(), skip_special_tokens=False)
+            byte_lp, nb = teacher.score(text)
+            byte_lp = byte_lp.to(dev)
+            for j, (a, e) in enumerate(rollout_token_byte_spans(blen, rolls[b])):
+                if e > a and e <= nb:
+                    rewards[b, j] = byte_lp[a:e].mean()
+
+        # GRPO: normalize within each prompt's group of G rollouts
+        rw = rewards.view(P, G, GL)
+        adv = ((rw - rw.mean(dim=(1, 2), keepdim=True))
+               / (rw.std(dim=(1, 2), keepdim=True) + 1e-6)).view(P * G, GL)
+
+        # ONE batched student forward+backward for the policy gradient
+        full = torch.cat([idx0, rolls], dim=1)                                 # [P*G, Pl+GL]
+        lp = student_logprobs(student, cfg, full, args.prompt_len)             # [P*G, GL] (grad)
+
+        # KL-to-reference anchor: penalize divergence of the trained policy from the
+        # frozen pre-OPD student. Without this, on-policy RL against the teacher reward
+        # collapses into degenerate high-reward text. k3 estimator (low-variance,
+        # non-negative) of KL(student || reference), per token.
+        with torch.no_grad():
+            ref_lp = student_logprobs(reference, cfg, full, args.prompt_len)   # [P*G, GL]
+        logr = (ref_lp - lp)
+        kl = (torch.exp(logr) - logr - 1.0).mean()                            # >= 0
+
+        pg = -(adv.detach() * lp).mean()
+        loss = pg + args.kl_beta * kl
+        loss.backward()
+        step_loss = float(pg.detach()); step_reward = float(rewards.mean()) * P
+        step_kl = float(kl.detach())
+
+        gn = torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
+        if torch.isfinite(gn):
+            opt.step()
+        else:
+            print(f"step {step}: non-finite grad, skip")
+
+        if step % 10 == 0 or step == 1:
+            dt = time.time() - t0
+            print(f"step {step:5d} | pg-loss {step_loss:+.4f} | teacher-reward/byte "
+                  f"{step_reward/args.prompts:+.3f} | kl {step_kl:.3f} | lr {lr:.2e} | {dt:.1f}s")
+        if args.ckpt_every and step % args.ckpt_every == 0:
+            T.save_ckpt(student, tok.vocab_size, step, args.out, opt)
+    T.save_ckpt(student, tok.vocab_size, args.steps, args.out, opt)
+    print("done.")
+
+
+if __name__ == "__main__":
+    main()

train_sft.bat

@@ -0,0 +1,16 @@
+@echo off
+REM Mycel-LM SFT: fine-tune the pretrained mycelium base (chkpt/quazimoto.pt) on the
+REM 4-way chat blend (ultrachat + ultrafeedback + UltraData + OpenThoughts).
+REM --init-ckpt loads the base; --resume continues from chkpt/quazimoto_sft.pt.
+REM   If you OOM: drop --batch to 4, or --block to 512.
+cd /d "%~dp0"
+python train_sft.py ^
+  --init-ckpt chkpt/quazimoto.pt ^
+  --out chkpt/quazimoto_sft.pt ^
+  --steps 4000 ^
+  --batch 8 ^
+  --block 1024 ^
+  --lr 2e-5 ^
+  --stream ^
+  --resume
+pause

train_sft.py

@@ -0,0 +1,235 @@
+"""
+train_sft.py -- supervised fine-tuning of Quazimoto-LM on UltraChat (chat SFT).
+
+Fine-tunes the pretrained base checkpoint on a blended chat-SFT mix (mlabonne
+ultrachat_200k_sft + ultrafeedback-sft + openbmb UltraData-SFT-2605/Knowledge +
+volcanos OpenThoughts2-1M-ShortThink reasoning), normalizing ChatML and ShareGPT
+schemas, rendered in the student's ChatML, with loss MASKED to assistant turns (std
+instruction SFT -- the model learns to produce assistant responses, not to parrot
+the user). Plain next-token cross-entropy -> stable, no reward model, no hacking.
+
+Masking detail: with a byte-merge tokenizer you must encode the whole conversation
+ONCE (encoding turn-by-turn would change byte-merges at the boundaries), then derive
+the mask by walking the atomic ChatML special tokens: loss is on for everything
+after an <|assistant|> marker up to and including the next <|im_end|>.
+
+Example:
+  python train_sft.py --init-ckpt chkpt/quazimoto.pt --steps 4000 --batch 8 \
+      --block 1024 --lr 2e-5 --stream --device cuda
+"""
+import argparse, math, os, time
+import numpy as np
+import torch
+
+from model import QuazimotoLM, QuazimotoConfig
+import train as T
+
+PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+ROLE_TOK = {"system": "<|system|>", "user": "<|user|>", "assistant": "<|assistant|>"}
+
+
+def make_renderer(tok):
+    vocab = tok.get_vocab()
+    asst_id = vocab.get("<|assistant|>")
+    imend_id = vocab.get("<|im_end|>")
+    if asst_id is None or imend_id is None:
+        raise ValueError("tokenizer missing ChatML special tokens (<|assistant|>/<|im_end|>)")
+
+    def render(messages, max_len):
+        """conversation -> (ids, loss_mask); loss_mask=1 on assistant response tokens."""
+        parts = []
+        for m in messages:
+            role = ROLE_TOK.get(m.get("role", "user"), "<|user|>")
+            parts.append(f"<|im_start|>{role}\n{m.get('content','')}<|im_end|>\n")
+        ids = tok.encode("".join(parts), add_special_tokens=False)[:max_len]
+        mask, in_asst = [], False
+        for tid in ids:
+            if tid == asst_id:                       # header marker -> no loss on it
+                in_asst = True; mask.append(0); continue
+            mask.append(1 if in_asst else 0)         # response tokens (incl the closing im_end)
+            if tid == imend_id and in_asst:
+                in_asst = False
+        return ids, mask
+    return render
+
+
+_ROLE_ALIAS = {"human": "user", "user": "user", "gpt": "assistant",
+               "assistant": "assistant", "model": "assistant", "bot": "assistant",
+               "system": "system"}
+
+
+def _messages_of(ex):
+    """Pull and NORMALIZE a conversation to [{role,content},...] from a dataset row,
+    tolerating both ChatML ({role,content}) and ShareGPT ({from,value}) schemas and
+    different field names (messages / conversations / conversation / chosen)."""
+    for k in ("messages", "conversations", "conversation", "chosen"):
+        v = ex.get(k)
+        if not (isinstance(v, list) and v and isinstance(v[0], dict)):
+            continue
+        out = []
+        for m in v:
+            if "content" in m:                      # ChatML {role, content}
+                role, content = m.get("role", "user"), m.get("content", "")
+            elif "value" in m:                      # ShareGPT {from, value}
+                role, content = m.get("from", "user"), m.get("value", "")
+            else:
+                continue
+            out.append({"role": _ROLE_ALIAS.get(str(role).lower(), "user"),
+                        "content": content or ""})
+        if out:
+            return out
+    return None
+
+
+def sft_stream(dataset, config, split):
+    """Infinite stream of chat conversations (list-of-messages) from one SFT dataset
+    (config may be None for single-config datasets)."""
+    from datasets import load_dataset
+    while True:
+        ds = load_dataset(dataset, name=config, split=split, streaming=True)
+        any_y = False
+        for ex in ds:
+            msgs = _messages_of(ex)
+            if msgs:
+                any_y = True; yield msgs
+        if not any_y:
+            raise RuntimeError(f"{dataset}:{config}:{split} yielded no conversations")
+
+
+def blended_convos(sources, seed=0):
+    """Blend several SFT datasets by per-conversation probability.
+    sources: list of (dataset, config, split, weight)."""
+    import random
+    rng = random.Random(seed)
+    gens = [sft_stream(d, c, s) for (d, c, s, _) in sources]
+    weights = [w for (_, _, _, w) in sources]
+    tot = sum(weights)
+    cum, acc = [], 0.0
+    for w in weights:
+        acc += w / tot; cum.append(acc)
+    print("SFT blend: " + " / ".join(f"{int(round(100*w/tot))}% {d}"
+                                     for (d, _, _, w) in sources))
+    while True:
+        r = rng.random()
+        yield next(gens[next(k for k, c in enumerate(cum) if r <= c)])
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--tok-dir", default=PKG_DIR)
+    p.add_argument("--init-ckpt", default=os.path.join(PKG_DIR, "chkpt", "quazimoto.pt"),
+                   help="pretrained base checkpoint to fine-tune from")
+    p.add_argument("--steps", type=int, default=4000)
+    p.add_argument("--batch", type=int, default=8)
+    p.add_argument("--block", type=int, default=1024)
+    p.add_argument("--lr", type=float, default=2e-5, help="SFT lr (lower than pretraining)")
+    p.add_argument("--warmup", type=int, default=100)
+    p.add_argument("--min-lr-frac", type=float, default=0.1)
+    p.add_argument("--weight-decay", type=float, default=0.0)
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--out", default=os.path.join(PKG_DIR, "chkpt", "quazimoto_sft.pt"))
+    p.add_argument("--ckpt-every", type=int, default=250)
+    p.add_argument("--resume", action="store_true")
+    p.add_argument("--stream", action="store_true", help="stream the SFT blend")
+    # blended SFT sources (default ~equal quarters)
+    p.add_argument("--ultrachat-frac", type=float, default=0.25)
+    p.add_argument("--ultrachat-dataset", default="mlabonne/ultrachat_200k_sft")
+    p.add_argument("--ultrachat-config", default=None)
+    p.add_argument("--ultrachat-split", default="train_sft")
+    p.add_argument("--feedback-frac", type=float, default=0.25)
+    p.add_argument("--feedback-dataset", default="mlabonne/ultrafeedback-sft")
+    p.add_argument("--feedback-config", default=None)
+    p.add_argument("--feedback-split", default="train")
+    p.add_argument("--ultradata-frac", type=float, default=0.25)
+    p.add_argument("--ultradata-dataset", default="openbmb/UltraData-SFT-2605")
+    p.add_argument("--ultradata-config", default="Knowledge")
+    p.add_argument("--ultradata-split", default="no_think")
+    p.add_argument("--thoughts-frac", type=float, default=0.25)
+    p.add_argument("--thoughts-dataset", default="volcanos/OpenThoughts2-1M-ShortThink")
+    p.add_argument("--thoughts-config", default=None)
+    p.add_argument("--thoughts-split", default="train")
+    p.add_argument("--seed", type=int, default=0)
+    args = p.parse_args()
+
+    torch.manual_seed(args.seed)
+    dev = args.device
+    tok = T.load_tokenizer(args.tok_dir)
+    pad = tok.pad_token_id or 0
+    render = make_renderer(tok)
+
+    ck = torch.load(args.init_ckpt, map_location=dev, weights_only=False)
+    cfg = QuazimotoConfig(**ck["family_config"])
+    model = QuazimotoLM(cfg); model.load_state_dict(ck["model"], strict=False)
+    model.to(dev)
+    print(f"SFT from {args.init_ckpt} (step {ck.get('step')}) | block {args.block} | UltraChat")
+
+    sources = [
+        (args.ultrachat_dataset, args.ultrachat_config, args.ultrachat_split, args.ultrachat_frac),
+        (args.feedback_dataset,  args.feedback_config,  args.feedback_split,  args.feedback_frac),
+        (args.ultradata_dataset, args.ultradata_config, args.ultradata_split, args.ultradata_frac),
+        (args.thoughts_dataset,  args.thoughts_config,  args.thoughts_split,  args.thoughts_frac),
+    ]
+    sources = [s for s in sources if s[3] > 0]
+    convos = blended_convos(sources, args.seed)
+
+    def sft_batch():
+        rows = []
+        while len(rows) < args.batch:
+            ids, mask = render(next(convos), args.block)
+            if sum(mask) > 0:                        # keep only convos with an assistant turn
+                rows.append((ids, mask))
+        L = max(len(r[0]) for r in rows)
+        x = torch.full((args.batch, L), pad, dtype=torch.long)
+        y = torch.full((args.batch, L), -1, dtype=torch.long)   # -1 = ignore
+        for j, (ids, mask) in enumerate(rows):
+            t = torch.tensor(ids)
+            x[j, :len(ids)] = t
+            # predict token i+1 at position i; loss only where target is assistant content
+            for i in range(len(ids) - 1):
+                if mask[i + 1]:
+                    y[j, i] = ids[i + 1]
+        return x.to(dev), y.to(dev)
+
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95),
+                            weight_decay=args.weight_decay)
+
+    start_step = 1
+    if args.resume and os.path.isfile(args.out):
+        r = torch.load(args.out, map_location=dev, weights_only=False)
+        model.load_state_dict(r["model"], strict=False)
+        if "optim" in r:
+            try: opt.load_state_dict(r["optim"])
+            except ValueError: pass
+        start_step = int(r.get("step", 0)) + 1
+        print(f"resumed SFT from {args.out} at step {r.get('step')} -> {start_step}")
+    if start_step > args.steps:
+        print(f"  nothing to do: already at {start_step-1} >= --steps {args.steps}."); return
+
+    model.train()
+    t0 = time.time()
+    for step in range(start_step, args.steps + 1):
+        for g in opt.param_groups:
+            g["lr"] = T.lr_at(step, args.lr, args.warmup, args.steps, args.min_lr_frac)
+        x, y = sft_batch()
+        _, loss, aux = model(x, y)
+        total = loss + sum(aux.values())
+        opt.zero_grad(set_to_none=True)
+        total.backward()
+        gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        if torch.isfinite(gn):
+            opt.step()
+        else:
+            print(f"step {step}: non-finite grad, skip"); continue
+        if step % 10 == 0 or step == 1:
+            dt = time.time() - t0
+            ntok = int((y != -1).sum())
+            print(f"step {step:5d} | sft-loss {loss.item():.3f} | ppl {math.exp(min(loss.item(),20)):.1f} "
+                  f"| asst-toks {ntok} | lr {opt.param_groups[0]['lr']:.2e} | {dt:.1f}s")
+        if args.ckpt_every and step % args.ckpt_every == 0:
+            T.save_ckpt(model, tok.vocab_size, step, args.out, opt)
+    T.save_ckpt(model, tok.vocab_size, args.steps, args.out, opt)
+    print("done ->", args.out)
+
+
+if __name__ == "__main__":
+    main()

visualize.py

@@ -0,0 +1,287 @@
+"""
+visualize.py -- live visualizer for Mycel-LM: watch the fungal colony grow in 3D as
+the model generates, token by token.
+
+Captures per generated token (via instrument.Recorder, which MycelBlock writes to):
+  * tip POSITIONS (all pos_dim axes) for one layer's colony -- the 3D scatter
+  * per-tip growth TRAIL over the growth steps -- the hyphae extending this token
+  * per-layer colony density + trait-station sites + attention + token stream
+
+Renders one self-contained HTML dashboard (data embedded -> opens via file://). The
+colony is a Three.js network you can drag-rotate / scroll-zoom: hyphal tips linked to
+their nearest neighbours as filaments (so it reads as a mycelial web, not loose dots),
+tips coloured by local density, with trait STATIONS as orange wire-spheres and faint
+grey growth trails. Three.js loads from a CDN (needs network the first time).
+
+Usage:
+  python visualize.py --prompt "the mycelium spreads" --tokens 50
+"""
+import argparse, json, os, sys, webbrowser
+import torch
+
+from model import QuazimotoLM, QuazimotoConfig
+import instrument
+
+PKG_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+def find_ckpt(path):
+    if path and os.path.isfile(path):
+        return path
+    folder = (os.path.dirname(path) if path else os.path.join(PKG_DIR, "chkpt")) or "."
+    if not os.path.isdir(folder):
+        return None
+    pts = [os.path.join(folder, f) for f in os.listdir(folder) if f.endswith(".pt")]
+    return max(pts, key=os.path.getmtime) if pts else None
+
+
+def load_tokenizer(tok_dir):
+    sys.path.insert(0, tok_dir)
+    from spike_tokenizer import SpikeTokenizer
+    return SpikeTokenizer(vocab_file=os.path.join(tok_dir, "tokenizer.json"))
+
+
+@torch.no_grad()
+def run_capture(model, cfg, tok, ids, n_tokens, temperature, top_k, device):
+    rec = instrument.Recorder(phase_layer=cfg.n_layer // 2, attn_layer=cfg.n_layer // 2)
+    instrument.set_rec(rec)
+    idx = torch.tensor([ids], device=device)
+    try:
+        for _ in range(n_tokens):
+            cond = idx[:, -cfg.block_size:]
+            rec.begin()
+            logits, _, _ = model(cond)
+            lg = torch.nan_to_num(logits[:, -1, :].float(), nan=0.0,
+                                  posinf=1e4, neginf=-1e4) / max(temperature, 1e-6)
+            if top_k:
+                v = torch.topk(lg, min(top_k, lg.size(-1))).values
+                lg = lg.masked_fill(lg < v[:, [-1]], float("-inf"))
+            probs = torch.softmax(lg, dim=-1)
+            nxt = (torch.argmax(lg, -1, keepdim=True) if temperature <= 1e-3
+                   else torch.multinomial(probs, 1))
+            top = torch.topk(torch.softmax(logits[:, -1].float(), -1), 5)
+            rec.end(token=int(nxt),
+                    char=tok.decode([int(nxt)], skip_special_tokens=False),
+                    top=[[int(t), round(float(p), 3)] for t, p in zip(top.indices[0], top.values[0])])
+            idx = torch.cat([idx, nxt], dim=1)
+    finally:
+        instrument.set_rec(None)
+    return rec.frames, idx[0].tolist()
+
+
+def main():
+    p = argparse.ArgumentParser(description="Mycel-LM 3D colony visualizer")
+    p.add_argument("--ckpt", default="")
+    p.add_argument("--tok_dir", default=PKG_DIR)
+    p.add_argument("--prompt", default="the mycelium spreads")
+    p.add_argument("--tokens", type=int, default=50)
+    p.add_argument("--temperature", type=float, default=0.0)
+    p.add_argument("--top_k", type=int, default=40)
+    p.add_argument("--device", default="cpu")
+    p.add_argument("--out", default=os.path.join(PKG_DIR, "viz.html"))
+    p.add_argument("--no_open", action="store_true")
+    args = p.parse_args()
+
+    path = find_ckpt(args.ckpt)
+    if path is None:
+        print("No checkpoint found; pass --ckpt or train one."); return
+    ckpt = torch.load(path, map_location=args.device, weights_only=False)
+    cfg = QuazimotoConfig(**ckpt["family_config"])
+    model = QuazimotoLM(cfg); model.load_state_dict(ckpt["model"], strict=False)
+    model.to(args.device).eval()
+    tok = load_tokenizer(args.tok_dir)
+    print(f"loaded {path} (step {ckpt.get('step')}) | capturing {args.tokens} tokens ...")
+
+    ids = tok.encode(args.prompt, add_special_tokens=False)
+    frames, _ = run_capture(model, cfg, tok, ids, args.tokens,
+                            args.temperature, args.top_k or None, args.device)
+
+    pl = cfg.n_layer // 2
+    blk = model.layers[pl].quaz
+    stations = blk.stations.anchors.detach().cpu().tolist() if getattr(blk, "use_stations", False) else []
+
+    data = {
+        "meta": {"ckpt": os.path.basename(path), "step": ckpt.get("step"),
+                 "n_layer": cfg.n_layer, "n_tips": cfg.n_tips, "pos_dim": cfg.mycel_pos_dim,
+                 "bound": cfg.osc_bound, "phase_layer": pl, "prompt": args.prompt,
+                 "stations": stations,
+                 "tip_rings": bool(getattr(cfg, "use_tip_rings", False)),
+                 "tip_ring_size": getattr(cfg, "tip_ring_size", 0)},
+        "prompt_chars": [tok.decode([t], skip_special_tokens=False) for t in ids],
+        "frames": frames,
+    }
+    html = HTML_TEMPLATE.replace("/*__DATA__*/", json.dumps(data))
+    with open(args.out, "w", encoding="utf-8") as f:
+        f.write(html)
+    print(f"wrote {args.out} ({os.path.getsize(args.out)//1024} KB)")
+    if not args.no_open:
+        webbrowser.open("file://" + os.path.abspath(args.out))
+
+
+HTML_TEMPLATE = r"""<!doctype html>
+<html><head><meta charset="utf-8"><title>Mycel-LM live</title>
+<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
+<script src="https://cdn.jsdelivr.net/npm/three@0.128.0/examples/js/controls/OrbitControls.js"></script>
+<style>
+  :root{--bg:#0d1117;--panel:#161b22;--ink:#e6edf3;--mut:#8b949e;--ac:#58a6ff;--hot:#f0883e;--spore:#3fb950}
+  *{box-sizing:border-box} body{margin:0;background:var(--bg);color:var(--ink);font:13px/1.4 ui-monospace,Menlo,Consolas,monospace}
+  header{padding:10px 16px;background:var(--panel);border-bottom:1px solid #30363d;display:flex;gap:16px;align-items:center;flex-wrap:wrap}
+  h1{font-size:15px;margin:0;color:var(--spore)} .mut{color:var(--mut)}
+  #wrap{display:grid;grid-template-columns:320px 1fr 300px;gap:12px;padding:12px}
+  .panel{background:var(--panel);border:1px solid #30363d;border-radius:8px;padding:10px}
+  .panel h2{font-size:12px;margin:0 0 8px;color:var(--mut);text-transform:uppercase;letter-spacing:.5px}
+  #stream{line-height:1.9;max-height:120px;overflow:auto}
+  .tk{padding:1px 2px;border-radius:3px;cursor:pointer;white-space:pre}
+  .tk.prompt{color:var(--mut)} .tk.cur{background:var(--spore);color:#000} .tk:hover{background:#30363d}
+  .ctrl{display:flex;gap:10px;align-items:center} input[type=range]{width:260px}
+  button{background:#21262d;color:var(--ink);border:1px solid #30363d;border-radius:6px;padding:4px 10px;cursor:pointer}
+  button:hover{border-color:var(--spore)} svg{display:block;width:100%;height:auto}
+  #colony3d{width:100%;height:440px;border-radius:6px;overflow:hidden}
+  .bar{height:10px;background:#21262d;border-radius:5px;overflow:hidden;margin:2px 0}.bar>div{height:100%}
+  table{border-collapse:collapse;width:100%}td{padding:1px 4px}.lbl{color:var(--mut);width:52px}
+  .grid{display:grid;gap:3px}.cell{height:15px;border-radius:2px}.small{font-size:11px;color:var(--mut)}
+</style></head><body>
+<header><h1>Mycel-LM &middot; colony 3D</h1><span class="mut" id="meta"></span>
+  <span style="flex:1"></span>
+  <div class="ctrl"><button id="play">▶ play</button><input type="range" id="seek" min="0" value="0"><span id="pos" class="mut"></span></div>
+</header>
+<div id="wrap">
+  <div style="display:flex;flex-direction:column;gap:12px">
+    <div class="panel"><h2>Token stream</h2><div id="stream"></div></div>
+    <div class="panel"><h2>Trait activity (this token)</h2><div id="traits"></div></div>
+    <div class="panel"><h2>Top predictions</h2><div id="top"></div></div>
+  </div>
+  <div style="display:flex;flex-direction:column;gap:12px">
+    <div class="panel"><h2>Colony &mdash; layer <span id="pl"></span> (drag rotate, scroll zoom; wire-spheres = trait stations)</h2>
+      <div id="colony3d"></div></div>
+    <div class="panel"><h2>Colony density per layer (bright = dense/clustered)</h2><div id="heat"></div></div>
+  </div>
+  <div style="display:flex;flex-direction:column;gap:12px">
+    <div class="panel"><h2>Attention &mdash; layer <span id="al"></span> (last token &rarr; context)</h2><div id="attn"></div></div>
+    <div class="panel"><h2>Notes</h2><div class="small" id="notes"></div></div>
+  </div>
+</div>
+<script>
+const D=/*__DATA__*/; const M=D.meta,F=D.frames,NL=M.n_layer,NT=M.n_tips,PD=M.pos_dim,BND=M.bound;
+let i=0,playing=false,timer=null; const $=id=>document.getElementById(id);
+$("meta").textContent=`${M.ckpt} · step ${M.step} · ${NT} tips · ${PD}D colony · ${NL} layers · "${M.prompt}"`;
+$("pl").textContent=M.phase_layer; $("al").textContent=M.phase_layer;
+const seek=$("seek"); seek.max=F.length-1;
+const hue2=v=>`hsl(${(1-Math.max(0,Math.min(1,v)))*140},70%,${30+45*Math.max(0,Math.min(1,v))}%)`;
+
+// hsl(green->orange) -> [r,g,b] in 0..1 for Three.js vertex colours (density mode)
+function colRGB(v){v=Math.max(0,Math.min(1,v));const h=(1-v)*140/360,s=0.7,l=0.32+0.42*v;
+  const a=s*Math.min(l,1-l),f=n=>{const k=(n+h*12)%12;return l-a*Math.max(-1,Math.min(k-3,9-k,1));};return[f(0),f(8),f(4)];}
+// ring PHASE -> [r,g,b]: hue = phase around the wheel, saturation/brightness = coherence r.
+// Synchronized tips share a hue (colony turns one colour); desynced tips are a rainbow.
+function phaseRGB(psi,r){const h=(psi/(2*Math.PI))+0.5,s=0.25+0.65*Math.max(0,Math.min(1,r)),l=0.5;
+  const a=s*Math.min(l,1-l),f=n=>{const k=(n+h*12)%12;return l-a*Math.max(-1,Math.min(k-3,9-k,1));};return[f(0),f(8),f(4)];}
+
+// ---- Three.js colony ----
+const KNN=3;                       // links per tip -> the mycelial web
+const MAXE=NT*KNN;                 // max filament segments
+let scene,camera,renderer,controls,geom,points,webGeom,web,trailGeom,trail;
+function init3d(){
+  const el=$("colony3d"),W=el.clientWidth||600,H=440;
+  scene=new THREE.Scene(); scene.background=new THREE.Color(0x0d1117);
+  camera=new THREE.PerspectiveCamera(55,W/H,0.1,1000); camera.position.set(BND*2.4,BND*1.7,BND*2.4);
+  renderer=new THREE.WebGLRenderer({antialias:true}); renderer.setSize(W,H); el.appendChild(renderer.domElement);
+  controls=new THREE.OrbitControls(camera,renderer.domElement); controls.enableDamping=true; controls.target.set(0,0,0);
+  scene.add(new THREE.LineSegments(new THREE.EdgesGeometry(new THREE.BoxGeometry(2*BND,2*BND,2*BND)),
+    new THREE.LineBasicMaterial({color:0x30363d})));
+  // filament web: each tip linked to its nearest neighbours -> looks like hyphae, not dots
+  webGeom=new THREE.BufferGeometry();
+  webGeom.setAttribute('position',new THREE.BufferAttribute(new Float32Array(MAXE*2*3),3));
+  webGeom.setAttribute('color',new THREE.BufferAttribute(new Float32Array(MAXE*2*3),3));
+  web=new THREE.LineSegments(webGeom,new THREE.LineBasicMaterial({vertexColors:true,transparent:true,opacity:0.55}));
+  scene.add(web);
+  // growth trails: each tip's path over the growth steps (the hypha extending this token)
+  const TS=(F[0]&&F[0].phases&&F[0].phases.traj)?F[0].phases.traj.length:0;
+  if(TS>1){trailGeom=new THREE.BufferGeometry();
+    trailGeom.setAttribute('position',new THREE.BufferAttribute(new Float32Array(NT*(TS-1)*2*3),3));
+    trail=new THREE.LineSegments(trailGeom,new THREE.LineBasicMaterial({color:0x8b949e,transparent:true,opacity:0.35}));
+    scene.add(trail);}
+  geom=new THREE.BufferGeometry();
+  geom.setAttribute('position',new THREE.BufferAttribute(new Float32Array(NT*3),3));
+  geom.setAttribute('color',new THREE.BufferAttribute(new Float32Array(NT*3),3));
+  points=new THREE.Points(geom,new THREE.PointsMaterial({size:BND*0.07,vertexColors:true}));
+  scene.add(points);
+  (M.stations||[]).forEach(s=>{const mesh=new THREE.Mesh(new THREE.SphereGeometry(BND*0.06,10,10),
+    new THREE.MeshBasicMaterial({color:0xf0883e,wireframe:true}));
+    mesh.position.set(s[0]||0,s[1]||0,s[2]||0); scene.add(mesh);});
+  window.addEventListener('resize',()=>{const w=el.clientWidth||600;camera.aspect=w/H;camera.updateProjectionMatrix();renderer.setSize(w,H);});
+  (function loop(){requestAnimationFrame(loop);controls.update();renderer.render(scene,camera);})();
+}
+function updateColony(f){
+  if(!geom||!f.phases)return;
+  const th=f.phases.theta,tp=f.phases.tip_psi,tr=f.phases.tip_r;
+  const pos=geom.attributes.position.array,col=geom.attributes.color.array;
+  const P=[]; for(let t=0;t<NT;t++)P.push([th[t*PD],th[t*PD+1],PD>2?th[t*PD+2]:0]);
+  let colr;
+  if(tp){                                  // colour by ring phase (synchronization view)
+    colr=P.map((_,t)=>phaseRGB(tp[t],tr?tr[t]:0.6));
+  }else{                                   // colour by local density (no rings)
+    const r2=(BND*0.25)**2, dens=P.map(a=>{let d=0;P.forEach(b=>{const dx=a[0]-b[0],dy=a[1]-b[1],dz=a[2]-b[2];if(dx*dx+dy*dy+dz*dz<r2)d++;});return d;});
+    const mx=Math.max(...dens,1); colr=dens.map(d=>colRGB(d/mx));
+  }
+  for(let t=0;t<NT;t++){pos[t*3]=P[t][0];pos[t*3+1]=P[t][1];pos[t*3+2]=P[t][2];
+    const c=colr[t];col[t*3]=c[0];col[t*3+1]=c[1];col[t*3+2]=c[2];}
+  geom.attributes.position.needsUpdate=true; geom.attributes.color.needsUpdate=true;
+  // rebuild the filament web: connect each tip to its KNN nearest neighbours
+  if(webGeom){const wp=webGeom.attributes.position.array,wc=webGeom.attributes.color.array;let e=0;
+    for(let a=0;a<NT&&e<MAXE;a++){
+      const nb=[]; for(let b=0;b<NT;b++){if(b===a)continue;
+        const dx=P[a][0]-P[b][0],dy=P[a][1]-P[b][1],dz=P[a][2]-P[b][2];nb.push([dx*dx+dy*dy+dz*dz,b]);}
+      nb.sort((x,y)=>x[0]-y[0]);
+      for(let n=0;n<Math.min(KNN,nb.length)&&e<MAXE;n++){const b=nb[n][1],ca=colr[a],cb=colr[b];
+        wp[e*6]=P[a][0];wp[e*6+1]=P[a][1];wp[e*6+2]=P[a][2];
+        wp[e*6+3]=P[b][0];wp[e*6+4]=P[b][1];wp[e*6+5]=P[b][2];
+        wc[e*6]=ca[0];wc[e*6+1]=ca[1];wc[e*6+2]=ca[2];wc[e*6+3]=cb[0];wc[e*6+4]=cb[1];wc[e*6+5]=cb[2];e++;}}
+    webGeom.setDrawRange(0,e*2);webGeom.attributes.position.needsUpdate=true;webGeom.attributes.color.needsUpdate=true;}
+  // growth trails: draw each tip's path across the captured growth steps
+  if(trail&&f.phases.traj){const tj=f.phases.traj,TS=tj.length,tp2=trailGeom.attributes.position.array;let g=0;
+    for(let t=0;t<NT;t++)for(let s=0;s<TS-1;s++){
+      const A=tj[s],Bp=tj[s+1];
+      tp2[g++]=A[t*PD];tp2[g++]=A[t*PD+1];tp2[g++]=PD>2?A[t*PD+2]:0;
+      tp2[g++]=Bp[t*PD];tp2[g++]=Bp[t*PD+1];tp2[g++]=PD>2?Bp[t*PD+2]:0;}
+    trailGeom.setDrawRange(0,NT*(TS-1)*2);trailGeom.attributes.position.needsUpdate=true;}
+}
+
+const stream=$("stream");
+D.prompt_chars.forEach(c=>{const s=document.createElement("span");s.className="tk prompt";s.textContent=esc(c);stream.appendChild(s);});
+F.forEach((f,k)=>{const s=document.createElement("span");s.className="tk gen";s.textContent=esc(f.char);s.onclick=()=>{i=k;render()};s.dataset.k=k;stream.appendChild(s);});
+function esc(c){return c.replace(/\n/g,"⏎").replace(/ /g,"·");}
+function heat(f){
+  let h=`<div class="grid" style="grid-template-columns:auto repeat(${NL},1fr)"><div class="small"></div>`;
+  for(let l=0;l<NL;l++)h+=`<div class="small" style="text-align:center">L${l}</div>`;
+  h+=`<div class="small">density</div>`; const mx=Math.max(...f.rings.map(r=>r.R[0]),1e-6);
+  for(let l=0;l<NL;l++){const s=f.rings[l]?f.rings[l].R[0]:0;h+=`<div class="cell" title="L${l} ${s.toFixed(3)}" style="background:${hue2(s/mx)}"></div>`;}
+  $("heat").innerHTML=h+"</div>";
+}
+function render(){
+  const f=F[i]; $("pos").textContent=`${i+1}/${F.length}`; seek.value=i;
+  document.querySelectorAll(".tk.gen").forEach(s=>s.classList.toggle("cur",+s.dataset.k===i));
+  updateColony(f); heat(f);
+  const qn=f.quaz_norm,mx=Math.max(...qn,1e-6);let L=qn.map((_,l)=>"grow L"+l),V=qn.map(v=>v/mx);
+  if(f.traits.hrm!=null){L.push("HRM");V.push(Math.min(1,f.traits.hrm/mx));}
+  if(f.traits.moe!=null){L.push("MoE");V.push(Math.min(1,f.traits.moe/mx));}
+  let th="<table>";V.forEach((v,k)=>{th+=`<tr><td class="lbl">${L[k]}</td><td style="width:100%"><div class="bar"><div style="width:${(v*100).toFixed(0)}%;background:${hue2(v)}"></div></div></td></tr>`;});
+  $("traits").innerHTML=th+"</table>";
+  let tp="<table>";f.top.forEach(([t,p])=>{tp+=`<tr><td class="lbl">${p.toFixed(2)}</td><td><div class="bar"><div style="width:${(p*100).toFixed(0)}%;background:var(--ac)"></div></div></td><td class="small">id ${t}</td></tr>`;});
+  $("top").innerHTML=tp+"</table>";
+  if(f.attn){const w=f.attn.w,m=Math.max(...w,1e-6);let h="<div style='display:flex;flex-wrap:wrap;gap:1px'>";
+    w.forEach((a,p)=>{h+=`<div title="pos ${p}: ${a.toFixed(3)}" style="width:8px;height:14px;background:${hue2(a/m)}"></div>`;});
+    $("attn").innerHTML=h+`</div><div class='small'>${w.length} context positions</div>`;}
+  else $("attn").innerHTML="<span class='small'>n/a</span>";
+  $("notes").innerHTML=`Tips: ${NT} · stations: ${(M.stations||[]).length}<br>Green = sparse growing edge, orange = dense core.<br>Filaments link each tip to its nearest neighbours (the mycelial web); faint grey trails are each tip's growth path this token.<br>Orange wire-spheres = trait stations. Drag to orbit, scroll to zoom.`;
+}
+seek.oninput=()=>{i=+seek.value;render()};
+$("play").onclick=()=>{playing=!playing;$("play").textContent=playing?"⏸ pause":"▶ play";if(playing)timer=setInterval(()=>{i=(i+1)%F.length;render();},400);else clearInterval(timer);};
+document.onkeydown=e=>{if(e.key==="ArrowRight"){i=Math.min(F.length-1,i+1);render();}if(e.key==="ArrowLeft"){i=Math.max(0,i-1);render();}};
+if(window.THREE){init3d();} else {$("colony3d").innerHTML="<div class='small' style='padding:20px'>Three.js failed to load (needs network for the CDN).</div>";}
+render();
+</script></body></html>"""
+
+
+if __name__ == "__main__":
+    main()