serve the full 2s arm library (single, stacked, routed, gain, memory and recurrent arms) with the best-measured arm as default; text-frame arms get their own frame on the specials craft

README.md

@@ -9,7 +9,7 @@ python_version: '3.12'
 app_file: app.py
 pinned: false
 license: mit
-short_description: Two byte-level AlephLLM crafts — arms and full splat
+short_description: Two byte-level AlephLLM crafts and their arm libraries
 ---
 
 # Beatrix — AlephLLM chat
@@ -23,9 +23,12 @@ adapters (3.2M–5.3M parameters), each trained on one exact frozen core.
 
 **mini-beatrix-2s** — 237.1M parameters, **full splat**: the signed
 address read runs in *every* block (v1 used it in three), governed from
-birth so the anchors never crowd. The active mission — still
-pretraining, no arms yet, and not yet through the phase that teaches a
-conversation frame, so **Completion** is the honest view of it.
+birth so the anchors never crowd. Its mission is complete (16.1 billion
+bytes); the closing anneal taught the conversation frame, so the core
+chats on its own. It carries its own **arm library** — single arms,
+stacks, routed dispatches and recurrent arms, each labelled with its
+measured score and how settled it is — and opens on the best-measured
+one.
 
 - **Chat** — conversation, in whichever frame the selected craft and arm
   were actually trained on.

app.py

@@ -12,8 +12,11 @@ own config, its own byte format, and its own arm library:
   mini-beatrix-2s  237.1M · FULL SPLAT (a governed multi-constellation
                    aleph read in every block) · the 0.8.0 SPECIALS format:
                    structural tokens live on the invalid-UTF-8 bytes, so
-                   the frame is unforgeable by content. Arms are read
-                   from the repo for both crafts — 2s has none yet.
+                   the frame is unforgeable by content. Its arm library
+                   holds several recipe types (single arms, stacks,
+                   routed dispatches, recurrent arms); arm_mount.py
+                   attaches any of them and detaches back to the bare
+                   core, and arms_catalog.json describes each one.
 
 Crafts are discovered from the repo at boot: a craft with no manifest or
 no checkpoint is skipped, and one that starts shipping appears on its
@@ -64,10 +67,7 @@ for _mod in list(sys.modules.values()):
 from huggingface_hub import HfApi, hf_hub_download
 from safetensors.torch import load_file
 
-import amoe
-from amoe.core.adapter import AdapterSpec, RelayPatchwork
-from amoe.io.checkpoint import load_anchor
-from amoe.runtime.attach import _per_block_state
+import arm_mount
 from geolip.alephllm.presets import AlephLMConfig
 from geolip.alephllm.model.alephlm import AlephLM
 from geolip.alephllm.amoe_bridge import CHAT_HEADER, USER_TAG, ASSISTANT_TAG
@@ -95,9 +95,8 @@ GEN_BUDGET_S = 95.0
 # for it are gated on manifest state (frame_trained) and switched
 # themselves off when anneal_mix went done. v1 stays one click away as
 # the detach exhibit: a core that does NOT chat plus a 3.2M arm that
-# makes it chat, removable bit-for-bit. Both crafts READ the arm index:
-# 2s has no arms today, and the day one is trained on a 2s core it
-# appears in the picker with no code change.
+# makes it chat, removable bit-for-bit. Both crafts READ the arm index;
+# 2s also reads arms_catalog.json, which overlays it.
 CRAFTS = [
     {"name": "mini-beatrix-1", "fmt": "v1tags", "wants_arms": True},
     {"name": "mini-beatrix-2s", "fmt": "specials", "wants_arms": True},
@@ -233,6 +232,20 @@ ALWAYS_OFFER = {"mini-beatrix-2s": [57607]}
 DEFAULT_ARM_NAME = {"mini-beatrix-1": "poly"}
 N_LISTED = 12          # newest N checkpoints offered in the dropdown
 
+# Space-side arm catalog (built from the training ledgers). For a craft it
+# lists, its rows REPLACE the hub index rows with the same path and add
+# the arms the hub index does not carry yet: stacks, routed dispatches,
+# recurrent arms. Each row names its mount recipe, frame, decode, status
+# and score, and at most one row per craft is marked the default — the
+# best-measured arm on that craft's default checkpoint.
+CATALOG_PATH = Path(__file__).with_name("arms_catalog.json")
+try:
+    CATALOG = json.load(open(CATALOG_PATH, encoding="utf-8"))["crafts"]
+except (OSError, ValueError, KeyError) as e:
+    print(f"[boot] no arm catalog ({type(e).__name__}) — hub index only",
+          flush=True)
+    CATALOG = {}
+
 
 def list_cores(files, craft):
     """{step: core_path} for every shipped checkpoint of one craft.
@@ -281,19 +294,43 @@ def load_arm_index(files, craft):
                                   else "behavior"),
                          "behavior": "", "spec": None, "sites": None,
                          "params": None})
+    over = (CATALOG.get(craft) or {}).get("rows") or []
+    if over:
+        mine = {r["path"] for r in over}
+        rows = [r for r in rows if r["path"] not in mine] + list(over)
+    # an arm is listed only when every file it mounts is on the hub:
+    # an index row can outlive its upload, and picking it would fail
+    # inside the GPU call instead of simply not being offered
+    present = set(files)
+    kept = [r for r in rows if all(p in present for p in _arm_files(r))]
+    if len(kept) < len(rows):
+        print(f"[boot] {craft}: {len(rows) - len(kept)} indexed arm(s) have "
+              "no file on the hub — not offered", flush=True)
     by_step = {}
-    for r in rows:
+    for r in kept:
         by_step.setdefault(int(r["step"]), []).append(r)
-    for s in by_step:                      # chat first, then turn-end
+    for s in by_step:   # catalog rank first, then chat, turn-end, behavior
         order = {"chat": 0, "turn-end": 1, "behavior": 2}
-        by_step[s].sort(key=lambda r: (order.get(r["kind"], 3),
+        by_step[s].sort(key=lambda r: (r.get("rank", 99),
+                                       order.get(r["kind"], 3),
                                        r.get("group", ""), r["name"]))
     return by_step
 
 
+def _arm_files(arm):
+    """The hub files one arm row mounts: its own checkpoint, or a stack's
+    members (a stack has no file of its own)."""
+    m = arm.get("mount") or {}
+    return list(m["members"]) if m.get("type") == "stack" else [arm["path"]]
+
+
 CHAT_STOPS = ("\nUser:", "\nBeatrix:", "User:", "Beatrix:")
 FRAME_TAG = {"chat": "chat transcript", "qa": "single-turn Q&A",
-             "raw": "raw text (no frame)"}
+             "text": "single-turn Q&A", "raw": "raw text (no frame)",
+             "specials": "specials frame"}
+# frames whose prompt is ONE plain-text question (no history): "qa" is the
+# v1 name and "text" the 2s ladder's name for the same bytes
+SINGLE_TURN = ("qa", "text", "raw")
 
 
 def _tmpl(arm):
@@ -311,6 +348,9 @@ def _tmpl(arm):
 
 
 def _arm_title(a):
+    if a.get("title"):
+        return a["title"] + (f" · {a['params']/1e6:.1f}M" if a.get("params")
+                             else "")
     who = f"{a['group']}/{a['name']}" if a.get("group") else a["name"]
     t = _tmpl(a)
     bits = [who, FRAME_TAG.get(t["frame"], t["frame"])]
@@ -324,11 +364,12 @@ def _arm_title(a):
     return " · ".join(bits)
 
 
-def _verify_anchor(local_arm, craft, step):
+def _verify_arm(local_arm, craft, step):
     """amoe's strict= check is inert for AlephLM (no model.config), so
-    provenance is asserted here from the anchor's own metadata."""
+    provenance is asserted here from the arm file's own metadata — every
+    format carries it (anchor, dispatch, relay-EMA, recurrent)."""
     try:
-        got = load_anchor(local_arm).meta.get("base_model_id")
+        got = arm_mount.base_model_id(local_arm)
     except Exception:
         return False
     want = f"alephllm/{craft}@step{step}"
@@ -339,19 +380,16 @@ def _verify_anchor(local_arm, craft, step):
     return True
 
 
-def _spec_of(ck):
-    """Read the adapter geometry from the anchor's OWN tensors. The arm
-    library holds more than one shape (n_slots 16 and 32, hidden 178 and
-    256), and amoe builds its modules from a spec before loading state —
-    assuming the default would raise on every wide arm."""
-    cb, proj = ck.adapters.get("0.addr.codebook"), ck.adapters.get(
-        "0.proj.weight")
-    cons = ck.adapters.get("0.consume.0.weight")
-    if cb is None or proj is None or cons is None:
-        return AdapterSpec()
-    d = int(cb.shape[1])
-    return AdapterSpec(n_slots=int(proj.shape[0] // d), K=int(cb.shape[0]),
-                       D=d, hidden=int(cons.shape[0]))
+def _mount(c, arm, step):
+    """Download, verify and attach one arm row to the BARE model. Returns
+    the handle, or None when any file fails provenance. Adapter geometry
+    is read from each file's own tensors (the library holds more than one
+    shape), so nothing is assumed from a default spec."""
+    local = [hf_hub_download(REPO, p) for p in _arm_files(arm)]
+    if not all(_verify_arm(p, c.name, step) for p in local):
+        return None
+    return arm_mount.mount(c.model, arm.get("mount") or {"type": "anchor"},
+                           local[0], local)
 
 
 class Craft:
@@ -454,42 +492,28 @@ class Craft:
         self.model.eval()
         self.core_params = self.model.param_count()
 
-        # attach ONCE so the adapter sites exist and ZeroGPU packs them;
-        # the template is THIS checkpoint's own default arm when it has
-        # one, so the common path needs no rebuild. Selecting an arm of
-        # another geometry rebuilds the site modules for that request.
-        self.wrappers, spec_now = [], None
-        tmpl = self.default_arm(self.step) or next(iter(self.arm_by_path),
-                                                   None)
-        if tmpl:
-            local = hf_hub_download(REPO, tmpl)
-            if _verify_anchor(local, self.name, self.arm_by_path[tmpl]["step"]):
-                ck = load_anchor(local)
-                spec_now = _spec_of(ck)
-                amoe.attach(self.model, ck, spec=spec_now, strict=False)
-                self.wrappers = [b for b in self.model.blocks
-                                 if hasattr(b, "adapter")]
-                print(f"[boot] {self.name}: adapter sites "
-                      f"{len(self.wrappers)} (template "
-                      f"{tmpl.split('/')[-1]}, n_slots={spec_now.n_slots}, "
-                      f"hidden={spec_now.hidden})", flush=True)
-
-        # core-key -> live-key map: attach wraps blocks, so a core
-        # checkpoint's "blocks.N.*" lands at "blocks.N.block.*".
+        # core-key -> live-tensor map, taken on the BARE model: a core
+        # reload always detaches the arm first, and detach restores these
+        # very block objects, so the map stays valid across every switch.
         self.live = dict(self.model.named_parameters())
         self.live.update(dict(self.model.named_buffers()))
-        self.keymap = {}
-        for k in sd:
-            if k in self.live:
-                self.keymap[k] = k
-            else:
-                alt = re.sub(r"^blocks\.(\d+)\.", r"blocks.\1.block.", k)
-                if alt in self.live:
-                    self.keymap[k] = alt
+        self.keymap = {k: k for k in sd if k in self.live}
+
+        # the default arm is mounted at import so the common first request
+        # needs no switch; every other arm mounts inside the GPU call.
+        # One handle at a time: an arm switch detaches back to the bare
+        # model before the next recipe attaches (recipes differ in shape —
+        # a stack, a routed dispatch and a recurrent arm share no sites).
+        self.handle, arm_now = None, None
+        want = self.default_arm(self.step)
+        if want:
+            self.handle = _mount(self, self.arm_by_path[want], self.step)
+            if self.handle is not None:
+                arm_now = want
+                print(f"[boot] {self.name}: default arm "
+                      f"{want.split('/')[-1]} mounted", flush=True)
         self.model.to("cuda")     # ZeroGPU packs here (standard pattern)
-        self.loaded = {"step": self.step,
-                       "arm_path": tmpl if self.wrappers else None,
-                       "spec": spec_now}
+        self.loaded = {"step": self.step, "arm_path": arm_now}
         print(f"[boot] {self.name}: {len(self.cores)} cores, "
               f"{len(self.arm_by_path)} arms, {self.core_params/1e6:.1f}M "
               f"core, ctx {self.cfg.context}, fmt {self.fmt}", flush=True)
@@ -512,6 +536,10 @@ class Craft:
         so nothing is misrepresented as a chat arm."""
         step = int(step)
         arms = self.arms.get(step, [])
+        if step == getattr(self, "step", None):
+            for a in arms:                  # the catalog's measured best
+                if a.get("default"):
+                    return a["path"]
         want = DEFAULT_ARM_NAME.get(self.name)
         if want and step == getattr(self, "step", None):
             for a in arms:
@@ -644,7 +672,25 @@ def _select(craft, step, arm_path):
     c = S(craft)
     step = _step(c, step)                   # stale/typed values fall back
     t0 = time.time()
+    # an arm belongs to exactly one core: a stale pick from another
+    # checkpoint (or another craft) is dropped rather than force-fitted
+    arm = _resolve_arm(c, arm_path, step)
+    want = arm["path"] if arm else None
+
+    # INVALIDATE FIRST, COMMIT LAST. This state survives the request (a
+    # ZeroGPU worker is reused, not forked per call), so a half-finished
+    # switch must never be recorded as a complete one. Any change of core
+    # or arm starts from the bare model: the handle is dropped and its
+    # bookkeeping cleared BEFORE detach, so if anything below raises, the
+    # next request finds "no arm loaded" and rebuilds rather than serving
+    # a stale mount under a new label. Core only = fully detached, so off
+    # is the untouched base model.
+    if c.handle is not None and (c.loaded["step"] != step
+                                 or c.loaded["arm_path"] != want):
+        h, c.handle, c.loaded["arm_path"] = c.handle, None, None
+        h.detach(verify=False)
     if c.loaded["step"] != step:
+        c.loaded["step"] = None
         sd = load_file(hf_hub_download(REPO, c.cores[step]))
         with torch.no_grad():
             for k, v in sd.items():
@@ -652,52 +698,16 @@ def _select(craft, step, arm_path):
                 if tgt is not None and tgt.shape == v.shape:
                     tgt.copy_(v.to(tgt.device, tgt.dtype))
         c.loaded["step"] = step
-
-    # an arm belongs to exactly one core: a stale pick from another
-    # checkpoint (or another craft) is dropped rather than force-fitted
-    arm = _resolve_arm(c, arm_path, step)
-    if arm is None:
-        arm_path = CORE_ONLY
-    arm_on = bool(c.wrappers) and bool(arm)
-    if arm_on and c.loaded.get("arm_path") != arm_path:
-        local = hf_hub_download(REPO, arm_path)
-        if _verify_anchor(local, c.name, step):   # provenance, every switch
-            ck = load_anchor(local)
-            spec = _spec_of(ck)
-            dev = next(c.model.parameters()).device
-            # INVALIDATE FIRST, COMMIT LAST. This state survives the
-            # request (a ZeroGPU worker is reused, not forked per call),
-            # so a partial load must never be recorded as a complete
-            # one: if the loop below raises, the wrappers hold a mix of
-            # two arms, and a loaded[arm_path] still naming the old arm
-            # would make the next request skip the reload and serve the
-            # chimera under the wrong provenance label.
-            prev_spec = c.loaded["spec"]
-            c.loaded["arm_path"] = None
-            if (spec.n_slots, spec.K, spec.D, spec.hidden) != (
-                    prev_spec.n_slots, prev_spec.K,
-                    prev_spec.D, prev_spec.hidden):
-                # different adapter geometry: rebuild the site modules
-                # for this request rather than forcing mismatched state
-                d_hidden = int(ck.adapters["0.proj.weight"].shape[1])
-                for w in c.wrappers:
-                    w.adapter = RelayPatchwork(d_hidden, spec).to(dev)
-                c.loaded["spec"] = spec
-            with torch.no_grad():
-                for i, w in enumerate(c.wrappers):
-                    st = _per_block_state(ck.adapters, i)
-                    if st:
-                        w.adapter.load_state_dict(
-                            {k: v.to(next(w.adapter.parameters()).device)
-                             for k, v in st.items()})
-            c.loaded["arm_path"] = arm_path       # commit only now
-        else:
-            arm_on, arm = False, None
-    for w in c.wrappers:
-        w.enabled = arm_on
+    if arm is not None and c.handle is None:
+        c.handle = _mount(c, arm, step)       # provenance, every switch
+        if c.handle is not None:
+            c.loaded["arm_path"] = want       # commit only now
+    arm_on = arm is not None and c.handle is not None
+    fp32 = arm_on and (arm.get("decode") or {}).get("precision") == "fp32"
     return {"craft": c.name, "checkpoint_step": step,
-            "arm": (arm_path.split("/")[-1] if arm_on else None),
+            "arm": (want.split("/")[-1] if arm_on else None),
             "arm_kind": (arm["kind"] if arm_on else "core-only"),
+            "precision": "fp32" if fp32 else "bf16",
             "switch_s": round(time.time() - t0, 2)}
 
 
@@ -716,22 +726,29 @@ def status_md(craft, step, arm_path=CORE_ONLY):
                     "`⟦SYS⟧ ⟦USER⟧ ⟦MODEL⟧ ⟦END⟧` frame a conversation. "
                     "They live on bytes UTF-8 can never produce, so "
                     "nothing you type can forge them.")
-    elif not c.wrappers:
-        arm_line = "**Arm:** adapter sites unavailable — core only"
     elif arm:
         sp = arm.get("spec") or {}
         t = _tmpl(arm)
-        dt, dp, _ = _decode_of(arm)
+        dt, dp, dr = _decode_of(arm)
+        kind = (arm.get("mount") or {}).get("type", "anchor")
         arm_line = (
             f"**Arm** `{arm['name']}`"
             + (f" ({arm['group']})" if arm.get("group") else "")
             + (f" · {arm['params']/1e6:.1f}M" if arm.get("params") else "")
             + (f" · {arm['sites']} sites" if arm.get("sites") else "")
-            + (f" · n{sp['n_slots']}/h{sp['hidden']}" if sp else "")
+            + (f" · n{sp['n_slots']}/h{sp['hidden']}"
+               if sp and kind == "anchor" else "")
+            + (f" · {MOUNT_TAG.get(kind, kind)}" if kind != "anchor" else "")
             + " · trained on THIS core"
+            + (f"  \n**{arm['status_word']}** · {arm['score']}"
+               if arm.get("score") else "")
             + f"  \n{FRAME_TAG.get(t['frame'], t['frame'])} · ends on "
               f"{t['turn_end']} · T={dt:g}/top-p {dp:g}"
+            + (f"/rep {dr:g}" if dr != 1.15 else "")
+            + (" · fp32" if (arm.get("decode") or {}).get("precision")
+               == "fp32" else "")
             + (f" · *{arm['behavior']}*" if arm.get("behavior") else "")
+            + (f"  \nℹ️ {arm['note']}" if arm.get("note") else "")
             + (f"  \n⚠️ {t['note']}" if t.get("note") else ""))
     else:
         arm_line = ("**Arm** none attached — raw core"
@@ -789,11 +806,17 @@ def status_md(craft, step, arm_path=CORE_ONLY):
          "conversational adapter"
          if arms and not any(a["kind"] == "chat" for a in arms)
          and (arm is None or arm["kind"] not in ("chat", "turn-end"))
-         else ""),
+         and not (c.fmt == "specials" and c.chats(step)) else ""),
+        ("ℹ️ this arm was trained on a **single-turn** frame, so Chat sends "
+         "one question at a time (no history) — that is the frame it "
+         "learned. *Core only* keeps the whole conversation"
+         if arm and _tmpl(arm)["frame"] in SINGLE_TURN
+         and c.fmt == "specials" else ""),
         ("ℹ️ this arm was trained on a **single-turn** frame, so Chat sends "
          "one question at a time (no history) — that is the frame it "
          "learned; the Completion tab is the plainer view"
-         if arm and _tmpl(arm)["frame"] == "qa" else ""),
+         if arm and _tmpl(arm)["frame"] in SINGLE_TURN
+         and c.fmt != "specials" else ""),
         ("⚠️ this is a **representation** arm with no turn convention: it "
          "was trained on raw lines, not conversations. Use Completion — "
          "Chat will not read as dialogue"
@@ -815,10 +838,27 @@ CORE_EX = [["Hello!"], ["Who are you?"], ["What can you do?"],
            ["How are you today?"], ["Tell me about books."]]
 
 
+MOUNT_TAG = {"stack": "stacked arms (always on)",
+             "routed": "routed dispatch", "relay_ema": "arm + byte memory",
+             "recurrent": "recurrent arm"}
+STATUS_WORD = {"best": "★ best trained", "settled": "replicated (two seeds)",
+               "one-seed": "one seed", "candidate": "candidate",
+               "refuted": "refuted", "control": "control row",
+               "superseded": "superseded", "in-progress": "still training",
+               "seed-unstable": "unstable across seeds",
+               "null": "no measured gain"}
+for _c in CATALOG.values():
+    for _r in _c.get("rows", []):
+        _r["status_word"] = STATUS_WORD.get(_r.get("status"), _r.get("status"))
+
+
 def _decode_of(arm):
+    """(temperature, top_p, repetition penalty) an arm asks for. Task arms
+    are measured greedy with no repetition penalty; the penalty slider's
+    1.15 default belongs to open chat on the bare core."""
     d = (arm or {}).get("decode") or {}
     return (float(d.get("temperature", 0.7)), float(d.get("top_p", 0.95)),
-            d.get("why", ""))
+            float(d.get("rep", 1.15)))
 
 
 def _examples_of(arm):
@@ -834,10 +874,11 @@ def on_arm(craft, step, arm_path, auto_decode, fam_examples):
     switchable — the checkboxes leave the controls alone when off."""
     c = S(craft)
     arm = _resolve_arm(c, arm_path, _step(c, step))
-    t, p, _ = _decode_of(arm)
+    t, p, r = _decode_of(arm)
     return (status_md(craft, step, arm_path),
             gr.update(value=t) if auto_decode else gr.update(),
             gr.update(value=p) if auto_decode else gr.update(),
+            gr.update(value=r) if auto_decode else gr.update(),
             gr.update(samples=_examples_of(arm)) if fam_examples
             else gr.update())
 
@@ -990,6 +1031,9 @@ def _stream(craft, ids, max_new, temperature, top_p, use_cache, arm_path,
     c = S(craft)
     model = c.model
     dev = next(model.parameters()).device
+    # arms decode in fp32, the precision every arm score was measured in;
+    # bf16 autocast stays the bare core's fast path
+    ac = dev.type == "cuda" and prov.get("precision") != "fp32"
     # The no-cache path recomputes the WHOLE context per byte, so on the
     # bigger craft it is minutes where the ZeroGPU budget is 120s. This
     # used to be a silent max_new = min(max_new, 96): the slider said 512
@@ -1011,13 +1055,14 @@ def _stream(craft, ids, max_new, temperature, top_p, use_cache, arm_path,
     label = (f"{prov['craft']} · step {prov['checkpoint_step']:,} · "
              + (f"arm `{prov['arm'].split('@')[0]}`" if prov["arm"]
                 else "core only")
+             + (" · fp32" if prov.get("precision") == "fp32" else "")
              + (f" · switched in {prov['switch_s']}s"
                 if prov["switch_s"] > 0.05 else ""))
+    getattr(c.handle, "reset_state", lambda: None)()   # fresh context
     with torch.no_grad():
         cache, prefill_s = None, 0.0
         if use_cache:
-            with torch.autocast("cuda", dtype=torch.bfloat16,
-                                enabled=dev.type == "cuda"):
+            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=ac):
                 logits, cache = model.prefill(tids)
             prefill_s = time.time() - t0
         t1 = time.time()
@@ -1031,8 +1076,7 @@ def _stream(craft, ids, max_new, temperature, top_p, use_cache, arm_path,
                 nxt = model._sample(_rep_penalty(logits, out, rep),
                                     temperature, top_p)
             else:
-                with torch.autocast("cuda", dtype=torch.bfloat16,
-                                    enabled=dev.type == "cuda"):
+                with torch.autocast("cuda", dtype=torch.bfloat16, enabled=ac):
                     fl, _ = model(tids[:, -c.cfg.context:])
                 nxt = model._sample(_rep_penalty(fl, out, rep),
                                     temperature, top_p)
@@ -1057,8 +1101,7 @@ def _stream(craft, ids, max_new, temperature, top_p, use_cache, arm_path,
                 return
             out.append(nid)
             if use_cache:
-                with torch.autocast("cuda", dtype=torch.bfloat16,
-                                    enabled=dev.type == "cuda"):
+                with torch.autocast("cuda", dtype=torch.bfloat16, enabled=ac):
                     logits = model.decode_step(nxt, cache)
             text = _render(out, c.fmt)
             cut = min((text.index(s) for s in stop if s in text),
@@ -1192,20 +1235,28 @@ def chat(craft, message, history, max_new, temperature, top_p, use_cache,
         return
     c = S(craft)
     step_i = _step(c, ckpt_step)
-    # On a plain-text craft the tags ARE the syntax, so a message
-    # containing "User:" has to be defanged. On a specials craft the
-    # frame lives on bytes text cannot contain, so the message needs no
-    # mangling — defanging it there would contradict the guarantee the
-    # status panel makes and quietly edit what the visitor typed.
-    clean = _clean_turn if c.fmt != "specials" else (
+    arm = _resolve_arm(c, arm_path, step_i)
+    tmpl = _tmpl(arm)
+    # The prompt is built in the frame the ARM was trained on. On the
+    # specials craft that is the specials frame for the bare core, but a
+    # single-turn text arm learned the plain-text tags — serving it the
+    # specials frame is a different tokenizer, and it answers as if never
+    # trained. (ARM TEMPLATE LAW: the template is the tokenizer.)
+    specials_frame = c.fmt == "specials" and (
+        arm is None or tmpl["frame"] not in SINGLE_TURN)
+    # Where plain-text tags ARE the syntax, a message containing "User:"
+    # has to be defanged. In the specials frame the structure lives on
+    # bytes text cannot contain, so the message needs no mangling —
+    # defanging it there would contradict the guarantee the status panel
+    # makes and quietly edit what the visitor typed.
+    clean = _clean_turn if not specials_frame else (
         lambda s: " ".join(str(s).split()))
     message = clean(message)
     history = [{"role": t["role"], "content": clean(_text(t["content"]))}
                for t in (history or [])]
-    tmpl = _tmpl(_resolve_arm(c, arm_path, step_i))
     stops, stop_ids = (), ()
     room = _room(c.cfg.context, max_new)
-    if c.fmt == "specials":
+    if specials_frame:
         # the frame's own terminators, as ids: END closes the turn, DOC
         # means she started a new document instead of answering. Built
         # to the SAME budget _fit enforces, so _fit never has to cut a
@@ -1213,6 +1264,8 @@ def chat(craft, message, history, max_new, temperature, top_p, use_cache,
         ids = _specials_chat_ids(history, message, c.cfg.context, room)
         stop_ids = (SP.END, SP.DOC, SP.USER)
     else:
+        if c.fmt == "specials":           # a frame byte is never a reply
+            stop_ids = (SP.END, SP.DOC, SP.USER)
         # the SELECTED arm's own frame and turn-end bytes, not one
         # hardcoded chat format: a Q&A arm never saw a transcript, and a
         # newline-pair arm's terminator is invisible to the "User:" stops
@@ -1235,7 +1288,7 @@ def chat(craft, message, history, max_new, temperature, top_p, use_cache,
                  "teaches the conversation frame — it has only ever read "
                  "documents. The Completion tab shows what it can do "
                  "today.)*")
-    elif tmpl["frame"] == "qa":
+    elif tmpl["frame"] in SINGLE_TURN:
         empty = ("*(this arm ended the turn immediately — it was trained "
                  "on single-turn task prompts, not conversation. Try one "
                  "of its examples below, or switch arms.)*")
@@ -1291,10 +1344,24 @@ address read runs in *every* block (v1 used it in three), with a
 governor holding the anchors apart from birth. Its mission is
 **complete** — 16.1 billion bytes over 61,422 steps, closing with a
 two-phase anneal that taught the conversation frame, so this core
-**chats on its own**: no arm is required, and none have been trained on
-it yet. Its format is specials-native — SYS/USER/MODEL/END are single
-reserved bytes rather than typed text, so the turn frame cannot be
-forged by anything you type.
+**chats on its own** — pick *Core only* for a multi-turn conversation.
+Its format is specials-native — SYS/USER/MODEL/END are single reserved
+bytes rather than typed text, so the turn frame cannot be forged by
+anything you type.
+
+**The 2s arm library.** Every arm trained on the final 2s core is
+selectable, each labelled with what it does, how settled it is
+(replicated across two seeds, one seed, candidate, refuted, control)
+and its measured score. The default is the best-measured arm: a
+rule-transfer arm that reads five if-then rules about words it has never
+seen and answers with only the final word, while still holding a
+conversation. Beyond single arms the library holds **stacks** (arms
+applied together, always on), **routed dispatches** (a learned signed
+router over several arms, some with a slowly varying per-arm gain) and
+**recurrent arms** (arms that carry state from byte to byte). These arms
+were trained on one question at a time, so Chat sends each message on its
+own with no history, and they decode greedy in fp32 exactly as they were
+measured.
 
 **mini-beatrix-1** — 112.5M parameters, pretrained on 15.3 billion raw
 bytes, carrying a **library of detachable arms**: small adapters
@@ -1405,14 +1472,15 @@ with gr.Blocks(title="Beatrix — AlephLLM chat") as demo:
         # auto-decode only fires on a change event, so a default arm that
         # asks for greedy was otherwise served sampled on first paint —
         # the exact setting its metadata says degenerates
-        _dt, _dp, _ = _decode_of(_C0.arm_by_path.get(_C0.default_arm(STEP)))
+        _dt, _dp, _dr = _decode_of(
+            _C0.arm_by_path.get(_C0.default_arm(STEP)))
         temperature = gr.Slider(0.0, 1.5, value=_dt, step=0.05,
                                 label="Temperature")
         top_p = gr.Slider(0.1, 1.0, value=_dp, step=0.05, label="Top-p")
         # byte-level: this damps whatever would CONTINUE a repeated
         # 8-byte phrase, not every byte already seen — on a 256-value
         # vocabulary the textbook form would just tax ordinary English
-        rep = gr.Slider(1.0, 2.0, value=1.15, step=0.05,
+        rep = gr.Slider(1.0, 2.0, value=_dr, step=0.05,
                         label="Repetition penalty",
                         info="damps bytes that would repeat a recent "
                              "8-byte phrase · 1.0 = off")
@@ -1420,8 +1488,8 @@ with gr.Blocks(title="Beatrix — AlephLLM chat") as demo:
     with gr.Accordion("Automation", open=False), gr.Row():
         auto_decode = gr.Checkbox(
             value=True, label="Auto decode per arm",
-            info="task arms switch to greedy on attach — they degenerate "
-                 "into punctuation when sampled off-distribution")
+            info="task arms switch to greedy with no repetition penalty on "
+                 "attach — the setting they were measured in")
         fam_examples = gr.Checkbox(
             value=True, label="Arm-family examples",
             info="offer prompts from the selected arm's own training "
@@ -1499,13 +1567,14 @@ with gr.Blocks(title="Beatrix — AlephLLM chat") as demo:
     # or nothing); decode and examples then follow the ARM, unless their
     # toggles say otherwise
     craft_pick.change(on_craft, [craft_pick, auto_decode, fam_examples],
-                      [ckpt, arm_pick, status_box, temperature, top_p,
+                      [ckpt, arm_pick, status_box, temperature, top_p, rep,
                        ex.dataset])
     ckpt.change(on_checkpoint, [craft_pick, ckpt, auto_decode, fam_examples],
-                [ckpt, arm_pick, status_box, temperature, top_p, ex.dataset])
+                [ckpt, arm_pick, status_box, temperature, top_p, rep,
+                 ex.dataset])
     arm_pick.change(on_arm,
                     [craft_pick, ckpt, arm_pick, auto_decode, fam_examples],
-                    [status_box, temperature, top_p, ex.dataset])
+                    [status_box, temperature, top_p, rep, ex.dataset])
     demo.load(js=RESIZE_JS)
 
 demo.launch()

arm_mount.py

@@ -0,0 +1,279 @@
+"""Arm mounting for the Beatrix Space.
+
+`mount(model, recipe, path)` makes the live model carry exactly one arm
+recipe and returns a handle whose `detach()` restores the bare blocks.
+Every recipe is served from the files it was trained into; nothing is
+re-derived at serve time.
+
+Recipe types (the `mount.type` field of an arm-index row):
+  anchor      one patchwork arm (an amoe anchor file)
+  stack       several anchor arms applied always-on, nested inner to
+              outer, with no mixer between them
+  routed      a dense signed dispatch over member arms; the checkpoint
+              carries the members and the router for every block
+  routed + gain
+              the same dispatch with a detached, bounded per-member gain
+              on the numerator, read per site or once at the first site
+              and shared by every site
+  relay_ema   the patchwork organ plus two fixed-decay causal averages
+              of its own read feeding its head
+  recurrent   a small recurrent cell per site fed by the attention read
+
+The relay_ema, recurrent and gain recipes carry state across cached
+decode steps; a fresh prefill starts that state over, and the uncached
+path treats every forward as a fresh context.
+
+Trained arms and their measurements:
+huggingface.co/AbstractPhil/alephllm-mini-beatrix-training
+(mini-beatrix-2s/arms/index.json and arms/dispatch_index.json).
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import amoe
+from amoe.binding.resolver import resolve
+from amoe.core.adapter import AdapterSpec, RelayPatchwork
+from amoe.core.dispatch import AnchorDispatch, BlockWithDispatch
+from amoe.io.checkpoint import load_anchor, load_dispatch
+from amoe.runtime.attach import AttachHandle, _probe
+
+MOUNT_TYPES = ("anchor", "stack", "routed", "relay_ema", "recurrent")
+
+
+# ------------------------------------------------------------- geometry
+def spec_from_state(adapters: dict, prefix: str | None = None) -> AdapterSpec:
+    """Adapter geometry read from the arm's own tensors (the arm library
+    holds more than one shape, and modules are built before state loads).
+    `prefix` defaults to the first block the arm actually carries."""
+    if prefix is None:
+        blocks = sorted(int(k.split(".")[0]) for k in adapters
+                        if k.endswith(".addr.codebook") and k.split(".")[0].isdigit())
+        prefix = f"{blocks[0]}." if blocks else "0."
+    cb = adapters.get(f"{prefix}addr.codebook")
+    proj = adapters.get(f"{prefix}proj.weight")
+    cons = adapters.get(f"{prefix}consume.0.weight")
+    if cb is None or proj is None or cons is None:
+        return AdapterSpec()
+    D = int(cb.shape[1])
+    return AdapterSpec(n_slots=int(proj.shape[0] // D), K=int(cb.shape[0]),
+                       D=D, hidden=int(cons.shape[0]))
+
+
+def _block_state(adapters: dict, i: int) -> dict:
+    pre = f"{i}."
+    return {k[len(pre):]: v for k, v in adapters.items() if k.startswith(pre)}
+
+
+def base_model_id(path: str) -> str | None:
+    """The core an arm file was trained on, from its own metadata."""
+    blob = torch.load(path, map_location="cpu", weights_only=False)
+    return (blob.get("meta") or {}).get("base_model_id")
+
+
+# ------------------------------------------------------------- handles
+class StackHandle:
+    """Nested always-on handles, detached outermost first."""
+
+    def __init__(self, handles):
+        self.handles = list(handles)
+
+    def reset_state(self):
+        for h in self.handles:
+            getattr(h, "reset_state", lambda: None)()
+
+    def detach(self, *, verify: bool = False):
+        for h in reversed(self.handles):
+            h.detach(verify=verify)
+
+
+# ------------------------------------------------------------- gain dispatch
+class GainDispatch(AnchorDispatch):
+    """AnchorDispatch with a detached, bounded per-member gain alpha on the
+    numerator (the denominator still sums every member, masked or not):
+
+        out = x + sum_a alpha_a(t) (sinh u_a / sum_all cosh u) (arm_a(x) - x)
+        s_a(t)  = cosh u_a / sum_all cosh u        (the member's share)
+        ubar_a  = causal EMA over bytes of s_a, rate rho, starting at 1/A
+        alpha_a = clamp(1 + gamma (A ubar_a - 1), a_min, a_max)
+
+    mode 'growth' lets the mean gain rise toward the cap; 'conserved'
+    renormalizes alpha to mean 1 over members at every byte.
+    coherent None reads alpha at every site; 'alpha' reads it once at the
+    first site and shares it; 'full' also shares the first site's signed
+    routing weight. Never a selector: no top-k, no softmax over members."""
+
+    def configure(self, *, mode, rho, gamma, a_min, a_max, coherent=None,
+                  role="ref", holder=None):
+        self.mode, self.rho, self.gamma = mode, rho, gamma
+        self.a_min, self.a_max = a_min, a_max
+        self.coherent, self.role = coherent, role
+        self.holder = holder if holder is not None else {}
+        self._state, self._call = None, "forward"
+        return self
+
+    def _alpha(self, ep, en, z):
+        with torch.no_grad():
+            s = ((ep + en) / z).float()
+            B, T, A = s.shape
+            st = (self._state if (self._call == "step" and self._state is not None
+                                  and self._state.shape[0] == B) else None)
+            if st is None:
+                st = torch.full((B, A), 1.0 / A, device=s.device, dtype=torch.float32)
+            if T == 1:
+                st = st + self.rho * (s[:, 0] - st)
+                ubar = st.unsqueeze(1)
+            elif T <= 4200:
+                q = 1.0 - self.rho
+                t = torch.arange(1, T + 1, device=s.device, dtype=torch.float64)
+                acc = torch.cumsum((q ** -t).view(1, T, 1) * s.double(), dim=1)
+                ubar = ((q ** t).view(1, T, 1)
+                        * (st.double().unsqueeze(1) + self.rho * acc)).float()
+                st = ubar[:, -1]
+            else:
+                ubar = torch.empty_like(s)
+                for i in range(T):
+                    st = st + self.rho * (s[:, i] - st)
+                    ubar[:, i] = st
+            self._state = st.clone() if self._call in ("prefill", "step") else None
+            alpha = (1.0 + self.gamma * (A * ubar - 1.0)).clamp(self.a_min, self.a_max)
+            if self.mode == "conserved":
+                alpha = alpha * (A / alpha.sum(dim=-1, keepdim=True))
+            return alpha.to(z.dtype)
+
+    def forward(self, x):
+        shared_route = self.coherent == "full" and self.role != "ref"
+        if not shared_route:
+            emb = F.normalize(x @ self.key_proj.T, dim=-1)
+            C = F.normalize(self.dispatch, dim=-1)
+            u = (emb @ C.T) / self.tau
+            m = u.abs().amax(dim=-1, keepdim=True)
+            ep, en = torch.exp(u - m), torch.exp(-u - m)
+            z = (ep + en).sum(dim=-1, keepdim=True)
+            g = (ep - en) / z
+        if self.coherent and self.role != "ref":
+            alpha = self.holder["alpha"]
+            if shared_route:
+                g = self.holder["g"]
+        else:
+            alpha = self._alpha(ep, en, z)
+            if self.coherent:
+                self.holder["alpha"], self.holder["g"] = alpha, g
+        out = x
+        for k, a in enumerate(self.anchors):
+            if self.enabled[k]:
+                out = out + (g[..., k:k + 1] * alpha[..., k:k + 1]) * (a(x) - x)
+        return out
+
+
+class GainBlock(BlockWithDispatch):
+    """Marks cached-decode calls so the gain state is carried across steps."""
+
+    def prefill(self, *args, **kwargs):
+        out, cache = self.block.prefill(*args, **kwargs)
+        self.disp._call = "prefill"
+        try:
+            return self.disp(out), cache
+        finally:
+            self.disp._call = "forward"
+
+    def step(self, *args, **kwargs):
+        out = self.block.step(*args, **kwargs)
+        self.disp._call = "step"
+        try:
+            if isinstance(out, tuple):
+                return (self.disp(out[0]),) + out[1:]
+            return self.disp(out)
+        finally:
+            self.disp._call = "forward"
+
+
+# ------------------------------------------------------------- mounts
+def _mount_anchor(model, path):
+    ck = load_anchor(path)
+    return amoe.attach(model, ck, spec=spec_from_state(ck.adapters), strict=False)
+
+
+def _mount_stack(model, paths):
+    handles = []
+    try:
+        for p in paths:
+            handles.append(_mount_anchor(model, p))
+    except Exception:
+        StackHandle(handles).detach()      # never leave half a stack attached
+        raise
+    return StackHandle(handles)
+
+
+def _mount_routed(model, path):
+    ck = load_dispatch(path)
+    b = resolve(model, None)
+    layers, d = list(b.layers(model)), b.hidden_size(model)
+    fingerprint = _probe(model, b)
+    names = list(ck.meta.get("anchors") or [])
+    A = len(names)
+    member0 = {k[len("anchors.0."):]: v for k, v in ck.dispatch[0].items()
+               if k.startswith("anchors.0.")}
+    spec = spec_from_state(member0, prefix="")
+    gain = ck.meta.get("alpha")
+    coherent = ck.meta.get("coherent_form")
+    holder, blocks = {}, []
+    for i, layer in enumerate(layers):
+        dev = next(layer.parameters()).device
+        members = nn.ModuleList([RelayPatchwork(d, spec) for _ in range(A)])
+        dp = AnchorDispatch(members, d)
+        dp.load_state_dict(ck.dispatch[i], strict=True)
+        dp.to(dev)
+        if gain:
+            dp.__class__ = GainDispatch
+            dp.configure(mode=ck.meta.get("alpha_mode") or "growth",
+                         rho=float(gain["rho"]), gamma=float(gain["gamma"]),
+                         a_min=float(gain["a_min"]), a_max=float(gain["a_max"]),
+                         coherent=coherent, role="ref" if i == 0 else "follow",
+                         holder=holder)
+            blocks.append(GainBlock(layer, dp))
+        else:
+            blocks.append(BlockWithDispatch(layer, dp))
+    b.set_layers(model, blocks)
+    return AttachHandle(model, b, layers, names, blocks, fingerprint)
+
+
+def _mount_relay_ema(model, path):
+    from amoe.arms.relay_ema import RelayEMASpec, attach_relay_ema
+    blob = torch.load(path, map_location="cpu", weights_only=False)
+    st, sp = blob["state"], blob.get("spec") or {}
+    spec = RelayEMASpec(base=spec_from_state(st), rho1=float(sp.get("rho1", 1 / 16)),
+                        rho2=float(sp.get("rho2", 1 / 64)))
+    name = (blob.get("meta") or {}).get("name", "relay_ema")
+    h = attach_relay_ema(model, name, spec, zero_bias=False)
+    with torch.no_grad():
+        for i, w in enumerate(h.wraps()):
+            w.arm.load_state_dict(_block_state(st, i), strict=True)
+    return h
+
+
+def _mount_recurrent(model, path):
+    from amoe.arms.recurrent_wide import RecurrentSpec, attach_recurrent
+    blob = torch.load(path, map_location="cpu", weights_only=False)
+    name = (blob.get("meta") or {}).get("name", "recurrent")
+    return attach_recurrent(model, name, RecurrentSpec(**blob["spec"]), state=blob["state"])
+
+
+def mount(model, recipe: dict, path: str, member_paths=None):
+    """Attach one arm recipe to the bare model. `path` is the local file of
+    the row's own checkpoint; `member_paths` are the local files of a
+    stack's members, inner first."""
+    kind = (recipe or {}).get("type", "anchor")
+    if kind == "anchor":
+        return _mount_anchor(model, path)
+    if kind == "stack":
+        return _mount_stack(model, member_paths)
+    if kind == "routed":
+        return _mount_routed(model, path)
+    if kind == "relay_ema":
+        return _mount_relay_ema(model, path)
+    if kind == "recurrent":
+        return _mount_recurrent(model, path)
+    raise ValueError(f"unknown mount type {kind!r}")

requirements.txt

@@ -4,4 +4,4 @@ safetensors>=0.4
 huggingface_hub>=0.24
 numpy
 geolip-alephllm @ git+https://github.com/AbstractEyes/alephllm@c6fff9fa3a9423a18d522635d8f2c666093d64c6
-amoe-lora @ git+https://github.com/AbstractEyes/amoe-lora@7e1baba9c0bb1cc38815e4424d969b80ae5a315d
+amoe-lora @ git+https://github.com/AbstractEyes/amoe-lora@e21baea08b065029f8bd1b80f7af15681e91ce3d