File size: 3,553 Bytes
906715b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """CPU-testable invariants on a tiny synthetic trunk (no downloads).
assert_toggle_law: with every anchor disabled, wrapped-model logits ==
base logits BIT-EXACT. assert_detach_bitexact: detach(verify=True)
round-trips. These are the package's contract with the research line's
toggle law (exp011: max|Δlogit| = 0.0).
"""
from __future__ import annotations
import torch
import torch.nn as nn
class _TinyBlock(nn.Module):
def __init__(self, d):
super().__init__()
self.lin = nn.Linear(d, d)
def forward(self, x):
return self.lin(x) # tensor output (tuple path tested too)
class _TupleBlock(_TinyBlock):
def forward(self, x):
return (self.lin(x), None)
class _TinyConfig:
model_type = "tiny"
hidden_size = 32
_name_or_path = "amoe/tiny-test"
class _Inner(nn.Module):
def __init__(self, d, L, tuple_blocks):
super().__init__()
cls = _TupleBlock if tuple_blocks else _TinyBlock
self.layers = nn.ModuleList([cls(d) for _ in range(L)])
class TinyTrunk(nn.Module):
"""model.model.layers shape -> resolves via generic_causal."""
def __init__(self, d=32, L=4, vocab=17, tuple_blocks=False):
super().__init__()
self.config = _TinyConfig()
self.model = _Inner(d, L, tuple_blocks)
self.emb = nn.Embedding(vocab, d)
self.head = nn.Linear(d, vocab)
def forward(self, input_ids):
h = self.emb(input_ids)
for layer in self.model.layers:
out = layer(h)
h = out[0] if isinstance(out, tuple) else out
class _O:
pass
o = _O()
o.logits = self.head(h)
return o
def _fresh_anchor_ckpt(d, L, name="test"):
from ..core.adapter import RelayPatchwork
from ..io.checkpoint import AnchorCheckpoint
torch.manual_seed(7)
state = {}
for i in range(L):
a = RelayPatchwork(d)
for k, v in a.state_dict().items():
state[f"{i}.{k}"] = v.clone()
return AnchorCheckpoint(state, {"name": name,
"base_model_id": "amoe/tiny-test"})
def assert_toggle_law(tuple_blocks=False) -> None:
from ..runtime.attach import attach
torch.manual_seed(3)
m = TinyTrunk(tuple_blocks=tuple_blocks)
ids = torch.arange(6).unsqueeze(0) % 5
with torch.no_grad():
base = m(ids).logits.clone()
h = attach(m, _fresh_anchor_ckpt(32, 4))
with h.all_off():
with torch.no_grad():
off = m(ids).logits
assert torch.equal(off, base), "toggle law violated (single)"
# dispatch path
m2 = TinyTrunk(tuple_blocks=tuple_blocks)
with torch.no_grad():
base2 = m2(ids).logits.clone()
h2 = attach(m2, [_fresh_anchor_ckpt(32, 4, "a"),
_fresh_anchor_ckpt(32, 4, "b")], dispatch="init")
with h2.all_off():
with torch.no_grad():
off2 = m2(ids).logits
assert torch.equal(off2, base2), "toggle law violated (dispatch)"
def assert_detach_bitexact() -> None:
from ..runtime.attach import attach
torch.manual_seed(5)
m = TinyTrunk()
h = attach(m, _fresh_anchor_ckpt(32, 4))
h.detach(verify=True) # raises on any non-bit-exact detach
def run_all() -> None:
assert_toggle_law(tuple_blocks=False)
assert_toggle_law(tuple_blocks=True)
assert_detach_bitexact()
print("amoe invariants: toggle law (tensor+tuple, single+dispatch) "
"and bit-exact detach — ALL GREEN")
if __name__ == "__main__":
run_all()
|