"""Unit tests: byte-tick decode label fix + opt-in truncated BPTT through the field (docs/DECODE_FIX_AND_BPTT.md). Decode fix (automatic in byte-tick, escape flag ``legacy_gen_labels`` / --legacy-gen-labels): * byte-tick generated commits are bit-identical (label + features + atomizer state) to ``encode_bytes`` of the same byte stream, in both pack modes; * the teacher-forced and generation paths give identical logits (max |dlogit| == 0); * ``legacy_gen_labels=True`` reproduces the pre-fix commit (commit f8a8918) exactly: labels, features, and a sampled generate_packets() byte string; * non-byte-tick (max_span > 1) commits are unchanged whatever the flag. Field BPTT (--field-bptt K, --field-leak-learnable): * off (default): multi-tick train+eval trace bit-identical to commit f8a8918 (no leak, scalar leak 0.3, leak spectrum), trainer checkpoint bit-identical too; * on: forward values (field, logits, loss) bit-identical to off at equal weights; eval/no_grad ticks never keep a graph; * on: d loss_t / d atom_{t-L} != 0 for L = 1..K-1, and == 0 across a window cut; off: == 0 for every L >= 1 (the ROOT_CAUSE measurement); * learnable leak: init == scalar leak (bit-identical trace), receives gradient under BPTT; * checkpoint roundtrip; old checkpoints load as off; trainer CLI default off, on, and CLI wins after --resume in both directions. """ from __future__ import annotations import importlib import shutil import subprocess import sys import tempfile import unittest from pathlib import Path import torch from src.atom_native import AtomNativeModel from src.io.atomizer import Atomizer ROOT = Path(__file__).resolve().parents[1] LEGACY_COMMIT = "f8a8918" TEXT = "Bonjour, comment ça va ? Il était une fois un petit atome. Qui es-tu ?\nOui, d'accord... " * 3 _LEG: dict = {} def _legacy(): """(AtomNativeModel, Atomizer) classes of the pre-change commit, or None.""" if "mods" in _LEG: return _LEG["mods"] try: tmp = Path(tempfile.mkdtemp(prefix="legacy_dfb_")) pkg = tmp / "legacy_src_dfb" pkg.mkdir() files = subprocess.run(["git", "-C", str(ROOT), "ls-tree", "-r", "--name-only", LEGACY_COMMIT, "src"], capture_output=True, text=True, check=True).stdout.split() for f in files: if not f.endswith(".py"): continue code = subprocess.run(["git", "-C", str(ROOT), "show", f"{LEGACY_COMMIT}:{f}"], capture_output=True, text=True, check=True).stdout dst = pkg / Path(f).relative_to("src") dst.parent.mkdir(parents=True, exist_ok=True) dst.write_text(code, encoding="utf-8") for d in [pkg, *[p for p in pkg.rglob("*") if p.is_dir()]]: (d / "__init__.py").touch(exist_ok=True) sys.path.insert(0, str(tmp)) try: an = importlib.import_module("legacy_src_dfb.atom_native") at = importlib.import_module("legacy_src_dfb.io.atomizer") finally: sys.path.remove(str(tmp)) _LEG["mods"] = (an.AtomNativeModel, at.Atomizer) except Exception: # pragma: no cover _LEG["mods"] = None return _LEG["mods"] def _build(cls=AtomNativeModel, atom_cls=Atomizer, **kw): torch.manual_seed(0) return cls( d_model=8, n_modes=8, n_atoms_max=32, max_payload_bytes=1, atomizer=atom_cls(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, **kw, ) def _run(model, n_train=40, n_eval=20): opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-3) packets = model.atomizer.encode(TEXT, reset=True) model.reset_state(reset_atomizer=False) model.train() trace = [] torch.manual_seed(5) for cur, tgt in list(zip(packets[:-1], packets[1:]))[:n_train]: opt.zero_grad(set_to_none=True) loss, _ = model.transition_loss(cur, tgt) loss.backward() opt.step() model.stabilize_dynamics_parameters() trace.append(loss.detach().clone()) trace.append(model.core.state.alpha.detach().clone()) model.eval() with torch.no_grad(): for cur, tgt in list(zip(packets[:-1], packets[1:]))[n_train:n_train + n_eval]: loss, _ = model.transition_loss(cur, tgt) trace.append(loss.detach().clone()) trace.append(model.core.state.alpha.detach().clone()) trace.extend(p.detach().clone() for p in model.parameters()) return trace def _eq(tc, a, b): tc.assertEqual(len(a), len(b)) for x, y in zip(a, b): tc.assertTrue(torch.equal(x, y)) class DecodeLabelFixTests(unittest.TestCase): def test_byte_tick_commit_equals_training_encode(self) -> None: raw = TEXT.encode("utf-8") for mode in ("linguistic", "eager"): tr = Atomizer(max_span_bytes=1, pack_mode=mode).encode_bytes(raw, reset=True, flush=False) gen_at = Atomizer(max_span_bytes=1, pack_mode=mode) gen = [gen_at.packet_from_payload(bytes([b])) for b in raw] self.assertEqual(len(tr), len(gen)) for a, b in zip(tr, gen): self.assertEqual((a.payload, a.start, a.boundary, a.level), (b.payload, b.start, b.boundary, b.level)) self.assertTrue(torch.equal(a.features, b.features)) self.assertEqual({p.boundary for p in gen} - {"max_span", "newline", "punctuation", "whitespace"}, set()) if mode == "linguistic": self.assertEqual({p.boundary for p in gen}, {"max_span", "newline"}) def test_legacy_flag_reproduces_old_labels_exactly(self) -> None: leg = _legacy() if leg is None: self.skipTest("legacy commit not available") _, OldAtomizer = leg raw = TEXT.encode("utf-8") for mode in ("linguistic", "eager"): new = Atomizer(max_span_bytes=1, pack_mode=mode); new.legacy_gen_labels = True old = OldAtomizer(max_span_bytes=1, pack_mode=mode) for b in raw: a, o = new.packet_from_payload(bytes([b])), old.packet_from_payload(bytes([b])) self.assertEqual((a.boundary, a.start, a.level), (o.boundary, o.start, o.level)) self.assertTrue(torch.equal(a.features, o.features)) # the fix does change labels (else the test above would prove nothing) fixed = Atomizer(max_span_bytes=1) self.assertEqual(fixed.packet_from_payload(b" ").boundary, "max_span") legacy = Atomizer(max_span_bytes=1); legacy.legacy_gen_labels = True self.assertEqual(legacy.packet_from_payload(b" ").boundary, "whitespace") def test_non_byte_tick_unchanged(self) -> None: leg = _legacy() if leg is None: self.skipTest("legacy commit not available") _, OldAtomizer = leg for flag in (False, True): new = Atomizer(max_span_bytes=8); new.legacy_gen_labels = flag old = OldAtomizer(max_span_bytes=8) for pl in (b"abc ", b"Oui,", b"x", b" ", b"\n", b"mot"): a, o = new.packet_from_payload(pl), old.packet_from_payload(pl) self.assertEqual(a.boundary, o.boundary) self.assertTrue(torch.equal(a.features, o.features)) def test_teacher_forced_equals_generation_path(self) -> None: m = _build(field_leak=0.3) m.eval() text = "Utilisateur: Qui es-tu ?\nAssistant: Je suis Atom, un petit modèle, d'accord ?\n" raw = text.encode() with torch.no_grad(): m.reset_state(reset_atomizer=True) pk = m.atomizer.encode_bytes(raw, reset=True) tf = [m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pk[:-1]] m.reset_state(reset_atomizer=True) n0 = 30 pp = m.atomizer.encode_bytes(raw[:n0], reset=True) gen = [m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pp[:-1]] cur = pp[-1] for b in raw[n0:]: gen.append(m.forward_packet(cur)["surface"]["byte_logits"].clone()) cur = m.atomizer.packet_from_payload(bytes([b])) self.assertEqual(len(tf), len(gen)) self.assertEqual(max(float((a - b).abs().max()) for a, b in zip(tf, gen)), 0.0) # legacy labels do diverge (sanity) m.atomizer.legacy_gen_labels = True with torch.no_grad(): m.reset_state(reset_atomizer=True) pp = m.atomizer.encode_bytes(raw[:n0], reset=True) g2 = [m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pp[:-1]] cur = pp[-1] for b in raw[n0:]: g2.append(m.forward_packet(cur)["surface"]["byte_logits"].clone()) cur = m.atomizer.packet_from_payload(bytes([b])) self.assertGreater(max(float((a - b).abs().max()) for a, b in zip(tf, g2)), 0.0) def test_generate_legacy_flag_bit_identical_to_old_code(self) -> None: leg = _legacy() if leg is None: self.skipTest("legacy commit not available") OldModel, OldAtomizer = leg new = _build(field_leak=0.3) old = _build(cls=OldModel, atom_cls=OldAtomizer, field_leak=0.3) for x, y in zip(new.parameters(), old.parameters()): self.assertTrue(torch.equal(x, y)) new.atomizer.legacy_gen_labels = True for prompt in ("Bonjour, comment ça va ?", "Il était une fois"): for seed in (0, 1, 2): torch.manual_seed(seed) a = new.generate_packets(prompt, max_packets=60, temperature=0.8, top_k=12, deterministic=False, max_length=60) torch.manual_seed(seed) b = old.generate_packets(prompt, max_packets=60, temperature=0.8, top_k=12, deterministic=False, max_length=60) self.assertEqual(a, b) def _atom_grads_by_lag(model, n=24, lags=(1, 2, 3, 5, 8)): """||d loss_T / d atom.r_{T-L}|| on a train-mode tick sequence (no updates).""" model.train() model.surface.last_atom_dropout = 0.0 stash = [] orig = model.compiler.forward def hooked(features, atom_count=0): atom, op, conf = orig(features, atom_count) atom.r.retain_grad(); atom.E.retain_grad() stash.append(atom) return atom, op, conf model.compiler.forward = hooked try: pk = model.atomizer.encode(TEXT, reset=True) model.reset_state(reset_atomizer=False) for i in range(n - 1): model.forward_packet(pk[i]) loss, _ = model.transition_loss(pk[n - 1], pk[n]) loss.backward() finally: model.compiler.forward = orig out = {} for L in (0, *lags): a = stash[n - 1 - L] g = 0.0 for t in (a.r, a.E): if t.grad is not None: g += float(t.grad.norm()) out[L] = g return out class FieldBpttTests(unittest.TestCase): def test_default_off(self) -> None: m = _build() self.assertEqual(m.field_bptt, 0) self.assertIsNone(m.field_leak_logit) self.assertEqual(m._bptt_checkpoint(), {}) def test_off_bit_identical_to_legacy_commit(self) -> None: leg = _legacy() if leg is None: self.skipTest("legacy commit not available") OldModel, OldAtomizer = leg for kw in ({}, {"field_leak": 0.3}, {"field_leak_spectrum": (0.02, 0.3)}): _eq(self, _run(_build(**kw)), _run(_build(cls=OldModel, atom_cls=OldAtomizer, **kw))) def test_on_forward_values_identical(self) -> None: for kw in ({"field_leak": 0.3}, {"field_leak": 0.05}, {"field_leak_spectrum": (0.02, 0.3)}, {}): a, b = _build(**kw), _build(field_bptt=16, **kw) pk = a.atomizer.encode(TEXT, reset=True) outs = [] for m in (a, b): m.train(); m.surface.last_atom_dropout = 0.0 m.reset_state(reset_atomizer=False) torch.manual_seed(3) tr = [] for cur, tgt in list(zip(pk[:-1], pk[1:]))[:30]: loss, info = m.transition_loss(cur, tgt) tr += [loss.detach().clone(), info["surface"]["byte_logits"].detach().clone(), m.core.state.alpha.detach().clone()] outs.append(tr) _eq(self, outs[0], outs[1]) self.assertIsNotNone(b._bptt_alpha) self.assertTrue(b._bptt_alpha.requires_grad) def test_eval_never_keeps_graph(self) -> None: m = _build(field_bptt=8, field_leak=0.05) pk = m.atomizer.encode(TEXT, reset=True) m.eval() with torch.no_grad(): for p in pk[:10]: m.forward_packet(p) self.assertIsNone(m._bptt_alpha) m.train() with torch.no_grad(): m.forward_packet(pk[10]) self.assertIsNone(m._bptt_alpha) def test_gradient_reaches_past_atoms_only_when_on(self) -> None: off = _atom_grads_by_lag(_build(field_leak=0.05)) self.assertGreater(off[0], 0.0) for L in (1, 2, 3, 5, 8): self.assertEqual(off[L], 0.0, f"lag {L} off") on = _atom_grads_by_lag(_build(field_leak=0.05, field_bptt=64)) for L in (0, 1, 2, 3, 5, 8): self.assertGreater(on[L], 0.0, f"lag {L} on") def test_window_cut_stops_gradient(self) -> None: m = _build(field_leak=0.05, field_bptt=4) m.train(); m.surface.last_atom_dropout = 0.0 stash = [] orig = m.compiler.forward def hooked(features, atom_count=0): atom, op, conf = orig(features, atom_count) atom.r.retain_grad(); stash.append(atom) return atom, op, conf m.compiler.forward = hooked pk = m.atomizer.encode(TEXT, reset=True) m.reset_state(reset_atomizer=False) for i in range(6): m.forward_packet(pk[i]) m.detach_field_graph() # window boundary after tick 6 for i in range(6, 9): m.forward_packet(pk[i]) loss, _ = m.transition_loss(pk[9], pk[10]) loss.backward() m.compiler.forward = orig self.assertIsNotNone(stash[7].r.grad) self.assertGreater(float(stash[7].r.grad.norm()), 0.0) for a in stash[:6]: self.assertTrue(a.r.grad is None or float(a.r.grad.norm()) == 0.0) def test_learnable_leak(self) -> None: # init equals the scalar leak (bit-identical forward trace, no BPTT, eval ticks) a, b = _build(field_leak=0.3), _build(field_leak=0.3, field_leak_learnable=True) lam = float(b.learnable_leak_rates().detach()[0]) self.assertAlmostEqual(lam, 0.3, places=6) b2 = _build(field_leak=lam) pk = a.atomizer.encode(TEXT, reset=True) tr = [] for m in (b2, b): m.eval(); m.reset_state(reset_atomizer=False) with torch.no_grad(): tr.append([m.forward_packet(p)["surface"]["byte_logits"].clone() for p in pk[:20]]) for x, y in zip(*tr): self.assertTrue(torch.allclose(x, y, atol=1e-6)) # receives gradient under BPTT c = _build(field_leak=0.05, field_leak_learnable=True, field_bptt=32) c.train() c.reset_state(reset_atomizer=False) tot = 0 for cur, tgt in list(zip(pk[:-1], pk[1:]))[:20]: loss, _ = c.transition_loss(cur, tgt); tot = tot + loss tot.backward() self.assertIsNotNone(c.field_leak_logit.grad) self.assertGreater(float(c.field_leak_logit.grad.abs().sum()), 0.0) with self.assertRaises(ValueError): _build(field_leak_spectrum=(0.02, 0.3), field_leak_learnable=True) def test_checkpoint_roundtrip(self) -> None: with tempfile.TemporaryDirectory() as t: m = _build(field_leak=0.05, field_bptt=64, field_leak_learnable=True) with torch.no_grad(): m.field_leak_logit.add_(torch.linspace(-1, 1, 8).view(8, 1)) p = Path(t) / "on.pt" m.save(p) f = _build() f.load(p) self.assertEqual(f.field_bptt, 64) self.assertTrue(torch.equal(f.field_leak_logit.detach(), m.field_leak_logit.detach())) # off ckpt: no new keys; loads as off into an "on" model p2 = Path(t) / "off.pt" _build().save(p2) blob = torch.load(p2, weights_only=False) self.assertNotIn("field_bptt", blob) self.assertNotIn("field_leak_logit", blob) g = _build(field_bptt=8, field_leak_learnable=True) g.load(p2) self.assertEqual(g.field_bptt, 0) self.assertIsNone(g.field_leak_logit) class TrainerCliTests(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", "10", "--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) -> dict: data = tmp / "corpus.txt" if not data.exists(): data.write_text(TEXT * 20, encoding="utf-8") cmd = [sys.executable, str(root / "tools" / "run_atom_native.py"), "--data", str(data), "--output-dir", str(tmp / name.replace(".pt", "_dir")), "--checkpoint-name", name, *self.BASE, *extra] subprocess.run(cmd, cwd=root, env={**self.ENV, "PYTHONPATH": str(root)}, check=True, capture_output=True, text=True) return torch.load(tmp / name.replace(".pt", "_dir") / name, map_location="cpu", weights_only=False) def test_trainer_off_bit_identical_to_legacy_commit(self) -> None: with tempfile.TemporaryDirectory() as t: tmp = Path(t) legacy_root = tmp / "legacy_repo" try: legacy_root.mkdir() arch = subprocess.run(["git", "-C", str(ROOT), "archive", LEGACY_COMMIT, "src", "tools"], capture_output=True, check=True).stdout subprocess.run(["tar", "-x", "-C", str(legacy_root)], input=arch, check=True) except Exception: self.skipTest("legacy commit not available") # (a) defaults with efference off: bit-identical; (b) efference on + # --legacy-gen-labels: bit-identical (the fix changes efference commits by design) self._compare(self._train(tmp, "new.pt", "--efference-every", "0"), self._train(tmp, "old.pt", "--efference-every", "0", root=legacy_root)) self._compare(self._train(tmp, "new2.pt", "--efference-every", "2", "--legacy-gen-labels"), self._train(tmp, "old2.pt", "--efference-every", "2", root=legacy_root)) def _compare(self, new, old): if True: 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())) 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_cli_on_and_wins_after_resume(self) -> None: with tempfile.TemporaryDirectory() as t: tmp = Path(t) a = self._train(tmp, "a.pt", "--efference-every", "0") self.assertNotIn("field_bptt", a) b = self._train(tmp, "b.pt", "--efference-every", "0", "--field-bptt", "8", "--field-leak", "0.05", "--field-leak-learnable") self.assertEqual(b["field_bptt"], 8) self.assertIn("field_leak_logit", b) # bptt ckpt resumed without the flags -> off c = self._train(tmp, "c.pt", "--efference-every", "0", "--resume", str(tmp / "b_dir" / "b.pt")) self.assertNotIn("field_bptt", c) self.assertNotIn("field_leak_logit", c) # plain ckpt resumed with the flag -> on; learnable leak kept from the ckpt when resumed with it d = self._train(tmp, "d.pt", "--efference-every", "0", "--resume", str(tmp / "a_dir" / "a.pt"), "--field-bptt", "16") self.assertEqual(d["field_bptt"], 16) e = self._train(tmp, "e.pt", "--efference-every", "0", "--resume", str(tmp / "b_dir" / "b.pt"), "--field-bptt", "8", "--field-leak-learnable", "--steps", "1") self.assertEqual(e["field_bptt"], 8) self.assertIn("field_leak_logit", e) # legacy gen labels flag runs (efference path) and is recorded only when on f = self._train(tmp, "f.pt", "--legacy-gen-labels") self.assertNotIn("field_bptt", f) if __name__ == "__main__": unittest.main()