ATOM (atom-ai): code snapshot (100fae0 + batched-train), README with measured numbers, curated checkpoints (part 4)
b9301f1 verified | """B parallel byte streams over a (possibly multi-GB) memory-mapped corpus. | |
| docs/BATCHED_TRAINING.md. Bytes only: no tokenizer, no vocabulary. The shard | |
| files (sorted glob, same order as ``resolve_shard_paths``) are viewed as ONE | |
| circular byte string, exactly the continuum ``StreamingPacketSource`` produces | |
| with ``--loop-shards`` (shard ends and the epoch wrap are plain byte | |
| transitions; in byte-tick mode the shard-end flush emits nothing). | |
| Stream b sits at an absolute, ever-increasing position ``pos[b]`` (the corpus | |
| byte is ``pos mod total``). Each ``next()`` returns, for every stream, the | |
| current byte, its 320-d byte-tick packet features and the target (next) byte, | |
| then advances every stream by one byte. | |
| Features are the ``Atomizer`` byte-tick features (``max_span_bytes == 1``) | |
| computed vectorised: everything that depends only on the byte comes from a | |
| 256-row table built with the scalar ``Atomizer`` code itself; the stream- | |
| dependent parts are the phase (``pos mod phase_period``), ``previous_length`` | |
| (0 on the first packet of a stream context, then 1), the context similarity and | |
| the 16-value rolling context sketch (EMA 0.85, low nibble). From a fresh | |
| context this equals ``Atomizer.encode_bytes`` feature for feature (tested). | |
| Resumable: ``state_dict()`` holds the corpus signature, the positions and the | |
| per-stream feature context (exact resume). | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from pathlib import Path | |
| from typing import Sequence | |
| import numpy as np | |
| import torch | |
| from .atomizer import Atomizer | |
| # Column layout of the 320 features: hist 0:256, first 256:272, last 272:288, | |
| # stats 288:304, context 304:320. | |
| _STATS = 288 | |
| _COL_SIN, _COL_COS, _COL_PREV, _COL_SIM = _STATS + 10, _STATS + 11, _STATS + 12, _STATS + 13 | |
| _CTX = 304 | |
| class ByteCorpus: | |
| """Ordered shard files as one circular, memory-mapped byte string.""" | |
| def __init__(self, paths: Sequence[str | Path]) -> None: | |
| self.paths = [Path(p).resolve() for p in paths] | |
| if not self.paths: | |
| raise ValueError("ByteCorpus needs at least one file") | |
| self.sizes = [p.stat().st_size for p in self.paths] | |
| if any(s == 0 for s in self.sizes): | |
| empty = [str(p) for p, s in zip(self.paths, self.sizes) if s == 0] | |
| raise ValueError(f"empty shard(s): {empty}") | |
| self.maps = [np.memmap(p, dtype=np.uint8, mode="r") for p in self.paths] | |
| self.starts = np.cumsum([0] + self.sizes[:-1]).astype(np.int64) | |
| self.total = int(sum(self.sizes)) | |
| if self.total < 2: | |
| raise ValueError("corpus must hold at least 2 bytes") | |
| def signature(self) -> dict: | |
| return {"paths": [str(p) for p in self.paths], "sizes": list(self.sizes), "total": self.total} | |
| def read(self, pos: int, n: int) -> np.ndarray: | |
| """n bytes starting at absolute position ``pos`` (wraps around the corpus).""" | |
| out = np.empty(n, dtype=np.uint8) | |
| g = int(pos) % self.total | |
| done = 0 | |
| while done < n: | |
| i = int(np.searchsorted(self.starts, g, side="right") - 1) | |
| off = g - int(self.starts[i]) | |
| take = min(n - done, self.sizes[i] - off) | |
| out[done:done + take] = self.maps[i][off:off + take] | |
| done += take | |
| g = (g + take) % self.total | |
| return out | |
| def byte_tick_feature_table(pack_mode: str = "linguistic") -> torch.Tensor: | |
| """[256, 320] static byte-tick features (context-free part), built with the | |
| scalar Atomizer: stream-dependent columns are overwritten per tick.""" | |
| at = Atomizer(max_span_bytes=1, pack_mode=pack_mode) | |
| rows = [] | |
| for b in range(256): | |
| at.reset() | |
| packets = at.encode_bytes(bytes([b]), reset=True, flush=True) | |
| assert len(packets) == 1 | |
| rows.append(packets[0].features) | |
| return torch.stack(rows) | |
| class BatchedByteStream: | |
| """B byte streams with vectorised byte-tick features. | |
| ``offsets``: starting absolute positions (default: evenly spaced over the | |
| corpus). ``warmup``: bytes before each non-zero offset replayed through the | |
| feature context only (0.85^64 ≈ 3e-5, so 64 bytes ≈ a full replay). | |
| """ | |
| def __init__(self, corpus: ByteCorpus, batch_size: int, *, offsets: Sequence[int] | None = None, | |
| chunk: int = 4096, pack_mode: str = "linguistic", phase_period: float = 256.0, | |
| warmup: int = 64, device: str | torch.device = "cpu") -> None: | |
| self.corpus = corpus | |
| self.B = int(batch_size) | |
| self.chunk = max(1, int(chunk)) | |
| self.pack_mode = pack_mode | |
| self.phase_period = float(phase_period) | |
| self.device = torch.device(device) | |
| if offsets is None: | |
| offsets = [(b * corpus.total) // self.B for b in range(self.B)] | |
| if len(offsets) != self.B: | |
| raise ValueError("need one offset per stream") | |
| self.table = byte_tick_feature_table(pack_mode).to(self.device) | |
| # sin/cos of the phase, built with math.* exactly like Atomizer._make_features. | |
| per = int(self.phase_period) | |
| if float(per) != self.phase_period: | |
| raise ValueError("phase_period must be an integer") | |
| ph = [2.0 * math.pi * (p % self.phase_period) / self.phase_period for p in range(per)] | |
| self.sin_t = torch.tensor([math.sin(x) for x in ph], dtype=torch.float32, device=self.device) | |
| self.cos_t = torch.tensor([math.cos(x) for x in ph], dtype=torch.float32, device=self.device) | |
| self.pos = np.asarray([int(o) for o in offsets], dtype=np.int64) | |
| self.ctx = torch.zeros(self.B, 16, dtype=torch.float32, device=self.device) | |
| self.prev_len = torch.zeros(self.B, dtype=torch.float32, device=self.device) | |
| self._buf: torch.Tensor | None = None | |
| self._buf_i = 0 | |
| self._ph = None # device copy of pos mod phase_period (no per-tick host->device copy) | |
| if warmup > 0: | |
| self._warmup(int(warmup)) | |
| # --------------------------------------------------------------- context | |
| def _warmup(self, n: int) -> None: | |
| start = self.pos.copy() | |
| for b in range(self.B): | |
| k = int(min(n, start[b])) | |
| if k <= 0: | |
| continue | |
| raw = torch.from_numpy(self.corpus.read(int(start[b]) - k, k).astype(np.int64)) | |
| for v in raw.tolist(): | |
| self._ctx_update_one(b, v) | |
| self._buf = None | |
| def _ctx_update_one(self, b: int, v: int) -> None: | |
| g = torch.zeros(16, dtype=torch.float32, device=self.device) | |
| g[v & 15] = 1.0 | |
| self.ctx[b].mul_(0.85).add_(0.15 * g) | |
| self.prev_len[b] = 1.0 | |
| def _refill(self) -> None: | |
| n = self.chunk + 1 | |
| arr = np.stack([self.corpus.read(int(p), n) for p in self.pos]) | |
| self._buf = torch.from_numpy(arr.astype(np.int64)).to(self.device) | |
| self._buf_i = 0 | |
| # ------------------------------------------------------------------ next | |
| def next(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: | |
| if self._buf is None or self._buf_i >= self.chunk: | |
| self._refill() | |
| i = self._buf_i | |
| cur = self._buf[:, i] | |
| tgt = self._buf[:, i + 1] | |
| feats = self.table.index_select(0, cur).clone() | |
| if self._ph is None: | |
| self._ph = torch.from_numpy(self.pos % int(self.phase_period)).to(self.device) | |
| ph = self._ph | |
| feats[:, _COL_SIN] = self.sin_t[ph] | |
| feats[:, _COL_COS] = self.cos_t[ph] | |
| feats[:, _COL_PREV] = self.prev_len | |
| low = cur & 15 | |
| dot = self.ctx.gather(1, low.view(-1, 1)).squeeze(1).double() | |
| denom = self.ctx.norm(dim=1).double() | |
| sim = torch.where(denom > 1e-12, dot / denom.clamp_min(1e-300), torch.zeros_like(dot)) | |
| feats[:, _COL_SIM] = ((sim + 1.0) / 2.0).float() | |
| feats[:, _CTX:] = self.ctx | |
| grouped = torch.zeros_like(self.ctx).scatter_(1, low.view(-1, 1), 1.0) | |
| self.ctx = self.ctx.mul(0.85).add(0.15 * grouped) | |
| self.prev_len = torch.ones_like(self.prev_len) | |
| self.pos += 1 | |
| self._ph = (ph + 1) % int(self.phase_period) | |
| self._buf_i += 1 | |
| return cur, feats, tgt | |
| __next__ = next | |
| def __iter__(self): | |
| return self | |
| # ------------------------------------------------------------ checkpoint | |
| def state_dict(self) -> dict: | |
| return { | |
| "corpus": self.corpus.signature(), "B": self.B, "pack_mode": self.pack_mode, | |
| "phase_period": self.phase_period, "pos": self.pos.tolist(), | |
| "ctx": self.ctx.detach().cpu().clone(), "prev_len": self.prev_len.detach().cpu().clone(), | |
| } | |
| def load_state_dict(self, sd: dict, *, strict_corpus: bool = True) -> None: | |
| if int(sd["B"]) != self.B: | |
| raise ValueError(f"stream state has B={sd['B']}, loader has B={self.B}") | |
| sig = self.corpus.signature() | |
| if strict_corpus and (sd["corpus"]["sizes"] != sig["sizes"]): | |
| raise ValueError("corpus changed since the checkpoint (sizes differ); pass strict_corpus=False to accept") | |
| self.pos = np.asarray(sd["pos"], dtype=np.int64) | |
| self.ctx = sd["ctx"].to(self.device).clone() | |
| self.prev_len = sd["prev_len"].to(self.device).clone() | |
| self._ph = None | |
| self._buf = None | |
| self._buf_i = 0 | |