Add MP-01 (9500) 100M base: weights, config, tokenizer, model code, model card and eval results

README.md

@@ -0,0 +1,108 @@
+---
+license: other
+license_name: to-be-decided
+language:
+- en
+library_name: vbq
+pipeline_tag: text-generation
+tags:
+- base-model
+- pretraining
+- decoder-only
+- small-language-model
+---
+
+# MP-01 (9500) 100M base
+
+MP-01 (9500) 100M base is a 100M-parameter decoder-only language model from Manprot, grown from a 75M base and pretrained further on new English data.
+This is the **base model** (no chat or instruction tuning). It continues text; it does not follow instructions.
+
+- **Author:** Surelynotvain (organisation: Manprot)
+- **Checkpoint:** training step 9500 of 34341 (expansion run `vbq-100m-expand`)
+- **Parameters:** 100.6M (tied input/output embeddings)
+
+## Architecture
+
+| | |
+|---|---|
+| Layers | 18 |
+| Model width (d_model) | 640 |
+| Attention heads / KV heads | 10 / 5 (grouped-query attention, QK-norm) |
+| MLP hidden | 1664 (SwiGLU) |
+| Vocabulary | 32768 (byte-level BPE, `tokenizer.json`) |
+| Context length | 2048 tokens (RoPE) |
+| Normalisation | RMSNorm |
+
+This is a custom architecture, not a `transformers` model. Load it with the `vbq/` code in this repository.
+
+## Training
+
+- Started from a 75M-parameter base, grown to 18 layers, then continued on new English data.
+- Optimizer: Muon. Peak learning rate 1.2e-3, batch 262,144 tokens, weight decay 0.1, z-loss 1e-4.
+- Sequence length 1,024 tokens for the first part, 2,048 from step 24,038 (not reached yet).
+- Data mixture of this stage (share of tokens): FineWeb-Edu 34%, Wikipedia 24%, Cosmopedia 18%,
+  DCLM 12%, FineMath 5%, Stack Exchange 3%, Python code 3%.
+- Expansion steps completed when this checkpoint was saved: 9500 of 34341
+  (about 1.88B tokens in this run).
+
+## Evaluation
+
+Scored on this repository's own harness (`scripts/eval_checkpoint.py`), zero-shot.
+
+Benchmark tasks use 500 questions each (the same sample every time). With 500 questions per task, a
+difference under about 2 points is within noise.
+
+| Task | Score |
+|---|---:|
+| arc_challenge | 27.6 |
+| arc_easy | 52.8 |
+| hellaswag | 29.8 |
+| openbookqa | 27.6 |
+| piqa | 62.4 |
+| sciq | 71.0 |
+| winogrande | 51.6 |
+| **Average (7 tasks)** | **46.11** |
+| **BLiMP (67 paradigms, grammatical vs ungrammatical minimal pairs)** | **76.88** |
+
+Validation loss (held-out text, nats per token): **2.7780** overall.
+
+| Source | Val loss |
+|---|---:|
+| fineweb_edu | 2.941 |
+| wikipedia | 2.714 |
+| cosmopedia | 2.183 |
+| dclm | 3.387 |
+| stackexchange | 2.566 |
+| finemath | 3.402 |
+| code_python | 1.823 |
+
+## Limitations
+
+- Base model: it has no instruction or chat tuning and can produce factually wrong or repetitive text.
+- Small size (100M parameters) and small benchmarks; these numbers are indicative only.
+- Trained on web and educational text that may contain bias and errors.
+- The training data is not released with this model.
+
+## Loading
+
+```python
+import torch
+from tokenizers import Tokenizer
+from vbq.model import VBQ, VBQConfig
+from safetensors.torch import load_file
+import json
+
+cfg = VBQConfig(**json.load(open("config.json"))["model_cfg"])
+model = VBQ(cfg)
+model.load_state_dict(load_file("model.safetensors"))  # the tied output weight is stored once
+model.eval()
+tok = Tokenizer.from_file("tokenizer.json")
+```
+
+## Repository contents
+
+- `model.safetensors`: weights (bf16)
+- `config.json`: architecture and training configuration
+- `tokenizer.json`: tokenizer
+- `vbq/`: model code
+- `eval_results.json`: full evaluation output for this checkpoint

config.json

@@ -0,0 +1,24 @@
+{
+  "model_type": "vbq",
+  "architectures": [
+    "VBQ"
+  ],
+  "model_cfg": {
+    "vocab_size": 32768,
+    "d_model": 640,
+    "n_layers": 18,
+    "n_heads": 10,
+    "n_kv_heads": 5,
+    "mlp_hidden": 1664,
+    "max_seq_len": 2048,
+    "rope_theta": 10000.0,
+    "norm_eps": 1e-05,
+    "qk_norm": true,
+    "tie_embeddings": true,
+    "n_experts": 0,
+    "capacity_factor": 1.25
+  },
+  "train_step": 9500,
+  "total_steps": 34341,
+  "display_name": "MP-01 (9500) 100M base"
+}

eval_results.json

@@ -0,0 +1,99 @@
+{
+  "ckpt": "step_009500.pt",
+  "step": 9500,
+  "val": {
+    "fineweb_edu": 2.9408107101917267,
+    "wikipedia": 2.713771253824234,
+    "cosmopedia": 2.1831925809383392,
+    "dclm": 3.386944741010666,
+    "stackexchange": 2.5659963488578796,
+    "finemath": 3.4021530747413635,
+    "code_python": 1.8225234597921371,
+    "overall": 2.7780119873583318
+  },
+  "bench": {
+    "arc_challenge": 27.6,
+    "arc_easy": 52.8,
+    "hellaswag": 29.8,
+    "openbookqa": 27.6,
+    "piqa": 62.4,
+    "sciq": 71.0,
+    "winogrande": 51.6
+  },
+  "bench_avg": 46.114285714285714,
+  "blimp": {
+    "per_paradigm": {
+      "adjunct_island": 82.9,
+      "anaphor_gender_agreement": 92.3,
+      "anaphor_number_agreement": 97.3,
+      "animate_subject_passive": 75.9,
+      "animate_subject_trans": 83.8,
+      "causative": 75.4,
+      "complex_NP_island": 50.3,
+      "coordinate_structure_constraint_complex_left_branch": 53.5,
+      "coordinate_structure_constraint_object_extraction": 75.9,
+      "determiner_noun_agreement_1": 98.2,
+      "determiner_noun_agreement_2": 97.8,
+      "determiner_noun_agreement_irregular_1": 90.6,
+      "determiner_noun_agreement_irregular_2": 95.9,
+      "determiner_noun_agreement_with_adj_2": 93.3,
+      "determiner_noun_agreement_with_adj_irregular_1": 90.1,
+      "determiner_noun_agreement_with_adj_irregular_2": 89.2,
+      "determiner_noun_agreement_with_adjective_1": 93.8,
+      "distractor_agreement_relational_noun": 84.0,
+      "distractor_agreement_relative_clause": 77.5,
+      "drop_argument": 76.5,
+      "ellipsis_n_bar_1": 83.8,
+      "ellipsis_n_bar_2": 87.6,
+      "existential_there_object_raising": 79.5,
+      "existential_there_quantifiers_1": 97.3,
+      "existential_there_quantifiers_2": 23.3,
+      "existential_there_subject_raising": 85.5,
+      "expletive_it_object_raising": 78.4,
+      "inchoative": 63.4,
+      "intransitive": 73.4,
+      "irregular_past_participle_adjectives": 90.8,
+      "irregular_past_participle_verbs": 92.5,
+      "irregular_plural_subject_verb_agreement_1": 88.8,
+      "irregular_plural_subject_verb_agreement_2": 91.9,
+      "left_branch_island_echo_question": 45.0,
+      "left_branch_island_simple_question": 54.8,
+      "matrix_question_npi_licensor_present": 44.7,
+      "npi_present_1": 55.7,
+      "npi_present_2": 60.9,
+      "only_npi_licensor_present": 62.0,
+      "only_npi_scope": 47.4,
+      "passive_1": 88.4,
+      "passive_2": 86.0,
+      "principle_A_c_command": 63.5,
+      "principle_A_case_1": 100.0,
+      "principle_A_case_2": 95.5,
+      "principle_A_domain_1": 96.2,
+      "principle_A_domain_2": 74.7,
+      "principle_A_domain_3": 57.6,
+      "principle_A_reconstruction": 36.3,
+      "regular_plural_subject_verb_agreement_1": 92.4,
+      "regular_plural_subject_verb_agreement_2": 91.7,
+      "sentential_negation_npi_licensor_present": 96.8,
+      "sentential_negation_npi_scope": 68.9,
+      "sentential_subject_island": 30.2,
+      "superlative_quantifiers_1": 91.8,
+      "superlative_quantifiers_2": 80.3,
+      "tough_vs_raising_1": 64.3,
+      "tough_vs_raising_2": 86.2,
+      "transitive": 86.9,
+      "wh_island": 70.5,
+      "wh_questions_object_gap": 75.6,
+      "wh_questions_subject_gap": 90.7,
+      "wh_questions_subject_gap_long_distance": 81.2,
+      "wh_vs_that_no_gap": 97.1,
+      "wh_vs_that_no_gap_long_distance": 96.4,
+      "wh_vs_that_with_gap": 46.1,
+      "wh_vs_that_with_gap_long_distance": 24.8
+    },
+    "paradigm_avg": 76.88059701492537,
+    "overall": 76.88059701492537,
+    "pairs": 67000
+  },
+  "seconds": 124
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:621d2129dd34133fa4fcb623db7ff05d94755e0271eba4473e7625cdf8e84030
+size 201265144

vbq/__init__.py

@@ -0,0 +1,7 @@
+"""VBQ — a 75M-parameter decoder-only language model by Voussoir AI."""
+import os
+
+# Flash / mem-efficient SDPA kernels on RDNA3 are gated behind this flag in PyTorch's ROCm build.
+# They are ~10x faster than the math fallback here, so enable them before torch is imported.
+os.environ.setdefault("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "1")
+os.environ.setdefault("PYTORCH_HIP_ALLOC_CONF", "expandable_segments:True")

vbq/chat.py

@@ -0,0 +1,65 @@
+"""Shared ChatML tokens, assistant supervision and context budgeting."""
+import os
+from functools import lru_cache
+from pathlib import Path
+from tokenizers import Tokenizer
+
+ROOT = Path(__file__).resolve().parent.parent
+SPECIALS = ('<|endoftext|>', '<|pad|>', '<|im_start|>', '<|im_end|>', '<|system|>', '<|user|>', '<|assistant|>', '<|reserved|>')
+STOP_IDS = {0, 3}
+
+
+@lru_cache(maxsize=1)
+def tokenizer():
+    return Tokenizer.from_file(str(ROOT / os.environ.get('VBQ_TOKENIZER', 'data/tokenizer/vbq-tokenizer.json')))
+
+
+def render(messages, add_generation_prompt=False):
+    ids, mask = [], []
+    tok = tokenizer()
+    for m in messages:
+        role, content = m['role'], m['content']
+        if role not in ('system', 'user', 'assistant') or not isinstance(content, str):
+            raise ValueError('messages must contain system/user/assistant roles and string content')
+        if any(s in content for s in SPECIALS):
+            raise ValueError('message content cannot contain reserved tokens')
+        header = tok.encode(f'<|im_start|>{role}\n', add_special_tokens=False).ids
+        body = tok.encode(content, add_special_tokens=False).ids
+        tail = tok.encode('\n', add_special_tokens=False).ids
+        ids.extend(header + body + [3] + tail)
+        mask.extend([0] * len(header) + [int(role == 'assistant')] * (len(body) + 1) + [0] * len(tail))
+    if add_generation_prompt:
+        header = tok.encode('<|im_start|>assistant\n', add_special_tokens=False).ids
+        ids.extend(header); mask.extend([0] * len(header))
+    return ids, mask
+
+
+def fit_history(messages, budget):
+    history = [dict(m) for m in messages]
+    if not history or history[-1]['role'] != 'user':
+        raise ValueError('chat must end with a user message')
+    if any(m['role'] == 'system' for m in history[1:]):
+        raise ValueError('system message must be first')
+    start = int(history[0]['role'] == 'system')
+    while True:
+        ids, _ = render(history, add_generation_prompt=True)
+        if len(ids) <= budget: return history, ids
+        if len(history) <= start + 1:
+            raise ValueError('system prompt and latest message exceed the context budget')
+        end = start + 1
+        while end < len(history) - 1 and history[end]['role'] != 'user': end += 1
+        del history[start:end]
+
+
+def stream_reply(model, messages, max_new=256, temperature=0.7, top_p=0.9, top_k=50, rep_penalty=1.1):
+    import torch
+    _, ids = fit_history(messages, model.c.max_seq_len - max_new)
+    out, sent = [], ''
+    for t in model.generate_cached(torch.tensor([ids], device=next(model.parameters()).device), max_new_tokens=max_new,
+                                   temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=rep_penalty, eos_id=STOP_IDS):
+        if t in STOP_IDS: break
+        out.append(t); text = tokenizer().decode(out)
+        if text.endswith('�'): continue
+        yield text[len(sent):]; sent = text
+    text = tokenizer().decode(out)
+    if len(text) > len(sent): yield text[len(sent):]

vbq/data.py

@@ -0,0 +1,92 @@
+"""Deterministic, resumable batch loader over the packed token corpus.
+
+The corpus is one flat uint16 array in which every source already holds exactly its mixture share.
+It is cut into 2048-token blocks and visited once in a fixed random order, so a single pass
+reproduces the intended mix with no repetition. Step N always sees the same blocks (resume-safe).
+"""
+from __future__ import annotations
+
+import json
+import queue
+import threading
+from pathlib import Path
+
+import numpy as np
+import torch
+
+BLOCK = 2048
+
+
+class BlockLoader:
+    def __init__(self, data_dir: Path, seed: int = 1337, corpus_json: str = "corpus.json", train_file: str = "train.bin"):
+        self.dir = Path(data_dir)
+        self.meta = json.loads((self.dir / corpus_json).read_text())
+        self.mm = np.memmap(self.dir / train_file, dtype=np.uint16, mode="r")
+        self.n_blocks = (len(self.mm) - 1) // BLOCK
+        self.perm = np.random.default_rng(seed).permutation(self.n_blocks).astype(np.int64)
+
+    def batch(self, step: int, blocks_per_step: int, seq_len: int) -> torch.Tensor:
+        """Returns int64 [blocks_per_step * (BLOCK // seq_len), seq_len + 1]: inputs and shifted targets in one array."""
+        ids = self.perm[(np.arange(blocks_per_step) + step * blocks_per_step) % self.n_blocks]
+        out = np.empty((blocks_per_step, BLOCK // seq_len, seq_len + 1), dtype=np.int64)
+        for i, b in enumerate(ids):
+            w = self.mm[b * BLOCK : b * BLOCK + BLOCK + 1].astype(np.int64)
+            for j in range(BLOCK // seq_len):
+                out[i, j] = w[j * seq_len : j * seq_len + seq_len + 1]
+        return torch.from_numpy(out.reshape(-1, seq_len + 1))
+
+
+class StagedLoader:
+    """first loader for steps < from_step, then a second corpus from its own beginning (opened when first needed)."""
+
+    def __init__(self, first: BlockLoader, from_step: int, second_dir: Path, seed: int):
+        self.first, self.from_step, self.second_dir, self.seed = first, from_step, Path(second_dir), seed
+        self.second, self.n_blocks, self.meta = None, first.n_blocks, first.meta
+
+    def ready(self) -> bool:
+        return (self.second_dir / "corpus.json").exists()
+
+    def batch(self, step: int, blocks_per_step: int, seq_len: int) -> torch.Tensor:
+        if step < self.from_step: return self.first.batch(step, blocks_per_step, seq_len)
+        if self.second is None: self.second = BlockLoader(self.second_dir, seed=self.seed)
+        return self.second.batch(step - self.from_step, blocks_per_step, seq_len)
+
+
+class Prefetcher:
+    """Background thread that prepares upcoming steps' batches."""
+
+    def __init__(self, loader: BlockLoader, start_step: int, blocks_per_step: int, seq_len_fn, depth: int = 3):
+        self.q: queue.Queue = queue.Queue(maxsize=depth)
+        self.loader, self.bps, self.seq_len_fn = loader, blocks_per_step, seq_len_fn
+        self.step = start_step
+        self.stop = False
+        self.t = threading.Thread(target=self._run, daemon=True)
+        self.t.start()
+
+    def _run(self):
+        s = self.step
+        while not self.stop:
+            T = self.seq_len_fn(s)
+            b = self.loader.batch(s, self.bps, T).pin_memory()
+            self.q.put((s, T, b))
+            s += 1
+
+    def get(self):
+        return self.q.get()
+
+
+class ValSet:
+    """Fixed held-out token windows, overall and per source."""
+
+    def __init__(self, data_dir: Path, seq_len: int = 2048, per_source_windows: int = 64):
+        meta = json.loads((Path(data_dir) / "corpus.json").read_text())
+        mm = np.memmap(Path(data_dir) / "val.bin", dtype=np.uint16, mode="r")
+        self.sets: dict[str, torch.Tensor] = {}
+        for name, s in meta["sources"].items():
+            n = (s["val_tokens"] - 1) // seq_len
+            if n <= 0:
+                continue
+            take = min(n, per_source_windows)
+            starts = np.linspace(0, n - 1, take).astype(int)
+            arr = np.stack([mm[s["val_start"] + i * seq_len : s["val_start"] + i * seq_len + seq_len + 1].astype(np.int64) for i in starts])
+            self.sets[name] = torch.from_numpy(arr)

vbq/evals.py

@@ -0,0 +1,88 @@
+"""Zero-shot multiple-choice evaluation (lm-evaluation-harness conventions) that fits next to a training job."""
+from __future__ import annotations
+
+import json
+import random
+import re
+from pathlib import Path
+
+import torch
+
+EVAL_DIR = Path(__file__).resolve().parent.parent / "data" / "evals"
+# chance level, for the dashboard
+CHANCE = {"hellaswag": 25.0, "arc_easy": 25.0, "arc_challenge": 25.0, "piqa": 50.0, "sciq": 25.0, "openbookqa": 25.0, "winogrande": 50.0}
+METRIC = {"hellaswag": "acc_norm", "arc_easy": "acc", "arc_challenge": "acc_norm", "piqa": "acc_norm", "sciq": "acc_norm", "openbookqa": "acc_norm", "winogrande": "acc"}
+
+
+def load_tasks(limit: int | None = None, seed: int = 0) -> dict[str, list[dict]]:
+    tasks = {}
+    for f in sorted(EVAL_DIR.glob("*.json")):
+        items = json.loads(f.read_text())
+        if limit and len(items) > limit:
+            items = random.Random(seed).sample(items, limit)
+        tasks[f.stem] = items
+    return tasks
+
+
+@torch.no_grad()
+def loglik_batch(model, pairs, device, pad_id=1, max_ctx=480, bs=48):
+    """pairs: list of (ctx_ids, cont_ids). Returns (sum logprob, n_tokens) per pair."""
+    res = [None] * len(pairs)
+    order = sorted(range(len(pairs)), key=lambda i: len(pairs[i][0]) + len(pairs[i][1]))
+    w = model.out_weight
+    for s in range(0, len(order), bs):
+        idxs = order[s : s + bs]
+        seqs = []
+        for i in idxs:
+            c, k = pairs[i]
+            c = c[-max_ctx:]
+            seqs.append((c + k, len(c), len(k)))
+        L = max(len(x[0]) for x in seqs)
+        inp = torch.full((len(seqs), L), pad_id, dtype=torch.long)
+        for r, (ids, _, _) in enumerate(seqs):
+            inp[r, : len(ids)] = torch.tensor(ids)
+        inp = inp.to(device)
+        with torch.autocast(device.type, torch.bfloat16):
+            h = model.hidden(inp)
+        rows, cols, tgts = [], [], []
+        for r, (ids, nc, nk) in enumerate(seqs):
+            for j in range(nk):
+                rows.append(r); cols.append(nc + j - 1); tgts.append(ids[nc + j])
+        hs = h[torch.tensor(rows, device=device), torch.tensor(cols, device=device)]
+        logits = (hs @ w.to(hs.dtype).t()).float()
+        lp = torch.log_softmax(logits, -1).gather(1, torch.tensor(tgts, device=device)[:, None]).squeeze(1).cpu()
+        p = 0
+        for r, i in enumerate(idxs):
+            nk = seqs[r][2]
+            res[i] = (float(lp[p : p + nk].sum()), nk); p += nk
+    return res
+
+
+def evaluate(model, tok, tasks: dict[str, list[dict]], device, pad_id=1, bs=48) -> dict[str, dict]:
+    model.eval()
+    out = {}
+    for name, items in tasks.items():
+        pairs, spans, labels, chs = [], [], [], []
+        for it in items:
+            start = len(pairs)
+            if "alt_ctx" in it:  # WinoGrande partial scoring: same continuation, two candidate contexts
+                tail = tok.encode(it["choices"][0]).ids
+                pairs += [(tok.encode(it["ctx"]).ids, tail), (tok.encode(it["alt_ctx"]).ids, tail)]
+                labels.append(it["wino_label"]); chs.append([it["choices"][0]] * 2)
+            else:
+                cids = tok.encode(it["ctx"]).ids
+                for ch in it["choices"]:
+                    pairs.append((cids, tok.encode(ch).ids or [0]))
+                labels.append(it["label"]); chs.append(it["choices"])
+            spans.append((start, len(pairs) - start))
+        ll = loglik_batch(model, pairs, device, pad_id=pad_id, bs=bs)
+        acc = accn = 0
+        for (st, n), lab, ch in zip(spans, labels, chs):
+            raw = [ll[st + k][0] for k in range(n)]
+            norm = [raw[k] / max(1, len(ch[k].encode())) for k in range(n)]
+            acc += int(max(range(n), key=raw.__getitem__) == lab)
+            accn += int(max(range(n), key=norm.__getitem__) == lab)
+        out[name] = {"acc": 100 * acc / len(items), "acc_norm": 100 * accn / len(items), "n": len(items)}
+        out[name]["score"] = out[name][METRIC.get(name, "acc")]
+    model.train()
+    return out

vbq/fastce.py

@@ -0,0 +1,117 @@
+"""Fused cross-entropy: same maths as FusedLinearCE, but the per-chunk softmax/grad chain is one compiled kernel set.
+
+On a 288 GB/s card the eager version streams ~5 GB of fp32 logits-sized tensors per chunk; here the bf16 logits are
+read twice and the bf16 gradient written once.
+"""
+import torch
+
+
+def _chunk(logits, tgt, z_loss: float, inv_n: float):
+    lg = logits.float()
+    lse = torch.logsumexp(lg, dim=-1)
+    tl = lg.gather(1, tgt[:, None]).squeeze(1)
+    loss = (lse - tl).sum()
+    p = torch.exp(lg - lse[:, None])
+    if z_loss:
+        loss = loss + z_loss * (lse * lse).sum()
+        p = p * (1.0 + 2.0 * z_loss * lse[:, None])
+    cols = torch.arange(lg.shape[1], device=lg.device)[None, :]
+    p = torch.where(cols == tgt[:, None], p - 1.0, p)
+    return loss, (p * inv_n).to(torch.bfloat16)
+
+
+chunk_compiled = torch.compile(_chunk, dynamic=False)
+
+
+class FastLinearCE(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, h, w, targets, chunk: int, z_loss: float):
+        N = h.shape[0]
+        wb = w.to(h.dtype)
+        grad_h = torch.empty_like(h)
+        grad_w = torch.zeros_like(w, dtype=torch.float32)
+        total = torch.zeros((), device=h.device, dtype=torch.float32)
+        for i in range(0, N, chunk):
+            hc, tc = h[i : i + chunk], targets[i : i + chunk]
+            loss_i, pb = chunk_compiled(hc @ wb.t(), tc, z_loss, 1.0 / N)
+            total += loss_i
+            grad_h[i : i + chunk] = pb @ wb
+            grad_w.add_((pb.t() @ hc).float())
+        ctx.save_for_backward(grad_h, grad_w)
+        return total / N
+
+    @staticmethod
+    def backward(ctx, g):
+        grad_h, grad_w = ctx.saved_tensors
+        return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None
+
+
+def _chunk_multi(logits, tgt, z_loss: float, inv_n: float):
+    """Patch-level loss: one distribution per position scored against all k tokens of the next patch.
+    tgt [n, k]; loss = sum_k CE(logits, tgt[:, k]); grad = k*softmax - sum_k onehot."""
+    lg = logits.float()
+    k = tgt.shape[1]
+    lse = torch.logsumexp(lg, dim=-1)
+    loss = (k * lse - lg.gather(1, tgt).sum(1)).sum()
+    p = k * torch.exp(lg - lse[:, None])
+    if z_loss:
+        loss = loss + k * z_loss * (lse * lse).sum()
+        p = p * (1.0 + 2.0 * z_loss * lse[:, None])
+    cols = torch.arange(lg.shape[1], device=lg.device)[None, :]
+    for j in range(k):   # unrolled at trace time and fused: no logits-sized one-hot buffer
+        p = p - (cols == tgt[:, j : j + 1]).to(p.dtype)
+    return loss, (p * inv_n).to(torch.bfloat16)
+
+
+chunk_multi = torch.compile(_chunk_multi, dynamic=False)
+
+
+class PatchLinearCE(torch.autograd.Function):
+    """FastLinearCE for patch-level training: h [N, d], targets [N, k]; returns the mean per-token CE."""
+    @staticmethod
+    def forward(ctx, h, w, targets, chunk: int, z_loss: float):
+        N, k = targets.shape
+        wb = w.to(h.dtype)
+        grad_h = torch.empty_like(h)
+        grad_w = torch.zeros_like(w, dtype=torch.float32)
+        total = torch.zeros((), device=h.device, dtype=torch.float32)
+        for i in range(0, N, chunk):
+            hc, tc = h[i : i + chunk], targets[i : i + chunk]
+            loss_i, pb = chunk_multi(hc @ wb.t(), tc, z_loss, 1.0 / (N * k))
+            total += loss_i
+            grad_h[i : i + chunk] = pb @ wb
+            grad_w.add_((pb.t() @ hc).float())
+        ctx.save_for_backward(grad_h, grad_w)
+        return total / (N * k)
+
+    @staticmethod
+    def backward(ctx, g):
+        grad_h, grad_w = ctx.saved_tensors
+        return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None
+
+
+# Variable assistant counts and CE tails must not specialize each compiled call.
+chunk_sft = torch.compile(_chunk, dynamic=True)
+
+
+class SFTLinearCE(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, h, w, targets, chunk: int, z_loss: float):
+        N = h.shape[0]
+        wb = w.to(h.dtype)
+        grad_h = torch.empty_like(h)
+        grad_w = torch.zeros_like(w, dtype=torch.float32)
+        total = torch.zeros((), device=h.device, dtype=torch.float32)
+        for i in range(0, N, chunk):
+            hc, tc = h[i : i + chunk], targets[i : i + chunk]
+            loss_i, pb = chunk_sft(hc @ wb.t(), tc, z_loss, torch.tensor(1.0 / N, device=h.device))
+            total += loss_i
+            grad_h[i : i + chunk] = pb @ wb
+            grad_w.add_((pb.t() @ hc).float())
+        ctx.save_for_backward(grad_h, grad_w)
+        return total / N
+
+    @staticmethod
+    def backward(ctx, g):
+        grad_h, grad_w = ctx.saved_tensors
+        return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None

vbq/model.py

@@ -0,0 +1,357 @@
+"""VBQ transformer: Llama-style decoder (RMSNorm, RoPE, GQA, QK-norm, SwiGLU, tied embeddings)."""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, asdict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@dataclass
+class VBQConfig:
+    vocab_size: int = 32768
+    d_model: int = 640
+    n_layers: int = 12
+    n_heads: int = 10
+    n_kv_heads: int = 5
+    mlp_hidden: int = 1664
+    max_seq_len: int = 2048
+    rope_theta: float = 10000.0
+    norm_eps: float = 1e-5
+    qk_norm: bool = True
+    tie_embeddings: bool = True
+    # Mixture of experts (0 = dense SwiGLU). Top-1 routing, Switch-style.
+    n_experts: int = 0
+    capacity_factor: float = 1.25
+
+    @property
+    def head_dim(self) -> int:
+        return self.d_model // self.n_heads
+
+    def n_params(self) -> int:
+        d, hd = self.d_model, self.head_dim
+        attn = d * d * 2 + d * (self.n_kv_heads * hd) * 2
+        mlp = 3 * d * self.mlp_hidden * max(1, self.n_experts) + (d * self.n_experts if self.n_experts else 0)
+        per_layer = attn + mlp + 2 * d + (2 * hd if self.qk_norm else 0)
+        emb = self.vocab_size * d * (1 if self.tie_embeddings else 2)
+        return emb + self.n_layers * per_layer + d
+
+    def to_dict(self):
+        return asdict(self)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return F.rms_norm(x, (x.shape[-1],), self.weight.to(x.dtype), self.eps)
+
+
+def build_rope(seq_len: int, head_dim: int, theta: float, device=None):
+    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device, dtype=torch.float32) / head_dim))
+    t = torch.arange(seq_len, device=device, dtype=torch.float32)
+    freqs = torch.outer(t, inv)
+    return freqs.cos(), freqs.sin()  # [T, hd/2]
+
+
+def apply_rope(x, cos, sin):
+    # x: [B, H, T, hd]; half-split rotation
+    x1, x2 = x.chunk(2, dim=-1)
+    cos = cos[None, None, : x.shape[2]].to(x.dtype)
+    sin = sin[None, None, : x.shape[2]].to(x.dtype)
+    return torch.cat((x1 * cos - x2 * sin, x2 * cos + x1 * sin), dim=-1)
+
+
+class Attention(nn.Module):
+    def __init__(self, c: VBQConfig):
+        super().__init__()
+        self.c = c
+        hd = c.head_dim
+        self.wq = nn.Linear(c.d_model, c.n_heads * hd, bias=False)
+        self.wk = nn.Linear(c.d_model, c.n_kv_heads * hd, bias=False)
+        self.wv = nn.Linear(c.d_model, c.n_kv_heads * hd, bias=False)
+        self.wo = nn.Linear(c.n_heads * hd, c.d_model, bias=False)
+        if c.qk_norm:
+            self.q_norm = RMSNorm(hd, c.norm_eps)
+            self.k_norm = RMSNorm(hd, c.norm_eps)
+
+    def forward(self, x, cos, sin):
+        B, T, _ = x.shape
+        c, hd = self.c, self.c.head_dim
+        q = self.wq(x).view(B, T, c.n_heads, hd).transpose(1, 2)
+        k = self.wk(x).view(B, T, c.n_kv_heads, hd).transpose(1, 2)
+        v = self.wv(x).view(B, T, c.n_kv_heads, hd).transpose(1, 2)
+        if c.qk_norm:
+            q, k = self.q_norm(q), self.k_norm(k)
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        if c.n_kv_heads != c.n_heads:
+            rep = c.n_heads // c.n_kv_heads
+            k = k[:, :, None].expand(B, c.n_kv_heads, rep, T, hd).reshape(B, c.n_heads, T, hd)
+            v = v[:, :, None].expand(B, c.n_kv_heads, rep, T, hd).reshape(B, c.n_heads, T, hd)
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.wo(y.transpose(1, 2).reshape(B, T, c.n_heads * hd))
+
+
+class MLP(nn.Module):
+    def __init__(self, c: VBQConfig):
+        super().__init__()
+        self.gate_up = nn.Linear(c.d_model, 2 * c.mlp_hidden, bias=False)
+        self.down = nn.Linear(c.mlp_hidden, c.d_model, bias=False)
+
+    def forward(self, x):
+        g, u = self.gate_up(x).chunk(2, dim=-1)
+        return self.down(F.silu(g) * u)
+
+
+class MoE(nn.Module):
+    """Top-1 mixture of SwiGLU experts (Switch Transformer routing).
+
+    Fixed shapes (compile-friendly): every expert gets a buffer of
+    capacity = ceil(capacity_factor * tokens / experts) slots while training
+    (tokens past it skip the MLP and keep only the residual); at eval every
+    token fits. The chosen expert's output is scaled by p / p.detach(), which
+    is 1 in value, so a model upcycled from identical copies computes exactly
+    what the dense model did, while the router still gets gradients.
+    Returns (y, aux) with aux = Switch load-balancing loss (1.0 = balanced).
+    """
+
+    def __init__(self, c: VBQConfig):
+        super().__init__()
+        self.c = c
+        E, d, h = c.n_experts, c.d_model, c.mlp_hidden
+        self.router = nn.Linear(d, E, bias=False)
+        self.gate_up = nn.Parameter(torch.empty(E, 2 * h, d))
+        self.down = nn.Parameter(torch.empty(E, d, h))
+        nn.init.normal_(self.gate_up, 0.0, 0.02)
+        nn.init.normal_(self.down, 0.0, 0.02)
+
+    def forward(self, x):
+        shape = x.shape
+        x = x.reshape(-1, shape[-1])
+        N, d = x.shape
+        E = self.c.n_experts
+        probs = F.softmax(self.router(x).float(), dim=-1)          # [N, E]
+        gate, idx = probs.max(-1)                                  # [N]
+        onehot = F.one_hot(idx, E)                                 # [N, E]
+        pos = ((onehot.cumsum(0) - 1) * onehot).sum(-1)            # slot within its expert
+        C = math.ceil(self.c.capacity_factor * N / E) if self.training else N
+        keep = pos < C
+        slot = torch.where(keep, idx * C + pos, torch.full_like(pos, E * C))  # overflow -> dump slot
+        buf = x.new_zeros(E * C + 1, d).index_put((slot,), x)
+        hbuf = buf[: E * C].view(E, C, d)
+        g, u = torch.bmm(hbuf, self.gate_up.to(x.dtype).transpose(1, 2)).chunk(2, dim=-1)
+        y = torch.bmm(F.silu(g) * u, self.down.to(x.dtype).transpose(1, 2))   # [E, C, d]
+        y = torch.cat([y.reshape(E * C, d), y.new_zeros(1, d)])[slot]
+        scale = (gate / gate.detach()) * keep
+        y = y * scale.to(y.dtype).unsqueeze(-1)
+        aux = E * (onehot.float().mean(0) * probs.mean(0)).sum()
+        return y.reshape(shape), aux
+
+
+class Block(nn.Module):
+    def __init__(self, c: VBQConfig):
+        super().__init__()
+        self.attn_norm = RMSNorm(c.d_model, c.norm_eps)
+        self.attn = Attention(c)
+        self.mlp_norm = RMSNorm(c.d_model, c.norm_eps)
+        self.mlp = MoE(c) if c.n_experts else MLP(c)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.attn_norm(x), cos, sin)
+        return x + self.ffn(self.mlp_norm(x))[0]
+
+    def ffn(self, x):
+        """(output, load-balancing loss) for dense and MoE blocks alike."""
+        y = self.mlp(x)
+        return y if isinstance(y, tuple) else (y, x.new_zeros((), dtype=torch.float32))
+
+
+class FusedLinearCE(torch.autograd.Function):
+    """Cross-entropy over a huge vocab without ever materialising full logits.
+
+    Gradients w.r.t. hidden states and the (tied) output matrix are computed chunk by chunk
+    during the forward pass, so the backward pass is a cheap scale. Saves ~4 GB at 32k vocab.
+    """
+
+    @staticmethod
+    def forward(ctx, h, w, targets, chunk: int, z_loss: float):
+        N = h.shape[0]
+        wb = w.to(h.dtype)
+        grad_h = torch.empty_like(h)
+        grad_w = torch.zeros_like(w, dtype=torch.float32)
+        total = torch.zeros((), device=h.device, dtype=torch.float32)
+        for i in range(0, N, chunk):
+            hc, tc = h[i : i + chunk], targets[i : i + chunk]
+            logits = (hc @ wb.t()).float()
+            lse = torch.logsumexp(logits, dim=-1)
+            tgt = logits.gather(1, tc[:, None]).squeeze(1)
+            total += (lse - tgt).sum()
+            p = torch.exp(logits - lse[:, None])
+            if z_loss:
+                total += z_loss * (lse * lse).sum()
+                p.mul_(1.0 + 2.0 * z_loss * lse[:, None])
+            p.scatter_add_(1, tc[:, None], torch.full_like(tc[:, None], -1, dtype=p.dtype))
+            p.mul_(1.0 / N)
+            pb = p.to(h.dtype)
+            grad_h[i : i + chunk] = pb @ wb
+            grad_w += (pb.t() @ hc).float()
+        ctx.save_for_backward(grad_h, grad_w)
+        return total / N
+
+    @staticmethod
+    def backward(ctx, g):
+        grad_h, grad_w = ctx.saved_tensors
+        return grad_h * g.to(grad_h.dtype), grad_w * g, None, None, None
+
+
+class VBQ(nn.Module):
+    def __init__(self, c: VBQConfig):
+        super().__init__()
+        self.c = c
+        self.tok_emb = nn.Embedding(c.vocab_size, c.d_model)
+        self.blocks = nn.ModuleList(Block(c) for _ in range(c.n_layers))
+        self.final_norm = RMSNorm(c.d_model, c.norm_eps)
+        self.lm_head = None if c.tie_embeddings else nn.Linear(c.d_model, c.vocab_size, bias=False)
+        cos, sin = build_rope(c.max_seq_len, c.head_dim, c.rope_theta)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+        self.apply(self._init)
+        for n, p in self.named_parameters():  # scaled init on residual output projections
+            if n.endswith("wo.weight") or n.endswith("down.weight"):
+                nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * c.n_layers))
+
+    @staticmethod
+    def _init(m):
+        if isinstance(m, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(m.weight, 0.0, 0.02)
+
+    @property
+    def out_weight(self):
+        return self.tok_emb.weight if self.lm_head is None else self.lm_head.weight
+
+    def hidden(self, idx):
+        return self.hidden_aux(idx)[0]
+
+    def hidden_aux(self, idx):
+        """Final hidden states and the mean MoE load-balancing loss (0 for dense)."""
+        return self._blocks(self.tok_emb(idx))
+
+    def hidden_patch(self, idx, k: int):
+        """Patch-level training (Shao et al. 2024): idx [B, P*k]; each group of k tokens enters as the mean of its
+        token embeddings, so the transformer runs on P positions. Same weights as token-level training."""
+        B, L = idx.shape
+        return self._blocks(self.tok_emb(idx).view(B, L // k, k, -1).mean(2))
+
+    def _blocks(self, x):
+        aux = x.new_zeros((), dtype=torch.float32)
+        for blk in self.blocks:
+            x = x + blk.attn(blk.attn_norm(x), self.rope_cos, self.rope_sin)
+            y, a = blk.ffn(blk.mlp_norm(x))
+            x = x + y
+            aux = aux + a
+        return self.final_norm(x), aux / len(self.blocks)
+
+    def forward(self, idx, targets=None, z_loss: float = 0.0, ce_chunk: int = 4096):
+        h = self.hidden(idx)
+        if targets is None:
+            return (h @ self.out_weight.to(h.dtype).t()).float()
+        loss = FusedLinearCE.apply(h.reshape(-1, h.shape[-1]), self.out_weight, targets.reshape(-1), ce_chunk, z_loss)
+        return loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens=128, temperature=0.8, top_k=50, top_p=0.95, eos_id=None, repetition_penalty=1.0):
+        self.eval()
+        for _ in range(max_new_tokens):
+            ctx = idx[:, -self.c.max_seq_len :]
+            logits = self(ctx)[:, -1]
+            if repetition_penalty != 1.0:
+                for tok in set(ctx[0].tolist()):
+                    logits[0, tok] = logits[0, tok] / repetition_penalty if logits[0, tok] > 0 else logits[0, tok] * repetition_penalty
+            if temperature <= 0:
+                nxt = logits.argmax(-1, keepdim=True)
+            else:
+                logits = logits / temperature
+                if top_k:
+                    v, _ = torch.topk(logits, min(top_k, logits.shape[-1]))
+                    logits[logits < v[:, [-1]]] = -float("inf")
+                probs = F.softmax(logits, dim=-1)
+                if top_p < 1.0:
+                    sp, si = torch.sort(probs, descending=True)
+                    cum = sp.cumsum(-1)
+                    sp[cum - sp > top_p] = 0
+                    probs = torch.zeros_like(probs).scatter_(1, si, sp)
+                    probs = probs / probs.sum(-1, keepdim=True)
+                nxt = torch.multinomial(probs, 1)
+            idx = torch.cat([idx, nxt], dim=1)
+            yield int(nxt.item())
+            if eos_id is not None and int(nxt.item()) == eos_id:
+                break
+
+    @torch.no_grad()
+    def generate_cached(self, idx, max_new_tokens=128, temperature=0.8, top_k=50, top_p=0.95, eos_id=None, repetition_penalty=1.0):
+        """KV-cached sampling. Yields token ids one by one. idx: [1, T] prompt."""
+        self.eval()
+        c = self.c
+        idx = idx[:, -(c.max_seq_len - 1):]
+        stop_ids = set() if eos_id is None else ({eos_id} if isinstance(eos_id, int) else set(eos_id))
+        caches: list = [None] * c.n_layers
+        seen = set(idx[0].tolist())
+
+        def forward(ids, pos):
+            T = ids.shape[1]
+            x = self.tok_emb(ids)
+            cos, sin = self.rope_cos[pos : pos + T], self.rope_sin[pos : pos + T]
+            for i, blk in enumerate(self.blocks):
+                a = blk.attn
+                h = blk.attn_norm(x)
+                q = a.wq(h).view(1, T, c.n_heads, c.head_dim).transpose(1, 2)
+                k = a.wk(h).view(1, T, c.n_kv_heads, c.head_dim).transpose(1, 2)
+                v = a.wv(h).view(1, T, c.n_kv_heads, c.head_dim).transpose(1, 2)
+                if c.qk_norm:
+                    q, k = a.q_norm(q), a.k_norm(k)
+                q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+                if caches[i] is not None:
+                    k, v = torch.cat([caches[i][0], k], 2), torch.cat([caches[i][1], v], 2)
+                caches[i] = (k, v)
+                rep = c.n_heads // c.n_kv_heads
+                kk = k.repeat_interleave(rep, dim=1) if rep > 1 else k
+                vv = v.repeat_interleave(rep, dim=1) if rep > 1 else v
+                y = F.scaled_dot_product_attention(q, kk, vv, is_causal=(T > 1))
+                x = x + a.wo(y.transpose(1, 2).reshape(1, T, c.n_heads * c.head_dim))
+                x = x + blk.ffn(blk.mlp_norm(x))[0]
+            return (self.final_norm(x[:, -1]) @ self.out_weight.to(x.dtype).t()).float()
+
+        logits = forward(idx, 0)
+        pos = idx.shape[1]
+        for _ in range(max_new_tokens):
+            if pos >= c.max_seq_len:
+                break
+            lg = logits[0].clone()
+            if repetition_penalty != 1.0:
+                for tok in seen:
+                    lg[tok] = lg[tok] / repetition_penalty if lg[tok] > 0 else lg[tok] * repetition_penalty
+            if temperature <= 0:
+                nxt = int(lg.argmax())
+            else:
+                lg = lg / temperature
+                if top_k:
+                    v, _ = torch.topk(lg, min(top_k, lg.numel()))
+                    lg[lg < v[-1]] = -float("inf")
+                probs = F.softmax(lg, dim=-1)
+                if top_p < 1.0:
+                    sp, si = torch.sort(probs, descending=True)
+                    sp[sp.cumsum(-1) - sp > top_p] = 0
+                    probs = torch.zeros_like(probs).scatter_(0, si, sp)
+                    probs = probs / probs.sum()
+                nxt = int(torch.multinomial(probs, 1))
+            yield nxt
+            if nxt in stop_ids:
+                break
+            seen.add(nxt)
+            logits = forward(torch.tensor([[nxt]], device=idx.device), pos)
+            pos += 1

vbq/optim.py

@@ -0,0 +1,55 @@
+"""Muon (Newton-Schulz orthogonalised momentum) for 2-D hidden matrices (and 3-D stacks of MoE expert matrices); scale matched to AdamW RMS (Moonlight)."""
+import torch
+
+
+@torch.no_grad()
+def newton_schulz(G: torch.Tensor, steps: int = 5) -> torch.Tensor:
+    a, b, c = 3.4445, -4.7750, 2.0315
+    X = G.bfloat16()
+    transposed = X.shape[0] > X.shape[1]
+    if transposed:
+        X = X.T
+    X = X / (X.norm() + 1e-7)
+    for _ in range(steps):
+        A = X @ X.T
+        X = a * X + (b * A + c * (A @ A)) @ X
+    return X.T if transposed else X
+
+
+class Muon(torch.optim.Optimizer):
+    def __init__(self, params, lr=1e-3, momentum=0.95, weight_decay=0.1, split=None):
+        # split: dict param-id -> number of equal row-chunks to orthogonalise independently (fused gate/up)
+        super().__init__(params, dict(lr=lr, momentum=momentum, weight_decay=weight_decay))
+        self.split = split or {}
+        self.bad = None   # device counter of updates dropped for being non-finite (read at log time, no sync here)
+
+    @torch.no_grad()
+    def step(self):
+        for g in self.param_groups:
+            for p in g["params"]:
+                if p.grad is None:
+                    continue
+                st = self.state[p]
+                if "buf" not in st:
+                    st["buf"] = torch.zeros_like(p.grad)
+                buf = st["buf"]
+                buf.lerp_(p.grad, 1 - g["momentum"])
+                upd = p.grad.lerp(buf, g["momentum"])  # nesterov
+                # 3-D = stacked MoE experts: each expert's matrix is orthogonalised on its own
+                mats = list(upd.unbind(0)) if upd.ndim == 3 else [upd]
+                done = []
+                for m in mats:
+                    outs = []
+                    for ch in m.chunk(self.split.get(id(p), 1), dim=0):
+                        o = newton_schulz(ch)
+                        outs.append(o * (0.2 * max(ch.shape) ** 0.5))
+                    done.append(torch.cat(outs, 0) if len(outs) > 1 else outs[0])
+                o = torch.stack(done) if upd.ndim == 3 else done[0]
+                # a rare non-finite Newton-Schulz result (seen on the RX 7600 with finite grads) would poison the
+                # weight for good: drop this parameter's update for the step instead
+                ok = torch.isfinite(o).all()
+                if self.bad is None: self.bad = torch.zeros((), device=o.device, dtype=torch.int64)
+                self.bad += (~ok).long()
+                o = torch.where(ok, o, torch.zeros((), device=o.device, dtype=o.dtype))
+                p.mul_(1 - g["lr"] * g["weight_decay"])
+                p.add_(o.to(p.dtype), alpha=-g["lr"])

vbq/sources.py

@@ -0,0 +1,171 @@
+"""Streaming adapters for each pretraining source. Each yields cleaned text documents."""
+from __future__ import annotations
+
+import io
+import json
+import random
+import re
+from pathlib import Path
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.parquet as pq
+import yaml
+from huggingface_hub import HfApi, hf_hub_download
+
+ROOT = Path(__file__).resolve().parent.parent
+RAW = ROOT / "data" / "raw"
+_BLANKS = re.compile(r"\n{4,}")
+
+
+def load_cfg(path=None):
+    return yaml.safe_load(open(path or ROOT / "configs" / "data.yaml"))
+
+
+def _clean(t: str) -> str:
+    t = t.replace("\r\n", "\n").replace("\x00", "")
+    return _BLANKS.sub("\n\n\n", t).strip()
+
+
+class Source:
+    def __init__(self, name: str, cfg: dict, glob: dict):
+        self.name, self.cfg, self.glob = name, cfg, glob
+        self.repo = cfg["repo"]
+        self.min_chars = cfg.get("min_chars", glob.get("min_chars", 200))
+        self.quota = int(glob["total_tokens"] * cfg["weight"])
+        self.per_shard_cap = cfg.get("per_shard_cap_tokens")
+        self.api = HfApi()
+
+    # -- file selection (deterministic shuffle so a partial run is still a uniform sample) --
+    def list_shards(self) -> list[str]:
+        files = [f.path for f in self.api.list_repo_tree(self.repo, repo_type="dataset", path_in_repo=self.cfg["prefix"].rstrip("/"), recursive=True)
+                 if hasattr(f, "size") and f.path.endswith((".parquet", ".jsonl.zst"))]
+        files.sort()
+        random.Random(self.glob["seed"]).shuffle(files)
+        return files
+
+    def download(self, path: str) -> Path:
+        return Path(hf_hub_download(self.repo, path, repo_type="dataset", local_dir=RAW / self.name))
+
+    def docs(self, local: Path):
+        raise NotImplementedError
+
+    def _emit(self, texts):
+        for t in texts:
+            if t is None:
+                continue
+            t = _clean(t)
+            if len(t) >= self.min_chars:
+                yield t
+
+
+class ParquetSource(Source):
+    text_col = "text"
+    extra_cols: list[str] = []
+
+    def mask(self, tbl: pa.Table):
+        return None
+
+    def format(self, tbl: pa.Table):
+        return tbl[self.text_col].to_pylist()
+
+    def docs(self, local: Path):
+        pf = pq.ParquetFile(local)
+        cols = [self.text_col] + self.extra_cols
+        for batch in pf.iter_batches(batch_size=4096, columns=cols):
+            tbl = pa.Table.from_batches([batch])
+            m = self.mask(tbl)
+            if m is not None:
+                tbl = tbl.filter(m)
+            yield from self._emit(self.format(tbl))
+
+
+class FineWebEdu(ParquetSource):
+    extra_cols = ["metadata"]
+
+    def mask(self, tbl):
+        score = pc.struct_field(tbl["metadata"], "score")
+        return pc.greater_equal(score, self.cfg["min_edu_score"])
+
+
+class Cosmopedia(ParquetSource):
+    pass
+
+
+class FineMath(ParquetSource):
+    extra_cols = ["score", "language", "language_score"]
+
+    def mask(self, tbl):
+        return pc.and_(pc.equal(tbl["language"], "en"), pc.greater_equal(tbl["language_score"], 0.9))
+
+
+class Wikipedia(ParquetSource):
+    extra_cols = ["title"]
+
+    def format(self, tbl):
+        out = []
+        for title, text in zip(tbl["title"].to_pylist(), tbl["text"].to_pylist()):
+            if title.startswith(("List of", "Lists of", "Index of")) or len(text) < self.cfg.get("min_chars", 2000):
+                continue
+            out.append(f"# {title}\n\n{text}")
+        return out
+
+
+class DCLM(ParquetSource):
+    extra_cols = ["fasttext_score", "language", "language_score"]
+
+    def list_shards(self):
+        base = "filtered/OH_eli5_vs_rw_v2_bigram_200k_train/fasttext_openhermes_reddit_eli5_vs_rw_v2_bigram_200k_train/processed_data"
+        rng = random.Random(self.glob["seed"])
+        out = [f"{base}/global-shard_{g:02d}_of_10/local-shard_{l}_of_10/shard_{n:08d}_processed.parquet"
+               for g in range(1, 11) for l in range(10) for n in range(0, 260)]
+        rng.shuffle(out)
+        return out[:self.cfg.get("max_shards", 200)]  # far more than needed; builder stops at quota
+
+    def mask(self, tbl):
+        m = pc.and_(pc.greater_equal(tbl["fasttext_score"], self.cfg["min_fasttext"]),
+                    pc.greater_equal(tbl["language_score"], self.cfg["min_lang_score"]))
+        return pc.and_(m, pc.equal(tbl["language"], "en"))
+
+
+class StackExchange(ParquetSource):
+    text_col = "ThreadText"
+    extra_cols = ["Score", "AnswerCount"]
+
+    def list_shards(self):
+        out = []
+        for site in self.cfg["sites"]:
+            d = site if "." in site else f"{site}.stackexchange.com"
+            if site == "stackoverflow":
+                d = "stackoverflow.com"
+            try:
+                fs = [f.path for f in self.api.list_repo_tree(self.repo, repo_type="dataset", path_in_repo=d) if f.path.endswith(".parquet")]
+            except Exception:
+                continue
+            off, n = self.cfg.get("shard_offset", 0), self.cfg.get("shards_per_site", 1)
+            out.extend(sorted(fs)[off:off + n])  # few shards per site; per_shard_cap keeps the mix balanced
+        random.Random(self.glob["seed"]).shuffle(out)
+        return out
+
+    def mask(self, tbl):
+        return pc.and_(pc.greater_equal(tbl["Score"], self.cfg["min_score"]), pc.greater_equal(tbl["AnswerCount"], 1))
+
+
+class CodePython(Source):
+    def docs(self, local: Path):
+        import zstandard
+        with open(local, "rb") as fh, zstandard.ZstdDecompressor().stream_reader(fh) as r:
+            for line in io.TextIOWrapper(r, encoding="utf-8"):
+                o = json.loads(line)
+                if o.get("score", 0) < self.cfg["min_score"] or o.get("length_bytes", 0) < self.cfg["min_bytes"]:
+                    continue
+                yield from self._emit([o["text"]])
+
+
+REGISTRY = {"fineweb_edu": FineWebEdu, "dclm": DCLM, "cosmopedia": Cosmopedia, "finemath": FineMath,
+            "wikipedia": Wikipedia, "stackexchange": StackExchange, "code_python": CodePython}
+
+
+def get_source(name: str, cfg: dict | None = None) -> Source:
+    cfg = cfg or load_cfg()
+    return REGISTRY[name](name, cfg["sources"][name], cfg)