ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """Run and document the atom-native ATOM experiment. | |
| Usage: | |
| PYTHONPATH=. .venv/bin/python tools/run_atom_native.py \ | |
| --data /tmp/wiki.train.raw --steps 500 \ | |
| --output-dir checkpoints/atom_native_500 | |
| Streaming full-dataset ingest (no giant packet list in RAM): | |
| PYTHONPATH=. .venv/bin/python tools/run_atom_native.py \ | |
| --stream --data-glob 'data/shards/*.txt' --steps 50000 \ | |
| --output-dir checkpoints/atom_native_stream | |
| The script intentionally uses the raw local corpus and never calls Hugging | |
| Face or the legacy GPT-2 tokenizer. It trains one transition per atom packet. | |
| Episodes still pick random corpus starts in non-stream mode, but by default the | |
| toroidal field is NOT wiped every short window (see ``--episode-length`` / | |
| ``--no-episode-reset``) so dialogue-scale persistence can accumulate. Use | |
| ``--episode-reset`` to restore the old hard reset-every-episode behaviour. | |
| In ``--stream`` mode the continuum is preferred: shards are read online and | |
| ``--no-episode-reset`` keeps the field alive across the infinite shard loop. | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import atexit | |
| import fcntl | |
| import json | |
| import os | |
| import random | |
| import time | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from src.atom_native import DEFAULT_ENERGY_DECAY_BOUNDS, AtomNativeModel | |
| from src.io.atomizer import Atomizer | |
| from src.io.stream_corpus import StreamingPacketSource, resolve_shard_paths | |
| # Advisory flock so two run_atom_native.py processes do not hammer one box by accident. | |
| # Default path (cwd-relative): checkpoints/atom-ai-train.lock | |
| # Fallback if checkpoints/ cannot be created: /tmp/atom-ai-train.lock | |
| DEFAULT_TRAIN_LOCK_PATH = Path("checkpoints") / "atom-ai-train.lock" | |
| FALLBACK_TRAIN_LOCK_PATH = Path("/tmp/atom-ai-train.lock") | |
| def acquire_train_lock(lock_path: Path | None = None): | |
| """Non-blocking exclusive flock; fail-fast if another train holds it. | |
| Returns an open file object that must stay open for the lock lifetime. | |
| Writes this process PID into the lock file for diagnostics. | |
| """ | |
| candidates: list[Path] = [] | |
| if lock_path is not None: | |
| candidates.append(Path(lock_path)) | |
| else: | |
| candidates.append(DEFAULT_TRAIN_LOCK_PATH) | |
| candidates.append(FALLBACK_TRAIN_LOCK_PATH) | |
| last_error: Exception | None = None | |
| for path in candidates: | |
| try: | |
| path.parent.mkdir(parents=True, exist_ok=True) | |
| fh = open(path, "a+", encoding="utf-8") | |
| except OSError as exc: | |
| last_error = exc | |
| continue | |
| try: | |
| fcntl.flock(fh.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) | |
| except BlockingIOError: | |
| fh.seek(0) | |
| holder = fh.read().strip() or "unknown" | |
| fh.close() | |
| raise SystemExit( | |
| f"error: another train holds the advisory lock\n" | |
| f" lock path: {path.resolve()}\n" | |
| f" holder: {holder}\n" | |
| f" this PID: {os.getpid()}\n" | |
| f"Pass --allow-parallel-train to opt out (not recommended on one box)." | |
| ) from None | |
| except OSError as exc: | |
| fh.close() | |
| last_error = exc | |
| continue | |
| fh.seek(0) | |
| fh.truncate() | |
| fh.write(f"pid={os.getpid()}\n") | |
| fh.flush() | |
| os.fsync(fh.fileno()) | |
| def _release(handle=fh, lock=path) -> None: | |
| try: | |
| if not handle.closed: | |
| fcntl.flock(handle.fileno(), fcntl.LOCK_UN) | |
| handle.close() | |
| except (OSError, ValueError): | |
| pass | |
| atexit.register(_release) | |
| print(f"train lock acquired: {path.resolve()} (pid={os.getpid()})") | |
| return fh | |
| raise SystemExit( | |
| f"error: could not create/acquire train lock " | |
| f"(tried {[str(p) for p in candidates]}): {last_error}" | |
| ) | |
| PROMPTS = ["The future of AI is", "Bonjour, je m'appelle", "Valkyria Chronicles III is"] | |
| def seed_all(seed: int) -> None: | |
| random.seed(seed) | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| def finite_model(model: AtomNativeModel) -> bool: | |
| return all(torch.isfinite(parameter).all().item() for parameter in model.parameters()) | |
| def transition_metrics( | |
| model: AtomNativeModel, | |
| current, | |
| target, | |
| train: bool, | |
| efference: bool = False, | |
| free_run_golds=None, | |
| free_run_weight: float = 1.0, | |
| ) -> tuple[float, dict]: | |
| if train: | |
| if free_run_golds is not None: | |
| loss, info = model.free_run_aux_loss( | |
| current, list(free_run_golds), weight=float(free_run_weight) | |
| ) | |
| elif efference: | |
| loss, info = model.efference_loss(current, target) | |
| else: | |
| loss, info = model.transition_loss(current, target) | |
| if not torch.isfinite(loss): | |
| raise FloatingPointError("non-finite atom-native loss") | |
| loss.backward() | |
| slow_head = getattr(model, "slow_readout", None) | |
| if slow_head is None: | |
| grad_norm = torch.nn.utils.clip_grad_norm_(list(model.trainable_parameters), 1.0) | |
| else: | |
| # --slow-mode-readout: clip the new head on its own so its (fresh, | |
| # large) gradients never rescale the control's parameter updates. | |
| head_ids = {id(p) for p in slow_head.parameters()} | |
| grad_norm = torch.nn.utils.clip_grad_norm_( | |
| [p for p in model.trainable_parameters if id(p) not in head_ids], 1.0 | |
| ) | |
| head_norm = torch.nn.utils.clip_grad_norm_( | |
| [p for p in slow_head.parameters() if p.requires_grad], 1.0 | |
| ) | |
| info = {**info, "slow_head_grad_norm": float(head_norm)} | |
| if not torch.isfinite(torch.as_tensor(grad_norm)): | |
| raise FloatingPointError("non-finite atom-native gradient") | |
| return float(loss.detach().item()), {**info, "grad_norm": float(grad_norm)} | |
| with torch.no_grad(): | |
| loss, info = model.transition_loss(current, target) | |
| return float(loss.item()), info | |
| def _bptt_meta(args) -> dict: | |
| """Config keys for the opt-in levers, present only when on (off = files as before).""" | |
| out = {} | |
| if int(getattr(args, "field_bptt", 0) or 0) > 0: | |
| out["field_bptt"] = int(args.field_bptt) | |
| if getattr(args, "field_leak_learnable", False): | |
| out["field_leak_learnable"] = True | |
| if float(getattr(args, "field_leak_slow_block", 0.0) or 0.0) > 0: | |
| out["field_leak_slow_block"] = {"frac": float(args.field_leak_slow_block), | |
| "cap": float(args.field_leak_slow_cap), "gain": args.field_leak_slow_gain} | |
| if getattr(args, "legacy_gen_labels", False): | |
| out["legacy_gen_labels"] = True | |
| if int(getattr(args, "stream_skip_bytes", 0) or 0) > 0: | |
| out["stream_skip_bytes"] = int(args.stream_skip_bytes) | |
| return out | |
| def bptt_window_step(model, optimizer, acc_loss, n_in_window): | |
| """Backprop the accumulated window loss (mean over its ticks), clip, step, | |
| and cut the field graph at the window boundary.""" | |
| loss = acc_loss / float(max(n_in_window, 1)) | |
| if not torch.isfinite(loss): | |
| raise FloatingPointError("non-finite atom-native loss (bptt window)") | |
| loss.backward() | |
| grad_norm = torch.nn.utils.clip_grad_norm_(list(model.trainable_parameters), 1.0) | |
| if not torch.isfinite(torch.as_tensor(grad_norm)): | |
| raise FloatingPointError("non-finite atom-native gradient (bptt window)") | |
| optimizer.step() | |
| optimizer.zero_grad(set_to_none=True) | |
| model.detach_field_graph() | |
| return float(grad_norm) | |
| def main() -> None: | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument( | |
| "--data", | |
| default=None, | |
| help="single corpus file (optional if --data-glob is set)", | |
| ) | |
| parser.add_argument( | |
| "--data-glob", | |
| default=None, | |
| help="glob of shard files, e.g. 'data/shards/*.txt' (ordered lexicographically)", | |
| ) | |
| parser.add_argument( | |
| "--stream", | |
| action="store_true", | |
| help="stream shards online (required for large corpora / globs); " | |
| "do not encode the whole file into one packet list", | |
| ) | |
| parser.add_argument( | |
| "--chunk-bytes", | |
| type=int, | |
| default=1 << 20, | |
| help="UTF-8-safe read chunk size in stream mode (default 1048576)", | |
| ) | |
| parser.add_argument( | |
| "--loop-shards", | |
| action=argparse.BooleanOptionalAction, | |
| default=True, | |
| help="in stream mode, loop shards forever for an infinite continuum (default: loop)", | |
| ) | |
| parser.add_argument( | |
| "--stream-buffer", | |
| type=int, | |
| default=2048, | |
| help="ring-buffer size of recent packets for stream-mode validation (default 2048)", | |
| ) | |
| parser.add_argument( | |
| "--stream-prefetch", | |
| action=argparse.BooleanOptionalAction, | |
| default=True, | |
| help="in --stream mode, prefetch next byte chunk on a background thread " | |
| "(default: on; use --no-stream-prefetch to disable)", | |
| ) | |
| parser.add_argument( | |
| "--stream-skip-packets", | |
| type=int, | |
| default=0, | |
| help="in --stream mode, skip N packet transitions before training " | |
| "(default 0; use k*steps when chaining resumed chunks so each " | |
| "chunk sees fresh corpus instead of re-training the prefix)", | |
| ) | |
| parser.add_argument( | |
| "--stream-skip-bytes", | |
| type=int, | |
| default=0, | |
| help="opt-in (default 0 = off): in --stream mode start the FIRST pass at this byte offset " | |
| "of the concatenated shards (seek, no replay; next codepoint start). In byte-tick " | |
| "(--max-span-bytes 1) 1 byte = 1 packet, so this replaces a long --stream-skip-packets " | |
| "replay on large corpora (atomizer state starts fresh at the offset)", | |
| ) | |
| parser.add_argument( | |
| "--allow-parallel-train", | |
| action="store_true", | |
| default=False, | |
| help="opt out of the default advisory train lock " | |
| f"(default lock: {DEFAULT_TRAIN_LOCK_PATH})", | |
| ) | |
| parser.add_argument( | |
| "--train-lock-path", | |
| default=None, | |
| help="override advisory lock path " | |
| f"(default: {DEFAULT_TRAIN_LOCK_PATH}; fallback {FALLBACK_TRAIN_LOCK_PATH})", | |
| ) | |
| parser.add_argument("--output-dir", default="checkpoints/atom_native_chat") | |
| parser.add_argument("--checkpoint-name", default="atom_native.pt") | |
| parser.add_argument("--resume", default=None, help="optional checkpoint to resume") | |
| parser.add_argument("--steps", type=int, default=500) | |
| parser.add_argument( | |
| "--episode-length", | |
| type=int, | |
| default=512, | |
| help="packets per random corpus window (default 512; was 64)", | |
| ) | |
| parser.add_argument( | |
| "--episode-reset", | |
| action=argparse.BooleanOptionalAction, | |
| default=False, | |
| help="if set, wipe toroidal/atomizer state at each episode boundary " | |
| "(default: keep field across windows for dialogue continuity)", | |
| ) | |
| parser.add_argument("--d-model", type=int, default=8) | |
| parser.add_argument("--n-modes", type=int, default=8) | |
| parser.add_argument("--n-atoms-max", type=int, default=128) | |
| parser.add_argument("--max-span-bytes", type=int, default=32, | |
| help="Atomizer/surface max span bytes (default 32; was 16)") | |
| parser.add_argument( | |
| "--atomizer-pack", | |
| choices=("linguistic", "eager"), | |
| default="linguistic", | |
| help="Atomizer packing: linguistic=longer WS/punct-aligned spans (FR dialogue); eager=flush every boundary (legacy)", | |
| ) | |
| parser.add_argument("--learning-rate", type=float, default=3e-4) | |
| parser.add_argument( | |
| "--surface-learning-rate", | |
| type=float, | |
| default=None, | |
| help="optional higher LR for AtomSurfaceHead only (param group); " | |
| "default: same as --learning-rate", | |
| ) | |
| parser.add_argument( | |
| "--field-leak", | |
| type=float, | |
| default=0.0, | |
| help="per-tick field leak λ in [0,1): α ← (1-λ)·α before each atom " | |
| "injection (contractive recency memory; keeps α off the RMS cap). " | |
| "Default 0.0 = exact legacy behavior. CLI wins after --resume.", | |
| ) | |
| parser.add_argument( | |
| "--field-leak-spectrum", | |
| type=float, | |
| nargs=2, | |
| default=None, | |
| metavar=("LMIN", "LMAX"), | |
| help="opt-in multi-timescale leak: per-mode λ_m geometric from LMIN (slow, " | |
| "modes with the least intrinsic decay, sin(mode_phase) high) to LMAX " | |
| "(fast); α_m ← (1-λ_m)·α_m before each injection. REPLACES --field-leak " | |
| "when set. Default off = exact current behavior. CLI wins after --resume " | |
| "(omit = off even if the checkpoint had one). See docs/FIELD_LEAK_SPECTRUM.md.", | |
| ) | |
| parser.add_argument( | |
| "--field-leak-spectrum-gain", | |
| choices=["none", "ema", "sqrt"], | |
| default="none", | |
| help="per-mode input gain with --field-leak-spectrum: none (g=1), ema " | |
| "(g=λ_m/LMAX: every mode is an EMA of the tick input, LMAX modes = " | |
| "scalar λ=LMAX), sqrt (g=sqrt(λ_m/LMAX)). Ignored when the spectrum is off.", | |
| ) | |
| parser.add_argument( | |
| "--field-leak-spectrum-slow-frac", | |
| type=float, | |
| default=1.0, | |
| help="with --field-leak-spectrum: fraction F in (0,1] of modes (those with the " | |
| "least intrinsic decay) that carry the LMIN->LMAX geometric ladder; the " | |
| "rest sit at LMAX. Default 1.0 = ladder over all modes.", | |
| ) | |
| parser.add_argument( | |
| "--slow-mode-readout", | |
| action=argparse.BooleanOptionalAction, | |
| default=False, | |
| help="opt-in: give the slow modes of --field-leak-spectrum (λ < LMAX) their own " | |
| "readout: Δlogits = W_o·tanh(W_t·ŝ), W_o zero-init (step-0 logits unchanged). " | |
| "Requires --field-leak-spectrum. Default off = exact current behavior. " | |
| "CLI wins after --resume. See docs/SLOW_MODE_READOUT.md.", | |
| ) | |
| parser.add_argument("--slow-readout-rank", type=int, default=64, | |
| help="rank r of the slow-mode readout trunk (default 64)") | |
| parser.add_argument( | |
| "--slow-readout-ahat", choices=["exclude", "mix"], default="exclude", | |
| help="with --slow-mode-readout: exclude (default) = the existing a_hat readout sees " | |
| "α with the slow modes zeroed; mix = slow modes stay in a_hat as well", | |
| ) | |
| parser.add_argument( | |
| "--slow-readout-lr", type=float, default=None, | |
| help="LR of the slow-mode readout param group (default: --surface-learning-rate)", | |
| ) | |
| parser.add_argument( | |
| "--slow-memory-loss", type=float, default=0.0, metavar="W", | |
| help="weight W of the slow-memory aux loss (needs --slow-mode-readout; default 0 = off)", | |
| ) | |
| parser.add_argument( | |
| "--slow-memory-lags", type=int, nargs="+", default=[16, 32, 64], metavar="K", | |
| help="memory lags K for --slow-memory-loss (default 16 32 64)", | |
| ) | |
| parser.add_argument( | |
| "--slow-memory-target", choices=["bag", "byte"], default="bag", | |
| help="bag (default): byte histogram of lags [K/2, K) from the slow head; " | |
| "byte: the single byte at lag K", | |
| ) | |
| parser.add_argument( | |
| "--field-max-rms", | |
| type=float, | |
| default=3.0, | |
| help="RMS saturate cap (default 3.0; prior chat runs used 4.0). " | |
| "Tighter cap reduces saturate wipe while keeping signal.", | |
| ) | |
| parser.add_argument( | |
| "--energy-decay-min", | |
| type=float, | |
| default=DEFAULT_ENERGY_DECAY_BOUNDS[0], | |
| help=f"always-on energy_decay floor (default {DEFAULT_ENERGY_DECAY_BOUNDS[0]})", | |
| ) | |
| parser.add_argument( | |
| "--energy-decay-max", | |
| type=float, | |
| default=DEFAULT_ENERGY_DECAY_BOUNDS[1], | |
| help=f"always-on energy_decay ceiling (default {DEFAULT_ENERGY_DECAY_BOUNDS[1]})", | |
| ) | |
| parser.add_argument("--seed", type=int, default=20260913) | |
| parser.add_argument( | |
| "--atom-flush-every", | |
| type=int, | |
| default=0, | |
| help="clear structural atom list every N steps (0=never, default). " | |
| "0 retains the multi-atom payload ring (MERGE keeps last ~128B) so " | |
| "α-gated copy-bias sees prompt-distinct bytes; set >0 only to bound cost.", | |
| ) | |
| parser.add_argument( | |
| "--log-every", | |
| type=int, | |
| default=25, | |
| help="print/record trajectory every N steps", | |
| ) | |
| parser.add_argument( | |
| "--enable-merge", | |
| action=argparse.BooleanOptionalAction, | |
| default=True, | |
| help="MERGE phase-coherent energetic atoms into heavier structure (default: on)", | |
| ) | |
| parser.add_argument( | |
| "--merge-coherence-threshold", | |
| type=float, | |
| default=0.45, | |
| help="phase*energy coherence gate for MERGE (default 0.45; mph~0.5 regime)", | |
| ) | |
| parser.add_argument( | |
| "--merge-energy-floor", | |
| type=float, | |
| default=0.08, | |
| help="min mean-E / geo-mean energy floor for MERGE pairs (default 0.08)", | |
| ) | |
| parser.add_argument( | |
| "--field-loss-weight", | |
| type=float, | |
| default=0.05, | |
| help="weight for persistence probe (reconstruct packet features from α); 0 disables", | |
| ) | |
| parser.add_argument( | |
| "--field-contrast-weight", | |
| type=float, | |
| default=0.1, | |
| help="weight for α contrastive hinge on surface logits; 0 disables", | |
| ) | |
| parser.add_argument( | |
| "--field-contrast-margin", | |
| type=float, | |
| default=0.55, | |
| help="max allowed cosine(true_α_logits, null/shuffled_α_logits) before hinge", | |
| ) | |
| parser.add_argument( | |
| "--field-ignorance-weight", | |
| type=float, | |
| default=0.08, | |
| help="weight for field-ignorance hinge L_ign=ReLU(cos(l(a),l(sg[a_bar]))-m); 0 disables", | |
| ) | |
| parser.add_argument( | |
| "--field-ignorance-margin", | |
| type=float, | |
| default=0.85, | |
| help="margin m for field-ignorance cosine hinge (default 0.85)", | |
| ) | |
| parser.add_argument( | |
| "--field-ignorance-ablate-shared", | |
| action=argparse.BooleanOptionalAction, | |
| default=False, | |
| help="for L_ign ℓ(α_bar) only: pass atom_r/persistence as None (zeros); " | |
| "ℓ(α) keeps live shared state (default: off)", | |
| ) | |
| parser.add_argument( | |
| "--field-ignorance-prompt-bank", | |
| action=argparse.BooleanOptionalAction, | |
| default=False, | |
| help="build α_bar from fixed distinct prompts (probe-style), not train ticks " | |
| "(default: off)", | |
| ) | |
| parser.add_argument( | |
| "--field-obligatory-hard", | |
| action="store_true", | |
| help="hard-v2 obligatory: mix floor=1.0 + freeze non-alpha bypass; logits=frozen_α+scale*α_proj (implies readout)", | |
| ) | |
| parser.add_argument( | |
| "--payload-copy", | |
| action=argparse.BooleanOptionalAction, | |
| default=True, | |
| help="allow the α-local payload copy-bias branch on the hard CE path " | |
| "(default: on; use --no-payload-copy for the speech line so " | |
| "train matches chat payload_copy=False)", | |
| ) | |
| parser.add_argument( | |
| "--last-atom-readout", | |
| action=argparse.BooleanOptionalAction, | |
| default=False, | |
| help="byte-tick last-atom readout + dentate 2-gram on the hard path " | |
| "(requires --max-span-bytes 1; default: off)", | |
| ) | |
| parser.add_argument( | |
| "--last-atom-scale", | |
| type=float, | |
| default=1.0, | |
| help="fixed scale on last_atom logits (hard mix): logits += S * la_byte " | |
| "(default 1.0 = current behavior; ignored when --last-atom-scale-adaptive)", | |
| ) | |
| parser.add_argument( | |
| "--last-atom-scale-adaptive", | |
| action=argparse.BooleanOptionalAction, | |
| default=False, | |
| help="opt-in: S_eff = clamp(α_rms/(la_rms+eps), S_min, S_max) each forward " | |
| "(detached RMS; overrides fixed --last-atom-scale; default OFF)", | |
| ) | |
| parser.add_argument( | |
| "--last-atom-scale-min", | |
| type=float, | |
| default=0.0, | |
| help="adaptive clamp floor for S_eff (default 0.0 = true RMS match; " | |
| "0.05 floor falsified under free-run-aux)", | |
| ) | |
| parser.add_argument( | |
| "--last-atom-scale-max", | |
| type=float, | |
| default=1.0, | |
| help="adaptive clamp ceiling for S_eff (default 1.0)", | |
| ) | |
| parser.add_argument( | |
| "--last-atom-dropout", | |
| type=float, | |
| default=0.0, | |
| help="opt-in: training only, per byte-tick position zero the last_atom " | |
| "logit contribution with prob P (no inverted scaling) so CE must come " | |
| "from the field path; ignored at eval/generation (default 0 = current)", | |
| ) | |
| parser.add_argument( | |
| "--field-hard-couple-scale", | |
| type=float, | |
| default=0.0, | |
| help="opt-in: under hard obligatory, add S * field_hard_couple_byte(" | |
| "LN([mean‖std‖persist‖atom.r])) into hard logits " | |
| "(default 0 = OFF; dedicated W, not frozen field_byte_skip)", | |
| ) | |
| parser.add_argument( | |
| "--legacy-gen-labels", | |
| action="store_true", | |
| help="byte-tick decode fix escape hatch (docs/DECODE_FIX_AND_BPTT.md): " | |
| "commit generated/efference bytes with the pre-fix labels " | |
| "(whitespace/punctuation). Default off = train-time labels.", | |
| ) | |
| parser.add_argument( | |
| "--field-bptt", | |
| type=int, | |
| default=0, | |
| help="truncated BPTT through the toroidal field: keep the α graph and a " | |
| "differentiable atom injection within windows of K ticks; one " | |
| "backward + optimizer step per window (0=off = exact legacy " | |
| "per-tick detach + per-step update)", | |
| ) | |
| parser.add_argument( | |
| "--field-leak-learnable", | |
| action="store_true", | |
| help="learnable per-mode leak λ_m = sigmoid(logit_m), init --field-leak " | |
| "(needs --field-bptt to receive gradient; off = scalar --field-leak)", | |
| ) | |
| parser.add_argument( | |
| "--field-leak-slow-block", | |
| type=float, | |
| default=0.0, | |
| metavar="F", | |
| help="opt-in long memory (default 0 = off; needs --field-leak-learnable): a contiguous " | |
| "block of round(F*n_modes) modes around m=n/2 whose learnable leak is " | |
| "lambda = floor + CAP*sigmoid(logit) (floor keeps every mode stable), init sigmoid=.25, " | |
| "with EMA input gain lambda/ref (docs/MIX_DATA_LONG_MEMORY.md)", | |
| ) | |
| parser.add_argument("--field-leak-slow-cap", type=float, default=0.02, | |
| help="max learnable leak above the stability floor for slow-block modes (default .02)") | |
| parser.add_argument("--field-leak-slow-gain", choices=("ema", "none"), default="ema", | |
| help="slow-block input gain: ema = lambda_m/ref (default), none = 1") | |
| parser.add_argument( | |
| "--efference-every", | |
| type=int, | |
| default=0, | |
| help="efference copy (bridge #9): every N train steps, commit the " | |
| "model's own predicted byte then CE on gold target (0=off; " | |
| "sandbox speech line used 10)", | |
| ) | |
| parser.add_argument( | |
| "--free-run-aux-every", | |
| type=int, | |
| default=0, | |
| help="multi-tick free-run aux (L_roll): every N steps replace CE with " | |
| "H-step free-run CE using generate commit path (0=OFF default; " | |
| "does not change live recipe until opted in)", | |
| ) | |
| parser.add_argument( | |
| "--free-run-aux-horizon", | |
| type=int, | |
| default=4, | |
| help="horizon H for --free-run-aux-every (default 4)", | |
| ) | |
| parser.add_argument( | |
| "--free-run-aux-weight", | |
| type=float, | |
| default=1.0, | |
| help="scale on mean free-run CE when free-run aux fires (default 1.0)", | |
| ) | |
| parser.add_argument( | |
| "--printable-aux-weight", | |
| type=float, | |
| default=0.0, | |
| help="aux weight rewarding printable UTF-8 mass on hard logits (0 disables; try ~0.08)", | |
| ) | |
| parser.add_argument( | |
| "--field-next-packet-weight", | |
| type=float, | |
| default=0.0, | |
| help="aux weight: predict next packet features from α+atoms (0 disables; try ~0.05)", | |
| ) | |
| parser.add_argument( | |
| "--field-obligatory-readout", | |
| action=argparse.BooleanOptionalAction, | |
| default=False, | |
| help="main CE path: α-only residual into surface logits with non-zero floor " | |
| "(cannot bypass α via persist/atom_r alone; default: off)", | |
| ) | |
| parser.add_argument( | |
| "--slow-every", | |
| type=int, | |
| default=1, | |
| help="run abstraction/consolidation every N ticks if field RMS stable (1=always)", | |
| ) | |
| parser.add_argument( | |
| "--slow-rms-rel-tol", | |
| type=float, | |
| default=0.15, | |
| help="relative RMS change threshold for dual-clock slow path", | |
| ) | |
| args = parser.parse_args() | |
| if not args.allow_parallel_train: | |
| acquire_train_lock( | |
| Path(args.train_lock_path) if args.train_lock_path else None | |
| ) | |
| else: | |
| print("train lock skipped (--allow-parallel-train)") | |
| if not args.data and not args.data_glob: | |
| raise SystemExit("error: provide --data and/or --data-glob") | |
| if args.data_glob and not args.stream: | |
| raise SystemExit("error: --data-glob requires --stream (refuse full in-RAM encode of a shard glob)") | |
| if (args.energy_decay_min is None) != (args.energy_decay_max is None): | |
| raise ValueError("provide both --energy-decay-min and --energy-decay-max") | |
| energy_decay_bounds = ( | |
| (float(args.energy_decay_min), float(args.energy_decay_max)) | |
| if args.energy_decay_min is not None | |
| else DEFAULT_ENERGY_DECAY_BOUNDS | |
| ) | |
| output_dir = Path(args.output_dir) | |
| output_dir.mkdir(parents=True, exist_ok=True) | |
| stream_mode = bool(args.stream) | |
| shard_paths = resolve_shard_paths(args.data, args.data_glob) | |
| train_packets = None | |
| validation_packets = None | |
| packets = None | |
| raw = b"" | |
| text = "" | |
| stream_source: StreamingPacketSource | None = None | |
| atomizer = Atomizer(max_span_bytes=args.max_span_bytes, pack_mode=args.atomizer_pack) | |
| if stream_mode: | |
| stream_source = StreamingPacketSource( | |
| atomizer, | |
| shard_paths, | |
| chunk_bytes=args.chunk_bytes, | |
| loop=args.loop_shards, | |
| buffer_size=args.stream_buffer, | |
| reset_atomizer=True, | |
| prefetch=bool(args.stream_prefetch), | |
| ) | |
| # Peek one transition to fail fast on tiny/empty shards. | |
| try: | |
| first_pair = stream_source.next_transition() | |
| except StopIteration as exc: | |
| raise ValueError("stream source produced no packet transitions") from exc | |
| # Re-create source so training sees a clean continuum from the start. | |
| stream_source = StreamingPacketSource( | |
| Atomizer(max_span_bytes=args.max_span_bytes, pack_mode=args.atomizer_pack), | |
| shard_paths, | |
| chunk_bytes=args.chunk_bytes, | |
| loop=args.loop_shards, | |
| buffer_size=args.stream_buffer, | |
| reset_atomizer=True, | |
| prefetch=bool(args.stream_prefetch), | |
| **({"start_offset": int(args.stream_skip_bytes)} if int(args.stream_skip_bytes) > 0 else {}), | |
| ) | |
| del first_pair | |
| if int(args.stream_skip_bytes) > 0: | |
| print(f"stream skip bytes: first pass starts at byte {int(args.stream_skip_bytes)}") | |
| skip = max(0, int(args.stream_skip_packets)) | |
| if skip: | |
| print(f"stream skip: advancing {skip} packet transitions...") | |
| for _ in range(skip): | |
| try: | |
| next(stream_source) | |
| except StopIteration as exc: | |
| raise RuntimeError( | |
| f"stream exhausted while skipping {skip} packets; " | |
| "enable --loop-shards or reduce --stream-skip-packets" | |
| ) from exc | |
| print(f"stream skip done: packets_seen={stream_source.packets_seen}") | |
| data_bytes = sum(p.stat().st_size for p in shard_paths) | |
| metadata_lengths_placeholder = True | |
| else: | |
| if len(shard_paths) != 1 or args.data_glob: | |
| raise SystemExit("non-stream mode supports a single --data file only; use --stream for globs") | |
| data_path = shard_paths[0] | |
| raw = data_path.read_bytes() | |
| text = raw.decode("utf-8") | |
| packets = atomizer.encode_bytes(raw) | |
| if len(packets) < 4: | |
| raise ValueError("the corpus must yield at least four atom packets") | |
| if b"".join(packet.payload for packet in packets) != raw: | |
| raise AssertionError("atomization is not lossless") | |
| train_packets = packets[:-64] if len(packets) > 96 else packets[:-4] | |
| validation_packets = packets[len(train_packets) :] | |
| if len(train_packets) < 3 or len(validation_packets) < 2: | |
| raise ValueError("not enough packets for train/validation split") | |
| data_bytes = len(raw) | |
| metadata_lengths_placeholder = False | |
| seed_all(args.seed) | |
| model = AtomNativeModel( | |
| d_model=args.d_model, | |
| n_modes=args.n_modes, | |
| n_atoms_max=args.n_atoms_max, | |
| max_payload_bytes=args.max_span_bytes, | |
| atomizer=Atomizer(max_span_bytes=args.max_span_bytes, pack_mode=args.atomizer_pack), | |
| field_max_rms=args.field_max_rms, | |
| field_leak=float(args.field_leak), | |
| energy_decay_bounds=energy_decay_bounds, | |
| enable_merge=bool(args.enable_merge), | |
| merge_coherence_threshold=float(args.merge_coherence_threshold), | |
| merge_energy_floor=float(args.merge_energy_floor), | |
| slow_every=int(args.slow_every), | |
| field_loss_weight=float(args.field_loss_weight), | |
| field_contrast_weight=float(args.field_contrast_weight), | |
| field_contrast_margin=float(args.field_contrast_margin), | |
| field_ignorance_weight=float(args.field_ignorance_weight), | |
| field_ignorance_margin=float(args.field_ignorance_margin), | |
| field_ignorance_ablate_shared=bool(args.field_ignorance_ablate_shared), | |
| field_ignorance_prompt_bank=bool(args.field_ignorance_prompt_bank), | |
| field_obligatory_readout=bool(args.field_obligatory_readout) or bool(args.field_obligatory_hard), | |
| field_obligatory_hard=bool(args.field_obligatory_hard), | |
| last_atom_readout=bool(args.last_atom_readout), | |
| printable_aux_weight=float(args.printable_aux_weight), | |
| field_next_packet_weight=float(args.field_next_packet_weight), | |
| slow_rms_rel_tol=float(args.slow_rms_rel_tol), | |
| ) | |
| training_state = None | |
| start_step = 0 | |
| if args.resume: | |
| training_state = model.load(args.resume) | |
| start_step = int((training_state or {}).get("step", 0) or 0) | |
| migrated = (training_state or {}).get("legacy_surface_migrated", False) | |
| span_mig = (training_state or {}).get("span_head_migrated", False) | |
| span_info = (training_state or {}).get("span_head_migrate_info") or {} | |
| dyn_load = (training_state or {}).get("dynamics_on_load") or {} | |
| print( | |
| f"resumed weights from {args.resume} " | |
| f"(recorded step={start_step}, legacy_surface_migrated={migrated}, " | |
| f"span_head_migrated={span_mig})" | |
| ) | |
| if span_mig: | |
| print( | |
| f"span-head pad-migrate: old_max={span_info.get('old_max_payload')} -> " | |
| f"new_max={span_info.get('new_max_payload')} " | |
| f"padded={len(span_info.get('padded_keys') or [])} tensors " | |
| f"(field/core kept; new rows only)" | |
| ) | |
| print( | |
| f"atomizer pack_mode={args.atomizer_pack} max_span_bytes={args.max_span_bytes}" | |
| ) | |
| if dyn_load: | |
| print( | |
| f"dynamics on load: energy_decay=" | |
| f"{dyn_load.get('energy_decay_before')} -> {dyn_load.get('energy_decay')} " | |
| f"(repaired={dyn_load.get('energy_decay_repaired')}, " | |
| f"bounds={dyn_load.get('energy_decay_bounds')})" | |
| ) | |
| # CLI wins over ckpt config for obligatory readout (smoke / migrate). | |
| model.field_obligatory_readout = bool(args.field_obligatory_readout) | |
| model.surface.field_obligatory_readout = bool(args.field_obligatory_readout) | |
| # CLI wins for printable aux (mechanism under test). | |
| model.printable_aux_weight = float(args.printable_aux_weight) | |
| model.field_next_packet_weight = float(args.field_next_packet_weight) | |
| # CLI wins for field persistence / contrast weights (same pattern as | |
| # printable_aux / next_packet). Without this, model.load() restores | |
| # ckpt 0.0 and drops construct CLI e.g. --field-loss-weight 0.05. | |
| model.field_loss_weight = float(args.field_loss_weight) | |
| model.field_contrast_weight = float(args.field_contrast_weight) | |
| if float(args.field_next_packet_weight) > 0: | |
| print( | |
| f"field_next_packet_weight={float(args.field_next_packet_weight)} " | |
| "(predict next packet features from α+atoms)" | |
| ) | |
| # CLI wins for MERGE gates (operational threshold fix under hard-v2+). | |
| model.enable_merge = bool(args.enable_merge) | |
| model.set_merge_thresholds( | |
| coherence_threshold=float(args.merge_coherence_threshold), | |
| energy_floor=float(args.merge_energy_floor), | |
| ) | |
| print( | |
| f"MERGE gates: enable={bool(args.enable_merge)} " | |
| f"coherence_thr={float(args.merge_coherence_threshold):.3f} " | |
| f"energy_floor={float(args.merge_energy_floor):.3f}" | |
| ) | |
| if getattr(args, "field_obligatory_hard", False): | |
| model.field_obligatory_hard = True | |
| model.field_obligatory_readout = True | |
| model.surface.set_obligatory_hard(True) | |
| print( | |
| "field_obligatory_hard=ON hard-v2+ " | |
| "(mix_floor=1.0, freeze bypass, logits=frozen+cap*(α_MLP+payload_prod))" | |
| ) | |
| if float(args.printable_aux_weight) > 0: | |
| print(f"printable_aux_weight={float(args.printable_aux_weight)} (reward printable mass on hard logits)") | |
| if args.field_obligatory_readout: | |
| print( | |
| "field_obligatory_readout=ON " | |
| f"(migrated={bool((training_state or {}).get('field_obligatory_migrated'))})" | |
| ) | |
| elif args.field_obligatory_readout: | |
| model.field_obligatory_readout = True | |
| model.surface.field_obligatory_readout = True | |
| print("field_obligatory_readout=ON (fresh)") | |
| # CLI wins for train-time payload copy-bias (after load, which restores | |
| # the checkpoint value for provenance). Speech line: --no-payload-copy. | |
| model.surface.payload_enabled = bool(args.payload_copy) | |
| print(f"payload_copy={bool(args.payload_copy)} (train-time copy-bias branch)") | |
| print(f"last_atom_readout={bool(args.last_atom_readout)} (byte-tick 2-gram + dentate)") | |
| # CLI wins for last_atom scale (default 1.0 / adaptive OFF = no behavior change). | |
| model.surface.last_atom_scale = float(args.last_atom_scale) | |
| model.surface.last_atom_scale_adaptive = bool(args.last_atom_scale_adaptive) | |
| model.surface.last_atom_scale_min = float(args.last_atom_scale_min) | |
| model.surface.last_atom_scale_max = float(args.last_atom_scale_max) | |
| # CLI wins for last_atom dropout (default 0 = no behavior change; train only). | |
| if not 0.0 <= float(args.last_atom_dropout) <= 1.0: | |
| raise SystemExit("error: --last-atom-dropout must be in [0, 1]") | |
| model.surface.last_atom_dropout = float(args.last_atom_dropout) | |
| print( | |
| f"last_atom_dropout={float(args.last_atom_dropout)} " | |
| "(train-only per-position zeroing of la_byte; 0=off/current)" | |
| ) | |
| # CLI wins for field leak (default 0 = no-op = exact legacy dynamics). | |
| if not 0.0 <= float(args.field_leak) < 1.0: | |
| raise SystemExit("error: --field-leak must be in [0, 1)") | |
| model.field_leak = float(args.field_leak) | |
| print( | |
| f"field_leak={float(args.field_leak)} " | |
| "(per-tick α*=(1-λ) before injection; 0=off/current)" | |
| ) | |
| # CLI wins for the multi-timescale leak spectrum too (omitted = off). | |
| try: | |
| model.set_field_leak_spectrum( | |
| tuple(args.field_leak_spectrum) if args.field_leak_spectrum else None, | |
| gain=args.field_leak_spectrum_gain, | |
| slow_frac=args.field_leak_spectrum_slow_frac, | |
| ) | |
| except ValueError as exc: | |
| raise SystemExit(f"error: --field-leak-spectrum: {exc}") | |
| if model.field_leak_spectrum is not None: | |
| _lam = model.field_leak_rates() | |
| print( | |
| f"field_leak_spectrum={model.field_leak_spectrum} gain={model.field_leak_spectrum_gain} " | |
| f"slow_frac={model.field_leak_spectrum_slow_frac} " | |
| f"(per-mode λ min={float(_lam.min()):.4f} max={float(_lam.max()):.4f}; " | |
| f"replaces field_leak={float(args.field_leak)})" | |
| ) | |
| # CLI wins for the slow-mode readout too (omitted = off, head dropped). | |
| try: | |
| if args.slow_mode_readout: | |
| model.set_slow_mode_readout( | |
| True, rank=int(args.slow_readout_rank), ahat=args.slow_readout_ahat, | |
| memory_weight=float(args.slow_memory_loss), | |
| memory_lags=tuple(args.slow_memory_lags), | |
| memory_target=args.slow_memory_target, | |
| ) | |
| else: | |
| if float(args.slow_memory_loss) > 0: | |
| raise ValueError("--slow-memory-loss needs --slow-mode-readout") | |
| model.set_slow_mode_readout(False) | |
| except ValueError as exc: | |
| raise SystemExit(f"error: --slow-mode-readout: {exc}") | |
| if model.slow_readout is not None: | |
| print( | |
| f"slow_mode_readout=ON slow_modes={model._slow_idx.tolist()} rank={model.slow_readout_rank} " | |
| f"ahat={model.slow_readout_ahat} memory_weight={model.slow_memory_weight} " | |
| f"lags={list(model.slow_memory_lags)} target={model.slow_memory_target} " | |
| f"({sum(p.numel() for p in model.slow_readout.parameters())} params)" | |
| ) | |
| # CLI wins for field→hard-decode couple (default 0 = no behavior change). | |
| model.surface.field_hard_couple_scale = float(args.field_hard_couple_scale) | |
| if float(args.field_hard_couple_scale) > 0: | |
| print( | |
| f"field_hard_couple_scale={float(args.field_hard_couple_scale)} " | |
| "(hard logits += S * field_hard_couple_byte(full feats))" | |
| ) | |
| else: | |
| print("field_hard_couple_scale=0 (OFF — hard path unchanged)") | |
| if bool(args.last_atom_scale_adaptive): | |
| print( | |
| f"last_atom_scale=adaptive " | |
| f"(S_eff=clamp(α_rms/(la_rms+eps), " | |
| f"{float(args.last_atom_scale_min)}, {float(args.last_atom_scale_max)}); " | |
| f"fixed S={float(args.last_atom_scale)} ignored)" | |
| ) | |
| else: | |
| print( | |
| f"last_atom_scale={float(args.last_atom_scale)} " | |
| "(1.0=off/current; hard-path la_byte *= S)" | |
| ) | |
| model._efference_every = max(0, int(args.efference_every)) | |
| print(f"efference_every={model._efference_every} (0=off)") | |
| # CLI wins for the byte-tick decode fix escape hatch (runtime-only flag). | |
| model.atomizer.legacy_gen_labels = bool(args.legacy_gen_labels) | |
| if args.legacy_gen_labels: | |
| print("legacy_gen_labels=ON (pre-fix whitespace/punctuation commit labels)") | |
| # CLI wins for truncated BPTT / learnable leak (after load, both directions). | |
| if int(args.field_bptt) < 0: | |
| raise SystemExit("error: --field-bptt must be >= 0") | |
| model.set_field_bptt(int(args.field_bptt)) | |
| if args.field_leak_learnable: | |
| if model.field_leak_logit is None: | |
| try: | |
| model.set_field_leak_learnable(True, init=float(args.field_leak)) | |
| except ValueError as exc: | |
| raise SystemExit(f"error: --field-leak-learnable: {exc}") | |
| _ll = model.learnable_leak_rates().detach() | |
| print(f"field_leak_learnable=ON per-mode λ min={float(_ll.min()):.4f} max={float(_ll.max()):.4f}") | |
| else: | |
| model.set_field_leak_learnable(False) | |
| # CLI wins for the slow block too (after load, both directions; same config = kept). | |
| _sb_f = float(args.field_leak_slow_block or 0.0) | |
| _sb_cur = getattr(model, "field_leak_slow_block", None) | |
| if _sb_f > 0: | |
| if not args.field_leak_learnable: | |
| raise SystemExit("error: --field-leak-slow-block needs --field-leak-learnable") | |
| _same = _sb_cur is not None and (_sb_cur["frac"], _sb_cur["cap"], _sb_cur["gain"]) == ( | |
| _sb_f, float(args.field_leak_slow_cap), args.field_leak_slow_gain) | |
| if not _same: | |
| model.set_field_leak_slow_block(None) | |
| model.set_field_leak_slow_block(_sb_f, cap=float(args.field_leak_slow_cap), | |
| gain=args.field_leak_slow_gain, ref=float(args.field_leak)) | |
| _ll = model.learnable_leak_rates().detach().reshape(-1) | |
| print(f"field_leak_slow_block=ON modes={model.field_leak_slow_block['modes']} " | |
| f"cap={model.field_leak_slow_block['cap']} gain={model.field_leak_slow_block['gain']} " | |
| f"slow λ={[round(float(_ll[i]), 4) for i in model.field_leak_slow_block['modes']]}") | |
| elif _sb_cur is not None: | |
| model.set_field_leak_slow_block(None) | |
| if model.field_bptt > 0: | |
| print(f"field_bptt={model.field_bptt} (truncated BPTT window; one optimizer step per window)") | |
| model._free_run_aux_every = max(0, int(args.free_run_aux_every)) | |
| model._free_run_aux_horizon = max(1, int(args.free_run_aux_horizon)) | |
| model._free_run_aux_weight = float(args.free_run_aux_weight) | |
| print( | |
| f"free_run_aux_every={model._free_run_aux_every} " | |
| f"horizon={model._free_run_aux_horizon} " | |
| f"weight={model._free_run_aux_weight} (0=off)" | |
| ) | |
| surface_lr = ( | |
| float(args.surface_learning_rate) | |
| if args.surface_learning_rate is not None | |
| else float(args.learning_rate) | |
| ) | |
| surface_ids = {id(p) for p in model.surface.parameters() if p.requires_grad} | |
| slow_ids = ( | |
| {id(p) for p in model.slow_readout.parameters()} if model.slow_readout is not None else set() | |
| ) | |
| surface_params = [p for p in model.trainable_parameters if id(p) in surface_ids] | |
| other_params = [ | |
| p for p in model.trainable_parameters if id(p) not in surface_ids and id(p) not in slow_ids | |
| ] | |
| param_groups = [ | |
| {"params": other_params, "lr": float(args.learning_rate)}, | |
| {"params": surface_params, "lr": surface_lr}, | |
| ] | |
| if slow_ids: | |
| slow_lr = float(args.slow_readout_lr) if args.slow_readout_lr is not None else surface_lr | |
| param_groups.append( | |
| {"params": [p for p in model.slow_readout.parameters() if p.requires_grad], "lr": slow_lr} | |
| ) | |
| print(f"optimizer group: slow_readout lr={slow_lr}") | |
| optimizer = torch.optim.AdamW(param_groups, weight_decay=0.01) | |
| print( | |
| f"optimizer groups: field/other lr={args.learning_rate} " | |
| f"({len(other_params)} tensors), surface lr={surface_lr} " | |
| f"({len(surface_params)} tensors); " | |
| f"energy_decay_bounds={energy_decay_bounds}; field_max_rms={args.field_max_rms}" | |
| ) | |
| if stream_mode: | |
| print( | |
| f"stream mode: {len(shard_paths)} shard(s), chunk_bytes={args.chunk_bytes}, " | |
| f"loop_shards={args.loop_shards}, stream_buffer={args.stream_buffer}, " | |
| f"stream_prefetch={bool(args.stream_prefetch)}" | |
| ) | |
| if training_state and "optimizer" in training_state: | |
| if training_state.get("legacy_surface_migrated"): | |
| print("optimizer restore skipped: surface LayerNorm migration changed bias scale") | |
| elif args.surface_learning_rate is not None: | |
| print("optimizer restore skipped: surface LR param groups differ from checkpoint") | |
| else: | |
| try: | |
| optimizer.load_state_dict(training_state["optimizer"]) | |
| print("restored optimizer state") | |
| except Exception as exc: | |
| print(f"optimizer restore skipped: {exc}") | |
| model.train() | |
| if metadata_lengths_placeholder: | |
| metadata = { | |
| "stream": True, | |
| "data_path": str(args.data) if args.data else None, | |
| "data_glob": args.data_glob, | |
| "shard_paths": [str(p) for p in shard_paths], | |
| "data_bytes": data_bytes, | |
| "data_characters": None, | |
| "atomizer_version": atomizer.VERSION, | |
| "feature_dim": atomizer.feature_dim, | |
| "packet_count": None, | |
| "train_packet_count": None, | |
| "validation_packet_count": None, | |
| "mean_packet_bytes": None, | |
| "max_packet_bytes": None, | |
| "min_packet_bytes": None, | |
| "bytes_seen": 0, | |
| "config": { | |
| "d_model": args.d_model, | |
| "n_modes": args.n_modes, | |
| "n_atoms_max": args.n_atoms_max, | |
| "max_span_bytes": args.max_span_bytes, | |
| "episode_length": args.episode_length, | |
| "episode_reset": bool(args.episode_reset), | |
| "atom_flush_every": args.atom_flush_every, | |
| "learning_rate": args.learning_rate, | |
| "surface_learning_rate": surface_lr, | |
| "steps": args.steps, | |
| "start_step": start_step, | |
| "seed": args.seed, | |
| "surface_alphabet": 256, | |
| "field_max_rms": args.field_max_rms, | |
| "energy_decay_bounds": energy_decay_bounds, | |
| "stream": True, | |
| "chunk_bytes": args.chunk_bytes, | |
| "loop_shards": bool(args.loop_shards), | |
| "stream_buffer": args.stream_buffer, | |
| "stream_prefetch": bool(args.stream_prefetch), | |
| "stream_skip_packets": max(0, int(args.stream_skip_packets)), | |
| "last_atom_readout": bool(args.last_atom_readout), | |
| "last_atom_scale": float(args.last_atom_scale), | |
| "last_atom_scale_adaptive": bool(args.last_atom_scale_adaptive), | |
| "last_atom_scale_min": float(args.last_atom_scale_min), | |
| "last_atom_scale_max": float(args.last_atom_scale_max), | |
| "field_hard_couple_scale": float(args.field_hard_couple_scale), | |
| "last_atom_dropout": float(args.last_atom_dropout), | |
| "field_leak": float(args.field_leak), | |
| "field_leak_spectrum": list(args.field_leak_spectrum) if args.field_leak_spectrum else None, | |
| "field_leak_spectrum_gain": ( | |
| args.field_leak_spectrum_gain if args.field_leak_spectrum else "none" | |
| ), | |
| "field_leak_spectrum_slow_frac": ( | |
| float(args.field_leak_spectrum_slow_frac) if args.field_leak_spectrum else 1.0 | |
| ), | |
| "slow_mode_readout": bool(args.slow_mode_readout), | |
| "slow_readout_rank": int(args.slow_readout_rank), | |
| "slow_readout_ahat": args.slow_readout_ahat, | |
| "slow_readout_lr": args.slow_readout_lr, | |
| "slow_memory_loss": float(args.slow_memory_loss), | |
| "slow_memory_lags": list(args.slow_memory_lags), | |
| "slow_memory_target": args.slow_memory_target, | |
| "efference_every": max(0, int(args.efference_every)), | |
| "free_run_aux_every": max(0, int(args.free_run_aux_every)), | |
| "free_run_aux_horizon": max(1, int(args.free_run_aux_horizon)), | |
| "free_run_aux_weight": float(args.free_run_aux_weight), | |
| "enable_merge": bool(args.enable_merge), | |
| "payload_copy": bool(args.payload_copy), | |
| "merge_coherence_threshold": float(args.merge_coherence_threshold), | |
| "merge_energy_floor": float(args.merge_energy_floor), | |
| "field_loss_weight": float(args.field_loss_weight), | |
| "field_contrast_weight": float(args.field_contrast_weight), | |
| "field_ignorance_weight": float(args.field_ignorance_weight), | |
| "field_ignorance_ablate_shared": bool(args.field_ignorance_ablate_shared), | |
| "field_ignorance_prompt_bank": bool(args.field_ignorance_prompt_bank), | |
| "field_obligatory_readout": bool(args.field_obligatory_readout) or bool(getattr(args, "field_obligatory_hard", False)), | |
| "field_obligatory_hard": bool(getattr(args, "field_obligatory_hard", False)), | |
| "printable_aux_weight": float(args.printable_aux_weight), | |
| "field_next_packet_weight": float(args.field_next_packet_weight), | |
| "slow_every": int(args.slow_every), | |
| **_bptt_meta(args), | |
| }, | |
| } | |
| else: | |
| lengths = [len(packet.payload) for packet in packets] | |
| metadata = { | |
| "stream": False, | |
| "data_path": str(args.data), | |
| "data_glob": None, | |
| "shard_paths": [str(p) for p in shard_paths], | |
| "data_bytes": data_bytes, | |
| "data_characters": len(text), | |
| "atomizer_version": atomizer.VERSION, | |
| "feature_dim": atomizer.feature_dim, | |
| "packet_count": len(packets), | |
| "train_packet_count": len(train_packets), | |
| "validation_packet_count": len(validation_packets), | |
| "mean_packet_bytes": sum(lengths) / len(lengths), | |
| "max_packet_bytes": max(lengths), | |
| "min_packet_bytes": min(lengths), | |
| "config": { | |
| "d_model": args.d_model, | |
| "n_modes": args.n_modes, | |
| "n_atoms_max": args.n_atoms_max, | |
| "max_span_bytes": args.max_span_bytes, | |
| "episode_length": args.episode_length, | |
| "episode_reset": bool(args.episode_reset), | |
| "atom_flush_every": args.atom_flush_every, | |
| "learning_rate": args.learning_rate, | |
| "surface_learning_rate": surface_lr, | |
| "steps": args.steps, | |
| "start_step": start_step, | |
| "seed": args.seed, | |
| "surface_alphabet": 256, | |
| "field_max_rms": args.field_max_rms, | |
| "energy_decay_bounds": energy_decay_bounds, | |
| "stream": False, | |
| "enable_merge": bool(args.enable_merge), | |
| "payload_copy": bool(args.payload_copy), | |
| "last_atom_readout": bool(args.last_atom_readout), | |
| "last_atom_scale": float(args.last_atom_scale), | |
| "last_atom_scale_adaptive": bool(args.last_atom_scale_adaptive), | |
| "last_atom_scale_min": float(args.last_atom_scale_min), | |
| "last_atom_scale_max": float(args.last_atom_scale_max), | |
| "field_hard_couple_scale": float(args.field_hard_couple_scale), | |
| "last_atom_dropout": float(args.last_atom_dropout), | |
| "field_leak": float(args.field_leak), | |
| "field_leak_spectrum": list(args.field_leak_spectrum) if args.field_leak_spectrum else None, | |
| "field_leak_spectrum_gain": ( | |
| args.field_leak_spectrum_gain if args.field_leak_spectrum else "none" | |
| ), | |
| "field_leak_spectrum_slow_frac": ( | |
| float(args.field_leak_spectrum_slow_frac) if args.field_leak_spectrum else 1.0 | |
| ), | |
| "slow_mode_readout": bool(args.slow_mode_readout), | |
| "slow_readout_rank": int(args.slow_readout_rank), | |
| "slow_readout_ahat": args.slow_readout_ahat, | |
| "slow_readout_lr": args.slow_readout_lr, | |
| "slow_memory_loss": float(args.slow_memory_loss), | |
| "slow_memory_lags": list(args.slow_memory_lags), | |
| "slow_memory_target": args.slow_memory_target, | |
| "efference_every": max(0, int(args.efference_every)), | |
| "free_run_aux_every": max(0, int(args.free_run_aux_every)), | |
| "free_run_aux_horizon": max(1, int(args.free_run_aux_horizon)), | |
| "free_run_aux_weight": float(args.free_run_aux_weight), | |
| "merge_coherence_threshold": float(args.merge_coherence_threshold), | |
| "merge_energy_floor": float(args.merge_energy_floor), | |
| "field_loss_weight": float(args.field_loss_weight), | |
| "field_contrast_weight": float(args.field_contrast_weight), | |
| "field_ignorance_weight": float(args.field_ignorance_weight), | |
| "field_ignorance_ablate_shared": bool(args.field_ignorance_ablate_shared), | |
| "field_ignorance_prompt_bank": bool(args.field_ignorance_prompt_bank), | |
| "field_obligatory_readout": bool(args.field_obligatory_readout) or bool(getattr(args, "field_obligatory_hard", False)), | |
| "field_obligatory_hard": bool(getattr(args, "field_obligatory_hard", False)), | |
| "printable_aux_weight": float(args.printable_aux_weight), | |
| "field_next_packet_weight": float(args.field_next_packet_weight), | |
| "slow_every": int(args.slow_every), | |
| **_bptt_meta(args), | |
| }, | |
| } | |
| (output_dir / "atomization_stats.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") | |
| trajectory: list[dict] = [] | |
| all_losses: list[float] = [] | |
| efference_steps = 0 | |
| free_run_aux_steps = 0 | |
| field_rms_before_values: list[float] = [] | |
| field_rms_after_values: list[float] = [] | |
| field_scales: list[float] = [] | |
| field_norm_values: list[float] = [] | |
| slow_mem_values: list[float] = [] | |
| start_time = time.perf_counter() | |
| episode = -1 | |
| episode_start = 0 | |
| # Soft cap: if atoms approach n_atoms_max, reset once to avoid unbounded growth | |
| # when --no-episode-reset keeps the field alive across windows. | |
| atom_soft_cap = max(16, int(args.n_atoms_max * 0.9)) | |
| bptt_acc = None | |
| bptt_n = 0 | |
| bptt_windows = 0 | |
| bptt_last_grad = float("nan") | |
| dynamics_state = model.stabilize_dynamics_parameters() if int(getattr(model, "field_bptt", 0) or 0) else None | |
| for step in range(1, args.steps + 1): | |
| current_episode = (step - 1) // args.episode_length | |
| free_run_golds = None | |
| do_free_run = bool(args.free_run_aux_every) and step % int(args.free_run_aux_every) == 0 | |
| horizon = max(1, int(args.free_run_aux_horizon)) if do_free_run else 1 | |
| if stream_mode: | |
| # Prefer continuous stream. Episode windows only force a field wipe | |
| # when --episode-reset is set; otherwise the field persists. | |
| if args.episode_reset and current_episode != episode: | |
| episode = current_episode | |
| model.reset_state(reset_atomizer=True) | |
| model.train() | |
| elif step == 1: | |
| episode = current_episode | |
| model.reset_state(reset_atomizer=True) | |
| model.train() | |
| else: | |
| episode = current_episode | |
| try: | |
| # Prefer iterator protocol (avoids extra Python method frame). | |
| current_packet, target_packet = next(stream_source) | |
| if do_free_run: | |
| golds = [target_packet] | |
| for _ in range(horizon - 1): | |
| _, g = next(stream_source) | |
| golds.append(g) | |
| free_run_golds = golds | |
| except StopIteration as exc: | |
| raise RuntimeError( | |
| f"stream exhausted at step {step}/{args.steps}; " | |
| "enable --loop-shards for infinite continuum" | |
| ) from exc | |
| else: | |
| if current_episode != episode: | |
| episode = current_episode | |
| if args.episode_reset or step == 1: | |
| model.reset_state(reset_atomizer=True) | |
| else: | |
| # Continuity-preserving compromise: at each window boundary, clear | |
| # the structural atom list (keeps step cost bounded) but KEEP the | |
| # toroidal field alpha + consolidation/abstraction persistence so | |
| # dialogue-scale memory is not wiped every 64/512 ticks. | |
| model.core.atoms = type(model.core.atoms)() | |
| model.core.energy_history = [] | |
| if len(model.core.atoms) >= atom_soft_cap: | |
| pass # already cleared | |
| model.train() | |
| max_start = max(0, len(train_packets) - args.episode_length - 2) | |
| episode_start = random.randint(0, max_start) if max_start > 0 else 0 | |
| index = episode_start + ((step - 1) % args.episode_length) | |
| if index >= len(train_packets) - 1: | |
| index = index % (len(train_packets) - 1) | |
| current_packet = train_packets[index] | |
| target_packet = train_packets[index + 1] | |
| if do_free_run: | |
| golds = [] | |
| for h in range(horizon): | |
| gi = index + 1 + h | |
| if gi >= len(train_packets): | |
| break | |
| golds.append(train_packets[gi]) | |
| if golds: | |
| free_run_golds = golds | |
| bptt_k = int(getattr(model, "field_bptt", 0) or 0) | |
| if not bptt_k: | |
| optimizer.zero_grad(set_to_none=True) | |
| # Free-run aux supersedes 1-tick efference when both would fire. | |
| do_eff = ( | |
| free_run_golds is None | |
| and bool(args.efference_every) | |
| and step % int(args.efference_every) == 0 | |
| ) | |
| if do_eff: | |
| efference_steps += 1 | |
| if do_free_run and free_run_golds is not None: | |
| free_run_aux_steps += 1 | |
| if bptt_k: | |
| # --field-bptt K: accumulate K tick losses, one backward through the | |
| # field graph, one optimizer step, cut the graph (window boundary). | |
| if free_run_golds is not None: | |
| loss_t, info = model.free_run_aux_loss( | |
| current_packet, list(free_run_golds), weight=float(args.free_run_aux_weight) | |
| ) | |
| elif do_eff: | |
| loss_t, info = model.efference_loss(current_packet, target_packet) | |
| else: | |
| loss_t, info = model.transition_loss(current_packet, target_packet) | |
| if not torch.isfinite(loss_t): | |
| raise FloatingPointError("non-finite atom-native loss") | |
| bptt_acc = loss_t if bptt_n == 0 else bptt_acc + loss_t | |
| bptt_n += 1 | |
| loss = float(loss_t.detach().item()) | |
| episode_ends = bool(args.episode_reset) and step % args.episode_length == 0 | |
| if bptt_n >= bptt_k or step == args.steps or episode_ends: | |
| bptt_last_grad = bptt_window_step(model, optimizer, bptt_acc, bptt_n) | |
| bptt_acc = None | |
| bptt_n = 0 | |
| bptt_windows += 1 | |
| dynamics_state = model.stabilize_dynamics_parameters() | |
| info = {**info, "grad_norm": bptt_last_grad} | |
| else: | |
| loss, info = transition_metrics( | |
| model, | |
| current_packet, | |
| target_packet, | |
| train=True, | |
| efference=do_eff, | |
| free_run_golds=free_run_golds, | |
| free_run_weight=float(args.free_run_aux_weight), | |
| ) | |
| optimizer.step() | |
| if args.atom_flush_every and step % args.atom_flush_every == 0: | |
| model.core.atoms = type(model.core.atoms)() | |
| model.core.energy_history = [] | |
| if not bptt_k: | |
| dynamics_state = model.stabilize_dynamics_parameters() | |
| model.core.training_loss = loss | |
| all_losses.append(loss) | |
| field_rms_before_values.append(info["field_rms_before"]) | |
| field_rms_after_values.append(info["field_rms_after"]) | |
| field_scales.append(info["field_scale"]) | |
| field_norm_values.append(info["field_norm"]) | |
| if "slow_memory_loss" in info: | |
| slow_mem_values.append(float(info["slow_memory_loss"])) | |
| global_step = start_step + step | |
| if step == 1 or step % args.log_every == 0 or step == args.steps: | |
| record = { | |
| "step": global_step, | |
| "run_step": step, | |
| "episode": episode, | |
| "loss": loss, | |
| "total_perplexity": float(np.exp(min(loss, 30.0))), | |
| "byte_loss": info["byte_loss"], | |
| "byte_perplexity": float(np.exp(min(info["byte_loss"], 30.0))), | |
| "length_loss": info["length_loss"], | |
| "field_norm": info["field_norm"], | |
| "field_rms_before": info["field_rms_before"], | |
| "field_rms_after": info["field_rms_after"], | |
| "field_scale": info["field_scale"], | |
| "grad_norm": info.get("grad_norm"), | |
| "n_atoms": info["n_atoms"], | |
| "efference": do_eff, | |
| "free_run_aux": bool(info.get("free_run_aux", False)), | |
| "free_run_horizon": int(info.get("free_run_horizon", 0) or 0), | |
| "merge_count": info.get("merge_count", 0), | |
| "merge_count_total": info.get("merge_count_total", 0), | |
| "slow_tick": info.get("slow_tick", True), | |
| "field_contrast_loss": info.get("field_contrast_loss", 0.0), | |
| "field_persist_loss": info.get("field_persist_loss", 0.0), | |
| "field_ignorance_loss": info.get("field_ignorance_loss", 0.0), | |
| "printable_aux_loss": info.get("printable_aux_loss", 0.0), | |
| "printable_mass_mean": info.get("printable_mass_mean", float("nan")), | |
| "field_next_packet_loss": info.get("field_next_packet_loss", 0.0), | |
| "max_phase_coherence": info.get("max_phase_coherence", 0.0), | |
| "n_pairs_above_energy_floor": info.get("n_pairs_above_energy_floor", 0), | |
| "energy_decay": dynamics_state["energy_decay"], | |
| "coupling_scale": dynamics_state["coupling_scale"], | |
| "phase_sync": dynamics_state["phase_sync"], | |
| } | |
| if "slow_memory_loss" in info: | |
| record["slow_memory_loss"] = info["slow_memory_loss"] | |
| record["slow_head_grad_norm"] = info.get("slow_head_grad_norm") | |
| if stream_mode and stream_source is not None: | |
| record["bytes_seen"] = stream_source.bytes_seen | |
| record["packets_seen"] = stream_source.packets_seen | |
| record["shard_index"] = stream_source.shard_index | |
| record["stream_epoch"] = stream_source.epoch | |
| trajectory.append(record) | |
| extra = "" | |
| if stream_mode and stream_source is not None: | |
| extra = ( | |
| f" bytes={stream_source.bytes_seen} " | |
| f"pkts={stream_source.packets_seen} " | |
| f"shard={stream_source.shard_index}" | |
| ) | |
| print( | |
| f"step={global_step:07d} (run {step}/{args.steps}) " | |
| f"loss={loss:.5f} byte_ppl={record['byte_perplexity']:.3f} " | |
| f"atoms={record['n_atoms']} merges={record['merge_count_total']} " | |
| f"field={record['field_norm']:.4f} " | |
| f"rms={record['field_rms_after']:.4f} scale={record['field_scale']:.4f} " | |
| f"fcos={info.get('field_logit_cos_zero', float('nan')):.3f} " | |
| f"ign={info.get('field_ignorance_loss', 0.0):.4f} " | |
| f"prn={info.get('printable_mass_mean', float('nan')):.3f} " | |
| f"nxp={info.get('field_next_packet_loss', 0.0):.3f} " | |
| f"mph={info.get('max_phase_coherence', 0.0):.3f}" | |
| f"{extra}" | |
| ) | |
| if not finite_model(model): | |
| raise FloatingPointError(f"non-finite parameter at step {step}") | |
| elapsed = time.perf_counter() - start_time | |
| if stream_mode and stream_source is not None: | |
| metadata["bytes_seen"] = stream_source.bytes_seen | |
| metadata["packets_seen"] = stream_source.packets_seen | |
| metadata["stream_epoch"] = stream_source.epoch | |
| metadata["validation_packet_count"] = len(stream_source.validation_packets()) | |
| (output_dir / "atomization_stats.json").write_text( | |
| json.dumps(metadata, indent=2), encoding="utf-8" | |
| ) | |
| checkpoint_path = output_dir / args.checkpoint_name | |
| model.save( | |
| checkpoint_path, | |
| extra_state={ | |
| "step": start_step + args.steps, | |
| "run_steps": args.steps, | |
| "optimizer": optimizer.state_dict(), | |
| "trajectory": trajectory, | |
| "seed": args.seed, | |
| "episode_length": args.episode_length, | |
| "episode_reset": bool(args.episode_reset), | |
| "atom_flush_every": args.atom_flush_every, | |
| "stream": stream_mode, | |
| "bytes_seen": (stream_source.bytes_seen if stream_source else data_bytes), | |
| }, | |
| ) | |
| # Validation: non-stream uses held-out tail; stream uses ring-buffer samples. | |
| model.reset_state(reset_atomizer=True) | |
| model.eval() | |
| validation_records = [] | |
| if stream_mode and stream_source is not None: | |
| val_pairs = stream_source.validation_transitions() | |
| # Cap validation work on huge rings. | |
| if len(val_pairs) > 256: | |
| val_pairs = val_pairs[-256:] | |
| else: | |
| val_pairs = list(zip(validation_packets[:-1], validation_packets[1:])) | |
| with torch.no_grad(): | |
| for current, target in val_pairs: | |
| loss, info = transition_metrics(model, current, target, train=False) | |
| validation_records.append({"loss": loss, **info}) | |
| if validation_records: | |
| validation_loss = float(np.mean([item["loss"] for item in validation_records])) | |
| validation_byte_loss = float(np.mean([item["byte_loss"] for item in validation_records])) | |
| else: | |
| validation_loss = float("nan") | |
| validation_byte_loss = float("nan") | |
| # Reload the final checkpoint before generation, then compare deterministic | |
| # samples so a restart cannot silently change the atom-native state path. | |
| reloaded = AtomNativeModel( | |
| d_model=args.d_model, | |
| n_modes=args.n_modes, | |
| n_atoms_max=args.n_atoms_max, | |
| max_payload_bytes=args.max_span_bytes, | |
| atomizer=Atomizer(max_span_bytes=args.max_span_bytes, pack_mode=args.atomizer_pack), | |
| field_max_rms=args.field_max_rms, | |
| energy_decay_bounds=energy_decay_bounds, | |
| enable_merge=bool(args.enable_merge), | |
| slow_every=int(args.slow_every), | |
| field_loss_weight=float(args.field_loss_weight), | |
| field_contrast_weight=float(args.field_contrast_weight), | |
| field_contrast_margin=float(args.field_contrast_margin), | |
| field_ignorance_weight=float(args.field_ignorance_weight), | |
| field_ignorance_margin=float(args.field_ignorance_margin), | |
| field_ignorance_ablate_shared=bool(args.field_ignorance_ablate_shared), | |
| field_ignorance_prompt_bank=bool(args.field_ignorance_prompt_bank), | |
| field_obligatory_readout=bool(args.field_obligatory_readout) or bool(args.field_obligatory_hard), | |
| field_obligatory_hard=bool(args.field_obligatory_hard), | |
| last_atom_readout=bool(args.last_atom_readout), | |
| printable_aux_weight=float(args.printable_aux_weight), | |
| slow_rms_rel_tol=float(args.slow_rms_rel_tol), | |
| ) | |
| reloaded.load(checkpoint_path) | |
| reload_state = { | |
| "n_atoms": len(reloaded.core.atoms), | |
| "field_norm": float(reloaded.core.state.alpha.detach().norm().item()), | |
| "time": float(reloaded.core.state.t.item()), | |
| "finite": bool(torch.isfinite(reloaded.core.state.alpha).all().item()), | |
| } | |
| generation = {} | |
| for prompt in PROMPTS: | |
| torch.manual_seed(args.seed + 1) | |
| source_bytes = model.generate_packets( | |
| prompt, | |
| max_packets=8, | |
| max_length=64, | |
| temperature=0.8, | |
| top_k=5, | |
| deterministic=False, | |
| ) | |
| torch.manual_seed(args.seed + 1) | |
| reload_bytes = reloaded.generate_packets( | |
| prompt, | |
| max_packets=8, | |
| max_length=64, | |
| temperature=0.8, | |
| top_k=5, | |
| deterministic=False, | |
| ) | |
| generation[prompt] = { | |
| "temperature": 0.8, | |
| "top_k": 5, | |
| "max_length": 20, | |
| "bytes_hex": source_bytes.hex(), | |
| "text": source_bytes.decode("utf-8", errors="replace"), | |
| "reload_text": reload_bytes.decode("utf-8", errors="replace"), | |
| "reload_match": source_bytes == reload_bytes, | |
| } | |
| metrics = { | |
| **metadata, | |
| "train": { | |
| "steps": args.steps, | |
| "efference_every": max(0, int(args.efference_every)), | |
| "free_run_aux_every": max(0, int(args.free_run_aux_every)), | |
| "free_run_aux_horizon": max(1, int(args.free_run_aux_horizon)), | |
| "free_run_aux_weight": float(args.free_run_aux_weight), | |
| "efference_steps": efference_steps, | |
| **({"field_bptt": int(model.field_bptt), "bptt_windows": bptt_windows} | |
| if int(getattr(model, "field_bptt", 0) or 0) else {}), | |
| **({"field_leak_learnable_rates": [float(x) for x in model.learnable_leak_rates().detach().reshape(-1)]} | |
| if getattr(model, "field_leak_logit", None) is not None else {}), | |
| "free_run_aux_steps": free_run_aux_steps, | |
| "free_run_aux_every": max(0, int(args.free_run_aux_every)), | |
| "free_run_aux_horizon": max(1, int(args.free_run_aux_horizon)), | |
| "free_run_aux_weight": float(args.free_run_aux_weight), | |
| "last_atom_scale": float(args.last_atom_scale), | |
| "last_atom_scale_adaptive": bool(args.last_atom_scale_adaptive), | |
| "last_atom_scale_min": float(args.last_atom_scale_min), | |
| "last_atom_scale_max": float(args.last_atom_scale_max), | |
| "field_hard_couple_scale": float(args.field_hard_couple_scale), | |
| "last_atom_dropout": float(args.last_atom_dropout), | |
| "field_leak": float(args.field_leak), | |
| "field_leak_spectrum": list(args.field_leak_spectrum) if args.field_leak_spectrum else None, | |
| "field_leak_spectrum_gain": ( | |
| args.field_leak_spectrum_gain if args.field_leak_spectrum else "none" | |
| ), | |
| "field_leak_spectrum_slow_frac": ( | |
| float(args.field_leak_spectrum_slow_frac) if args.field_leak_spectrum else 1.0 | |
| ), | |
| "slow_mode_readout": bool(args.slow_mode_readout), | |
| "slow_readout_ahat": args.slow_readout_ahat, | |
| "slow_readout_lr": args.slow_readout_lr, | |
| "slow_memory_loss": float(args.slow_memory_loss), | |
| "slow_memory_lags": list(args.slow_memory_lags), | |
| "slow_memory_target": args.slow_memory_target, | |
| "slow_memory_loss_mean": ( | |
| float(np.mean(slow_mem_values)) if slow_mem_values else None | |
| ), | |
| "slow_memory_loss_last_1000": ( | |
| float(np.mean(slow_mem_values[-1000:])) if slow_mem_values else None | |
| ), | |
| "start_step": start_step, | |
| "total_steps": start_step + args.steps, | |
| "episode_length": args.episode_length, | |
| "episode_reset": bool(args.episode_reset), | |
| "atom_flush_every": args.atom_flush_every, | |
| "elapsed_seconds": elapsed, | |
| "transitions_per_second": args.steps / max(elapsed, 1e-9), | |
| "first_loss": all_losses[0], | |
| "final_loss": all_losses[-1], | |
| "final_total_perplexity": float(np.exp(min(all_losses[-1], 30.0))), | |
| "finite_losses": all(np.isfinite(all_losses)), | |
| "finite_parameters": finite_model(model), | |
| "nan_or_inf": not all(np.isfinite(all_losses)) or not finite_model(model), | |
| "field_rms_before_max": max(field_rms_before_values), | |
| "field_rms_after_max": max(field_rms_after_values), | |
| "field_scale_min": min(field_scales), | |
| "field_norm_max": max(field_norm_values), | |
| "field_limited_steps": sum(scale < 0.999999 for scale in field_scales), | |
| "field_limited_fraction": sum(scale < 0.999999 for scale in field_scales) / len(field_scales), | |
| "final_dynamics": dynamics_state, | |
| "final_atoms_before_validation": len(reloaded.core.atoms), | |
| "stream": stream_mode, | |
| "bytes_seen": (stream_source.bytes_seen if stream_source else data_bytes), | |
| }, | |
| "validation": { | |
| "transitions": len(validation_records), | |
| "loss": validation_loss, | |
| "total_perplexity": float(np.exp(min(validation_loss, 30.0))) if np.isfinite(validation_loss) else None, | |
| "byte_loss": validation_byte_loss, | |
| "byte_perplexity": float(np.exp(min(validation_byte_loss, 30.0))) if np.isfinite(validation_byte_loss) else None, | |
| "source": "stream_ring" if stream_mode else "held_out_tail", | |
| }, | |
| "checkpoint": str(checkpoint_path), | |
| "reload": reload_state, | |
| "generation": generation, | |
| } | |
| (output_dir / "loss_trajectory.json").write_text(json.dumps(trajectory, indent=2), encoding="utf-8") | |
| (output_dir / "run_metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8") | |
| print(json.dumps(metrics, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |