ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """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() | |