AbstractPhil's picture
mini-beatrix-2.5s: the 2s core (bit-identical) packaged with its arm library — 13 mountable rows over 11 weight files, a self-contained arm runtime verified bit-exact against amoe-lora, and an automodel that loads under current transformers
fc9a663 verified
Raw History Blame Contribute Delete
23.3 kB
"""arms.py — the detachable arm runtime for mini-beatrix-2.5s.
An *arm* is a small aleph-addressed patch head placed after every decoder
block of one exact frozen core. It is not a low-rank matrix: each site
projects the block output into `n_slots` slots of dimension D, reads them
through a K-atom aleph address (a closed-form signed read, never a
selector), and adds a gated patch back into the residual stream:
slots = proj(x).view(B, n, n_slots, D)
m_hat = sum_k sinh(u_k) A_k / sum_k cosh(u_k), u = (x_hat . A) / tau
x = x + sigmoid(gate) * consume(m_hat)
`mount()` makes the live core carry exactly one recipe and returns a
handle; `handle.detach(verify=True)` restores the untouched blocks and
asserts the restored logits are BIT-EXACT with a pre-attach fingerprint.
That exactness is the point of the design: an arm is a guest, never a
migration of the core.
Recipes carried by this package:
anchor one arm, always on
stack several anchor arms applied always-on, nested inner to outer,
with no mixer between them (the ungated composition)
routed a dense signed dispatch over member arms, optionally with a
detached bounded per-member gain read from a causal usage EMA
THE DAMPING LAW (do not "fix" it): the dispatch denominator sums cosh
over ALL members, including masked ones, and masking a member removes
its delta WITHOUT renormalizing. Masked is not solo.
This file is self-contained (torch + safetensors). It is a vendored,
verified re-expression of the amoe-lora runtime the arms were trained
with (github.com/AbstractEyes/amoe-lora): the packaged arms were checked
to produce bit-identical logits through both paths before release.
"""
from __future__ import annotations
import contextlib
import json
import os
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F
ARM_TYPES = ("anchor", "stack", "routed")
# --------------------------------------------------------------- geometry
@dataclass
class AdapterSpec:
n_slots: int = 16
K: int = 64
D: int = 4
tau: float = 0.1
hidden: int = 178
gate_init: float = -3.0
class AlephAddress(nn.Module):
"""m_hat(x) = sum_k sinh(u_k) A_k / sum_k cosh(u_k), max-|u| stabilized.
`home` snapshots the codebook at init and is required in every arm
state (the drift gauge)."""
def __init__(self, K: int = 64, D: int = 4, tau: float = 0.1):
super().__init__()
self.K, self.D, self.tau = K, D, tau
self.codebook = nn.Parameter(F.normalize(torch.randn(K, D), dim=-1))
self.register_buffer("home", self.codebook.detach().clone())
def m_hat(self, x: torch.Tensor) -> torch.Tensor:
A = F.normalize(self.codebook, dim=-1)
u = (F.normalize(x, dim=-1) @ A.transpose(-1, -2)) / self.tau
m = u.abs().amax(dim=-1, keepdim=True)
ep, en = torch.exp(u - m), torch.exp(-u - m)
return ((ep - en) @ A) / (ep + en).sum(dim=-1, keepdim=True)
@torch.no_grad()
def drift(self) -> float:
"""Mean angular drift (radians) of the codebook from `home`."""
A = F.normalize(self.codebook, dim=-1)
H = F.normalize(self.home, dim=-1)
return float(torch.arccos((A * H).sum(-1).clamp(-1.0, 1.0)).mean())
class SquaredReLU(nn.Module):
def forward(self, x):
return F.relu(x) ** 2
class RelayPatchwork(nn.Module):
"""One arm site: projection -> aleph address -> squared-ReLU patch head,
added through a sigmoid gate."""
def __init__(self, d: int, spec: AdapterSpec | None = None):
super().__init__()
s = spec or AdapterSpec()
self.spec = s
self.n_slots = s.n_slots
self.proj = nn.Linear(d, s.n_slots * s.D, bias=False)
self.addr = AlephAddress(s.K, s.D, s.tau)
self.consume = nn.Sequential(
nn.Linear(s.n_slots * s.D, s.hidden), SquaredReLU(),
nn.LayerNorm(s.hidden), nn.Linear(s.hidden, d))
self.gate = nn.Parameter(torch.tensor(float(s.gate_init)))
def forward(self, x):
B, n, _ = x.shape
slots = self.proj(x).view(B, n, self.n_slots, self.spec.D)
feats = self.addr.m_hat(slots).reshape(B, n, -1)
return x + torch.sigmoid(self.gate) * self.consume(feats)
def spec_from_state(state: dict, prefix: str | None = None) -> AdapterSpec:
"""Arm geometry read from the arm's own tensors (the library holds more
than one shape, and modules are built before the state loads)."""
if prefix is None:
blocks = sorted(int(k.split(".")[0]) for k in state
if k.endswith(".addr.codebook") and k.split(".")[0].isdigit())
prefix = f"{blocks[0]}." if blocks else "0."
cb = state.get(f"{prefix}addr.codebook")
proj = state.get(f"{prefix}proj.weight")
cons = state.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]))
# --------------------------------------------------------------- wrappers
class BlockWithAdapter(nn.Module):
"""Wraps one decoder block. `enabled=False` returns the block output
untouched — the single-arm half of the toggle law."""
def __init__(self, block: nn.Module, adapter: RelayPatchwork):
super().__init__()
self.block = block
self.adapter = adapter
self.enabled = True
def forward(self, *args, **kwargs):
out = self.block(*args, **kwargs)
if not self.enabled:
return out
if isinstance(out, tuple):
return (self.adapter(out[0]),) + out[1:]
return self.adapter(out)
def prefill(self, *args, **kwargs):
out, cache = self.block.prefill(*args, **kwargs)
return (self.adapter(out) if self.enabled else out), cache
def step(self, *args, **kwargs):
out = self.block.step(*args, **kwargs)
if not self.enabled:
return out
if isinstance(out, tuple):
return (self.adapter(out[0]),) + out[1:]
return self.adapter(out)
class AnchorDispatch(nn.Module):
"""Dense signed sinh/cosh dispatch over frozen member arms.
out = x + sum_a (sinh u_a / sum_all cosh u) (arm_a(x) - x)
No top-k, no softmax over members, no load balancing. The denominator
sums over every member, masked or not (the damping law)."""
def __init__(self, anchors: nn.ModuleList, d: int, emb: int = 64,
tau: float = 0.1):
super().__init__()
self.anchors = anchors
self.tau = tau
A = len(anchors)
self.register_buffer("key_proj", torch.zeros(emb, d))
self.dispatch = nn.Parameter(torch.zeros(A, emb))
self.enabled = [True] * A
def forward(self, x):
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)
w = ep - en
z = (ep + en).sum(dim=-1, keepdim=True) # ALL members (law)
out = x
for k, a in enumerate(self.anchors):
if self.enabled[k]:
out = out + (w[..., k:k + 1] / z) * (a(x) - x)
return out
class GainDispatch(AnchorDispatch):
"""AnchorDispatch with a detached bounded per-member gain on the
numerator (the denominator is untouched, the signs are untouched):
s_a(t) = cosh u_a / sum_all cosh u (the member's share)
ubar_a = causal EMA of s_a over bytes, rate rho, start 1/A
alpha_a = clamp(1 + gamma (A ubar_a - 1), a_min, a_max)
out = x + sum_a alpha_a (sinh u_a / sum_all cosh u)(arm_a(x) - x)
mode 'growth' lets the mean gain rise toward the cap; 'conserved'
renormalizes alpha to mean 1 over members at every byte. `coherent`
reads alpha once at the first site and shares it ('alpha'), and with
'full' shares the first site's signed routing weight too."""
def configure(self, *, mode="growth", rho=1 / 16, gamma=1.5, a_min=0.5,
a_max=2.0, coherent=None, role="ref", holder=None):
self.mode, self.rho, self.gamma = mode, float(rho), float(gamma)
self.a_min, self.a_max = float(a_min), float(a_max)
self.coherent, self.role = coherent, role
self.holder = {} if holder is None else holder
self._state, self._call = None, "forward"
return self
def reset_state(self):
self._state = None
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 BlockWithDispatch(nn.Module):
def __init__(self, block: nn.Module, disp: AnchorDispatch):
super().__init__()
self.block = block
self.disp = disp
def forward(self, *args, **kwargs):
out = self.block(*args, **kwargs)
if isinstance(out, tuple):
return (self.disp(out[0]),) + out[1:]
return self.disp(out)
def prefill(self, *args, **kwargs):
out, cache = self.block.prefill(*args, **kwargs)
self.disp._call = getattr(self.disp, "_call", "forward")
if isinstance(self.disp, GainDispatch):
self.disp._call = "prefill"
try:
return self.disp(out), cache
finally:
self.disp._call = "forward"
return self.disp(out), cache
def step(self, *args, **kwargs):
out = self.block.step(*args, **kwargs)
if isinstance(self.disp, GainDispatch):
self.disp._call = "step"
try:
return self.disp(out[0]) if isinstance(out, tuple) else self.disp(out)
finally:
self.disp._call = "forward"
if isinstance(out, tuple):
return (self.disp(out[0]),) + out[1:]
return self.disp(out)
# --------------------------------------------------------------- handles
def _probe(aleph) -> torch.Tensor:
"""Deterministic logits fingerprint on a tiny fixed input."""
dev = next(aleph.parameters()).device
ids = torch.arange(8, dtype=torch.long, device=dev).unsqueeze(0) % 7 + 1
was_training = aleph.training
aleph.eval()
with torch.no_grad():
out = aleph(ids)
if was_training:
aleph.train()
logits = out[0] if isinstance(out, (tuple, list)) else out.logits
return logits.detach().float().cpu()
class ArmHandle:
"""One mounted recipe. `detach(verify=True)` restores the core and
proves the restoration is bit-exact."""
def __init__(self, aleph, original, names, wrapped, fingerprint, card=None):
self.aleph = aleph
self._original = original
self.names = list(names)
self._blocks = wrapped
self._fingerprint = fingerprint
self.card = card or {}
# -- masking (the toggle law) --------------------------------------
def set_mask(self, mask: dict) -> None:
enabled = [bool(mask.get(n, True)) for n in self.names]
for b in self._blocks:
if isinstance(b, BlockWithDispatch):
b.disp.enabled = list(enabled)
else:
b.enabled = enabled[0]
def enable(self, *names):
self.set_mask({n: (n in names or not names) for n in self.names})
def disable(self, *names):
self.set_mask({n: n not in names for n in self.names})
def _snapshot(self):
return [list(b.disp.enabled) if isinstance(b, BlockWithDispatch) else b.enabled
for b in self._blocks]
def _restore(self, snap):
for b, s in zip(self._blocks, snap):
if isinstance(b, BlockWithDispatch):
b.disp.enabled = list(s)
else:
b.enabled = s
@contextlib.contextmanager
def only(self, *names):
prev = self._snapshot()
self.set_mask({n: n in names for n in self.names})
try:
yield self
finally:
self._restore(prev)
@contextlib.contextmanager
def all_off(self):
prev = self._snapshot()
self.set_mask({n: False for n in self.names})
try:
yield self
finally:
self._restore(prev)
# -- decode state ---------------------------------------------------
def reset_state(self):
"""Forget the carried gain state (a fresh context starts over)."""
for b in self._blocks:
if isinstance(b, BlockWithDispatch):
getattr(b.disp, "reset_state", lambda: None)()
# -- detach ---------------------------------------------------------
def detach(self, *, verify: bool = True):
self.aleph.blocks = self._original
if verify:
post = _probe(self.aleph)
if not torch.equal(post, self._fingerprint):
raise RuntimeError(
"detach verification FAILED: the restored core's logits "
"are not bit-exact with the pre-mount fingerprint")
return self.aleph
class StackHandle:
"""Nested always-on handles, detached outermost first."""
def __init__(self, handles, names=None, card=None):
self.handles = list(handles)
self.names = list(names or [h.names[0] for h in self.handles])
self.card = card or {}
def reset_state(self):
for h in self.handles:
h.reset_state()
def set_mask(self, mask: dict) -> None:
for h, n in zip(self.handles, self.names):
h.set_mask({h.names[0]: bool(mask.get(n, True))})
@contextlib.contextmanager
def all_off(self):
prev = [h._snapshot() for h in self.handles]
for h in self.handles:
h.set_mask({h.names[0]: False})
try:
yield self
finally:
for h, p in zip(self.handles, prev):
h._restore(p)
def detach(self, *, verify: bool = True):
for h in reversed(self.handles):
h.detach(verify=verify)
return self.handles[0].aleph if self.handles else None
# --------------------------------------------------------------- loading
def load_arm_file(path: str):
"""-> (tensors, meta). Reads the packaged safetensors arm; also accepts
an amoe .pt anchor/dispatch checkpoint for local experiments."""
if str(path).endswith(".safetensors"):
from safetensors.torch import load_file
from safetensors import safe_open
with safe_open(path, framework="pt") as f:
meta = dict(f.metadata() or {})
card = json.loads(meta.get("arm", "{}"))
return load_file(path), card
blob = torch.load(path, map_location="cpu", weights_only=True)
if "adapters" in blob:
return blob["adapters"], dict(blob.get("meta") or {})
if "dispatch" in blob: # flatten to "<block>.<key>"
flat = {f"{i}.{k}": v for i, st in enumerate(blob["dispatch"])
for k, v in st.items()}
return flat, dict(blob.get("meta") or {})
raise ValueError(f"unrecognized arm file: {path}")
def _per_block(state: dict, i: int) -> dict:
pre = f"{i}."
return {k[len(pre):]: v for k, v in state.items() if k.startswith(pre)}
def _target_blocks(aleph):
return list(aleph.blocks)
# --------------------------------------------------------------- mounting
def mount_anchor(aleph, state: dict, name: str = "arm", card=None) -> ArmHandle:
layers = _target_blocks(aleph)
d = int(aleph.cfg.d_model)
fingerprint = _probe(aleph)
spec = spec_from_state(state)
blocks, wrapped = [], []
for i, layer in enumerate(layers):
st = _per_block(state, i)
if not st:
blocks.append(layer)
continue
p = next(layer.parameters())
a = RelayPatchwork(d, spec)
a.load_state_dict(st, strict=True)
a.to(device=p.device, dtype=p.dtype)
w = BlockWithAdapter(layer, a)
blocks.append(w)
wrapped.append(w)
original = aleph.blocks
aleph.blocks = nn.ModuleList(blocks)
return ArmHandle(aleph, original, [name], wrapped, fingerprint, card)
def mount_routed(aleph, state: dict, names, gain=None, mode="growth",
coherent=None, card=None) -> ArmHandle:
layers = _target_blocks(aleph)
d = int(aleph.cfg.d_model)
fingerprint = _probe(aleph)
A = len(names)
member0 = {k[len("anchors.0."):]: v for k, v in _per_block(state, 0).items()
if k.startswith("anchors.0.")}
spec = spec_from_state(member0, prefix="")
holder, blocks, wrapped = {}, [], []
for i, layer in enumerate(layers):
st = _per_block(state, i)
p = next(layer.parameters())
members = nn.ModuleList([RelayPatchwork(d, spec) for _ in range(A)])
emb = int(st["key_proj"].shape[0])
dp = AnchorDispatch(members, d, emb=emb)
dp.load_state_dict(st, strict=True)
dp.to(device=p.device, dtype=p.dtype)
if gain:
dp.__class__ = GainDispatch
dp.configure(mode=mode, rho=gain["rho"], gamma=gain["gamma"],
a_min=gain["a_min"], a_max=gain["a_max"],
coherent=coherent, role="ref" if i == 0 else "follow",
holder=holder)
w = BlockWithDispatch(layer, dp)
blocks.append(w)
wrapped.append(w)
original = aleph.blocks
aleph.blocks = nn.ModuleList(blocks)
return ArmHandle(aleph, original, list(names), wrapped, fingerprint, card)
def mount(aleph, card: dict, files) -> "ArmHandle | StackHandle":
"""Attach one recipe. `card` is an arms/index.json row; `files` maps the
row's weight keys to local paths (a stack passes one file per member,
inner first)."""
kind = (card.get("mount") or {}).get("type", "anchor")
if kind == "anchor":
state, _ = load_arm_file(files[0])
return mount_anchor(aleph, state, card.get("name", "arm"), card)
if kind == "stack":
members = (card.get("mount") or {}).get("members") or []
handles = []
try:
for f, n in zip(files, members or [f"member{i}" for i in range(len(files))]):
state, _ = load_arm_file(f)
handles.append(mount_anchor(aleph, state, n, card))
except Exception:
StackHandle(handles).detach(verify=False) # never leave half a stack on
raise
return StackHandle(handles, members, card)
if kind == "routed":
state, _ = load_arm_file(files[0])
m = card.get("mount") or {}
return mount_routed(aleph, state, m.get("members") or [],
gain=m.get("gain"), mode=m.get("gain_mode", "growth"),
coherent=m.get("coherent"), card=card)
raise ValueError(f"unknown mount type {kind!r} (packaged types: {ARM_TYPES})")
# --------------------------------------------------------------- frames
CHAT_HEADER = ("A conversation with Beatrix, a small byte-level "
"language model still in pretraining.\n")
USER_TAG = "User: "
ASSISTANT_TAG = "Beatrix:"
TEXT_STOPS = ("\nUser:", "\nBeatrix:", "User:", "Beatrix:", "\n\n")
SPECIALS = {"DOC": 0xFF, "USER": 0xFE, "MODEL": 0xFD, "END": 0xFC, "SYS": 0xFB}
def render_text_frame(prompt: str, history=None) -> bytes:
"""The plain text frame every packaged arm was trained in:
header + 'User: ...\\n' + 'Beatrix:' — the reply follows the tag."""
out = CHAT_HEADER
for turn in (history or []):
tag = USER_TAG if turn.get("role") == "user" else ASSISTANT_TAG + " "
out += f"{tag}{turn.get('content', '')}\n"
out += f"{USER_TAG}{prompt}\n{ASSISTANT_TAG}"
return out.encode("utf-8")
def render_specials_frame(prompt: str, system: str | None = None,
history=None) -> bytes:
"""The core's own annealed frame: [SYS] .. [END] [USER] .. [END] [MODEL].
Encoded text can never contain a special, so the frame is unforgeable."""
S = SPECIALS
ids = []
if system:
ids += [S["SYS"]] + list(system.encode("utf-8")) + [S["END"]]
for turn in (history or []):
tag = S["USER"] if turn.get("role") == "user" else S["MODEL"]
ids += [tag] + list(str(turn.get("content", "")).encode("utf-8")) + [S["END"]]
ids += [S["USER"]] + list(prompt.encode("utf-8")) + [S["END"], S["MODEL"]]
return bytes(ids)
def cut_text_reply(raw: bytes) -> bytes:
"""Cut a text-frame continuation at the first stop string."""
best = len(raw)
for s in TEXT_STOPS:
j = raw.find(s.encode("utf-8"))
if j != -1:
best = min(best, j)
return raw[:best]
def cut_specials_reply(raw: bytes) -> bytes:
j = raw.find(bytes([SPECIALS["END"]]))
return raw if j == -1 else raw[:j]