"""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"]) @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available") 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()