ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """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 | |
| def device(self) -> torch.device: | |
| return self.model.core.state.alpha.device | |
| 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, | |
| } | |