code/fac_bed.py

"""fac_bed.py — FAC (factorized absolute code) experiment bed.  #TAG:fac #TAG:loss_campaign #TAG:autoregressive #TAG:sign_code

THE QUESTION: can a cosh-Bregman pull onto a FROZEN +/-1 code table replace the
trained CE readout on the certified byte-LM operating point — and does swapping
only-the-loss revive the certified addr_head collapse (the P4 cell)?

Built ON TOP of tools/ar_differentiation_bed.py (as-run certified reference —
imported, never modified, never copied). The certified operating point:
wikitext-2-raw bytes, block 256, batch 32, d=192, 4 layers, 2000 steps, pure
Adam lr 3e-4 wd=0. Incumbent control: addr_msl64 (3-seed certified mean bpb
2.4685; timeline 2026-07-09 W6).

THE FAC CONSTRUCTION
  feats  = the certified addr_msl64 pre-readout read: P=64 parallel D=4 slots
           through the shared K=64 aleph codebook, concat M_hat -> 256 dims.
           (The spec sketch guessed P=16/64-dim; the BED CODE is authoritative:
           ByteLM 'addr_msl<P>' parses P from the name -> addr_msl64 is 64
           slots x D=4 = 256 dims. Recorded as smoke S0.)
  R      = FIXED frame (64 x 256), orthonormal ROWS, torch.linalg.qr on a
           crc32-seeded gaussian; registered buffer, NEVER trained. Gauge-fixed
           by construction — a learnable frame reproduces either the tied-M_hat
           starvation or the L-PS1 moving-target failure.
  s      = F.normalize(feats, dim=-1) @ R.T          (B, T, 64);  |s_k| <= 1
  v      = s / t_loss                                 (t_loss 0.3; sweep .1/.3/1)
  C      = FROZEN code table in {-1,+1}^{256 x 64}. Two constructions:
             fac_ecc — sign of an iid crc32-seeded gaussian (injective w.h.p.);
             fac_lsh — sign of the same gaussian box-filtered (width 3) along
             the byte-value axis, so numerically adjacent bytes share more bits
             (a semantic-adjacency proxy for the byte alphabet; bytes carry no
             pretrained embedding here, so a fixed random frame IS the SimHash).
  LOSS   L = mean( cosh( clamp(v - C[y]*mu, -4, 4) ) - 1 ),  mu = 1.0 (S4).
           This is the residual-form cosh-Bregman D_Phi(v - C[y]*mu, 0) with
           Phi = sum cosh: D_Phi(r,0) = cosh(r) - cosh(0) - sinh(0)*r
           = cosh(r) - 1 exactly (smoke S2; the target-anchored D_Phi(v, v*)
           is a DIFFERENT function — coincides only at v* = 0 — S2 reports the
           gap so the naming is honest).
  EVAL   score(y') = s @ C[y'].T; bpb-of-record = log_softmax over the 256
           scores (raw scores as logits — softmax temperature NOT calibrated;
           caveat recorded in every ledger line). Decoded-token accuracy and
           code-collision rate (Hamming <= 2 over distinct byte pairs) ride
           along. Decode argmax is a READOUT only — no selection event in any
           gradient path (aleph rider).

ARMS (--arm)
  ce               certified addr_msl64 incumbent — bed's own model + CE.
  ce_fixedcode     SAME init (same seed, same RNG draw order), readout replaced
                   by the frozen C: logits = s @ C.T, trained with CE. Zero
                   trainable readout params — THE param-matched control.
  fac_lsh/fac_ecc  the FAC loss with each code table.
  fac_none         FAC loss, s from a plain Linear(d,64) on the trunk hidden
                   state; the aleph head is present but UNREAD (gradient-dead).
                   Isolates the geometry.
  p4_addr_head_ce  the certified COLLAPSE configuration, verbatim: ByteLM
                   'addr_head' (raw hidden -> single-slot address, K=32,
                   tau=0.1, coefficients->logits) + CE. Expected: bpb ~5.66,
                   usage_ppl ~1.88/64.
  p4_addr_head_fac the SAME model (identical params; the coeff->logit head is
                   computed but gradient-dead) with the FAC loss on the signed
                   coefficient vector, embedded isometrically by a 64x32
                   orthonormal-column frame. ONLY THE LOSS DIFFERS.
                   Prediction: usage_ppl >= 8/64 while win|cos| stays > 0.99.

LEDGER  JSON lines, one file per run: tools/fac_runs/{arm}_s{seed}_t{steps}.jsonl
  bpb (partitioned), decoded_acc, train-curve summary, collision_rate,
  sign_fidelity (exp015 gauge, L-078), anchor_drift mean + binding_fraction +
  usage_ppl + win|cos| (via the bed's own vitals -> geolip_vitals), gate stats
  (none of these arms carry gates — recorded as null), wall, peak_mem, params
  (+ delta vs ce — the S9 honesty line), seeds from crc32(arm:seed).

RIDERS  pure Adam wd=0 ONLY; fp32, TF32 off; cuda memory fraction 0.73;
  crc32 seeds never hash(); GPU-only verdict runs; data_root OUTSIDE the repo;
  Colab-cell-safe (paste-ahead imports, parse_known_args, no __file__ logic).

Terminal:  python tools/fac_bed.py                    # smoke battery (default)
           python tools/fac_bed.py --smoke            # same
           python tools/fac_bed.py --bench            # ~20-step throughput probe
           python tools/fac_bed.py --arm fac_lsh --seed 0        # verdict run
Colab:     paste geolip_vitals.py, ar_differentiation_bed.py, then this file
           (smokes auto-run); train_arm("fac_lsh", seed=0) in the next cell.
"""
from __future__ import annotations
import json
import math
import os
import sys
import time
import types
import zlib

import torch
import torch.nn as nn
import torch.nn.functional as F

# ---------------------------------------------------------------- environment
def _repo_root():
    d = os.path.abspath(os.getcwd())
    while True:
        if os.path.exists(os.path.join(d, "MANIFEST.md")):
            return d
        p = os.path.dirname(d)
        if p == d:
            return os.getcwd()
        d = p

ROOT = _repo_root()
_TOOLS = os.path.join(ROOT, "tools")
if os.path.isdir(_TOOLS) and _TOOLS not in sys.path:
    sys.path.insert(0, _TOOLS)

if "ByteLM" in globals() and "AlephAddress" in globals():   # Colab paste-ahead
    bed = types.SimpleNamespace(
        ByteLM=globals()["ByteLM"], AlephAddress=globals()["AlephAddress"],
        _wikitext_bytes=globals()["_wikitext_bytes"],
        _batch=globals()["_batch"], VOCAB=globals()["VOCAB"])
else:
    import ar_differentiation_bed as bed                    # certified, read-only

torch.backends.cuda.matmul.allow_tf32 = False               # pin_precision law
torch.backends.cudnn.allow_tf32 = False
DEV = "cuda" if torch.cuda.is_available() else "cpu"
if DEV == "cuda":
    torch.cuda.set_per_process_memory_fraction(0.73)        # WDDM standing cap

DATA_ROOT = os.environ.get("GEOLIP_DATA", "./data")
RUNS_DIR = (os.path.join(_TOOLS, "fac_runs") if os.path.isdir(_TOOLS)
            else os.path.abspath("./fac_runs"))

CODE_BITS = 64        # width of the absolute code (matches the K=64 aleph book)
T_LOSS = 0.3          # loss temperature (knob; sweep {0.1, 0.3, 1.0})
MU = 1.0              # target margin in v-units (S4 decision: KEPT at 1.0)
CLAMP = 4.0           # residual clamp — cosh(4) ~ 27.3, fp32-safe by construction


def seed_for(name: str) -> int:
    """crc32, never hash() — PYTHONHASHSEED nondeterminism is a recorded law."""
    return zlib.crc32(name.encode("utf-8")) & 0x7FFFFFFF


# ------------------------------------------------------------- frozen buffers
def orthonormal_frame(out_dim: int, in_dim: int, seed: int) -> torch.Tensor:
    """(out_dim, in_dim) fixed frame, QR on a crc32-seeded gaussian, fp64 then
    fp32, canonical sign fix. NEVER trained (registered as a buffer).
      in_dim >= out_dim: orthonormal ROWS  (R R^T = I)  -> |s_k| <= 1 for unit feats
      in_dim <  out_dim: orthonormal COLS  (R^T R = I)  -> isometric embed,
                         ||s|| = ||feats_hat||, |s_k| <= ||row_k|| <= 1."""
    g = torch.Generator().manual_seed(seed)
    n, m = max(out_dim, in_dim), min(out_dim, in_dim)
    G = torch.randn(n, m, generator=g, dtype=torch.float64)
    Q, Rq = torch.linalg.qr(G)
    sgn = torch.where(torch.diagonal(Rq) >= 0, 1.0, -1.0)
    Q = Q * sgn.unsqueeze(0)
    R = Q.T if in_dim >= out_dim else Q
    return R.float().contiguous()


def build_code(kind: str) -> torch.Tensor:
    """FROZEN code table C in {-1,+1}^(256 x 64), crc32-seeded, buffer-only.
      ecc — sign of iid gaussian rows: injective w.h.p., no byte structure.
      lsh — sign of the gaussian box-filtered (width 3, replicate-padded) along
            the byte-value axis: adjacent byte values share ~73% of bits
            (corr 2/3 -> sign agreement 1 - arccos(2/3)/pi) — the
            semantic-adjacency proxy for the byte alphabet."""
    g = torch.Generator().manual_seed(seed_for(f"fac:C:{kind}"))
    G = torch.randn(bed.VOCAB, CODE_BITS, generator=g, dtype=torch.float64)
    if kind == "lsh":
        Gp = torch.cat([G[:1], G, G[-1:]], dim=0)          # replicate pad
        G = (Gp[:-2] + Gp[1:-1] + Gp[2:]) / 3.0            # width-3 box filter
    elif kind != "ecc":
        raise ValueError(f"unknown code table '{kind}' (ecc|lsh)")
    return torch.where(G >= 0, 1.0, -1.0).float().contiguous()


@torch.no_grad()
def code_collision_rate(C: torch.Tensor, thresh: int = 2) -> float:
    """Fraction of distinct byte pairs whose codes are within Hamming <= thresh."""
    ham = (C.shape[1] - C @ C.t()) / 2
    iu = torch.triu_indices(C.shape[0], C.shape[0], offset=1)
    return float((ham[iu[0], iu[1]] <= thresh).float().mean())


@torch.no_grad()
def sign_fidelity(book: torch.Tensor, n: int = 3000, seed: int = 0) -> float:
    """Spearman(sign-code Hamming, true angle) over random pairs on S^(D-1),
    using the book as the LSH frame — the exp015 content gauge (L-078),
    reimplemented verbatim-in-spirit (importing exp015 drags its module state)."""
    A = F.normalize(book.detach().float().cpu(), dim=-1)
    g = torch.Generator().manual_seed(seed)
    D = A.shape[1]
    x = F.normalize(torch.randn(n, D, generator=g), dim=-1)
    y = F.normalize(torch.randn(n, D, generator=g), dim=-1)
    ang = torch.arccos((x * y).sum(-1).clamp(-1, 1))
    ham = (torch.sign(x @ A.T) != torch.sign(y @ A.T)).float().mean(-1)
    ra = ang.argsort().argsort().float()
    rb = ham.argsort().argsort().float()
    ra = (ra - ra.mean()) / ra.std()
    rb = (rb - rb.mean()) / rb.std()
    return round(float((ra * rb).mean()), 4)


def margin_audit(R: torch.Tensor, t_loss: float, mu: float) -> dict:
    """The S4 reachability chain, as data: per-axis max |s_k| = ||row_k|| <= 1;
    v = s/t_loss so per-axis reachable |v_k| = ||row_k||/t_loss; mu must sit
    strictly inside. (The JOINT target ||C[y]*mu|| = mu*8 exceeds the reachable
    ||v|| <= 1/t_loss ball — by design the loss is a per-axis margin pull, not
    an attainable minimum; recorded, not a failure.)"""
    rn = R.norm(dim=-1)
    lo, hi = float(rn.min()), float(rn.max())
    return {"row_norm_min": round(lo, 4), "row_norm_max": round(hi, 4),
            "v_reach_min": round(lo / t_loss, 4),
            "v_reach_max": round(hi / t_loss, 4), "t_loss": t_loss, "mu": mu,
            "per_axis_reachable": bool(mu < lo / t_loss)}


# --------------------------------------------------------------------- model
class FacModel(nn.Module):
    """Wraps the CERTIFIED bed ByteLM — subclass-free reuse, zero copied math.
    mode:
      'addr' — lm = ByteLM('addr_msl64'); lm.head := Identity, so the certified
               forward RETURNS the pre-readout feats itself (64 slots x D4 =
               256). Same seed => bit-identical init to the ce arm everywhere
               except the (removed) readout.
      'none' — lm = ByteLM('addr_msl64') kept whole (aleph present, UNREAD,
               gradient-dead); trunk walked via lm's own submodules; feats =
               plain Linear(d, 64) on the hidden state (orthogonal init,
               bias-free — the bed's head_proj idiom).
      'p4'   — lm = ByteLM('addr_head') UNMODIFIED (the certified collapse
               configuration: K=32, tau=0.1, coefficients->logits). The signed
               coefficient vector is captured by a forward PRE-HOOK on lm.head,
               so the two P4 cells have IDENTICAL parameters and identical
               forward compute — only the loss differs (the coeff->logit head
               is trainable-but-gradient-dead under FAC).
    The FAC read: s = normalize(feats) @ R.T (fixed frame), scores = s @ C.T."""

    def __init__(self, mode: str, code: str = "ecc", d: int = 192,
                 layers: int = 4, block: int = 256, t_loss: float = T_LOSS,
                 mu: float = MU):
        super().__init__()
        self.mode, self.code_kind = mode, code
        self.t_loss, self.mu = t_loss, mu
        if mode == "addr":
            self.lm = bed.ByteLM("addr_msl64", d=d, layers=layers, block=block)
            in_dim = self.lm.n_slots * 4                    # 256, read off the arm
            self.lm.head = nn.Identity()                    # readout removed
        elif mode == "none":
            self.lm = bed.ByteLM("addr_msl64", d=d, layers=layers, block=block)
            assert not self.lm.trigram and not self.lm.use_relay
            self.proj_none = nn.Linear(d, CODE_BITS, bias=False)
            nn.init.orthogonal_(self.proj_none.weight)
            in_dim = CODE_BITS
        elif mode == "p4":
            self.lm = bed.ByteLM("addr_head", d=d, layers=layers, block=block)
            in_dim = self.lm.head_addr.K                    # 32
            self._w = None
            self.lm.head.register_forward_pre_hook(self._grab)
        else:
            raise ValueError(f"unknown mode '{mode}'")
        self.in_dim = in_dim
        self.register_buffer("R", orthonormal_frame(
            CODE_BITS, in_dim, seed_for(f"fac:R:{in_dim}")))
        self.register_buffer("C", build_code(code))

    def _grab(self, module, inputs):                        # p4 pre-hook
        self._w = inputs[0]

    def _trunk(self, idx):
        """The bed's trunk, walked via the wrapped submodules (mode 'none' only
        — the certified forward would read the aleph, which this arm forbids)."""
        lm = self.lm
        x = lm.emb(idx) + lm.pos[:, : idx.shape[1]]
        for b in lm.blocks:
            x = b(x)
        h = lm.nf(x)
        lm._last_h = h.detach()                             # vitals hookup
        return h

    def feats(self, idx):
        if self.mode == "addr":
            return self.lm(idx)                             # head=Identity -> feats
        if self.mode == "none":
            return self.proj_none(self._trunk(idx))
        _ = self.lm(idx)                                    # head computed, unused
        return self._w

    def address(self, idx):
        return F.normalize(self.feats(idx), dim=-1) @ self.R.t()

    def forward(self, idx):
        return self.address(idx)

    @torch.no_grad()
    def vitals(self) -> dict:
        return self.lm.vitals()                             # the bed's own readouts


# ---------------------------------------------------------------- loss + eval
def fac_loss(s, y, C, t_loss: float = T_LOSS, mu: float = MU):
    """The FAC objective: residual-form cosh-Bregman pull onto the frozen code.
    L = mean(cosh(clamp(s/t_loss - C[y]*mu, -4, 4)) - 1)  — S2/S3/S5/S6."""
    v = s / t_loss
    r = (v - C[y] * mu).clamp(-CLAMP, CLAMP)
    return (torch.cosh(r) - 1.0).mean()


def eval_scores(kind: str, model, x):
    """Per-token scores on the shared 256-way scale. ce -> the model's own
    logits; fixed-code kinds -> s @ C.T (raw scores as logits; softmax
    temperature NOT calibrated — the recorded caveat)."""
    if kind == "ce":
        return model(x)
    s = model.address(x)
    return s @ model.C.t()


def compute_loss(kind: str, model, x, y):
    """THE single loss site of this bed (mirrors the bed's one F.cross_entropy)."""
    if kind in ("ce", "ce_fixedcode"):
        logits = eval_scores(kind, model, x)
        return F.cross_entropy(logits.reshape(-1, bed.VOCAB), y.reshape(-1))
    s = model.address(x)
    return fac_loss(s, y, model.C, model.t_loss, model.mu)


@torch.no_grad()
def evaluate(kind: str, model, va, batch, block, device, g, n_batches=20):
    model.eval()
    nll = acc = 0.0
    for _ in range(n_batches):
        xv, yv = bed._batch(va, batch, block, device, g)
        sc = eval_scores(kind, model, xv)
        nll += float(F.cross_entropy(sc.reshape(-1, bed.VOCAB), yv.reshape(-1)))
        acc += float((sc.argmax(-1) == yv).float().mean())
    model.train()
    return nll / n_batches / math.log(2), acc / n_batches


# ---------------------------------------------------------------------- arms
ARMS = {
    "ce":               dict(kind="ce", mode=None,
                             note="certified addr_msl64 incumbent (3-seed 2.4685)"),
    "ce_fixedcode":     dict(kind="ce_fixedcode", mode="addr",
                             note="frozen-C readout, CE — THE param-matched control"),
    "fac_lsh":          dict(kind="fac", mode="addr", code="lsh"),
    "fac_ecc":          dict(kind="fac", mode="addr", code="ecc"),
    "fac_none":         dict(kind="fac", mode="none",
                             note="aleph UNREAD (gradient-dead) — geometry isolation"),
    "p4_addr_head_ce":  dict(kind="ce", mode="p4ce",
                             note="certified collapse config + CE (expect ~5.66 bpb, ppl ~1.88/64)"),
    "p4_addr_head_fac": dict(kind="fac", mode="p4",
                             note="same params, FAC loss (predict ppl >= 8/64, win|cos| > .99)"),
}


def build_model(arm: str, code: str = "ecc", d: int = 192, layers: int = 4,
                block: int = 256, t_loss: float = T_LOSS, mu: float = MU):
    spec = ARMS[arm]
    if arm == "ce":
        return bed.ByteLM("addr_msl64", d=d, layers=layers, block=block)
    if arm == "p4_addr_head_ce":
        return bed.ByteLM("addr_head", d=d, layers=layers, block=block)
    return FacModel(spec["mode"], code=spec.get("code", code), d=d,
                    layers=layers, block=block, t_loss=t_loss, mu=mu)


def _trainable(model) -> int:
    return sum(p.numel() for p in model.parameters() if p.requires_grad)


@torch.no_grad()
def model_vitals(model) -> dict:
    lm = model.lm if isinstance(model, FacModel) else model
    return {"vitals": model.vitals(),
            "sign_fidelity": sign_fidelity(lm.head_addr.codebook),
            "gates": None}                       # no gates in any FAC arm


# ---------------------------------------------------------------------- train
def train_arm(arm: str, seed: int = 0, steps: int = 2000, batch: int = 32,
              block: int = 256, d: int = 192, layers: int = 4,
              device: str = "cuda", data_root: str | None = None,
              t_loss: float = T_LOSS, mu: float = MU, code: str | None = None,
              eval_every: int = 500, save: bool = True, lr: float = 3e-4):
    """Verdict run — GPU only, pure Adam wd=0, ledger JSONL under tools/fac_runs."""
    if device == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("Verdict runs are GPU-only (never CPU-train for accuracy).")
    spec = ARMS[arm]
    code = code or spec.get("code", "ecc")
    data_root = data_root or DATA_ROOT
    base = seed_for(f"{arm}:{seed}")             # crc32(arm + seed index)
    torch.manual_seed(base)
    g = torch.Generator().manual_seed(base)
    tr, va = bed._wikitext_bytes(data_root)
    print(f"data ready: train {tr.numel():,} bytes, val {va.numel():,} bytes",
          flush=True)
    model = build_model(arm, code=code, d=d, layers=layers, block=block,
                        t_loss=t_loss, mu=mu).to(device)
    with torch.random.fork_rng():                # S9 honesty line, same init seed
        torch.manual_seed(base)
        ce_ref = bed.ByteLM("addr_msl64", d=d, layers=layers, block=block)
        ce_trainable = _trainable(ce_ref)
        del ce_ref
    n_train = _trainable(model)
    audit = (margin_audit(model.R, t_loss, mu)
             if isinstance(model, FacModel) else None)
    if audit is not None and not audit["per_axis_reachable"]:
        print(f"WARNING: mu={mu} not per-axis reachable at t_loss={t_loss} "
              f"({audit}) — the margin never engages; consider mu<="
              f"{0.8 * audit['v_reach_min']:.2f}", flush=True)
    os.makedirs(RUNS_DIR, exist_ok=True)
    led = os.path.join(RUNS_DIR, f"{arm}_s{seed}_t{steps}.jsonl")
    def emit(obj, first=False):
        with open(led, "w" if first else "a", encoding="utf-8") as f:
            f.write(json.dumps(obj) + "\n")
    emit({"event": "config", "arm": arm, "seed": seed, "base_seed": base,
          "steps": steps, "batch": batch, "block": block, "d": d,
          "layers": layers, "lr": lr, "t_loss": t_loss, "mu": mu,
          "code": (code if isinstance(model, FacModel) else None),
          "margin_audit": audit, "params_trainable": n_train,
          "params_trainable_ce_ref": ce_trainable,
          "param_delta_vs_ce": ce_trainable - n_train,
          "readout_note": {
              "ce": "trained Linear(256,256) readout",
              "ce_fixedcode": "frozen C readout — ZERO trainable readout params",
              "fac_lsh": "frozen C readout — ZERO trainable readout params",
              "fac_ecc": "frozen C readout — ZERO trainable readout params",
              "fac_none": "frozen C readout; aleph head params present but UNREAD",
              "p4_addr_head_ce": "trained Linear(32,256) coeff->logit head",
              "p4_addr_head_fac": "identical params to the ce cell; coeff->logit "
                                  "head computed but gradient-dead under FAC",
          }[arm], "partition_note": "bpb uses raw s@C.T scores as logits "
          "(log_softmax); softmax temperature NOT calibrated — caveat",
          "torch": torch.__version__, "device": device,
          "note": spec.get("note", "")}, first=True)
    if device == "cuda":
        torch.cuda.reset_peak_memory_stats()
    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=0.0)
    kind, curve, t0, bpb, acc = spec["kind"], [], time.time(), None, None
    for step in range(1, steps + 1):
        x, y = bed._batch(tr, batch, block, device, g)
        loss = compute_loss(kind, model, x, y)
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()
        if step == 1 or step % 100 == 0:
            curve.append([step, round(float(loss), 5)])
        if step % eval_every == 0 or step == steps:
            bpb, acc = evaluate(kind, model, va, batch, block, device, g)
            vit = model_vitals(model)
            emit({"event": "eval", "step": step, "bpb": round(bpb, 4),
                  "decoded_acc": round(acc, 4), **vit})
            print(f"[{arm} s{seed}] step {step} bpb={bpb:.4f} acc={acc:.4f} "
                  f"vitals={vit['vitals']}", flush=True)
    losses = [c[1] for c in curve]
    final = {"event": "final", "arm": arm, "seed": seed, "steps": steps,
             "t_loss": t_loss, "mu": mu,
             "code": (code if isinstance(model, FacModel) else None),
             "bpb": round(bpb, 4), "decoded_acc": round(acc, 4),
             "collision_rate": (code_collision_rate(model.C.cpu())
                                if isinstance(model, FacModel) else None),
             "train_curve": {"first": losses[0], "final": losses[-1],
                             "min": min(losses), "every100": curve},
             **model_vitals(model),
             "wall_s": round(time.time() - t0, 1),
             "peak_mem_gb": (round(torch.cuda.max_memory_allocated() / 2**30, 3)
                             if device == "cuda" else 0.0),
             "params_trainable": n_train,
             "param_delta_vs_ce": ce_trainable - n_train,
             "mu_note": "mu in v-units; per-axis reachability audited at config; "
                        "joint target ||C[y]*mu||=8 exceeds ||v||<=1/t_loss — "
                        "per-axis margin pull by design"}
    emit(final)
    print(json.dumps(final), flush=True)
    if save:
        ck = os.path.join(data_root, "fac_ckpts")
        os.makedirs(ck, exist_ok=True)
        path = os.path.join(ck, f"{arm}_s{seed}_t{steps}.pt")
        torch.save({"arm": arm, "seed": seed, "steps": steps, "bpb": bpb,
                    "state_dict": {k: v.cpu() for k, v in
                                   model.state_dict().items()}}, path)
        print(f"saved specimen: {path}", flush=True)
    return final


# ---------------------------------------------------------------------- bench
def bench(arm: str = "fac_lsh", steps: int = 20, warmup: int = 5,
          batch: int = 32, block: int = 256, device: str = "cuda",
          data_root: str | None = None):
    """~20-step throughput probe at the full operating point -> min/2000 steps.
    Falls back to synthetic random bytes if the parquet cache is unreachable
    (identical compute — the batch path indexes a flat uint8 tensor either way)."""
    if device == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("bench is a GPU probe")
    try:
        tr, _ = bed._wikitext_bytes(data_root or DATA_ROOT)
        src = "wikitext-2-raw"
    except Exception as e:
        g0 = torch.Generator().manual_seed(seed_for("fac:bench:data"))
        tr = torch.randint(0, 256, (2_000_000,), generator=g0,
                           dtype=torch.uint8)
        src = f"synthetic ({type(e).__name__})"
    torch.manual_seed(seed_for(f"{arm}:bench"))
    g = torch.Generator().manual_seed(seed_for(f"{arm}:bench"))
    model = build_model(arm).to(device)
    kind = ARMS[arm]["kind"]
    opt = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=0.0)
    if device == "cuda":
        torch.cuda.reset_peak_memory_stats()
    for _ in range(warmup):
        x, y = bed._batch(tr, batch, block, device, g)
        loss = compute_loss(kind, model, x, y)
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
    if device == "cuda":
        torch.cuda.synchronize()
    t0 = time.time()
    for _ in range(steps):
        x, y = bed._batch(tr, batch, block, device, g)
        loss = compute_loss(kind, model, x, y)
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
    if device == "cuda":
        torch.cuda.synchronize()
    sec = (time.time() - t0) / steps
    peak = (torch.cuda.max_memory_allocated() / 2**30) if device == "cuda" else 0.0
    print(f"BENCH [{arm}] data={src} {sec*1000:.1f} ms/step -> "
          f"{sec*2000/60:.1f} min / 2000 steps (+~{20*sec:.0f}s eval overhead); "
          f"peak {peak:.2f} GB", flush=True)
    return {"arm": arm, "ms_per_step": sec * 1000,
            "min_per_2000": sec * 2000 / 60, "peak_gb": peak, "data": src}


# ------------------------------------------------------------- smoke battery
RESULTS = []
def record(tid, name, ok, detail=""):
    RESULTS.append((tid, name, "PASS" if ok else "FAIL", detail))
    return ok


def run_smokes() -> bool:
    """FORMULA smokes only — shapes, gradients, identities, fp32 safety,
    causality, param honesty. NO training, ever (MANIFEST rider); the S3
    minimizer probe optimizes one free 64-vector, not a model."""
    del RESULTS[:]
    t0 = time.time()
    sd, sl, sb = 96, 2, 64                                  # small config

    # S0 — feats-dim truth: the certified addr_msl64 read is 64 slots x D4.
    torch.manual_seed(seed_for("fac:s0"))
    m_addr = FacModel("addr", "lsh", d=sd, layers=sl, block=sb).to(DEV)
    record("S0", "feats match the certified addr_msl64 read",
           m_addr.in_dim == 256 and m_addr.lm.n_slots == 64,
           "64 slots x D=4 = 256-dim feats (bed code authoritative; spec's "
           "'P=16' guess corrected) -> R is 64x256 orthonormal-rows")

    # S1 — gradient flow: FAC backward reaches the aleph codebook AND the slot
    # projection (the tied-M_hat failure would show a near-zero codebook grad).
    g1 = torch.Generator().manual_seed(seed_for("fac:s1"))
    seq = torch.randint(0, 256, (2, sb + 1), generator=g1)
    x, y = seq[:, :-1].to(DEV), seq[:, 1:].to(DEV)
    loss = fac_loss(m_addr.address(x), y, m_addr.C)
    loss.backward()
    g_cb = float(m_addr.lm.head_addr.codebook.grad.norm())
    g_pj = float(m_addr.lm.head_proj.weight.grad.norm())
    g_em = float(m_addr.lm.emb.weight.grad.norm())
    record("S1", "FAC gradient flow (codebook + slot proj + trunk)",
           g_cb > 0 and g_pj > 0 and g_em > 0
           and all(map(math.isfinite, (g_cb, g_pj, g_em))),
           "|g| codebook %.2e, head_proj %.2e, emb %.2e — all nonzero/finite"
           % (g_cb, g_pj, g_em))

    # S2 — Bregman identity. Phi = sum cosh. The implemented loss is the
    # RESIDUAL-form divergence D_Phi(r, 0) = cosh(r) - cosh(0) - sinh(0)*r
    # = cosh(r) - 1 ("up to the constant" = cosh(0)). The target-anchored
    # D_Phi(v, v*) is a different function (equal only at v* = 0) — its gap is
    # reported so the 'cosh-Bregman' name stays honest.
    g2 = torch.Generator().manual_seed(seed_for("fac:s2"))
    r = torch.empty(4096, dtype=torch.float64).uniform_(-3.9, 3.9, generator=g2)
    d_res = float((torch.cosh(r) - math.cosh(0.0) - math.sinh(0.0) * r
                   - (torch.cosh(r) - 1.0)).abs().max())
    vs = torch.where(torch.rand(4096, generator=g2) > 0.5, 1.0, -1.0).double()
    v = r + vs
    d_anchor = float((torch.cosh(v) - torch.cosh(vs) - torch.sinh(vs) * (v - vs)
                      - (torch.cosh(v - vs) - 1.0)).abs().max())
    record("S2", "Bregman identity (residual form, up to cosh(0))",
           d_res <= 1e-6,
           "residual-form dev %.1e; target-anchored D_Phi(v,c*mu) differs by "
           "up to %.2f — loss is D_Phi(v - C[y]mu, 0), coincides at v*=0"
           % (d_res, d_anchor))

    # S3 — minimizer identity: 200 Adam steps on a free 64-vector -> the code.
    C_ecc = build_code("ecc")
    g3 = torch.Generator().manual_seed(seed_for("fac:s3"))
    y0 = int(torch.randint(0, 256, (1,), generator=g3))
    target = C_ecc[y0] * MU
    vfree = nn.Parameter(torch.zeros(CODE_BITS))
    # pure Adam wd=0; beta2=0.9 so the second-moment memory (1000-step at the
    # default 0.999) cannot suppress late updates inside a 200-step anneal —
    # measured: default betas freeze the error at ~3e-3.
    opt3 = torch.optim.Adam([vfree], lr=1.0, betas=(0.9, 0.9), weight_decay=0.0)
    for _ in range(200):
        l3 = (torch.cosh((vfree - target).clamp(-CLAMP, CLAMP)) - 1.0).mean()
        opt3.zero_grad(set_to_none=True)
        l3.backward()
        opt3.step()
        for pg in opt3.param_groups:
            pg["lr"] *= 0.93                     # anneal; Adam alone orbits at lr
    err3 = float((vfree.detach() - target).abs().max())
    record("S3", "minimizer identity (free v -> C[y]*mu)",
           err3 < 1e-4, "||v - C[y]mu||_inf = %.1e after 200 Adam steps" % err3)

    # S4 — margin reachability chain. |s_k| <= ||R row_k|| <= 1 (R construction);
    # v = s/t_loss -> per-axis reachable |v_k| = ||row_k||/0.3 (=3.33 for the
    # orthonormal-row frames); mu = 1.0 sits strictly inside for ALL frames,
    # and a perfectly-aligned s pulls strictly toward the target on every axis.
    # DECISION: mu stays 1.0 (in v-units; the s-unit boundary worry dissolves
    # because the margin lives in v-space).
    frames = {n: orthonormal_frame(CODE_BITS, i, seed_for(f"fac:R:{i}"))
              for n, i in (("addr", 256), ("none", 64), ("p4", 32))}
    audits = {n: margin_audit(Rf, T_LOSS, MU) for n, Rf in frames.items()}
    ortho_dev = max(
        float((frames["addr"] @ frames["addr"].t()
               - torch.eye(CODE_BITS)).abs().max()),
        float((frames["none"] @ frames["none"].t()
               - torch.eye(CODE_BITS)).abs().max()),
        float((frames["p4"].t() @ frames["p4"]
               - torch.eye(32)).abs().max()))
    s_align = C_ecc[y0] / math.sqrt(CODE_BITS)              # unit, code-aligned
    pull = -torch.sinh((s_align / T_LOSS - C_ecc[y0] * MU).clamp(-CLAMP, CLAMP)
                       ) * C_ecc[y0]
    ok4 = (all(a["per_axis_reachable"] for a in audits.values())
           and all(a["row_norm_max"] <= 1.0 + 1e-5 for a in audits.values())
           and ortho_dev < 1e-5 and bool((pull > 0).all()))
    record("S4", "margin reachability (mu=1.0 KEPT, v-units)", ok4,
           "v-reach addr/none/p4 = %.2f/%.2f/%.2f > mu=1; ortho dev %.0e; "
           "aligned-s pull > 0 on 64/64 axes"
           % (audits["addr"]["v_reach_min"], audits["none"]["v_reach_min"],
              audits["p4"]["v_reach_min"], ortho_dev))

    # S5 — antipodal invariance: L(v, c) == L(-v, -c) bit-exact.
    g5 = torch.Generator().manual_seed(seed_for("fac:s5"))
    s5 = (torch.randn(4, 32, CODE_BITS, generator=g5) * 0.4).to(DEV)
    y5 = torch.randint(0, 256, (4, 32), generator=g5).to(DEV)
    C5 = C_ecc.to(DEV)
    la, lb = fac_loss(s5, y5, C5), fac_loss(-s5, y5, -C5)
    record("S5", "antipodal invariance L(v,c)==L(-v,-c)",
           bool(torch.equal(la, lb)),
           "bit-exact on %s: %.6f == %.6f" % (DEV, float(la), float(lb)))

    # S6 — fp32 overflow: finite loss AND gradient over the full reachable |v|
    # range (t_loss down to 0.1 -> |v| <= 10, swept to 12) with the clamp; the
    # zero-feats normalize edge is finite too.
    v6 = torch.linspace(-12.0, 12.0, 100001).requires_grad_(True)
    l6 = (torch.cosh((v6 - 1.0).clamp(-CLAMP, CLAMP)) - 1.0).sum()
    l6.backward()
    z = F.normalize(torch.zeros(3, CODE_BITS), dim=-1) @ frames["none"].t()
    ok6 = (bool(torch.isfinite(l6)) and bool(torch.isfinite(v6.grad).all())
           and bool(torch.isfinite(z).all()))
    record("S6", "fp32 safety across the reachable v-range",
           ok6, "cosh capped at cosh(4)=%.1f; grad finite on [-12,12]; "
           "zero-feats normalize edge finite" % math.cosh(CLAMP))

    # S7 — decode consistency: planted s = C[y]*t_loss*mu -> argmax score == y,
    # 1000 draws under fac_ecc (assert); lsh failure rate reported, not gated.
    g7 = torch.Generator().manual_seed(seed_for("fac:s7"))
    y7 = torch.randint(0, 256, (1000,), generator=g7)
    acc_ecc = float(((C_ecc[y7] * T_LOSS * MU) @ C_ecc.t()
                     ).argmax(-1).eq(y7).float().mean())
    C_lsh = build_code("lsh")
    acc_lsh = float(((C_lsh[y7] * T_LOSS * MU) @ C_lsh.t()
                     ).argmax(-1).eq(y7).float().mean())
    adj = float((C_lsh[:-1] == C_lsh[1:]).float().mean())
    record("S7", "decode consistency (planted code -> argmax)",
           acc_ecc == 1.0,
           "ecc 1000/1000; lsh fail rate %.4f (report-only); collisions@H<=2 "
           "ecc %.1e lsh %.1e; lsh adjacent-byte bit-share %.3f"
           % (1.0 - acc_lsh, code_collision_rate(C_ecc),
              code_collision_rate(C_lsh), adj))

    # S8 — causality: the bed's future-leak check, replicated on every FAC
    # read path (addr / none / p4): a future byte must not move past scores.
    leaks = {}
    for mode in ("addr", "none", "p4"):
        torch.manual_seed(seed_for(f"fac:s8:{mode}"))
        m8 = FacModel(mode, "ecc", d=sd, layers=sl, block=sb).to(DEV).eval()
        g8 = torch.Generator().manual_seed(seed_for("fac:s8:x"))
        x8 = torch.randint(0, 256, (2, sb), generator=g8).to(DEV)
        with torch.no_grad():
            a = (m8.address(x8) @ m8.C.t())[0, 10]
            x8b = x8.clone()
            x8b[0, 40] = (x8b[0, 40] + 7) % 256
            b = (m8.address(x8b) @ m8.C.t())[0, 10]
        leaks[mode] = float((a - b).abs().max())
    record("S8", "causality (no future leak, all FAC read paths)",
           all(v <= 1e-4 for v in leaks.values()),
           "max |dscore@t=10| after t=40 edit: " +
           ", ".join(f"{k} {v:.1e}" for k, v in leaks.items()))

    # S9 — toggle/param-match honesty at the full operating point (d=192, L=4):
    # ce_fixedcode == ce minus EXACTLY the readout table; zero trainable
    # readout params; the two P4 cells are parameter-IDENTICAL.
    with torch.random.fork_rng():
        torch.manual_seed(seed_for("fac:s9"))
        ce_ref = bed.ByteLM("addr_msl64")
        torch.manual_seed(seed_for("fac:s9"))
        fc = FacModel("addr", "ecc")
        torch.manual_seed(seed_for("fac:s9"))
        p4c = bed.ByteLM("addr_head")
        torch.manual_seed(seed_for("fac:s9"))
        p4f = FacModel("p4", "ecc")
    head_n = ce_ref.head.weight.numel() + ce_ref.head.bias.numel()
    tr_ce, tr_fc = _trainable(ce_ref), _trainable(fc)
    fc_readout = 0 if isinstance(fc.lm.head, nn.Identity) else -1
    record("S9", "param match (readout delta exact; P4 cells identical)",
           tr_ce - tr_fc == head_n == 65792 and fc_readout == 0
           and _trainable(p4c) == _trainable(p4f),
           "ce %s vs ce_fixedcode %s (delta %s == readout %s; fixed-code "
           "readout trainable=0); p4 pair %s == %s"
           % (f"{tr_ce:,}", f"{tr_fc:,}", f"{tr_ce - tr_fc:,}",
              f"{head_n:,}", f"{_trainable(p4c):,}", f"{_trainable(p4f):,}"))

    # ------------------------------------------------------------------ table
    wall = time.time() - t0
    peak = (torch.cuda.max_memory_allocated() / 2**30) if DEV == "cuda" else 0.0
    print("\nFAC FORMULA-SMOKE BATTERY  (%s, %.1fs, peak %.2f GB)"
          % (DEV, wall, peak))
    print("-" * 100)
    npass = nfail = 0
    for tid, name, st, detail in RESULTS:
        npass += st == "PASS"
        nfail += st == "FAIL"
        print("%-5s %-4s %-46s %s" % (tid, st, name[:46], detail))
    print("-" * 100)
    print("PASS %d  FAIL %d  SKIP %d" % (npass, nfail,
                                         len(RESULTS) - npass - nfail))
    return nfail == 0


def print_launch_matrix(steps: int = 2000):
    print("\nARM MATRIX (verdict runs — NOT launched by this bed; P4 cells first):")
    py = ".venv/Scripts/python.exe"
    for arm in ("p4_addr_head_ce", "p4_addr_head_fac"):
        print(f"  {py} tools/fac_bed.py --arm {arm} --seed 0")
    for arm in ("ce", "ce_fixedcode", "fac_lsh", "fac_ecc", "fac_none"):
        for seed in (0, 1, 2):
            print(f"  {py} tools/fac_bed.py --arm {arm} --seed {seed}")
    print("  # knobs: --steps N | --t_loss {0.1,0.3,1.0} | --mu M | --code {ecc,lsh}")


def _in_notebook() -> bool:
    try:
        get_ipython()  # type: ignore[name-defined]  # noqa: F821
        return True
    except NameError:
        return False


if __name__ == "__main__":
    if _in_notebook():
        _ok = run_smokes()
        print_launch_matrix()
        print("Notebook mode: train_arm('fac_lsh', seed=0) in the next cell (GPU).")
    else:
        import argparse
        ap = argparse.ArgumentParser()
        ap.add_argument("--arm", type=str, default=None, choices=sorted(ARMS))
        ap.add_argument("--seed", type=int, default=0)
        ap.add_argument("--steps", type=int, default=2000)
        ap.add_argument("--smoke", action="store_true")
        ap.add_argument("--bench", action="store_true")
        ap.add_argument("--t_loss", type=float, default=T_LOSS)
        ap.add_argument("--mu", type=float, default=MU)
        ap.add_argument("--code", type=str, default=None, choices=("ecc", "lsh"))
        ap.add_argument("--data_root", type=str, default=None)
        ap.add_argument("--device", type=str, default="cuda")
        a, _ = ap.parse_known_args()
        if a.bench:
            bench(steps=20, device=a.device, data_root=a.data_root)
        elif a.arm and not a.smoke:
            train_arm(a.arm, seed=a.seed, steps=a.steps, device=a.device,
                      data_root=a.data_root, t_loss=a.t_loss, mu=a.mu,
                      code=a.code)
        else:
            ok = run_smokes()
            print_launch_matrix(a.steps)
            sys.exit(0 if ok else 1)