Spaces:
Running
Running
| """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() | |