ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """Batched byte-tick training (docs/BATCHED_TRAINING.md). | |
| * the vectorised loader's features == Atomizer.encode_bytes (both pack modes), and its | |
| multi-shard circular stream == StreamingPacketSource(--loop-shards) transitions; | |
| * B=1 batched path == single-stream AtomNativeModel.transition_loss path (logits, loss, | |
| grads, params after AdamW steps, field/persistence state) over 300 ticks including BPTT | |
| window cuts, atom flushes and an episode reset (BPTT K=16 + learnable leak + last-atom | |
| dropout; and K=0 + scalar leak); | |
| * B=4 == 4 independent B=1 runs (per-stream logits/losses, batch grad == mean of the | |
| per-stream grads), also for a per-stream (masked) reset; | |
| * trainer checkpoint loads in the old loader / chat tool / old trainer, same parameter | |
| names+shapes as a run_atom_native.py checkpoint; exact resume; | |
| * unsupported options raise; CUDA == CPU when CUDA is available. | |
| """ | |
| from __future__ import annotations | |
| import copy | |
| import subprocess | |
| import sys | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| import torch | |
| from src.atom_batched import BatchedByteTick, BatchedUnsupported | |
| from src.atom_native import AtomNativeModel | |
| from src.io.atomizer import Atomizer | |
| from src.io.batched_stream import BatchedByteStream, ByteCorpus | |
| from src.io.stream_corpus import StreamingPacketSource | |
| ROOT = Path(__file__).resolve().parents[1] | |
| PY = sys.executable | |
| TEXT = ("Utilisateur: Bonjour, comment ça va ?\nAssistant: Très bien, merci ! Il était une fois " | |
| "un petit atome qui parlait français... Qui es-tu ? 42 élèves, 7 déjà.\n") * 4 | |
| TOL = 1e-5 | |
| torch.set_num_threads(1) | |
| def _model(d=8, K=16, leak=0.05, learn=True, drop=0.5, thr=0.2, **kw): | |
| torch.manual_seed(0) | |
| m = AtomNativeModel(d_model=d, n_modes=d, n_atoms_max=64, max_payload_bytes=1, | |
| atomizer=Atomizer(max_span_bytes=1), field_max_rms=3.0, field_obligatory_hard=True, | |
| enable_merge=False, last_atom_readout=True, field_loss_weight=0.05, field_leak=leak, | |
| field_bptt=K, field_leak_learnable=learn, **kw) | |
| m.surface.payload_enabled = False | |
| m.surface.last_atom_dropout = drop | |
| # lower threshold so the consolidation/persistence path is exercised in a short run | |
| m.core.consolidation.consolidation_threshold = thr | |
| return m | |
| def _corpus(tmp: Path, parts=(TEXT,)) -> ByteCorpus: | |
| paths = [] | |
| for i, t in enumerate(parts): | |
| p = tmp / f"c{i:02d}.txt" | |
| p.write_bytes(t.encode("utf-8")) | |
| paths.append(p) | |
| return ByteCorpus(paths) | |
| def _grads(model): | |
| return {n: (p.grad.detach().clone() if p.grad is not None else torch.zeros_like(p)) | |
| for n, p in model.named_parameters() if p.requires_grad} | |
| def _opt(model): | |
| return torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=2e-3, weight_decay=0.01) | |
| class LoaderTests(unittest.TestCase): | |
| def test_features_equal_atomizer(self): | |
| with tempfile.TemporaryDirectory() as d: | |
| corpus = _corpus(Path(d)) | |
| raw = TEXT.encode("utf-8") | |
| for mode in ("linguistic", "eager"): | |
| pk = Atomizer(max_span_bytes=1, pack_mode=mode).encode_bytes(raw, reset=True, flush=False) | |
| st = BatchedByteStream(corpus, 1, offsets=[0], chunk=37, pack_mode=mode) | |
| for i in range(len(pk) - 1): | |
| c, f, t = st.next() | |
| self.assertEqual((int(c), int(t)), (pk[i].payload[0], pk[i + 1].payload[0])) | |
| self.assertTrue(torch.equal(f[0], pk[i].features), f"{mode} byte {i}") | |
| def test_multishard_circular_equals_stream_source(self): | |
| with tempfile.TemporaryDirectory() as d: | |
| parts = (TEXT[:100], TEXT[100:170], TEXT[170:260]) | |
| corpus = _corpus(Path(d), parts) | |
| src = StreamingPacketSource(Atomizer(max_span_bytes=1), corpus.paths, chunk_bytes=32, loop=True, | |
| prefetch=False) | |
| st = BatchedByteStream(corpus, 1, offsets=[0], chunk=50) | |
| n = 2 * corpus.total + 17 # two epoch wraps | |
| for i in range(n): | |
| cur, tgt = next(src) | |
| c, f, t = st.next() | |
| self.assertEqual((int(c), int(t)), (cur.payload[0], tgt.payload[0]), i) | |
| self.assertTrue(torch.equal(f[0], cur.features), i) | |
| def test_offsets_and_state_roundtrip(self): | |
| with tempfile.TemporaryDirectory() as d: | |
| corpus = _corpus(Path(d)) | |
| a = BatchedByteStream(corpus, 3, chunk=16) | |
| self.assertEqual(a.pos.tolist(), [0, corpus.total // 3, 2 * corpus.total // 3]) | |
| for _ in range(40): | |
| a.next() | |
| sd = a.state_dict() | |
| b = BatchedByteStream(corpus, 3, chunk=16) | |
| b.load_state_dict(sd) | |
| for _ in range(60): | |
| x, y = a.next(), b.next() | |
| for u, v in zip(x, y): | |
| self.assertTrue(torch.equal(u, v)) | |
| def _legacy_vs_batched(tc, K, leak, learn, n=300, ep=150, fl=32, d=8): | |
| m = _model(d=d, K=K, leak=leak, learn=learn) | |
| mb = copy.deepcopy(m) | |
| with tempfile.TemporaryDirectory() as tmp: | |
| corpus = _corpus(Path(tmp)) | |
| pk = Atomizer(max_span_bytes=1).encode_bytes(TEXT.encode("utf-8"), reset=True, flush=False) | |
| run = BatchedByteTick(mb, 1, atom_flush_every=fl) | |
| st = BatchedByteStream(corpus, 1, offsets=[0]) | |
| oa, ob = _opt(m), _opt(mb) | |
| m.train(); mb.train() | |
| acc_a = acc_b = None | |
| k = 0 | |
| worst = {"logit": 0.0, "loss": 0.0, "grad": 0.0} | |
| n_windows = 0 | |
| persist_seen = 0.0 | |
| for tick in range(n): | |
| if tick % ep == 0: | |
| m.reset_state(reset_atomizer=True) | |
| run.reset_state() | |
| c, f, t = st.next() | |
| torch.manual_seed(1000 + tick) # same last-atom dropout draw on both paths | |
| la, info = m.transition_loss(pk[tick], pk[tick + 1]) | |
| torch.manual_seed(1000 + tick) | |
| out = run.step(c, f, t) | |
| worst["logit"] = max(worst["logit"], float((info["surface"]["byte_logits"].detach().reshape(-1) | |
| - out["byte_logits"].detach().reshape(-1)).abs().max())) | |
| worst["loss"] = max(worst["loss"], abs(float(la.detach()) - float(out["loss"].detach().mean()))) | |
| persist_seen = max(persist_seen, float(run.persistence.abs().max())) | |
| step_now = True | |
| if K > 0: | |
| acc_a = la if k == 0 else acc_a + la | |
| acc_b = out["loss"].mean() if k == 0 else acc_b + out["loss"].mean() | |
| k += 1 | |
| step_now = k >= K or (tick + 1) % ep == 0 or tick == n - 1 | |
| if step_now: | |
| (acc_a / k).backward() | |
| (acc_b / k).backward() | |
| else: | |
| la.backward() | |
| out["loss"].mean().backward() | |
| if step_now: | |
| ga, gb = _grads(m), _grads(mb) | |
| for name in ga: | |
| scale = float(ga[name].abs().max()) + 1e-12 | |
| worst["grad"] = max(worst["grad"], float((ga[name] - gb[name]).abs().max()) / scale) | |
| for mod, o in ((m, oa), (mb, ob)): | |
| torch.nn.utils.clip_grad_norm_([p for p in mod.parameters() if p.requires_grad], 1.0) | |
| o.step() | |
| o.zero_grad(set_to_none=True) | |
| mod.stabilize_dynamics_parameters() | |
| m.detach_field_graph() | |
| run.detach_field_graph() | |
| acc_a = acc_b = None | |
| k = 0 | |
| n_windows += 1 | |
| if (tick + 1) % fl == 0: | |
| m.core.atoms = type(m.core.atoms)() | |
| m.core.energy_history = [] | |
| run.flush_atoms() | |
| tc.assertLessEqual(worst["logit"], TOL, worst) | |
| tc.assertLessEqual(worst["loss"], TOL, worst) | |
| tc.assertLessEqual(worst["grad"], 1e-4, worst) | |
| tc.assertTrue(torch.allclose(m.core.state.alpha, run.alpha[0], atol=TOL)) | |
| tc.assertTrue(torch.allclose(m.core.consolidation.persistence, run.persistence[0], atol=TOL)) | |
| tc.assertGreater(persist_seen, 0.0, "consolidation path not exercised") | |
| for (na, pa), (nb, pb) in zip(m.named_parameters(), mb.named_parameters()): | |
| if pa.requires_grad: | |
| tc.assertTrue(torch.allclose(pa, pb, atol=1e-5), na) | |
| tc.assertGreater(n_windows, 2) | |
| return worst | |
| class SingleStreamEquivalenceTests(unittest.TestCase): | |
| def test_b1_matches_single_stream_bptt(self): | |
| w = _legacy_vs_batched(self, K=16, leak=0.05, learn=True) | |
| print(f"\n[B=1 vs single-stream, K=16 learnable leak] worst |dlogit|={w['logit']:.3g} " | |
| f"|dloss|={w['loss']:.3g} rel dgrad={w['grad']:.3g}") | |
| def test_b1_matches_single_stream_no_bptt(self): | |
| w = _legacy_vs_batched(self, K=0, leak=0.3, learn=False, n=200) | |
| print(f"\n[B=1 vs single-stream, K=0 scalar leak] worst |dlogit|={w['logit']:.3g} " | |
| f"|dloss|={w['loss']:.3g} rel dgrad={w['grad']:.3g}") | |
| class BatchIndependenceTests(unittest.TestCase): | |
| def _run_fixed_weights(self, masked_reset: bool): | |
| B, K, EP, FL, N = 4, 16, 100, 32, 170 | |
| base = _model(K=K) | |
| with tempfile.TemporaryDirectory() as tmp: | |
| corpus = _corpus(Path(tmp)) | |
| offs = [0, 57, 191, 333] | |
| gen = torch.Generator().manual_seed(7) | |
| drops = torch.rand(N, B, generator=gen) < 0.5 | |
| mb = copy.deepcopy(base) | |
| mb.train() | |
| big = BatchedByteTick(mb, B, atom_flush_every=FL) | |
| sb = BatchedByteStream(corpus, B, offsets=offs) | |
| singles = [] | |
| for b in range(B): | |
| ms = copy.deepcopy(base) | |
| ms.train() | |
| singles.append((ms, BatchedByteTick(ms, 1, atom_flush_every=FL), | |
| BatchedByteStream(corpus, 1, offsets=[offs[b]]))) | |
| acc_big = None | |
| accs = [None] * B | |
| k = 0 | |
| for tick in range(N): | |
| if tick % EP == 0: | |
| big.reset_state() | |
| for _, r, _ in singles: | |
| r.reset_state() | |
| if masked_reset and tick == 60: # reset stream 2 only | |
| big.reset_state(torch.tensor([False, False, True, False])) | |
| singles[2][1].reset_state() | |
| c, f, t = sb.next() | |
| o = big.step(c, f, t, dropout_mask=drops[tick]) | |
| acc_big = o["loss"].mean() if k == 0 else acc_big + o["loss"].mean() | |
| for b, (ms, r, s) in enumerate(singles): | |
| cs, fs, ts = s.next() | |
| os_ = r.step(cs, fs, ts, dropout_mask=drops[tick, b:b + 1]) | |
| self.assertLessEqual(float((os_["byte_logits"][0] - o["byte_logits"][b]).abs().max()), TOL) | |
| self.assertLessEqual(abs(float(os_["loss"][0]) - float(o["loss"][b])), TOL) | |
| accs[b] = os_["loss"].mean() if k == 0 else accs[b] + os_["loss"].mean() | |
| k += 1 | |
| if k >= K or (tick + 1) % EP == 0 or tick == N - 1: | |
| (acc_big / k).backward() | |
| big.detach_field_graph() | |
| for b, (ms, r, _) in enumerate(singles): | |
| (accs[b] / k).backward() | |
| r.detach_field_graph() | |
| acc_big, accs, k = None, [None] * B, 0 | |
| if (tick + 1) % FL == 0: | |
| big.flush_atoms() | |
| for _, r, _ in singles: | |
| r.flush_atoms() | |
| gbig = _grads(mb) | |
| gs = [_grads(ms) for ms, _, _ in singles] | |
| for name, g in gbig.items(): | |
| mean = sum(x[name] for x in gs) / B | |
| scale = float(mean.abs().max()) + 1e-12 | |
| self.assertLessEqual(float((g - mean).abs().max()) / scale, 1e-4, name) | |
| for b, (_, r, _) in enumerate(singles): | |
| self.assertTrue(torch.allclose(r.alpha[0], big.alpha[b], atol=TOL)) | |
| def test_b4_equals_four_b1(self): | |
| self._run_fixed_weights(masked_reset=False) | |
| def test_b4_masked_reset(self): | |
| self._run_fixed_weights(masked_reset=True) | |
| class UnsupportedTests(unittest.TestCase): | |
| def test_unsupported_raise(self): | |
| cases = { | |
| "merge": lambda m: setattr(m, "enable_merge", True), | |
| "payload": lambda m: setattr(m.surface, "payload_enabled", True), | |
| "efference": lambda m: setattr(m, "_efference_every", 2), | |
| "spectrum": lambda m: (m.set_field_leak_learnable(False), m.set_field_leak_spectrum((0.01, 0.3))), | |
| "contrast": lambda m: setattr(m, "field_contrast_weight", 0.1), | |
| "slow_every": lambda m: setattr(m, "slow_every", 4), | |
| "adaptive": lambda m: setattr(m.surface, "last_atom_scale_adaptive", True), | |
| } | |
| for name, fn in cases.items(): | |
| m = _model() | |
| fn(m) | |
| with self.assertRaises(BatchedUnsupported, msg=name): | |
| BatchedByteTick(m, 2, atom_flush_every=64) | |
| torch.manual_seed(0) | |
| span = AtomNativeModel(d_model=8, n_modes=8, max_payload_bytes=4, field_obligatory_hard=True, | |
| enable_merge=False) | |
| span.surface.payload_enabled = False | |
| with self.assertRaises(BatchedUnsupported): | |
| BatchedByteTick(span, 2, atom_flush_every=64) | |
| with self.assertRaises(BatchedUnsupported): | |
| BatchedByteTick(_model(), 2, atom_flush_every=0) | |
| TRAIN_ARGS = ["--d-model", "8", "--n-modes", "8", "--n-atoms-max", "32", "--threads", "1", | |
| "--episode-length", "120", "--atom-flush-every", "32", "--field-bptt", "16", "--log-every", "40"] | |
| class TrainerCheckpointTests(unittest.TestCase): | |
| def _train(self, tmp: Path, out: str, *extra: str) -> subprocess.CompletedProcess: | |
| cmd = [PY, "tools/train_batched.py", "--data-glob", str(tmp / "c*.txt"), "--output-dir", str(tmp / out), | |
| *TRAIN_ARGS, *extra] | |
| res = subprocess.run(cmd, cwd=ROOT, capture_output=True, text=True, env={"PYTHONPATH": str(ROOT), | |
| "PATH": "/usr/bin:/bin"}) | |
| self.assertEqual(res.returncode, 0, res.stdout[-2000:] + res.stderr[-3000:]) | |
| return res | |
| def test_checkpoint_roundtrip_into_old_code(self): | |
| with tempfile.TemporaryDirectory() as d: | |
| tmp = Path(d) | |
| _corpus(tmp) | |
| self._train(tmp, "b", "-B", "3", "--total-ticks", "100", "--ckpt-every", "100", "--samples") | |
| ck = tmp / "b" / "atom_native_step_100_batched.pt" | |
| self.assertTrue(ck.exists()) | |
| blob = torch.load(ck, map_location="cpu", weights_only=False) | |
| # old loader | |
| m = AtomNativeModel(d_model=8, n_modes=8, n_atoms_max=32, max_payload_bytes=1, | |
| field_obligatory_hard=True, last_atom_readout=True, enable_merge=False) | |
| m.load(ck) | |
| self.assertTrue(torch.equal(m.surface.state_dict()["alpha_byte_proj.fc1.weight"], | |
| blob["surface"]["alpha_byte_proj.fc1.weight"])) | |
| self.assertEqual(m.field_bptt, 16) | |
| self.assertIsNotNone(m.field_leak_logit) | |
| # chat tool on CPU | |
| sys.path.insert(0, str(ROOT / "tools")) | |
| from chat_atom_native import generate_reply, load_model | |
| cm = load_model(ck, device="cpu") | |
| for (n1, p1), (n2, p2) in zip(cm.named_parameters(), m.named_parameters()): | |
| self.assertEqual(n1, n2) | |
| self.assertTrue(torch.equal(p1, p2), n1) | |
| torch.manual_seed(0) | |
| reply = generate_reply(cm, "Bonjour", max_packets=20, max_length=20) | |
| self.assertIsInstance(reply, str) | |
| # same parameter names/shapes as an old-trainer checkpoint of this recipe | |
| res = subprocess.run( | |
| [PY, "tools/run_atom_native.py", "--stream", "--data-glob", str(tmp / "c*.txt"), "--loop-shards", | |
| "--chunk-bytes", "512", "--output-dir", str(tmp / "old"), "--checkpoint-name", "old.pt", | |
| "--steps", "4", "--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", | |
| "--efference-every", "0", "--allow-parallel-train", "--atom-flush-every", "32", | |
| "--field-loss-weight", "0.05", "--field-leak", "0.05", "--field-bptt", "16", | |
| "--field-leak-learnable", "--resume", str(ck)], | |
| cwd=ROOT, capture_output=True, text=True, env={"PYTHONPATH": str(ROOT), "PATH": "/usr/bin:/bin"}) | |
| self.assertEqual(res.returncode, 0, res.stdout[-2000:] + res.stderr[-3000:]) | |
| self.assertIn("restored optimizer state", res.stdout) | |
| old = torch.load(tmp / "old" / "old.pt", map_location="cpu", weights_only=False) | |
| for key in ("compiler", "surface", "field_probe", "field_next_probe"): | |
| self.assertEqual({k: tuple(v.shape) for k, v in blob[key].items()}, | |
| {k: tuple(v.shape) for k, v in old[key].items()}, key) | |
| for key in ("encoder", "dynamics", "state", "consolidation"): | |
| self.assertEqual({k: tuple(v.shape) for k, v in blob["core"][key].items() if torch.is_tensor(v)}, | |
| {k: tuple(v.shape) for k, v in old["core"][key].items() if torch.is_tensor(v)}, key) | |
| self.assertEqual(tuple(blob["field_leak_logit"].shape), tuple(old["field_leak_logit"].shape)) | |
| self.assertEqual(int(old["training"]["step"]), 104) | |
| def test_exact_resume(self): | |
| with tempfile.TemporaryDirectory() as d: | |
| tmp = Path(d) | |
| _corpus(tmp) | |
| common = ["-B", "2", "--total-ticks", "160", "--ckpt-every", "80", "--no-samples"] | |
| self._train(tmp, "a", *common) | |
| self._train(tmp, "b", *common, "--run-ticks", "80") | |
| self._train(tmp, "b", *common) | |
| a = torch.load(tmp / "a" / "atom_native_step_160_batched.pt", map_location="cpu", weights_only=False) | |
| b = torch.load(tmp / "b" / "atom_native_step_160_batched.pt", map_location="cpu", weights_only=False) | |
| for key in ("compiler", "surface", "field_probe"): | |
| for k in a[key]: | |
| self.assertTrue(torch.equal(a[key][k], b[key][k]), (key, k)) | |
| self.assertTrue(torch.equal(a["field_leak_logit"], b["field_leak_logit"])) | |
| ra, rb = a["training"]["batched"]["runtime"], b["training"]["batched"]["runtime"] | |
| self.assertTrue(torch.equal(ra["alpha"], rb["alpha"])) | |
| self.assertEqual(a["training"]["batched"]["stream"]["pos"], b["training"]["batched"]["stream"]["pos"]) | |
| class CudaTests(unittest.TestCase): | |
| def test_cuda_matches_cpu(self): | |
| torch.backends.cuda.matmul.allow_tf32 = False | |
| m = _model(K=16) | |
| mc = copy.deepcopy(m).cuda() | |
| with tempfile.TemporaryDirectory() as tmp: | |
| corpus = _corpus(Path(tmp)) | |
| r1, r2 = BatchedByteTick(m, 4, atom_flush_every=32), BatchedByteTick(mc, 4, atom_flush_every=32) | |
| s1 = BatchedByteStream(corpus, 4) | |
| s2 = BatchedByteStream(corpus, 4, device="cuda") | |
| m.train(); mc.train() | |
| drops = torch.rand(64, 4) < 0.5 | |
| acc1 = acc2 = 0 | |
| for tick in range(64): | |
| a, b = s1.next(), s2.next() | |
| o1 = r1.step(*a, dropout_mask=drops[tick]) | |
| o2 = r2.step(*b, dropout_mask=drops[tick]) | |
| self.assertLessEqual(float((o1["byte_logits"] - o2["byte_logits"].cpu()).abs().max()), 1e-3) | |
| acc1 = acc1 + o1["loss"].mean() | |
| acc2 = acc2 + o2["loss"].mean() | |
| if (tick + 1) % 32 == 0: | |
| r1.flush_atoms(); r2.flush_atoms() | |
| acc1.backward(); acc2.backward() | |
| for (n, p1), (_, p2) in zip(m.named_parameters(), mc.named_parameters()): | |
| if p1.grad is not None: | |
| scale = float(p1.grad.abs().max()) + 1e-9 | |
| self.assertLessEqual(float((p1.grad - p2.grad.cpu()).abs().max()) / scale, 1e-2, n) | |
| if __name__ == "__main__": | |
| unittest.main() | |