File size: 14,245 Bytes
b9301f1 | 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 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 | """Unit tests for docs/MIX_DATA_LONG_MEMORY.md (all levers opt-in; defaults bit-identical).
* Trainer with default flags, with the current BPTT recipe (--field-bptt --field-leak-learnable),
and in --stream mode with --stream-skip-packets: checkpoint (all tensors, config, loss
trajectory) bit-identical to the pre-change commit 507903a.
* --stream-skip-bytes / StreamingPacketSource(start_offset=): byte-exact resume at the offset
(single and multi shard, mid-codepoint -> next codepoint), later loops start at byte 0,
start_offset=0 identical to the old source.
* --field-leak-slow-block: off = plain sigmoid leak, no checkpoint key; on = slow 位 inside
[floor, floor+cap], init floor+cap/4, stable net retention for every allowed energy_decay,
other modes untouched, train(BPTT)/eval forward identical, gradient reaches the slow logits
and past atoms, EMA gain, checkpoint roundtrip, disable keeps 位, trainer CLI.
"""
from __future__ import annotations
import math
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path
import torch
from src.atom_native import AtomNativeModel, DEFAULT_ENERGY_DECAY_BOUNDS, field_leak_slow_block_modes
from src.io.atomizer import Atomizer
from src.io.stream_corpus import StreamingPacketSource
ROOT = Path(__file__).resolve().parents[1]
BASE_COMMIT = "507903a"
TEXT = "Bonjour, comment 莽a va ? Il 茅tait une fois un petit atome. Qui es-tu ?\nOui, d'accord... " * 3
def _model(learn=True, leak=0.05, n=8, d=8):
torch.manual_seed(3)
m = AtomNativeModel(d_model=d, n_modes=n, n_atoms_max=32, max_payload_bytes=1,
atomizer=Atomizer(max_span_bytes=1), field_max_rms=3.0, enable_merge=False,
field_loss_weight=0.05, field_obligatory_hard=True, last_atom_readout=True,
field_leak=leak, field_leak_learnable=learn)
m.surface.payload_enabled = False
return m
def _trace(m, text=TEXT[:120], train=False):
m.train(train)
m.reset_state(reset_atomizer=True)
pk = m.atomizer.encode_bytes(text.encode(), reset=True)
out = []
for c, t in zip(pk[:-1], pk[1:]):
loss, info = m.transition_loss(c, t)
out.append(float(loss.detach()))
return out, m.core.state.alpha.detach().clone()
class StreamSkipBytesTests(unittest.TestCase):
def _bytes(self, src, n):
return b"".join(next(src)[1].payload for _ in range(n))
def test_offset_resume_exact(self):
with tempfile.TemporaryDirectory() as t:
a, b = Path(t) / "a.txt", Path(t) / "b.txt"
a.write_bytes(("茅" * 50 + "abc\n" * 30).encode()); b.write_bytes(("xyz 莽a " * 40).encode())
cat = a.read_bytes() + b.read_bytes()
for off in (0, 7, 100, len(a.read_bytes()) - 1, len(a.read_bytes()) + 5):
s = StreamingPacketSource(Atomizer(max_span_bytes=1), [a, b], chunk_bytes=64, loop=True,
prefetch=False, start_offset=off)
x = next(s)
got = x[0].payload + x[1].payload + self._bytes(s, 60)
o = off
while o < len(cat) and 0x80 <= cat[o] <= 0xBF:
o += 1 # moved to the next codepoint start
self.assertEqual(got, cat[o:o + len(got)], f"offset {off}")
# mid-codepoint: '茅' = 2 bytes; offset 1 -> starts at byte 2
s = StreamingPacketSource(Atomizer(max_span_bytes=1), [a], chunk_bytes=64, loop=True,
prefetch=False, start_offset=1)
self.assertEqual(next(s)[0].payload, cat[2:3])
def test_second_epoch_from_zero_and_default_identical(self):
with tempfile.TemporaryDirectory() as t:
a = Path(t) / "a.txt"; a.write_bytes(b"0123456789" * 5)
s = StreamingPacketSource(Atomizer(max_span_bytes=1), [a], chunk_bytes=16, loop=True,
prefetch=False, start_offset=45)
x = next(s)
seq = x[0].payload + x[1].payload + self._bytes(s, 9)
self.assertEqual(seq, b"56789" + b"012345")
s0 = StreamingPacketSource(Atomizer(max_span_bytes=1), [a], chunk_bytes=16, loop=True, prefetch=False)
s1 = StreamingPacketSource(Atomizer(max_span_bytes=1), [a], chunk_bytes=16, loop=True, prefetch=False,
start_offset=0)
for _ in range(120):
x, y = next(s0), next(s1)
self.assertEqual((x[0].payload, x[1].payload), (y[0].payload, y[1].payload))
self.assertTrue(torch.equal(x[0].features, y[0].features))
class SlowBlockTests(unittest.TestCase):
def test_off_default(self):
m = _model()
self.assertIsNone(m.field_leak_slow_block)
self.assertTrue(torch.equal(m.learnable_leak_rates(), torch.sigmoid(m.field_leak_logit)))
self.assertNotIn("field_leak_slow_block", m._bptt_checkpoint())
self.assertIsNone(m.slow_block_input_gain())
def test_modes_contiguous(self):
self.assertEqual(field_leak_slow_block_modes(32, 0.25), list(range(12, 20)))
self.assertEqual(field_leak_slow_block_modes(8, 0.25), [3, 4])
def test_rates_floor_cap_stability(self):
m = _model(n=32)
before = m.learnable_leak_rates().detach().clone()
m.set_field_leak_slow_block(0.25, cap=0.02)
lam = m.learnable_leak_rates().detach().reshape(-1)
modes = m.field_leak_slow_block["modes"]
floor = m._slow_block_floor.reshape(-1)
for i in range(32):
if i in modes:
self.assertAlmostEqual(float(lam[i]), float(floor[i]) + 0.02 * 0.25, places=6)
else:
self.assertEqual(float(lam[i]), float(before.reshape(-1)[i]))
# every allowed energy_decay: net retention (1-位_max)路exp(dt(sin 蠁 + ed - 1)) < 1 even at 位=floor
lo, hi = DEFAULT_ENERGY_DECAY_BOUNDS
for i in modes:
phi = 2 * math.pi * i / 32
for ed in (lo, hi):
self.assertLess((1 - float(floor[i])) * math.exp(0.01 * (math.sin(phi) + ed - 1)), 1.0)
with torch.no_grad():
m.field_leak_logit.fill_(50.0)
lam = m.learnable_leak_rates().detach().reshape(-1)
for i in modes:
self.assertLessEqual(float(lam[i]), float(floor[i]) + 0.02 + 1e-7)
def test_gain_and_train_eval_identical_and_grad(self):
m = _model(n=8)
m.set_field_leak_slow_block(0.25, cap=0.02)
m.set_field_bptt(64)
ev, a_ev = _trace(m, train=False)
m.surface.last_atom_dropout = 0.0
with torch.no_grad():
tr, a_tr = _trace(m, train=True)
self.assertEqual(ev, tr)
self.assertTrue(torch.equal(a_ev, a_tr))
# BPTT graph: loss at the end reaches the slow logits and an early atom
m.train(); m.reset_state(reset_atomizer=True)
pk = m.atomizer.encode_bytes(TEXT[:60].encode(), reset=True)
acc = 0
for c, t in zip(pk[:-1], pk[1:]):
acc = acc + m.transition_loss(c, t)[0]
acc.backward()
g = m.field_leak_logit.grad.reshape(-1)
self.assertTrue(all(float(g[i]) != 0.0 for i in m.field_leak_slow_block["modes"]))
# gain = 位/ref on slow modes, 1 elsewhere
gain = m.slow_block_input_gain().detach().reshape(-1)
lam = m.learnable_leak_rates().detach().reshape(-1)
for i in range(8):
exp = float(lam[i]) / 0.05 if i in m.field_leak_slow_block["modes"] else 1.0
self.assertAlmostEqual(float(gain[i]), exp, places=6)
def test_gain_none_equals_plain_learnable_with_same_rates(self):
a = _model(n=8); a.set_field_leak_slow_block(0.25, cap=0.02, gain="none")
b = _model(n=8)
with torch.no_grad():
lam = a.learnable_leak_rates().detach()
b.field_leak_logit.copy_(torch.log(lam / (1 - lam)))
ta, aa = _trace(a); tb, ab = _trace(b)
for x, y in zip(ta, tb):
self.assertAlmostEqual(x, y, places=4)
self.assertTrue(torch.allclose(aa, ab, atol=1e-5))
def test_checkpoint_roundtrip_and_disable(self):
with tempfile.TemporaryDirectory() as t:
m = _model(n=8); m.set_field_leak_slow_block(0.25, cap=0.03)
m.stabilize_dynamics_parameters() # load() applies the same repair
with torch.no_grad():
m.field_leak_logit.add_(torch.randn_like(m.field_leak_logit) * 0.3)
lam = m.learnable_leak_rates().detach().clone()
p = Path(t) / "x.pt"; m.save(p)
g = _model(n=8); g.load(p)
self.assertEqual(g.field_leak_slow_block["cap"], 0.03)
self.assertTrue(torch.equal(g.learnable_leak_rates().detach(), lam))
self.assertEqual(_trace(g)[0], _trace(m)[0])
g.set_field_leak_slow_block(None)
self.assertIsNone(g.field_leak_slow_block)
self.assertTrue(torch.allclose(g.learnable_leak_rates().detach(), lam, atol=1e-6))
# old-layout checkpoint (no slow key) loads with the block off
q = Path(t) / "y.pt"; _model(n=8).save(q)
h = _model(n=8); h.set_field_leak_slow_block(0.25); h.load(q)
self.assertIsNone(h.field_leak_slow_block)
class TrainerTests(unittest.TestCase):
ENV = {"OMP_NUM_THREADS": "1", "MKL_NUM_THREADS": "1", "PATH": "/usr/bin:/bin"}
BASE = ["--steps", "40", "--d-model", "8", "--n-modes", "8", "--n-atoms-max", "32",
"--max-span-bytes", "1", "--field-obligatory-hard", "--no-enable-merge", "--no-payload-copy",
"--last-atom-readout", "--allow-parallel-train", "--seed", "1", "--atom-flush-every", "16",
"--episode-length", "25", "--efference-every", "0", "--field-loss-weight", "0.05",
"--last-atom-dropout", "0.5", "--field-leak", "0.3"]
def _train(self, tmp: Path, name: str, *extra: str, root: Path = ROOT, stream=False) -> dict:
data = tmp / "corpus.txt"
if not data.exists():
data.write_text(TEXT * 20, encoding="utf-8")
src = ["--stream", "--data-glob", str(tmp / "corpus*.txt"), "--loop-shards", "--chunk-bytes", "512"] if stream \
else ["--data", str(data)]
cmd = [sys.executable, str(root / "tools" / "run_atom_native.py"), *src,
"--output-dir", str(tmp / name.replace(".pt", "_dir")), "--checkpoint-name", name, *self.BASE, *extra]
r = subprocess.run(cmd, cwd=root, env={**self.ENV, "PYTHONPATH": str(root)}, capture_output=True, text=True)
self.assertEqual(r.returncode, 0, r.stderr[-2000:])
return torch.load(tmp / name.replace(".pt", "_dir") / name, map_location="cpu", weights_only=False)
def _compare(self, new, old):
for part in ("compiler", "surface"):
for k, v in old[part].items():
self.assertTrue(torch.equal(v, new[part][k]), f"{part}.{k}")
self.assertTrue(torch.equal(old["core"]["state"]["alpha"], new["core"]["state"]["alpha"]))
for k, v in old["core"]["dynamics"].items():
self.assertTrue(torch.equal(v, new["core"]["dynamics"][k]))
self.assertEqual(set(old.keys()), set(new.keys()))
if "field_leak_logit" in old:
self.assertTrue(torch.equal(old["field_leak_logit"], new["field_leak_logit"]))
self.assertEqual(old["config"], new["config"])
self.assertEqual([r["loss"] for r in old["training"]["trajectory"]],
[r["loss"] for r in new["training"]["trajectory"]])
def test_defaults_bit_identical_to_base_commit(self):
with tempfile.TemporaryDirectory() as t:
tmp = Path(t); old_root = tmp / "base_repo"
try:
old_root.mkdir()
arch = subprocess.run(["git", "-C", str(ROOT), "archive", BASE_COMMIT, "src", "tools"],
capture_output=True, check=True).stdout
subprocess.run(["tar", "-x", "-C", str(old_root)], input=arch, check=True)
except Exception:
self.skipTest("base commit not available")
self._compare(self._train(tmp, "n1.pt"), self._train(tmp, "o1.pt", root=old_root))
recipe = ("--field-bptt", "8", "--field-leak", "0.05", "--field-leak-learnable")
self._compare(self._train(tmp, "n2.pt", *recipe), self._train(tmp, "o2.pt", *recipe, root=old_root))
st = ("--stream-skip-packets", "37", *recipe)
self._compare(self._train(tmp, "n3.pt", *st, stream=True),
self._train(tmp, "o3.pt", *st, root=old_root, stream=True))
def test_cli_slow_block_and_skip_bytes(self):
with tempfile.TemporaryDirectory() as t:
tmp = Path(t)
recipe = ("--field-bptt", "8", "--field-leak", "0.05", "--field-leak-learnable")
a = self._train(tmp, "a.pt", *recipe, "--field-leak-slow-block", "0.25", "--stream-skip-bytes", "100",
stream=True)
self.assertEqual(a["field_leak_slow_block"]["cap"], 0.02)
import json
cfg = json.loads((tmp / "a_dir" / "run_metrics.json").read_text())["config"]
self.assertEqual(cfg["stream_skip_bytes"], 100)
self.assertEqual(cfg["field_leak_slow_block"]["frac"], 0.25)
# resumed without the flag -> off (CLI wins); with it -> kept (logits not re-initialised)
b = self._train(tmp, "b.pt", *recipe, "--resume", str(tmp / "a_dir" / "a.pt"), stream=True)
self.assertNotIn("field_leak_slow_block", b)
c = self._train(tmp, "c.pt", *recipe, "--field-leak-slow-block", "0.25", "--steps", "1",
"--resume", str(tmp / "a_dir" / "a.pt"), stream=True)
self.assertIn("field_leak_slow_block", c)
n = self._train(tmp, "n.pt", *recipe, stream=True)
ncfg = json.loads((tmp / "n_dir" / "run_metrics.json").read_text())["config"]
self.assertNotIn("stream_skip_bytes", ncfg)
self.assertNotIn("field_leak_slow_block", ncfg)
self.assertNotIn("field_leak_slow_block", n)
if __name__ == "__main__":
unittest.main()
|