"""Online token-streaming LM loader with a prefetch queue. Streams the Phase-1 web mixture from the HF Hub, quality-filters + tokenizes in background threads, and serves packed ``[batch, seqlen]`` int32 windows -- a drop-in for ``cortex.data.ShardLoader`` but with NO on-disk pre-tokenized dataset (removes the uint32 storage wall). Pipeline (all threads release the GIL on the hot paths -- network I/O and the Rust tokenizer): reader thread/source -> stream + filter + tokenize -> per-source doc queue packer thread -> token-weighted interleave -> pack -> batch queue main (training) thread -> next() pops a ready [B,T] batch Every queue is bounded, so the producers block when the consumer is busy (backpressure) -- HF fetch + tokenization overlap the accelerator step without running unboundedly ahead. Cross-session data continuity: every session streams a FRESH shard order from a seed derived from ``session_seed`` (the resume step), with raw doc counters carried across sessions for visibility. Exact HF-state resume is deliberately NOT used: on datasets 5.0.0, ``load_state_dict()`` on a shuffled stream permanently breaks subsequent ``state_dict()`` tracking (returns zeroed states -- verified in isolation), so a session that loads a position cannot hand one off. Session-seeded reshuffles keep the data fresh deterministically (the failure that caused the CE-3.2 plateau was every session re-reading the SAME seed-0 prefix); overlap across sessions is negligible at corpus scale. loader = StreamingMixLoader(batch=16, seqlen=2048) # plug straight into run_training it = iter(loader); batch = next(it) # np.int32 [16, 2048] """ from __future__ import annotations import queue import threading import numpy as np from .webdata import default_sources, make_decontam, stream_texts from .tokenizer import CortexTokenizer def _mx(o): """Max numeric leaf -- 0 for None/empty: a quick 'has this stream state advanced' probe.""" if isinstance(o, dict): return max([_mx(v) for v in o.values()] + [0]) if isinstance(o, list): return max([_mx(v) for v in o] + [0]) if isinstance(o, (int, float)) and not isinstance(o, bool): return int(o) return 0 class StreamingMixLoader: def __init__(self, batch: int, seqlen: int, *, grad_accum: int = 1, sources=None, tokenizer=None, prefetch: int = 8, doc_queue: int = 256, shuffle_buf: int = 10_000, seed: int = 0, decontam_file: str | None = None, holdout_mod: int = 100, resume_states: dict | None = None, session_seed: int = 0, snapshot_every: int = 200): self.batch, self.seqlen, self.grad_accum = batch, seqlen, grad_accum # the CPU prefetch packs a whole macro-batch (G micro-batches) so the gather # overlaps the device step; G==1 -> plain [B,T] (drop-in for ShardLoader). self._need = grad_accum * batch * seqlen self._mshape = (batch, seqlen) if grad_accum == 1 else (grad_accum, batch, seqlen) self.sources = sources or default_sources() w = np.array([s.weight for s in self.sources], dtype=np.float64) self.weights = w / w.sum() self.tok = tokenizer or CortexTokenizer() self.shuffle_buf = shuffle_buf self.holdout_mod = holdout_mod # train skips the val partition (webdata.holdout_is_val) self._decontam = make_decontam(decontam_file) self.stats = [{"seen": 0, "kept": 0, "tok": 0} for _ in self.sources] try: import datasets as _hfds print(f"[stream_data] datasets {_hfds.__version__} | session_seed {session_seed}", flush=True) except Exception: pass rs = resume_states or {} n = len(self.sources) self._docs0 = [0] * n # docs consumed in PRIOR sessions (carried) self._seed_used = [0] * n # shuffle seed in effect (recorded in state()) self._restarts = [0] * n self._passes = [0] * n # passes completed per source (loop/exhaustion detector) for i, src in enumerate(self.sources): blob = rs.get(src.name) bdocs = int(blob.get("docs") or 0) if isinstance(blob, dict) and "hf" in blob else 0 self._docs0[i] = bdocs self._seed_used[i] = seed + 131 * i + 7919 * int(session_seed) print(f"[stream_data] {src.name}: fresh-shuffle (seed {self._seed_used[i]}) | " f"prior docs {bdocs}", flush=True) self._doc_qs = [queue.Queue(maxsize=doc_queue) for _ in self.sources] self._batch_q: queue.Queue = queue.Queue(maxsize=prefetch) self._stop = threading.Event() self._threads = [] for i, src in enumerate(self.sources): t = threading.Thread(target=self._reader, args=(i, src), name=f"reader-{src.name}", daemon=True) t.start() self._threads.append(t) pk = threading.Thread(target=self._packer, name="packer", daemon=True) pk.start() self._threads.append(pk) # --- dynamic mixture: cap a source once it has looped (exhausted its unique docs) --- EXHAUST_CAP = 0.07 # capped weight for a source on pass >=2 def _bump_pass(self, i): self._passes[i] += 1 if self._passes[i] == 2: # finished pass 1 -> now repeating print(f"[stream_data] {self.sources[i].name}: exhausted (looped) -> capping weight to " f"{self.EXHAUST_CAP:.0%}, upweighting fresh sources", flush=True) def _eff_weights(self): """Token weights with looped (exhausted) sources capped and the rest renormalized up.""" w = self.weights.copy() if any(p >= 2 for p in self._passes): for i in range(len(w)): if self._passes[i] >= 2: w[i] = min(w[i], self.EXHAUST_CAP) w = w / w.sum() return w # --- producers ----------------------------------------------------------- def _reader(self, i, src): n = 0 while not self._stop.is_set(): try: # each (re)start gets its own shard order: a transient-error restart must # never replay the prefix it already consumed this session. seed_i = self._seed_used[i] + 104_729 * self._restarts[i] self._restarts[i] += 1 for text in stream_texts(src, shuffle_buf=self.shuffle_buf, seed=seed_i, decontam=self._decontam, loop=True, stats=self.stats[i], holdout="train", holdout_mod=self.holdout_mod, on_ds=lambda ds, _i=i: self._bump_pass(_i)): if self._stop.is_set(): return ids = self.tok.encode(text, add_eot=True) while not self._stop.is_set(): try: self._doc_qs[i].put(ids, timeout=1.0) break except queue.Full: continue n += 1 except Exception as e: # transient HF/network error -> rebuild the stream if self._stop.is_set(): return print(f"[stream_data] reader {src.name}: {type(e).__name__}: {str(e)[:120]} -- restarting") def _packer(self): need = self._need tok_count = np.zeros(len(self.sources)) buf: list[int] = [] while not self._stop.is_set(): i = int(np.argmin(tok_count / self._eff_weights())) # most behind; looped sources are capped try: ids = self._doc_qs[i].get(timeout=1.0) except queue.Empty: continue buf.extend(ids) tok_count[i] += len(ids) self.stats[i]["tok"] += len(ids) while len(buf) >= need: chunk = np.asarray(buf[:need], dtype=np.int32).reshape(self._mshape) buf = buf[need:] while not self._stop.is_set(): try: self._batch_q.put(chunk, timeout=1.0) break except queue.Full: continue # --- consumer ------------------------------------------------------------ def __iter__(self): return self def __next__(self): while True: try: return self._batch_q.get(timeout=5.0) except queue.Empty: if self._stop.is_set(): raise StopIteration continue def ratios(self) -> dict: tot = sum(s["tok"] for s in self.stats) or 1 return {self.sources[i].name: self.stats[i]["tok"] / tot for i in range(len(self.sources))} def state(self) -> dict: """Per-source resume blob {seed, docs, hf:None} -> serialized in the checkpoint's resume.json. Only `docs` (always-advancing raw counter) carries across sessions; the next session derives a fresh shard order from its own session seed.""" return {self.sources[i].name: {"seed": self._seed_used[i], "docs": self._docs0[i] + self.stats[i]["seen"], "hf": None} for i in range(len(self.sources))} def position(self) -> dict: """Per-source progress: raw docs consumed (always grows -> loader-health signal).""" return {self.sources[i].name: f"docs={self._docs0[i] + self.stats[i]['seen']}" for i in range(len(self.sources))} def stop(self): self._stop.set()