"""Batched byte-tick runtime for ``AtomNativeModel`` (B independent byte streams). docs/BATCHED_TRAINING.md. This module adds **no parameter** and does not touch ``AtomNativeModel``: it re-implements the per-tick path of the byte-tick recipe --max-span-bytes 1 --field-obligatory-hard --no-enable-merge --no-payload-copy --last-atom-readout --efference-every 0 --slow-every 1 [--field-bptt K] [--field-leak L] [--field-leak-learnable] [--last-atom-dropout P] [--field-loss-weight W] --atom-flush-every N with a leading batch dimension, reading the weights of a wrapped single-stream model. Only the *runtime* state (field α, consolidation persistence, energy history, last byte) is batched, so a checkpoint written with ``model.save`` is the ordinary single-stream checkpoint (same parameter names and shapes) and is loaded by ``run_atom_native.py`` / ``chat_atom_native.py`` unchanged. Per stream b and tick t the computation is the one of ``AtomNativeModel._advance`` + ``transition_loss`` for that recipe (same ops, same order): * compile packet features -> (r, φ, ω, E, κ) (AtomCompiler.mapping) * α ← leak(α_{t-1}) + projection(atom) (graph kept under --field-bptt) * α ← RK4(α, drive = r) ; phase-sync Laplacian ; RMS cap (FieldAmplitudeController) * consolidation: stability of the living-atom energy history -> persistence EMA * hard readout: obl·α-MLP(â) + 0.3·frozen(â) + last-atom 2-gram readout (dropout) * loss = CE(length) + CE(next byte) + W·(1 − cos(field_probe(feats), features)) Paths of ``_advance`` that cannot change the loss or gradients of this recipe are skipped: MERGE (off), aggregation/abstraction (only feed the zeroed+frozen legacy decoder under hard), field interaction energy (diagnostic), α-bank snapshots (only used by contrast/ignorance losses, which must be 0). Everything else that is not supported raises ``BatchedUnsupported``. """ from __future__ import annotations import math import torch import torch.nn.functional as F from .atom_native import AtomNativeModel class BatchedUnsupported(ValueError): """A model/trainer option that the batched byte-tick runtime does not implement.""" def check_batched_supported(model: AtomNativeModel, *, efference_every: int = 0, free_run_aux_every: int = 0) -> None: """Raise ``BatchedUnsupported`` listing every unsupported option that is on.""" bad: list[str] = [] s = model.surface if int(model.max_payload_bytes) != 1 or int(model.atomizer.max_span_bytes) != 1: bad.append("max_span_bytes/max_payload_bytes != 1 (spans > 1 byte)") if not bool(model.field_obligatory_hard): bad.append("field_obligatory_hard off (only the hard readout is batched)") if bool(model.enable_merge): bad.append("enable_merge (MERGE)") if bool(getattr(s, "payload_enabled", True)): bad.append("payload copy (use --no-payload-copy)") if int(efference_every or 0) > 0 or int(getattr(model, "_efference_every", 0) or 0) > 0: bad.append("efference (--efference-every must be 0)") if int(free_run_aux_every or 0) > 0: bad.append("free-run aux") if getattr(model, "field_leak_spectrum", None) is not None: bad.append("field_leak_spectrum") if getattr(model, "slow_readout", None) is not None: bad.append("slow_mode_readout") if getattr(model, "field_leak_slow_block", None) is not None: bad.append("field_leak_slow_block") for name in ("field_contrast_weight", "field_ignorance_weight", "printable_aux_weight", "field_next_packet_weight", "slow_memory_weight"): if float(getattr(model, name, 0.0) or 0.0) != 0.0: bad.append(f"{name} != 0") if int(model.slow_every) != 1: bad.append("slow_every != 1") if bool(getattr(s, "last_atom_scale_adaptive", False)): bad.append("last_atom_scale_adaptive") if float(getattr(s, "field_hard_couple_scale", 0.0) or 0.0) > 0.0: bad.append("field_hard_couple_scale > 0") if bad: raise BatchedUnsupported("batched byte-tick mode does not support: " + "; ".join(bad)) class BatchedByteTick: """Runtime state + per-tick forward for B parallel byte streams. ``step(cur_bytes, features, target_bytes)`` advances every stream by one byte and returns per-stream losses (shape [B]); ``loss.mean()`` is the batch loss. With ``model.field_bptt > 0`` (and train mode + grad) α keeps its autograd graph across ticks until ``detach_field_graph()``. """ def __init__(self, model: AtomNativeModel, batch_size: int, *, atom_flush_every: int, efference_every: int = 0, free_run_aux_every: int = 0) -> None: check_batched_supported(model, efference_every=efference_every, free_run_aux_every=free_run_aux_every) if not 1 <= int(batch_size): raise ValueError("batch_size must be >= 1") cap = int(atom_flush_every or 0) if cap < 1: raise BatchedUnsupported("batched mode needs --atom-flush-every N >= 1 (bounded atom memory)") self.model = model self.B = int(batch_size) self.cap = cap st = model.core.state self.M = int(st.n_modes) self.D = int(st.d_model) self.reset_state() # ------------------------------------------------------------------ state @property def device(self) -> torch.device: return self.model.core.state.alpha.device @property def dtype(self) -> torch.dtype: return self.model.core.state.alpha.dtype def reset_state(self, mask: torch.Tensor | None = None) -> None: """Episode reset (``AtomNativeModel.reset_state``) of all streams, or of the streams where ``mask`` (bool [B]) is True.""" B, M, D, dev, dt = self.B, self.M, self.D, self.device, self.dtype if mask is None or not hasattr(self, "alpha"): self.alpha = torch.zeros(B, M, D, device=dev, dtype=dt) self.persistence = torch.zeros(B, D, device=dev, dtype=dt) self.persistence_strength = torch.zeros(B, device=dev, dtype=dt) self.e_buf = torch.zeros(B, self.cap, D, device=dev, dtype=dt) self.h_buf = torch.zeros(B, self.cap, D, device=dev, dtype=dt) self.last_byte = torch.zeros(B, dtype=torch.long, device=dev) self.n_atoms = [0] * B self.ticks = [0] * B return m = mask.to(device=dev, dtype=torch.bool) keep = (~m).to(dt) self.alpha = self.alpha * keep.view(B, 1, 1) self.persistence = self.persistence * keep.view(B, 1) self.persistence_strength = self.persistence_strength * keep self.e_buf = self.e_buf * keep.view(B, 1, 1) self.h_buf = self.h_buf * keep.view(B, 1, 1) self.last_byte = torch.where(m, torch.zeros_like(self.last_byte), self.last_byte) for b, flag in enumerate(m.tolist()): if flag: self.n_atoms[b] = 0 self.ticks[b] = 0 def flush_atoms(self, mask: torch.Tensor | None = None) -> None: """``--atom-flush-every``: clear the living atoms + energy history (field kept).""" if mask is None: self.n_atoms = [0] * self.B return for b, flag in enumerate(mask.tolist()): if flag: self.n_atoms[b] = 0 def detach_field_graph(self) -> None: """BPTT window boundary: keep α's value, drop its autograd history.""" self.alpha = self.alpha.detach() def state_dict(self) -> dict: return { "B": self.B, "cap": self.cap, "alpha": self.alpha.detach().cpu().clone(), "persistence": self.persistence.detach().cpu().clone(), "persistence_strength": self.persistence_strength.detach().cpu().clone(), # only the live prefix (after a flush n_atoms == 0: nothing to store) "e_buf": self.e_buf[:, :max(self.n_atoms)].detach().cpu().clone(), "h_buf": self.h_buf[:, :max(self.n_atoms)].detach().cpu().clone(), "last_byte": self.last_byte.detach().cpu().clone(), "n_atoms": list(self.n_atoms), "ticks": list(self.ticks), } def load_state_dict(self, sd: dict) -> None: if int(sd["B"]) != self.B or int(sd["cap"]) != self.cap: raise ValueError(f"batched state B/cap {sd['B']}/{sd['cap']} != {self.B}/{self.cap}") dev, dt = self.device, self.dtype self.alpha = sd["alpha"].to(device=dev, dtype=dt).clone() self.persistence = sd["persistence"].to(device=dev, dtype=dt).clone() self.persistence_strength = sd["persistence_strength"].to(device=dev, dtype=dt).clone() for name in ("e_buf", "h_buf"): buf = torch.zeros(self.B, self.cap, self.D, device=dev, dtype=dt) src = sd[name].to(device=dev, dtype=dt) buf[:, :src.shape[1]] = src setattr(self, name, buf) self.last_byte = sd["last_byte"].to(device=dev).clone() self.n_atoms = [int(x) for x in sd["n_atoms"]] self.ticks = [int(x) for x in sd["ticks"]] def sync_to_model(self, stream: int = 0) -> None: """Copy stream ``stream``'s field/persistence into the single-stream model runtime buffers (what a single-stream checkpoint stores as live state). The structural atom list is left empty (it is episode memory).""" m = self.model with torch.no_grad(): m.core.state.alpha.copy_(self.alpha[stream].detach()) m.core.state.t.fill_(float(self.ticks[stream]) * float(m.core.dynamics.dt)) m.core.consolidation.persistence.copy_(self.persistence[stream].detach()) m.core.consolidation.persistence_strength.copy_(self.persistence_strength[stream].detach()) m.core.atoms = type(m.core.atoms)() m.core.energy_history = [] m._bptt_alpha = None # ---------------------------------------------------------------- pieces def _const(self, ref: torch.Tensor) -> dict: """Per-(device, dtype) constant tensors (avoids re-creating them every tick).""" key = (ref.device, ref.dtype) cache = self.__dict__.setdefault("_const_cache", {}) if key not in cache: n, D = self.M, self.D st, dyn = self.model.core.state, self.model.core.dynamics.dynamics cache[key] = { "mode_phase": st.mode_phase.to(device=ref.device, dtype=ref.dtype).view(1, n, 1), "mode_index": torch.arange(n, device=ref.device, dtype=ref.dtype).view(1, n, 1), "dyn_phase": dyn.mode_phase.to(device=ref.device, dtype=ref.dtype).view(1, n, 1), "dim_phase": torch.arange(D, device=ref.device, dtype=ref.dtype).view(1, 1, D), "zeros_long": torch.zeros(self.B, dtype=torch.long, device=ref.device), } cache[key]["sin_phase"] = torch.sin(cache[key]["dyn_phase"]) return cache[key] def _to_dev(self, t: torch.Tensor) -> torch.Tensor: """Host->device copy without a host sync on CUDA (pinned, non-blocking).""" if self.device.type == "cuda": return t.pin_memory().to(self.device, non_blocking=True) return t.to(self.device) def _bptt_active(self) -> bool: m = self.model return bool(int(getattr(m, "field_bptt", 0) or 0) > 0 and m.training and torch.is_grad_enabled()) def _compile(self, features: torch.Tensor): c = self.model.compiler raw = c.mapping(features).view(features.shape[0], 8, c.d_model) r = F.normalize(raw[:, 0], dim=-1) phi = raw[:, 1] omega = F.softplus(raw[:, 2]) + 1e-3 energy = torch.sigmoid(raw[:, 3]) kappa = F.softplus(raw[:, 4]) return r, phi, omega, energy, kappa def _projection(self, r, phi, omega, E, kappa) -> torch.Tensor: """FractalSuperpositionState._projection per stream (no sum over atoms).""" n = self.M cst = self._const(r) mode_phase = cst["mode_phase"] phase = phi.mean(dim=-1, keepdim=True).unsqueeze(1) freq = omega.mean(dim=-1, keepdim=True).unsqueeze(1).abs() freq_angle = 2.0 * torch.pi * freq / (1.0 + freq) centre = (freq_angle / (2.0 * torch.pi)) * n mode_index = cst["mode_index"] delta = torch.remainder(mode_index - centre + n / 2, n) - n / 2 width = 1.0 + 0.25 * torch.sigmoid(freq) kernel = torch.exp(-(delta / width) ** 2) modulation = 0.5 + 0.5 * torch.cos(mode_phase - phase) strength = (E * kappa).mean(dim=-1, keepdim=True).unsqueeze(1) vectors = r.unsqueeze(1) * strength return vectors * kernel * modulation def _dyn(self, alpha: torch.Tensor, tok: torch.Tensor) -> torch.Tensor: """ToroidalDynamics.forward with a per-stream drive token (B, D).""" dyn = self.model.core.dynamics.dynamics left = torch.roll(alpha, shifts=1, dims=-2) right = torch.roll(alpha, shifts=-1, dims=-2) laplacian = left + right - 2.0 * alpha coupling = dyn.phase_sync * laplacian cst = self._const(alpha) phase = cst["dyn_phase"] rotated = cst["sin_phase"] * alpha dim_phase = cst["dim_phase"] t = tok.reshape(tok.shape[0], 1, self.D) drive = t * torch.cos(phase + 0.5 * t * dim_phase) drive = dyn.coupling_scale * drive decay = (dyn.energy_decay - 1.0) * alpha return coupling + rotated + drive + decay def _evolve(self, alpha: torch.Tensor, tok: torch.Tensor) -> torch.Tensor: dt = self.model.core.dynamics.dt k1 = self._dyn(alpha, tok) k2 = self._dyn(alpha + dt * 0.5 * k1, tok) k3 = self._dyn(alpha + dt * 0.5 * k2, tok) k4 = self._dyn(alpha + dt * k3, tok) return alpha + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4) def _cap(self, field: torch.Tensor): ctl = self.model.field_controller rms_before = torch.sqrt(field.pow(2).mean(dim=(1, 2)) + ctl.eps) if ctl.max_rms is None: return field, rms_before, torch.ones_like(rms_before) cst = self._const(field) if "max_rms" not in cst or cst["max_rms_val"] != ctl.max_rms: cst["max_rms"] = torch.full((), float(ctl.max_rms), device=field.device, dtype=field.dtype) cst["max_rms_val"] = ctl.max_rms scale = torch.minimum(torch.ones_like(rms_before), cst["max_rms"] / (rms_before + ctl.eps)) return field * scale.view(-1, 1, 1), rms_before, scale def _last_atom(self, last: torch.Tensor, prev: torch.Tensor, phi: torch.Tensor) -> torch.Tensor: la = self.model.surface.last_atom pair = torch.cat([la.byte_embed(last), la.byte_embed(prev)], dim=-1) h = F.relu(la.expand(pair)) keep = la.dentate_keep(h.shape[-1]) top_values, top_index = torch.topk(h, keep, dim=-1) sparse = torch.zeros_like(h).scatter(-1, top_index, top_values) z = la.contract(sparse) phase = torch.cat([torch.cos(phi), torch.sin(phi)], dim=-1) out = F.gelu(z + la.phase_proj(phase)) return la.head(out) def _mean_upto(self, buf: torch.Tensor, lo: list[int], hi: list[int], std: bool = False): """Per-stream mean (and population std) of buf[b, lo_b:hi_b] (hi_b > lo_b).""" if all(x == lo[0] for x in lo) and all(x == hi[0] for x in hi): sl = buf[:, lo[0]:hi[0]] mean = sl.mean(dim=1) return (mean, sl.std(dim=1, unbiased=False)) if std else mean idx = torch.arange(buf.shape[1], device=buf.device).view(1, -1) lo_t = torch.tensor(lo, device=buf.device).view(-1, 1) hi_t = torch.tensor(hi, device=buf.device).view(-1, 1) w = ((idx >= lo_t) & (idx < hi_t)).to(buf.dtype).unsqueeze(-1) cnt = (hi_t - lo_t).clamp_min(1).to(buf.dtype).view(-1, 1) mean = (buf * w).sum(dim=1) / cnt if not std: return mean var = (((buf - mean.unsqueeze(1)) ** 2) * w).sum(dim=1) / cnt return mean, var.sqrt() # ------------------------------------------------------------------ tick def step(self, cur_bytes: torch.Tensor, features: torch.Tensor, target_bytes: torch.Tensor, *, dropout_mask: torch.Tensor | None = None) -> dict: """One byte tick for every stream. ``cur_bytes``/``target_bytes``: long [B]; ``features``: [B, 320] packet features of the current byte. ``dropout_mask`` (bool [B], tests only) overrides the last-atom dropout draw.""" m = self.model B, dev = self.B, self.device cur_bytes = cur_bytes.to(dev) target_bytes = target_bytes.to(dev) features = features.to(device=dev, dtype=self.dtype) self.n_atoms = [n + 1 for n in self.n_atoms] self.ticks = [t + 1 for t in self.ticks] r, phi, omega, energy, kappa = self._compile(features) # ---- inject (leak, projection); graph kept across ticks under BPTT bptt = self._bptt_active() prev = self.alpha if bptt else self.alpha.detach() learn_lam = m.learnable_leak_rates() if getattr(m, "field_leak_logit", None) is not None else None if learn_lam is not None: prev = prev * (1.0 - learn_lam) if bptt else (prev * (1.0 - learn_lam)).detach() else: leak = float(getattr(m, "field_leak", 0.0) or 0.0) if leak > 0.0: prev = prev * (1.0 - leak) if bptt: alpha = prev + self._projection(r, phi, omega, energy, kappa) else: with torch.no_grad(): alpha = prev + self._projection(r, phi, omega, energy, kappa) # ---- living-atom energy buffer (memory atoms are detached) n_list = self.n_atoms if max(n_list) > self.cap: raise RuntimeError(f"living atoms exceed atom_flush_every={self.cap}; call flush_atoms()") lockstep = all(n == n_list[0] for n in n_list) if lockstep: pos = n_list[0] - 1 ar = slice(None) else: pos = self._to_dev(torch.tensor([n - 1 for n in n_list])) ar = torch.arange(B, device=dev) self.e_buf[ar, pos] = energy.detach() e_mean = self._mean_upto(self.e_buf, [0] * B, n_list) # ---- field dynamics alpha_new = self._evolve(alpha, r) ps = m.core.dynamics.dynamics.phase_sync left = torch.roll(alpha_new, 1, 1) right = torch.roll(alpha_new, -1, 1) alpha_new = alpha_new + 0.1 * ps * (left + right - 2.0 * alpha_new) alpha_new, rms_before, scale = self._cap(alpha_new) # ---- consolidation (slow path, every tick): streams with > 5 history entries hist_n = [n - 1 for n in n_list] act = [h > 5 for h in hist_n] if any(act): lo = [max(0, h - 10) if a else 0 for h, a in zip(hist_n, act)] hi = [h if a else 1 for h, a in zip(hist_n, act)] with torch.no_grad(): h_mean, h_std = self._mean_upto(self.h_buf, lo, hi, std=True) cv = h_std / h_mean.abs().clamp_min(1e-6) stability = torch.exp(-cv).clamp(0, 1) * e_mean.clamp(0, 1) score = e_mean.clamp(0, 1) * stability.clamp(0, 1) sel = score > m.core.consolidation.consolidation_threshold # [B, D] cnt = sel.sum(dim=1) upd = cnt > 0 if not all(act): upd = upd & self._to_dev(torch.tensor(act)) a_det = alpha_new.detach() incoming = (a_det * sel.unsqueeze(1).to(a_det.dtype)).sum(dim=(1, 2)) / ( cnt.clamp_min(1).to(a_det.dtype) * self.M) strength = sel.to(a_det.dtype).mean(dim=1) pd = m.core.consolidation.persistence_decay new_p = self.persistence * pd + (1 - pd) * incoming.view(B, 1) new_s = self.persistence_strength * pd + (1 - pd) * strength self.persistence = torch.where(upd.view(B, 1), new_p, self.persistence) self.persistence_strength = torch.where(upd, new_s, self.persistence_strength) self.h_buf[ar, pos] = e_mean # ---- hard readout s = m.surface flat = alpha_new.reshape(B, -1) a_only = flat @ s.alpha_jl.t() a_rms = a_only.pow(2).mean(dim=-1, keepdim=True).sqrt().clamp_min(1e-8) a_hat = a_only / a_rms frozen_byte = a_hat @ s.alpha_byte_frozen.t() frozen_len = a_hat @ s.alpha_length_frozen.t() train_byte = s.alpha_byte_proj(a_hat) train_len = s.alpha_length_proj(a_hat) obl = s.obligatory_scale() byte_logits = obl * train_byte + 0.3 * frozen_byte la_on = getattr(s, "last_atom", None) is not None and not bool(getattr(s, "last_atom_force_off", False)) drop = None if la_on: p_drop = float(getattr(s, "last_atom_dropout", 0.0) or 0.0) if dropout_mask is not None: drop = dropout_mask.detach().to("cpu", dtype=torch.bool).reshape(B) elif p_drop > 0.0 and m.training and torch.is_grad_enabled(): # same draw as the single-stream path (one uniform per stream, CPU RNG) drop = torch.ones(B, dtype=torch.bool) if p_drop >= 1.0 else torch.rand(B) < p_drop if drop is not None and bool(drop.all()): la_on = False # like zeros_like(la_byte): last-atom params get no grad at all if la_on: if lockstep: prev_b = self.last_byte if n_list[0] >= 2 else cur_bytes else: prev_b = torch.where(self._to_dev(torch.tensor([n >= 2 for n in n_list])), self.last_byte, cur_bytes) la_raw = self._last_atom(cur_bytes, prev_b, phi) la_scale = float(getattr(s, "last_atom_scale", 1.0)) la_byte = la_raw if la_scale == 1.0 else la_raw * la_scale # x*1.0 == x exactly if drop is not None and bool(drop.any()): la_byte = la_byte * self._to_dev((~drop).to(la_byte.dtype)).unsqueeze(-1) byte_logits = byte_logits + la_byte length_logits = obl * train_len + 0.3 * frozen_len # ---- losses (per stream) length_loss = F.cross_entropy(length_logits, self._const(length_logits)["zeros_long"], reduction="none") byte_loss = F.cross_entropy(byte_logits, target_bytes, reduction="none") loss = length_loss + byte_loss w = float(m.field_loss_weight) persist_loss = torch.zeros_like(byte_loss) if w > 0: mean = alpha_new.mean(dim=1) std = alpha_new.std(dim=1, unbiased=False) feats = s.field_feat_norm(torch.cat([mean, std, self.persistence, r], dim=-1)) pred = m.field_probe(feats) persist_loss = 1.0 - F.cosine_similarity(pred, features, dim=-1) loss = loss + w * persist_loss self.alpha = alpha_new if bptt else alpha_new.detach() self.last_byte = cur_bytes return { "loss": loss, "byte_loss": byte_loss, "length_loss": length_loss, "persist_loss": persist_loss, "byte_logits": byte_logits, "field": alpha_new, "field_rms_before": rms_before, "field_scale": scale, }