"""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()