Spaces:
Running on Zero
Running on Zero
The arms fitted on the final mini-beatrix-3 core, the lookup arm's versions, and the four control twins of 2s. On mini-beatrix-3's final checkpoint (step 245,674, still opening bare) the arm list now holds the arm sets (stages 1-4, stages 1-8, the caption set, the research sets; two seeds each), served one row per set with a box per arm so any subset runs; the sixteen solo arms; and every version of the lookup arm (softmax attention after blocks 2, 6, 12 and 20, born at zero), mounted through amoe 0.2.12 from the arm's own tensors, which are read from the checkpoint by range request when no arm-only file exists. The lookup rows get recall prompts in the Completion tab and an activity line per site; the set rows get each member stage's prompts. The run's own stage programs (the checkpoints to step 148,000) use the same member switch. New crafts: mini-beatrix-2s-control (the original softmax twin) and its re-runs -attn-fp16-r10k (still training), -attn-fp16 and -bf16 (stopped), decoded in fp32. Pin: amoe-lora 5c934cf -> 9b1942d (v0.2.12); geolip-alephllm unchanged.
70c3e58 verified | """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 | |
| program a co-trained stage-arm program (mini-beatrix-3): one anchor | |
| arm per curriculum stage, trained together with the core, | |
| nested in stage order, each member switchable on its own; | |
| the arm sets fitted on the finished core mount the same way | |
| lookup a zero-born causal softmax-attention arm after chosen | |
| blocks (amoe.arms.lookup), read from the lookup_* tensors | |
| of a trainer checkpoint or of an arm-only file | |
| The relay_ema, recurrent, gain and lookup 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; the | |
| mini-beatrix-3 stage arms ship beside every checkpoint as | |
| mini-beatrix-3/arms/<arm>_step<N>.safetensors). | |
| """ | |
| from __future__ import annotations | |
| import contextlib | |
| import re | |
| 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", | |
| "program", "lookup") | |
| # ------------------------------------------------------------- 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 _load(path: str, *, weights_only: bool): | |
| """torch.load by CONTENT, never by name. torch >= 2.13 reads any path | |
| ending in .safetensors as a safetensors file, and the mini-beatrix-3 | |
| arm files are torch archives under that extension (amoe's | |
| AnchorCheckpoint.save is torch.save): loaded by path they fail with | |
| 'header too large'. An open file handle skips the name check.""" | |
| with open(path, "rb") as fh: | |
| return torch.load(fh, map_location="cpu", weights_only=weights_only) | |
| def _anchor(path: str): | |
| """amoe's load_anchor, fed an open handle for the same reason.""" | |
| with open(path, "rb") as fh: | |
| return load_anchor(fh) | |
| try: # amoe >= 0.2.6: block anchors may be real safetensors files | |
| from amoe.io.safetensors_io import is_safetensors | |
| except ImportError: # pragma: no cover - older pins: torch files only | |
| def is_safetensors(_src): | |
| return False | |
| def base_model_id(path: str) -> str | None: | |
| """The core an arm file was trained on, from its own metadata, read in | |
| whichever format the file actually is: a real safetensors block anchor | |
| through amoe's load_anchor (which reads either anchor format by | |
| content), every other arm file (anchor, dispatch, relay-EMA, recurrent | |
| torch archives) as the torch archive it is.""" | |
| if is_safetensors(path): | |
| return _anchor(path).meta.get("base_model_id") | |
| blob = _load(path, 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) | |
| def stage_index(phase) -> int | None: | |
| """curriculum_s3 -> 3; anything else -> None.""" | |
| m = re.fullmatch(r"curriculum_s(\d+)", str(phase or "")) | |
| return int(m.group(1)) if m else None | |
| class ProgramHandle(StackHandle): | |
| """A co-trained stage-arm program: every member is an anchor arm that | |
| trained in the same optimizer steps as the core, and each is nested | |
| around the blocks in the order the trainer attached it (stage order). | |
| Switching a member off sets its block wrappers to pass their input | |
| through unchanged, so the switch is an exact removal of that member's | |
| write, and the others read exactly what they would read without it. | |
| Nothing renormalizes. `detach` restores the bare blocks.""" | |
| def __init__(self, handles, names, meta): | |
| super().__init__(handles) | |
| self.names = list(names) | |
| self.meta = list(meta) | |
| self.by_name = dict(zip(self.names, self.handles)) | |
| self.live = {n: True for n in self.names} | |
| def set_live(self, live): | |
| want = set(live) | |
| for n, h in self.by_name.items(): | |
| h.set_mask({n: n in want}) | |
| self.live[n] = n in want | |
| def masked(self, name): | |
| """One member off for the duration (the others as they are).""" | |
| with self.by_name[name].all_off(): | |
| yield self | |
| def detach(self, *, verify: bool = False): | |
| # every member back on first: each member's detach check compares | |
| # against the model as it stood when that member attached, and | |
| # every member inside it was on then | |
| self.set_live(self.names) | |
| super().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 = _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_program(model, paths): | |
| """Nest every member in the order the trainer attached it: curriculum | |
| stage order, read from each anchor's own `phase` (the listing order is | |
| the fallback). The order matters because each member reads the block | |
| output the members inside it have already written to.""" | |
| cks = [_anchor(p) for p in paths] | |
| order = sorted(range(len(cks)), key=lambda i: ( | |
| stage_index(cks[i].meta.get("phase")) is None, | |
| stage_index(cks[i].meta.get("phase")) or 0, i)) | |
| handles, names, meta = [], [], [] | |
| try: | |
| for i in order: | |
| ck = cks[i] | |
| name = ck.meta.get("name") or f"arm{i}" | |
| handles.append(amoe.attach(model, ck, | |
| spec=spec_from_state(ck.adapters), | |
| strict=False)) | |
| names.append(name) | |
| meta.append({"name": name, "phase": ck.meta.get("phase"), | |
| "lambda": ck.meta.get("lambda"), | |
| "n_params": ck.meta.get("n_params")}) | |
| except Exception: | |
| StackHandle(handles).detach() # never leave half a program on | |
| raise | |
| return ProgramHandle(handles, names, meta) | |
| def _mount_routed(model, path): | |
| with open(path, "rb") as fh: | |
| ck = load_dispatch(fh) | |
| 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 = _load(path, 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 = _load(path, weights_only=False) | |
| name = (blob.get("meta") or {}).get("name", "recurrent") | |
| return attach_recurrent(model, name, RecurrentSpec(**blob["spec"]), state=blob["state"]) | |
| # ------------------------------------------------------------- lookup arm | |
| LOOKUP_PREFIXES = ("lookup_arms.", "lookup_norms.") | |
| _ST_DTYPES = {"BF16": torch.bfloat16, "F16": torch.float16, | |
| "F32": torch.float32, "F64": torch.float64} | |
| def _st_tensor(meta, raw): | |
| t = torch.frombuffer(bytearray(raw), dtype=_ST_DTYPES[meta["dtype"]]) | |
| return t.reshape(meta["shape"]).clone() | |
| def fetch_lookup_tensors(repo: str, path: str, cache_dir: str) -> str: | |
| """The lookup arm's tensors out of a trainer checkpoint on the hub, by | |
| HTTP range reads: the safetensors header, then only the lookup_* | |
| tensors' bytes (about 34 MB of a 790 MB file per attention layer set). | |
| Written once to a local safetensors file under cache_dir (atomically, | |
| so two workers never read a half-written file) and reused after.""" | |
| import json | |
| import os | |
| import struct | |
| import tempfile | |
| from huggingface_hub import HfFileSystem | |
| from safetensors.torch import save_file | |
| out = os.path.join(cache_dir, path.replace("/", "__") + ".lookup.safetensors") | |
| if os.path.exists(out): | |
| return out | |
| os.makedirs(cache_dir, exist_ok=True) | |
| with HfFileSystem().open(f"{repo}/{path}", "rb") as fh: | |
| n = struct.unpack("<Q", fh.read(8))[0] | |
| header = json.loads(fh.read(n)) | |
| base = 8 + n | |
| spans = sorted((header[k]["data_offsets"][0], header[k]["data_offsets"][1], k) | |
| for k in header if k.startswith(LOOKUP_PREFIXES)) | |
| if not spans: | |
| raise KeyError(f"{path}: the checkpoint carries no lookup arm") | |
| lo, hi = spans[0][0], spans[-1][1] | |
| tensors = {} | |
| if hi - lo <= 2 * sum(e - b for b, e, _ in spans): # one read | |
| fh.seek(base + lo) | |
| blob = fh.read(hi - lo) | |
| for b, e, k in spans: | |
| tensors[k] = _st_tensor(header[k], blob[b - lo:e - lo]) | |
| else: | |
| for b, e, k in spans: | |
| fh.seek(base + b) | |
| tensors[k] = _st_tensor(header[k], fh.read(e - b)) | |
| fd, tmp = tempfile.mkstemp(dir=cache_dir, suffix=".part") | |
| os.close(fd) | |
| save_file(tensors, tmp, metadata={"source": f"{repo}/{path}", | |
| "keys": str(len(tensors))}) | |
| os.replace(tmp, out) | |
| return out | |
| def _mount_lookup(model, path, name="lookup"): | |
| """A lookup arm: zero-born causal softmax attention after chosen blocks, | |
| each layer behind its own LayerNorm (amoe.arms.lookup). The geometry is | |
| read from the keys (sites, depth, heads from the q/k gains) and checked | |
| against the bound trunk; the attention runs in fp32 under autocast, as | |
| it was trained. Cached decode keeps the arm's own key/value cache.""" | |
| from safetensors.torch import load_file | |
| from amoe.arms.lookup import attach_from_alephllm, from_alephllm_state | |
| sd = {k: v for k, v in load_file(path).items() | |
| if k.startswith(LOOKUP_PREFIXES)} | |
| _, sites, depth = from_alephllm_state(sd) | |
| b = resolve(model, None) | |
| n_blocks, d = len(list(b.layers(model))), b.hidden_size(model) | |
| if max(sites) >= n_blocks: | |
| raise ValueError(f"lookup arm sites {sites} outside a {n_blocks}-block trunk") | |
| w = sd[f"lookup_arms.{sites[0]}.qkv.weight"] | |
| if tuple(w.shape) != (3 * d, d): | |
| raise ValueError(f"lookup arm width {tuple(w.shape)} does not fit d={d}") | |
| return attach_from_alephllm(model, sd, name=name, attn_fp32=True) | |
| 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, or of a stage program's members.""" | |
| kind = (recipe or {}).get("type", "anchor") | |
| if kind == "lookup": | |
| return _mount_lookup(model, path, (recipe or {}).get("name", "lookup")) | |
| if kind == "anchor": | |
| return _mount_anchor(model, path) | |
| if kind == "stack": | |
| return _mount_stack(model, member_paths) | |
| if kind == "program": | |
| return _mount_program(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}") | |