sleep-form-rescue-count / reader /read_shuffle_cost.py
RonBuildsAI's picture
Rescue count: 40 seeds, plain vs sleep form
1272654 verified
Raw History Blame Contribute Delete
11.8 kB
# SPDX-License-Identifier: BSD-3-Clause-Clear
# Copyright (c) 2026 Ronald Lee Hampton. All rights reserved.
# No express or implied licenses to any party's patent rights are granted by this license.
#
# read_shuffle_cost.py: re-measures the shuffle cost of every shipped brain and re-counts the switchboards.
#
# Usage (from the package folder):
# python -I reader/read_shuffle_cost.py --brains brains --data data --expect results/per_brain_record.csv
# Needs: Python 3, numpy, torch, safetensors. CPU only, one thread, about a minute for the 80 brains. No network.
# Side effects: none; it prints only (redirect the output to keep it).
#
# THE MEASURE. Each brain is run on the 256 held-out windows (66 tokens each: two context words, then 64 targets), 16,384 positions in
# one batch. The reasoning core's quantized output (qk, one 32-number vector per position) is then shuffled across positions with
# numpy default_rng(0).permutation(16384), everything else held fixed, and the network is run again. The SHUFFLE COST is the mean, over
# the positions whose target is not one of the 40 excluded ids, of (cross-entropy shuffled minus cross-entropy as trained), in nats.
# THE LINE (fixed before the run): a core is a SWITCHBOARD if its shuffle cost is below 0.025 nats, else a CONTEXT WORKER.
# A RESCUE is a seed whose plain (ST-0) core is a switchboard and whose sleep-form (S-SLEEP) core is a context worker; a HARM is the
# reverse. THE BAR (fixed before the run): S-SLEEP ends with fewer switchboards than ST-0, AND at least 5 seeds are discordant, AND the
# exact one-sided sign test on the discordant seeds gives p at or below 0.05.
#
# This file holds only the forward pass the measure needs (no training code). Tensor names keep the run's own prefixes:
# tiny.* the routed base network: two word embeddings summed, a router to one of 48 experts (top 1, its score nudged by the pull
# 0.4 x (mean size minus the expert's size), sizes in "y"), each expert a two-layer block, a layer norm and a shared head;
# rc.* the reasoning core: written from the integer part's hidden state, one gated pass, quantized to whole numbers in -127..127;
# cc.* the small integer part with ternary weights: reads the two context words' ids as 12 bits each, and writes a gain on the
# base network's expert output from its own hidden state and the core's quantized output (zero gain at the start).
# cc.Wv and cc.bv, a two-output head the measure does not use, were removed from the shipped files (results/head_removed_record.csv).
# Every operation below is written in the same order as the run's own code so that the numbers can match it bit for bit on the same
# software and hardware; on other hardware or torch versions the last digits may differ (the --expect check reports by how much).
import sys, os, csv, glob, math, hashlib, argparse
import numpy as np
import torch
import torch.nn.functional as F
V, D, T, CTX, NBITS = 4096, 32, 64, 2, 12
S_PULL, RHO = 0.4, 0.5
LINE = 0.025
FORMS = ("ST-0", "S-SLEEP")
P_BAR, MIN_DISC = 0.05, 5
BITW = 2 ** torch.arange(NBITS)
def sha256_file(p):
h = hashlib.sha256()
with open(p, "rb") as f:
for blk in iter(lambda: f.read(1 << 20), b""):
h.update(blk)
return h.hexdigest()
def sign_p(r, n):
"""exact one-sided p of at least r rescues among n discordant seeds, each a fair coin."""
return float(sum(math.comb(n, k) for k in range(r, n + 1)) / 2 ** n) if n else 1.0
def tern(w):
return (w > 0.5).float() - (w < -0.5).float()
def ste_tern(w): # the training-time expression, kept so the floating-point result is the same
return w + (tern(w) - w)
def ste_round(x):
return x + (torch.round(x) - x)
def ste_floor(x):
return x + (torch.floor(x) - x)
def rms(x, w):
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * w
def msg(xs):
"""the two context words' token ids as 12 bits each, +1/-1, newest word first."""
parts = [((x[:, None] // BITW) % 2).float() * 2 - 1 for x in xs[:2]]
return torch.cat(parts, dim=-1)
def hidden(P, xs):
"""the integer part's hidden state: ternary weights, whole-number biases, clamps; every value a whole number."""
m = msg(xs)
a1 = m @ ste_tern(P["cc.W1"]).t() + ste_round(P["cc.b1"] * 1.0)
h1 = torch.clamp(a1, 0.0, 31.0)
a2 = h1 @ ste_tern(P["cc.W2"]).t() + ste_round(P["cc.b2"] * 16.0)
return torch.clamp(ste_floor(a2 / 4.0), 0.0, 127.0)
def pull_bias(P):
y = P["y"].numpy()
return torch.tensor(S_PULL * (y.mean() - y), dtype=torch.float32)
def forward(P, xs, bias, h2, perm=None):
"""logits at every position; with perm, the reasoning core's quantized output is shuffled across positions first."""
# the routed base network
x = F.embedding(xs[0], P["tiny.emb.0.weight"])
x = x + F.embedding(xs[1], P["tiny.emb.1.weight"])
r = F.linear(x, P["tiny.router.weight"], P["tiny.router.bias"])
score = r + bias
p = F.softmax(r, dim=-1)
sel = torch.argmax(score, dim=-1)
g = p.gather(1, sel[:, None]).squeeze(1)
h = F.gelu(torch.bmm(x[:, None, :], P["tiny.w1"][sel])).squeeze(1)
y = torch.bmm(h[:, None, :], P["tiny.w2"][sel]).squeeze(1)
contrib = g[:, None] * y
# the reasoning core, written from the integer part's hidden state
K0 = (1 - RHO) * torch.tanh(rms(h2 / 127.0, P["rc.nwc.weight"]) @ P["rc.wkc"])
z = rms(K0, P["rc.ns.weight"])
u = F.silu(z @ P["rc.w1"]) * (z @ P["rc.w3"])
K2 = RHO * K0 + (1 - RHO) * torch.tanh(u @ P["rc.wh"])
qk = torch.clamp(ste_round(K2 * 127.0), -127.0, 127.0)
if perm is not None:
qk = qk[perm] # THE SHUFFLE: the core's output moved across positions, nothing else changed
# the integer part's write: a gain on the expert output
c = h2 @ ste_tern(P["cc.Wcg"]).t()
c = c + qk @ ste_tern(P["cc.Wck"]).t()
gain = 2.0 * torch.sigmoid(c / 256.0) - 1.0
o = x + contrib * (1.0 + gain)
o = F.layer_norm(o, (D,), P["tiny.norm.weight"], P["tiny.norm.bias"], 1e-5)
return F.linear(o, P["tiny.head.weight"], P["tiny.head.bias"])
def heldout(data_dir):
win = np.load(os.path.join(data_dir, "heldout_windows_uint16.npy"), allow_pickle=False).astype(np.int64)
if win.shape != (256, T + CTX):
raise SystemExit("REFUSED: held-out windows have shape %s" % (win.shape,))
xs = [torch.from_numpy(np.ascontiguousarray(win[:, CTX - 1 - k:T + CTX - 1 - k])).reshape(-1) for k in range(CTX)]
tgt = torch.from_numpy(np.ascontiguousarray(win[:, CTX:])).reshape(-1)
ids = [int(l) for l in open(os.path.join(data_dir, "excluded_target_ids.txt"), encoding="utf-8") if l.strip() and not l.startswith("#")]
excl = torch.zeros(V, dtype=torch.bool)
excl[torch.tensor(ids)] = True
allowed = (~excl[tgt]).numpy()
return win, xs, tgt, allowed, ids
def main(a):
from safetensors.torch import load_file
torch.set_num_threads(1)
win, xs, tgt, allowed, ids = heldout(a.data)
perm = torch.from_numpy(np.random.default_rng(0).permutation(len(tgt)))
print("held-out windows %d, positions %d, scored positions %d, excluded ids %d" % (win.shape[0], len(tgt), int(allowed.sum()), len(ids)))
print("torch %s, numpy %s, threads %d" % (torch.__version__, np.__version__, torch.get_num_threads()))
paths = sorted(glob.glob(os.path.join(a.brains, "*.safetensors")))
per = {}
for path in paths:
name = os.path.basename(path)[:-len(".safetensors")]
form, s = name.split("_")[0], int(name.split("_")[1][1:])
P = load_file(path)
with torch.no_grad():
bias = pull_bias(P)
h2 = hidden(P, xs)
lg_on = forward(P, xs, bias, h2)
lg_sh = forward(P, xs, bias, h2, perm)
ce_on = F.cross_entropy(lg_on, tgt, reduction="none").double().numpy()
ce_sh = F.cross_entropy(lg_sh, tgt, reduction="none").double().numpy()
cost = float((ce_sh - ce_on)[allowed].mean())
per[(form, s)] = {"name": name, "cost": cost, "label": "SWITCHBOARD" if cost < LINE else "CONTEXT WORKER",
"ce_on": float(ce_on[allowed].mean()),
"logits_sha256": hashlib.sha256(lg_on.contiguous().numpy().tobytes()).hexdigest(),
"file_sha256": sha256_file(path)}
print("%-8s s%d shuffle cost %.6f %-14s (held-out loss as trained %.6f)" % (form, s, cost, per[(form, s)]["label"], per[(form, s)]["ce_on"]))
seeds = sorted({s for (_f, s) in per})
if not seeds or not all((f, s) in per for f in FORMS for s in seeds):
print("REFUSED: no brains, or a seed is missing one of its two brains")
return 2
sb = {f: [s for s in seeds if per[(f, s)]["label"] == "SWITCHBOARD"] for f in FORMS}
res = [s for s in seeds if per[("ST-0", s)]["label"] == "SWITCHBOARD" and per[("S-SLEEP", s)]["label"] == "CONTEXT WORKER"]
harm = [s for s in seeds if per[("ST-0", s)]["label"] == "CONTEXT WORKER" and per[("S-SLEEP", s)]["label"] == "SWITCHBOARD"]
n = len(res) + len(harm)
p = sign_p(len(res), n)
passed = len(sb["S-SLEEP"]) < len(sb["ST-0"]) and n >= MIN_DISC and p <= P_BAR
print("")
print("seeds %d (%d to %d); THE LINE: switchboard below %.3f nats of shuffle cost" % (len(seeds), seeds[0], seeds[-1], LINE))
for f in FORMS:
print(" %-8s switchboards %d of %d: %s" % (f, len(sb[f]), len(seeds), sb[f]))
print(" rescues %d %s; harms %d %s; discordant %d; exact one-sided sign test p %.5f" % (len(res), res, len(harm), harm, n, p))
print("THE BAR (S-SLEEP fewer, at least %d discordant, p <= %.2f): %s" % (MIN_DISC, P_BAR, "CONFIRMED" if passed else "NOT CONFIRMED"))
near = sorted((abs(v["cost"] - LINE), "%s s%d" % k, round(v["cost"], 6)) for k, v in per.items())[:3]
print(" nearest the line: %s" % [(k, c) for _d, k, c in near])
rc = 0
if a.expect:
rows = {(r["form"], int(r["seed"])): r for r in csv.DictReader(open(a.expect, encoding="utf-8", newline=""))}
keys = [k for k in per if k in rows]
lab_eq = sum(1 for k in keys if rows[k]["label"] == per[k]["label"])
file_eq = sum(1 for k in keys if rows[k]["safetensors_sha256"] == per[k]["file_sha256"])
cost_eq = sum(1 for k in keys if float(rows[k]["shuffle_cost"]) == per[k]["cost"])
ce_eq = sum(1 for k in keys if float(rows[k]["ce_on_allowed_mean"]) == per[k]["ce_on"])
lg_eq = sum(1 for k in keys if rows[k]["logits_sha256_on"] == per[k]["logits_sha256"])
dmax = max(abs(float(rows[k]["shuffle_cost"]) - per[k]["cost"]) for k in keys) if keys else float("nan")
cmax = max(abs(float(rows[k]["ce_on_allowed_mean"]) - per[k]["ce_on"]) for k in keys) if keys else float("nan")
print("")
print("AGAINST %s (%d rows):" % (os.path.basename(a.expect), len(rows)))
print(" brain files equal to the record's sha256: %d of %d" % (file_eq, len(per)))
print(" labels equal: %d of %d" % (lab_eq, len(per)))
print(" shuffle costs equal bit for bit: %d of %d; largest absolute difference %.3g nats" % (cost_eq, len(per), dmax))
print(" held-out loss as trained equal bit for bit: %d of %d; largest absolute difference %.3g nats" % (ce_eq, len(per), cmax))
print(" logits as trained equal bit for bit (sha256): %d of %d" % (lg_eq, len(per)))
if lab_eq != len(per) or file_eq != len(per) or len(rows) != len(per):
rc = 1
return rc
if __name__ == "__main__":
ap = argparse.ArgumentParser()
ap.add_argument("--brains", default="brains")
ap.add_argument("--data", default="data")
ap.add_argument("--expect", default="")
sys.exit(main(ap.parse_args()))