| """Deterministic streaming loader over packed uint16 shards. |
| |
| - One PackedStream per language: contiguous (seq_len+1)-token windows, |
| epoch-level window permutation (seeded), transparent epoch rollover with |
| epoch counting (matters for AR, whose pool may be smaller than the budget). |
| - MixedStream: bilingual token-level mixing. For every global sample slot a |
| stateless hash of (seed, slot) picks the language, so the mixture is exactly |
| reproducible, independent of world size, and resumable from a slot counter. |
| |
| Mixing mode (thesis decision, configurable in the run config): |
| token-level p=0.5 matches ATLAS's compute-matched framing: each language |
| contributes half the training *tokens*; under a starved tokenizer the |
| cross-script partner therefore sees less *content* -- that is exactly the |
| effect under audit. `content` mode instead sets p from the byte premiums so |
| content is matched and token counts differ. |
| """ |
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
| from ..paths import shard_dir |
|
|
|
|
| def _splitmix64(x: int) -> int: |
| x = (x + 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF |
| z = x |
| z = ((z ^ (z >> 30)) * 0xBF58476D1CE4E5B9) & 0xFFFFFFFFFFFFFFFF |
| z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & 0xFFFFFFFFFFFFFFFF |
| return z ^ (z >> 31) |
|
|
|
|
| class PackedStream: |
| def __init__(self, lang: str, tok_name: str, seq_len: int, seed: int): |
| self.lang, self.seq_len, self.seed = lang, seq_len, seed |
| d = shard_dir(lang, tok_name) |
| index = json.loads((d / "index.json").read_text()) |
| |
| |
| |
| |
| names = [name for name in sorted(index) if index[name] > 0] |
| self.paths = [d / name for name in names] |
| self.lens = np.array([index[name] for name in names], dtype=np.int64) |
| self.cum = np.concatenate([[0], np.cumsum(self.lens)]) |
| self.total_tokens = int(self.cum[-1]) |
| self.n_windows = self.total_tokens // (seq_len + 1) |
| if self.n_windows == 0: |
| raise RuntimeError(f"shards for {lang}/{tok_name} too small") |
| self._mm = [None] * len(self.paths) |
| self._epoch = -1 |
| self._perm = None |
| self.consumed = 0 |
|
|
| def _memmap(self, i: int): |
| if self._mm[i] is None: |
| self._mm[i] = np.memmap(self.paths[i], dtype=np.uint16, mode="r") |
| return self._mm[i] |
|
|
| def _window(self, w: int) -> np.ndarray: |
| start = w * (self.seq_len + 1) |
| i = int(np.searchsorted(self.cum, start, side="right") - 1) |
| out = np.empty(self.seq_len + 1, dtype=np.uint16) |
| filled = 0 |
| while filled < self.seq_len + 1: |
| off = start + filled - self.cum[i] |
| take = int(min(self.lens[i] - off, self.seq_len + 1 - filled)) |
| out[filled:filled + take] = self._memmap(i)[off:off + take] |
| filled += take |
| i += 1 |
| return out |
|
|
| def get(self, k: int) -> np.ndarray: |
| """k-th window in the global consumption order (deterministic).""" |
| epoch, pos = divmod(k, self.n_windows) |
| if epoch != self._epoch: |
| rng = np.random.default_rng([self.seed, epoch, 0xDA7A]) |
| self._perm = rng.permutation(self.n_windows) |
| self._epoch = epoch |
| return self._window(int(self._perm[pos])) |
|
|
| def next(self) -> np.ndarray: |
| x = self.get(self.consumed) |
| self.consumed += 1 |
| return x |
|
|
| @property |
| def epochs(self) -> float: |
| return self.consumed / self.n_windows |
|
|
| def state_dict(self): |
| return {"consumed": self.consumed} |
|
|
| def load_state_dict(self, s): |
| self.consumed = s["consumed"] |
|
|
|
|
| class MixedStream: |
| """Token-level language mixing across one or two PackedStreams.""" |
|
|
| def __init__(self, langs: list[str], tok_name: str, seq_len: int, seed: int, |
| probs: list[float] | None = None): |
| self.langs = langs |
| self.streams = {l: PackedStream(l, tok_name, seq_len, seed + j) |
| for j, l in enumerate(langs)} |
| self.probs = probs or [1.0 / len(langs)] * len(langs) |
| assert abs(sum(self.probs) - 1.0) < 1e-6 |
| self.seed = seed |
| self.slot = 0 |
|
|
| def _choose(self, slot: int) -> str: |
| if len(self.langs) == 1: |
| return self.langs[0] |
| u = _splitmix64(self.seed * 0x100000001 + slot) / 2**64 |
| acc = 0.0 |
| for l, p in zip(self.langs, self.probs): |
| acc += p |
| if u < acc: |
| return l |
| return self.langs[-1] |
|
|
| def next_batch(self, n: int) -> tuple[np.ndarray, dict[str, int]]: |
| """n windows -> array (n, seq_len+1); also per-language window counts.""" |
| rows, counts = [], {l: 0 for l in self.langs} |
| for _ in range(n): |
| l = self._choose(self.slot) |
| rows.append(self.streams[l].next()) |
| counts[l] += 1 |
| self.slot += 1 |
| return np.stack(rows), counts |
|
|
| def draw(self) -> tuple[str, int]: |
| """Advance one global slot; return (lang, occurrence-index) WITHOUT reading. |
| |
| Occurrence-index is the k-th window of that language in global order, so |
| stream.get(occ) is a pure deterministic lookup. All DDP ranks call draw() |
| for every slot (keeping counters identical) but only materialise their own |
| rows via streams[lang].get(occ) -- giving each rank a disjoint, reproducible |
| 1/world shard of the global batch. |
| """ |
| l = self._choose(self.slot) |
| occ = self.streams[l].consumed |
| self.streams[l].consumed += 1 |
| self.slot += 1 |
| return l, occ |
|
|
| def rank_batch(self, n_global: int, rank: int, world: int |
| ) -> tuple[np.ndarray, dict[str, int]]: |
| """n_global windows drawn globally; rows n_global/world owned by `rank`.""" |
| rows, counts = [], {l: 0 for l in self.langs} |
| for j in range(n_global): |
| l, occ = self.draw() |
| if j % world == rank: |
| rows.append(self.streams[l].get(occ)) |
| counts[l] += 1 |
| return np.stack(rows), counts |
|
|
| def skip_to(self, slot: int) -> None: |
| """Fast-forward mixture state without reading data (for resume).""" |
| assert slot >= self.slot |
| for s in range(self.slot, slot): |
| self.streams[self._choose(s)].consumed += 1 |
| self.slot = slot |
|
|
| def state_dict(self): |
| return {"slot": self.slot, |
| "streams": {l: s.state_dict() for l, s in self.streams.items()}} |
|
|
| def load_state_dict(self, s): |
| self.slot = s["slot"] |
| for l, st in s["streams"].items(): |
| self.streams[l].load_state_dict(st) |
|
|
| def stats(self) -> dict: |
| return {l: {"windows": s.consumed, "epochs": round(s.epochs, 4)} |
| for l, s in self.streams.items()} |
|
|
