Add HF streaming data pipeline and three research reports

- tools/prepare_hf_data.py: stream any of 13 registered HF datasets
  (FineWeb-Edu, DCLM, Cosmopedia, FineMath, OpenWebMath, Wikipedia, C4,
  SmolTalk, ...) or weighted mixtures/presets (smollm-web, textbooks,
  decay-anneal) straight into train.bin/val.bin — no full downloads
- SMALL_MODELS_RESEARCH.md: how sub-1B models get logical, useful text
  (data quality, deep-and-thin, distillation, test-time compute)
- INFERENCE_RESEARCH.md: Apple Silicon inference speed playbook tied to
  the .forge format (bandwidth math, fused-dequant GEMV, KV cache,
  residency sets, warmup, quant layouts) with a prioritized roadmap
- ARCHITECTURES.md: config matrix to train Llama/Mistral/Qwen/Gemma/
  DeepSeek/Kimi-class variants, with a 3-wave implementation plan
- README: training modes, .forge format, tools

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

simon-pierre boucher committed 5 days ago (Aug 5, 2026) parent 8b73bf1

Showing 5 changed files with +585 and −3

added ARCHITECTURES.md +68 −0

		`@@ -0,0 +1,68 @@`
	1	+<!-- Author: Simon-Pierre Boucher — contact@spboucher.ai -->
	2	+
	3	+# Matrice d'architectures — entraîner du « Mistral-like », « Qwen-like », « Kimi-like » avec Forge
	4	+### Rapport de recherche web (août 2026) + plan d'implémentation basé sur le code actuel
	5	+
	6	+---
	7	+
	8	+## Où en est Forge
	9	+
	10	+Base actuelle : pre-norm RMSNorm/LayerNorm · SwiGLU/GELU · RoPE complet ou positions apprises · MHA/GQA · embeddings liés · MoE top-k (softmax, renormalisé, aux loss) · QAT int8/ternaire · AdamW/Muon · cosine/WSD. Déjà ajoutés cette session : `qk_norm`, `final_softcap`, `scale_embeddings`, `n_shared_experts`.
	11	+
	12	+## Les deltas par famille (chiffres réels des configs HF)
	13	+
	14	+\| Famille \| Ce qui la distingue \| Statut Forge \|
	15	+\|---\|---\|---\|
	16	+\| Llama 3.2 (1B/3B) \| GQA, SwiGLU, tied ≤3B, rope_theta 500k, rope_scaling "llama3" (factor 32, blend low/high freq) \| ✅ sauf rope_scaling \|
	17	+\| Mistral 7B \| Sliding window 4096 partout ; Ministral : SWA entrelacé full/window \| ❌ SWA (kernel) \|
	18	+\| Qwen 2.5 \| Biais sur Q/K/V seulement, rope_theta 1e6, GQA 28H/4KV \| ❌ attention_bias \|
	19	+\| Qwen 3 (dense) \| Remplace le biais par QK-Norm, head_dim découplé (128), tied ≤4B \| ✅ qk_norm ; ❌ head_dim découplé \|
	20	+\| Qwen3-30B-A3B (MoE) \| 128 experts top-8, pas d'expert partagé, softmax renormalisé, moe_d_ff 768 ≠ d_ff \| ✅ sauf `moe_intermediate_size` séparé \|
	21	+\| Gemma 2 \| Local(4096)/global 1:1, softcap attn 50 / final 30, sandwich norm, GeGLU, head_dim 256, emb×√d, query_pre_attn_scalar \| ✅ softcap final + emb scale ; ❌ reste \|
	22	+\| Gemma 3 \| 5 local : 1 global, window 1024, softcap→QK-norm, rope_theta 10k local / 1M global \| ✅ qk_norm ; ❌ patterns par couche \|
	23	+\| DeepSeek V3 \| MLA (kv_lora 512, q_lora 1536, nope 128 + rope 64, v 128), 256 routed + 1 shared top-8, routage sigmoïde + équilibrage par biais sans aux loss (γ=0.001), routed_scaling 2.5, 3 premières couches denses, MTP λ=0.3→0.1 \| ✅ shared expert ; ❌ MLA, sigmoïde/noaux \|
	24	+\| Kimi K2 \| = V3 avec 64 têtes (÷2), 384 experts, couche 0 seule dense, routed_scaling 2.827, MuonClip (Muon + QK-Clip : si logit max > τ=100, W_q,W_k ← ×√(τ/S_max) par tête) \| ✅ Muon ; ❌ QK-Clip, reste \|
	25	+\| GPT-OSS \| Dense/sliding-128 alterné, attention sinks appris (biais par tête dans le dénominateur softmax), MoE 32/top-4, biais partout, YaRN \| ❌ sinks (petit patch kernel) \|
	26	+\| OLMo 2 \| Post-norm réordonné (`x = x + norm(attn(x))`) + QK-norm — la combinaison fait la stabilité \| ✅ qk_norm ; ❌ norm_placement \|
	27	+\| SmolLM3 \| Llama-like + NoPE 1 couche sur 4 (RoPE sauté) → long contexte 64k+ \| ❌ nope_layers (facile) \|
	28	+\| nanochat (Karpathy) \| ReLU², head délié zero-init, RMSNorm sans paramètre, qk-norm après RoPE, softcap 15, value embeddings, Muon \| ✅ Muon, softcap ; ❌ relu², reste \|
	29	+
	30	+## Plan d'implémentation priorisé (valeur ÷ coût, ancré dans le code actuel)
	31	+
	32	+Vague 1 — pure plomberie config (pas de kernel, ~1 jour)
	33	+1. `rope_scaling` type llama3 : précalculer les inv_freq modifiés (formule ci-dessous) — le kernel RoPE actuel calcule les freqs inline ; passer à une table de freqs précalculée le débloque ET prépare les theta par couche.
	34	+2. `nope_layers` (SmolLM3) : sauter `ops::rope` selon la couche — un `if` dans `attention.h` + liste en config.
	35	+3. `attention_bias` (Qwen2.5) : `Linear` supporte déjà le biais — juste le flag.
	36	+4. `head_dim` découplé de d_model/n_heads (Qwen3) : `head_dim` explicite en config, wq → n_heads*head_dim.
	37	+5. `moe_intermediate_size` séparé + `norm_topk_prob` optionnel + `routed_scaling_factor`.
	38	+6. `relu2` comme 3e activation (kernel elementwise trivial + backward).
	39	+7. `norm_placement: pre \| post_olmo \| sandwich` : recomposition dans `TransformerBlock`.
	40	+
	41	+Vague 2 — un patch du kernel d'attention (la plus grosse valeur kernel)
	42	+8. `sliding_window` + `layer_types` par couche (Mistral, Gemma 2/3, GPT-OSS) : dans les kernels flash (scalar + MMA) et la référence, le masque causal devient `q−W < k ≤ q` ; bonus perf réel : les tuiles KV hors fenêtre sont sautées (borne de boucle), pas juste masquées. Avec `rope_theta` par type de couche (Gemma 3).
	43	+9. `attn_logit_softcap` (Gemma 2) : `cap·tanh(s/cap)` sur les scores pré-softmax, dans le kernel (l'op `softcap` + son backward existent déjà côté élémentwise).
	44	+10. Attention sinks (GPT-OSS) : initialiser le (max, somme) du softmax online avec le logit sink appris par tête — quelques lignes dans le kernel flash.
	45	+
	46	+Vague 3 — investissements
	47	+11. Routage sigmoïde + équilibrage sans aux loss (V3/K2) : sélection top-k sur `s+b`, gates sur `s` seul ; mise à jour du biais ±γ selon la charge après chaque step (état dans l'optimiseur/trainer).
	48	+12. QK-Clip / MuonClip (K2) : monitorer le logit d'attention max par tête (le kernel flash garde déjà le max par ligne — l'exposer), rescaler W_q/W_k si > τ.
	49	+13. MLA (V3/K2) : le seul gros chantier kernel — d_k(192) ≠ d_v(128) + RoPE partiel sur une tranche de 64. Praticable d'abord en « décompressé » si le kernel accepte d_k≠d_v.
	50	+14. MTP (V3) : bloc supplémentaire + tête partagée, coût trainer modéré.
	51	+
	52	+## Formules clés
	53	+
	54	+RoPE scaling llama3 (par composante, old=8192) :
	55	+```
	56	+wavelen = 2π/inv_freq
	57	+si wavelen < old/high_freq_factor : inchangé
	58	+si wavelen > old/low_freq_factor : inv_freq/factor
	59	+sinon : blend lisse s=(old/wavelen − low)/(high − low) ; (1−s)·inv_freq/factor + s·inv_freq
	60	+```
	61	+
	62	+Routage V3 : `s_i = σ(u·e_i)` ; top-k sur `s_i + b_i` (b pour la sélection seulement) ; `g_i = s_i/Σ_topk s_j` ; `y = x + scaling·Σ g_i·E_i(x) + E_shared(x)` ; après chaque step `b_i ∓ γ` selon surcharge/sous-charge (γ=0.001).
	63	+
	64	+MLA (d=7168) : `c_kv = W_dkv·x` (→512) ; `k_rope = RoPE(W_kr·x)` (→64, partagé entre têtes) ; par tête : `k_h = [W_ukv,h·c_kv ; k_rope]` (192), `v_h = W_uv,h·c_kv` (128) ; query via `c_q` (1536) ; scale 1/√192.
	65	+
	66	+## Sources
	67	+
	68	+[Llama rope_scaling (HF rope_utils)](https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py) · [Mistral 7B](https://ar5iv.labs.arxiv.org/html/2310.06825) · [Ministral](https://mistral.ai/news/ministraux/) · [Qwen3](https://arxiv.org/html/2505.09388v1) · [Gemma 3](https://arxiv.org/html/2503.19786v1) · [DeepSeek-V3 config](https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/config.json) · [V3 report](https://arxiv.org/pdf/2412.19437) · [Kimi K2 config](https://huggingface.co/moonshotai/Kimi-K2-Instruct/blob/main/config.json) · [K2 report (MuonClip)](https://arxiv.org/pdf/2507.20534) · [GPT-OSS](https://arxiv.org/pdf/2508.10925) · [OLMo 2](https://arxiv.org/pdf/2501.00656) · [SmolLM3](https://github.com/huggingface/blog/blob/main/smollm3.md) · [nanochat gpt.py](https://github.com/karpathy/nanochat/blob/master/nanochat/gpt.py) · [Raschka — comparaison d'architectures](https://magazine.sebastianraschka.com/p/the-big-llm-architecture-comparison)

added INFERENCE_RESEARCH.md +71 −0

		`@@ -0,0 +1,71 @@`
	1	+<!-- Author: Simon-Pierre Boucher — contact@spboucher.ai -->
	2	+
	3	+# Booster la vitesse d'inférence sur Apple Silicon — avec le format `.forge`
	4	+### Rapport de recherche web, août 2026 — llama.cpp/Metal, MLX, BaseRT, GGUF/safetensors/Core ML, Ollama/Xet
	5	+
	6	+---
	7	+
	8	+## TL;DR — la loi qui gouverne tout
	9	+
	10	+Le décodage autorégressif (batch 1) est borné par la bande passante mémoire, pas par le calcul : latence/token ≈ (octets de poids + KV lus par token) ÷ bande passante. Budgets M-series : M4 120 GB/s · M5 153 · M4 Pro 273 · M5 Max ≈ 550 · M2/M3 Ultra ~800. Trois leviers, dans l'ordre :
	11	+
	12	+1. Réduire les octets → quantization 4-bit (~4× moins de trafic = ~4× plus de tok/s), KV f16, GQA
	13	+2. Saturer la bande passante → GEMV avec déquantification fusionnée dans la boucle interne, lectures coalescées
	14	+3. Tuer l'overhead par token → 1 command buffer/token, zéro allocation dans la boucle de décodage, kernels fusionnés
	15	+
	16	+Le prefill, lui, est borné par le calcul → `simdgroup_matrix` (déjà dans Forge) et, sur M5, les Neural Accelerators via `tensor_ops::matmul2d` (~4× le prefill de M4 dans MLX ; le décodage ne gagne que 19–27 % — bande passante toujours).
	17	+
	18	+Références de vitesse : Llama-3.1-8B Q4 ≈ 51–75 tok/s sur M4 Max ; BaseRT (meilleur runtime Metal publié) : Qwen3-0.6B Q4 à 464 tok/s sur M4 Pro, +15–56 % vs llama.cpp, principalement par fusion de kernels.
	19	+
	20	+---
	21	+
	22	+## 1. Le chemin de décodage rapide (kernels)
	23	+
	24	+- GEMV dequant-fusionné : le kernel matvec déquantifie les blocs 4-bit dans la boucle interne (jamais de poids f32 matérialisés). Design llama.cpp : chaque simdgroup possède une bande de lignes de sortie (N_R0 lignes/simdgroup, N_SG simdgroups/threadgroup, réglés par format) ; 32 lanes stride le long des blocs de la ligne ; réduction `simd_sum` finale. Intensité arithmétique ~2 ops/octet — pur streaming.
	25	+- Fusions à implémenter (chaque fusion = un lancement + un aller-retour mémoire en moins ; l'overhead de dispatch coûtait +17 % sur SmolLM2-360M avant fusion) : rmsnorm+matvec QKV, residual+norm, gate·up SwiGLU en un kernel, dequant+matvec (toujours), matvec logits + sampling top-k sur GPU (ne rapatrier que le token).
	26	+- Zéro overhead par token : un seul command buffer encodant toute la pile de couches ; jamais de `waitUntilCompleted` au milieu ; double-buffering (encoder le token N+1 pendant que N s'exécute) ; « the decode loop allocates zero bytes » (BaseRT). Spécialiser par `MTLFunctionConstantValues` (head_dim, n_heads, params quant cuits dans le binaire) — Forge fait déjà ça pour le matmul.
	27	+- Routage par M : la décision n°1 du dispatcher est mat-mat (prefill, M grand) vs mat-vec (décodage, M=1).
	28	+
	29	+## 2. KV cache
	30	+
	31	+- Pré-alloué à la longueur max au chargement, layout coalescé le long de l'axe temporel par tête (`[layer][head][seq][d]`), f16 par défaut. Kernel décodage dédié : une query, split-K le long du cache pour les longs contextes.
	32	+- GQA est la plus grosse économie de bande passante KV (déjà supporté par Forge via `n_kv_heads`).
	33	+- Paged attention sur Metal : existe en expérimental, personne ne le met en prod pour du batch-1 — cache contigu simple = plus rapide. KV quantifié q8/q4 : fragile sur Metal (bugs llama.cpp), pas prioritaire.
	34	+
	35	+## 3. Ce que la recherche valide dans `.forge` (et ce qu'elle ajoute)
	36	+
	37	+Notre choix central — tenseurs alignés page 16 KB + mmap + `newBuffer(bytesNoCopy)` — est exactement le pattern llama.cpp/Metal, en mieux : GGUF n'aligne qu'à 32 octets (obligeant des fenêtres page-alignées bricolées) et safetensors à 8 octets — c'est l'anti-pattern démontré (impossible de passer un tenseur à bytesNoCopy). Le sharding contenu-adressé + manifests ≈ le blob store d'Ollama et les xorbs de HF Xet — l'industrie a convergé sur notre design.
	38	+
	39	+Upgrades format à faire (issus de la recherche) :
	40	+
	41	+1. Ordonner les tenseurs dans l'ordre du forward pass dans les shards (embeddings → blk.0 … blk.N → head) : le streaming de couches devient de l'I/O séquentielle (prefetch du shard N+1 pendant le calcul de la couche N, ~+10 % façon AirLLM).
	42	+2. Warmup séquentiel au chargement : le demand-faulting à froid est lent (fautes 16 KB aléatoires) ; un passage `F_RDADVISE`/lecture séquentielle par shard sature le SSD (~7 GB/s) — llama.cpp mesure +13–14 %. Option `mlock` quand le modèle tient en RAM (l'éviction du page cache sous pression = 0.025 tok/s de thrash).
	43	+3. MTLResidencySet (macOS 15+) sur tous les buffers de poids, attaché à la queue — remplace `useResource`, évite les stalls d'éviction GPU.
	44	+4. Quantization planaire, pas interleavée : data 4-bit packée en uint32 + tenseurs `scales`/`biases` séparés (style MLX, groupe 64) plutôt que les blocs 18 octets de GGUF — garde les loads vectoriels 16 B alignés et les scales résidents en threadgroup memory. Ajouter un type LUT/palettisé (indices 4-bit + LUT f16 16 entrées, style Core ML). Réserver des IDs dtype fp8 (E4M3/E5M2 — footprint seulement : fp8 émulé à 0.94× f16 sur M4, natif sur M5+).
	45	+5. Layout canonique sur disque : la seule transformation on-disk qui paie est la pré-transposition pour le sens d'accès du matvec ; PAS de fichiers pré-swizzlés (llama.cpp a supprimé Q4_0_4_4 et repacke au chargement — mais sur Apple Silicon, préférer des kernels qui consomment le layout canonique, car repacker touche toutes les pages et tue le zéro-copie).
	46	+6. Intégrité en deux vitesses : BLAKE3 (7–16 GB/s, Merkle) comme adresse de contenu vérifiée au pull/push ; XXH3 (30 GB/s) par tenseur en footer pour un `--verify` opt-in fusionné au warmup (coût ~0). Jamais de hash sur le chemin de chargement par défaut — mmap veut dire qu'on n'a pas encore lu les octets. + GC mark-and-sweep manifests→shards (style Ollama) et un index repo `tensor → (shard, offset)` façon `safetensors.index.json`.
	47	+
	48	+## 4. À NE PAS faire (impasses documentées)
	49	+
	50	+- Kernels ternaires GPU : TQ1_0/TQ2_0 ciblent les CPU ; pas d'implémentation Metal, jugée difficile par les mainteneurs. Notre QAT ternaire reste un outil d'entraînement/recherche ; à l'inférence, servir en 4-bit groupé.
	51	+- Speculative decoding naïf : perte nette sur Metal à toutes les configs (le coût de dispatch draft/verify mange le gain). Ne marche que verify-fusionné dans un seul graphe (Recurrent Drafter d'Apple : 2.3×) — plus tard.
	52	+- ANE pour le décodage : Core ML impose 2–4× d'overhead, shapes statiques ; le GPU est 2–5× plus rapide. Niche : basse conso, ou prefill-ANE + decode-GPU.
	53	+
	54	+## 5. Roadmap Forge priorisée (impact ÷ effort)
	55	+
	56	+\| # \| Chantier \| Gain attendu \|
	57	+\|---\|---\|---\|
	58	+\| 1 \| KV cache (le `generate` actuel recompute tout le contexte à chaque token !) \| 10–100× sur la génération, avant tout le reste \|
	59	+\| 2 \| GEMV batch-1 f32→f16 + kernel matvec dédié (routage M=1) \| ~2× décodage \|
	60	+\| 3 \| Quant 4-bit planaire dans `.forge` + GEMV dequant-fusionné \| ~4× décodage vs f32 \|
	61	+\| 4 \| 1 command buffer/token, sampling GPU, zéro alloc \| +15–20 % \|
	62	+\| 5 \| Warmup séquentiel + ResidencySet + mlock opt-in au chargement `.forge` \| cold-start ÷10, pas de stalls \|
	63	+\| 6 \| Fusions (norm+matvec, SwiGLU, residual) \| +15–50 % (résultat BaseRT) \|
	64	+\| 7 \| Tenseurs ordre-forward + index repo + XXH3 footer dans `.forge` v2 \| streaming + intégrité \|
	65	+\| 8 \| M5 : `tensor_ops::matmul2d` pour le prefill (déjà exploré dans bench_precision) \| ~4× prefill/TTFT \|
	66	+
	67	+## Sources principales
	68	+
	69	+Kernels/runtime : [BaseRT](https://arxiv.org/html/2607.00501) · [BaseRT+M5](https://arxiv.org/html/2607.19438v1) · [llama.cpp Metal backend](https://deepwiki.com/ggml-org/llama.cpp/5.2-metal-backend-(apple)) · [MLX sur M5 (Apple ML)](https://machinelearning.apple.com/research/exploring-llms-mlx-m5) · [spec-decode Metal #23752](https://github.com/ggml-org/llama.cpp/issues/23752) · [MLX quantization](https://deepwiki.com/ml-explore/mlx/7-quantization) · [Rigel (simdgroup/fp8)](https://arxiv.org/abs/2606.12765) · [WWDC26 Metal tensors](https://developer.apple.com/videos/play/wwdc2026/330/)
	70	+
	71	+Format/chargement : [GGUF spec](https://github.com/ggml-org/ggml/blob/master/docs/gguf.md) · [safetensors](https://github.com/safetensors/safetensors) · [fastsafetensors](https://arxiv.org/html/2505.23072v1) · [repack au chargement PR #9921](https://github.com/ggml-org/llama.cpp/pull/9921) · [mmap (justine.lol)](https://justine.lol/mmap/) · [prefetch séquentiel #18758](https://github.com/ggml-org/llama.cpp/discussions/18758) · [Ollama blob store](https://deepwiki.com/ollama/ollama/2.4-storage-and-blob-transfer) · [HF Xet dedup](https://huggingface.co/docs/xet/en/deduplication) · [palettisation Core ML](https://apple.github.io/coremltools/docs-guides/source/opt-palettization-overview.html) · [BLAKE3/XXH3 vitesses](https://jolynch.github.io/posts/use_fast_data_algorithms/)

modified README.md +30 −3

		`@@ -219,20 +219,47 @@ bit-comparable with the GPU path, which is what the parity suite checks.`
219	219	\| path \| contents \|
220	220	\|:--\|:--\|
221	221	\| `src/core/` \| `Tensor` (shared-storage views), `Allocator` (bucketed MTLBuffer pool), `Device` (queue + pipeline cache), autograd tape \|
222		−\| `src/kernels/` \| 12 `.metal` files: GEMM (naive→tiled→simdgroup→`matmul2d`), flash attention (scalar + MMA), softmax, norms, elementwise, embedding, cross-entropy, AdamW \|
	222	+\| `src/kernels/` \| 14 `.metal` files: GEMM (naive→tiled→simdgroup→`matmul2d`), flash attention (scalar + MMA), softmax, norms, elementwise, embedding, cross-entropy, AdamW, fake-quant (QAT), MoE gating \|
223	223	\| `src/ops/` \| `cpu/` reference impls · `metal/` dispatch + batched `Stream` · `ops.cpp` autograd layer routing to either backend \|
224		−\| `src/nn/` \| `Module`, `Linear` (pruning-mask + quantization seams), attention (GQA + RoPE behind an interface), SwiGLU/GELU MLP, `Transformer` \|
225		−\| `src/train/` \| mmap dataloader, AdamW, warmup+cosine schedule, trainer, resumable checkpoints \|
	224	+\| `src/nn/` \| `Module`, `Linear` (pruning mask + int8/ternary QAT), attention (GQA + RoPE behind an interface), SwiGLU/GELU MLP or top-k MoE, `Transformer` \|
	225	+\| `src/train/` \| mmap dataloader, AdamW or Muon, warmup+cosine or WSD schedule, trainer, resumable checkpoints \|
226	226	\| `src/tokenizer/` \| byte-level BPE, verified identical to the Python encoder \|
	227	+\| `src/core/fmodel.` \| `.forge` weight format* — Apple-native, git-style (see below) \|
	228	+\| `tools/` \| `prepare_data.py` (TinyStories) · `prepare_hf_data.py` (any HF dataset/mixture, streamed) · `fmodel.py` (inspect / history / safetensors export) \|
227	229	\| `tests/` \| parity · gradcheck · overfit · tokenizer · 3 benchmarks \|
228	230
229	231	Model: decoder-only transformer · RMSNorm or LayerNorm · SwiGLU or GELU · RoPE
230	232	(interleaved pairs) or learned positions · GQA · tied embeddings — all from config.
231	233
	234	+Training modes (config-selected): optimizer `adamw` \| `muon` (Newton–Schulz
	235	+orthogonalized momentum on hidden matrices, AdamW on embeddings/head) · schedule
	236	+`cosine` \| `wsd` (flat plateau + 1−√ cooldown, extendable runs) · `quant`
	237	+`none` \| `int8` \| `ternary` (BitNet-style fake-quant each forward, STE backward,
	238	+f32 master weights) · `n_experts`/`moe_top_k` (softmax router, renormalized top-k
	239	+gates, differentiable load-balance loss; v1 computes experts densely). See
	240	+`configs/gpt-50m-{muon,deep,ternary,moe}.json`.
	241	+
232	242	Verified: 85 CPU↔Metal parity checks (≤1e-4, most bit-exact) · numerical gradient
233	243	checks on every parameterized op and a full transformer · single-batch overfit to
234	244	loss < 0.05 in 86 steps · exact checkpoint resume · BPE round-trip vs Python.
235	245
	246	+`.forge` weight format (`forge export --checkpoint ckpt.bin --out model.forge`):
	247	+a model repository rather than a file. Tensors are 16 KB-page-aligned inside
	248	+content-addressed shards capped at 95 MB (GitHub-pushable), and loading is
	249	+`mmap` + `newBuffer(bytesNoCopy)` — on unified memory the file-cache pages are
	250	+the GPU memory, so a multi-GB model loads in milliseconds with zero copies.
	251	+Saves are git-style commits: a tiny JSON manifest per save, and **only tensors
	252	+whose content hash changed since the parent are written** — repeated exports
	253	+cost only the delta. Store as f32 (zero-copy load) or f16/bf16 (half size).
	254	+`generate`/`eval` accept a `.forge` repo directly; `tools/fmodel.py` gives
	255	+`inspect`, `log` (history), and `to-safetensors` for PyTorch/HF interop.
	256	+
	257	+HF data pipeline (`tools/prepare_hf_data.py`): streams any of 13 registered
	258	+Hugging Face datasets (FineWeb-Edu, DCLM, Cosmopedia, FineMath, OpenWebMath,
	259	+Wikipedia, C4, SmolTalk, …) or a weighted mixture (`--mix fineweb-edu:0.6,dclm:0.4`,
	260	+or presets like `smollm-web`) straight into `train.bin`/`val.bin` — no
	261	+full-corpus downloads, pay only for the megabytes you keep.
	262	+
236	263	---
237	264
238	265	## 🗺️ Roadmap

added SMALL_MODELS_RESEARCH.md +168 −0

		`@@ -0,0 +1,168 @@`
	1	+<!-- Author: Simon-Pierre Boucher — contact@spboucher.ai -->
	2	+
	3	+# Rendre un modèle < 1B paramètres capable de texte logique et utile
	4	+### Rapport de recherche — état de l'art 2023–2026, orienté Forge
	5	+
	6	+Synthèse de recherche web (arXiv, HuggingFace, Meta AI, Microsoft Research, DeepMind, blogs d'ingénierie), août 2026. Trois axes couverts en parallèle : données, architecture/entraînement, post-training/inférence.
	7	+
	8	+---
	9	+
	10	+## TL;DR — la thèse en une phrase
	11	+
	12	+La taille n'est plus le facteur limitant sous 1B : c'est (1) la qualité/adéquation des données, (2) la forme du réseau et la distillation, et (3) le compute dépensé à l'inférence. Empilées, ces trois familles de techniques permettent aujourd'hui à des modèles 0.5–3B de battre des modèles 10× à 100× plus gros sur des tâches ciblées — et à des modèles de 7M–30M de produire du texte parfaitement cohérent dans un domaine restreint.
	13	+
	14	+Les preuves les plus spectaculaires :
	15	+
	16	+\| Résultat \| Petit modèle \| Bat \| Source \|
	17	+\|---\|---\|---\|---\|
	18	+\| Cohérence narrative \| TinyStories 28M \| GPT-2-XL 1.5B (54× plus gros) \| arXiv 2305.07759 \|
	19	+\| Code (HumanEval 45%) \| phi-1-small 350M \| modèles 10× plus gros \| arXiv 2306.11644 \|
	20	+\| MATH-500 (78.2 vs 71.4) \| Llama-3.2-3B + test-time search \| Llama-3.1-405B \| arXiv 2502.06703 \|
	21	+\| MATH-500 (76.4 vs 74.6) \| Qwen2.5-0.5B + TTS \| GPT-4o \| arXiv 2502.06703 \|
	22	+\| MATH-500 83.9% \| R1-Distill-Qwen-1.5B (SFT seul) \| GPT-4o, Claude 3.5 Sonnet \| DeepSeek-R1 \|
	23	+\| ARC-AGI-1 ~45% \| Samsung TRM 7M (récursion ×16) \| scores publiés de R1 / o3-mini \| arXiv 2510.04871 \|
	24	+
	25	+---
	26	+
	27	+## 1. Les données : le levier n°1 (et de loin)
	28	+
	29	+### 1.1 TinyStories — adapter la largeur des données à la capacité du modèle
	30	+Le résultat fondateur (Eldan & Li 2023) : des modèles de 1M à 33M paramètres (parfois un seul bloc transformer) génèrent des histoires fluides, grammaticales et logiques quand le corpus est restreint au vocabulaire d'un enfant de 3–4 ans (~480M tokens synthétiques GPT-3.5/4).
	31	+
	32	+- Les capacités émergent hiérarchiquement : grammaire d'abord, cohérence contextuelle ensuite, créativité/logique de l'intrigue en dernier (le plus coûteux en capacité).
	33	+- Largeur ↔ connaissances factuelles ; profondeur ↔ cohérence narrative.
	34	+- Pourquoi ça marche : sur du web générique, un petit modèle gaspille sa capacité en vocabulaire et connaissances du monde qu'il ne pourra jamais contenir. Restreindre le domaine libère toute la capacité pour la grammaire, la logique et la cohérence. → C'est exactement le régime des runs Forge actuels (vocab 4096, TinyStories) — le bon choix.
	35	+
	36	+### 1.2 phi — « Textbooks Are All You Need »
	37	+- phi-1 (1.3B) : 7B tokens seulement — code web filtré par classifieur + manuels synthétiques GPT-3.5 → 50.6% HumanEval. La version 350M atteint 45% avec le même pipeline : preuve directe que la qualité des données change la ligue d'un sous-1B.
	38	+- phi-1.5 (1.3B) : +20B tokens synthétiques « façon manuel scolaire » sur 20K sujets choisis → égale Llama-2-7B / Vicuna-13B sur le sens commun.
	39	+- Microsoft parle de régime « data-optimal » : à taille de modèle fixée, on optimise les données, pas le modèle.
	40	+
	41	+### 1.3 FineWeb-Edu / SmolLM2 — le filtrage par classifieur, l'astuce au meilleur ROI
	42	+- FineWeb-Edu : Llama-3-70B note 500K pages web (valeur éducative 0–5), un petit classifieur entraîné dessus filtre 15T → 1.3T tokens (seuil ≥3, >90% des tokens jetés) → +5 pts MMLU/ARC à compute égal.
	43	+- DCLM : un simple classifieur fastText gardant le top ~10% du web → +6 pts MMLU. Le filtrage par modèle (même minuscule) domine toutes les heuristiques.
	44	+- SmolLM2 (le playbook ouvert de référence) : mélange web optimal 60% FineWeb-Edu + 40% DCLM refiltré ; les modèles 135M/360M préfèrent un mélange uniforme de haute qualité (pas de curriculum), le 1.7B bénéficie d'un entraînement par étapes avec maths/code de haute qualité suréchantillonnés dans la phase de decay du LR.
	45	+- Constat clé pour les minuscules modèles : le contenu de niveau collège (simple mais correct) donne les meilleurs scores presque partout.
	46	+
	47	+### 1.4 Combien de tokens ? Chinchilla est mort pour les petits modèles
	48	+Le ratio 20 tokens/param est obsolète quand le coût d'inférence compte — on sur-entraîne massivement :
	49	+
	50	+\| Modèle \| Params \| Tokens \| Tok/param \|
	51	+\|---\|---\|---\|---\|
	52	+\| SmolLM2-135M \| 135M \| 2T \| ~15 000 \|
	53	+\| Qwen2.5-0.5B \| 0.5B \| 18T \| ~36 000 \|
	54	+\| Llama-3.2-1B \| 1B \| 9T (+ distillation) \| ~9 000 \|
	55	+
	56	+- La qualité continue de monter jusqu'à 10 000+ tokens/param.
	57	+- Contrainte de données : répéter un corpus curé de haute qualité est quasi gratuit jusqu'à ~4 epochs, rendements décroissants vers ~16, destructeur au-delà (Muennighoff, data-constrained scaling). → Recette : curer dur, puis répéter ≤4×, plutôt que diluer avec du médiocre.
	58	+
	59	+### 1.5 Données synthétiques (Cosmopedia) et curriculum
	60	+- Cosmopedia : 25B tokens générés par Mixtral, diversité assurée par graines web × (audience × style) — les deux leviers anti-effondrement : couverture des sujets par seeds et conditionnement d'audience (viser un jeune public = digestible par petit modèle).
	61	+- Curriculum fin = évidence mixte-à-négative (BabyLM : « largely unsuccessful »). Ce qui marche : la version grossière — mélanges par étapes avec suréchantillonnage maths/code/raisonnement dans la phase de decay (SmolLM2/3, phi).
	62	+
	63	+---
	64	+
	65	+## 2. Architecture & entraînement : la forme compte
	66	+
	67	+### 2.1 Deep-and-thin (MobileLLM, Meta ICML 2024)
	68	+Sous 1B, la profondeur bat la largeur :
	69	+
	70	+\| Modèle \| Couches \| d_model \| Zero-shot moyen \| vs baseline \|
	71	+\|---\|---\|---\|---\|---\|
	72	+\| MobileLLM-125M \| 30 \| 576 \| 46.3% \| OPT-125M : 42.6% \|
	73	+\| MobileLLM-350M \| 32 \| 960 \| 51.3% \| Pythia-410M : 46.6% \|
	74	+
	75	+Quatre choix empilés : SwiGLU + deep-and-thin + embeddings liés (à 125M, la table d'embedding ≈ 20% des params ; la lier libère du budget pour des couches) + GQA (ratio 3:1). Règle pratique 100–500M : d_model / n_layers ≈ 20–30 (GPT-2 était à ~60).
	76	+
	77	+→ Les configs Forge actuelles sont trop « larges » selon ce critère : gpt-50m est à 640/10 = 64 ; une variante MobileLLM-style serait ~20 couches × d_model 448.
	78	+
	79	+- Partage de blocs immédiat (MobileLLM-LS) : chaque bloc exécuté 2× de suite → +0.7–0.8% à zéro paramètre ajouté. Idéal sur Apple Silicon (mémoire unifiée : re-exécuter des poids chauds coûte peu, la RAM est la contrainte). C'est une feature de config (`share_pattern`), pas une refonte.
	80	+- Recursive Transformers relâchés (DeepMind 2024) : un bloc partagé bouclé K fois + LoRA par boucle — Gemma-1B récursif ≈ Gemma-2B (58.4 vs 58.6) après seulement 60B tokens de reprise.
	81	+
	82	+### 2.2 Distillation : le plus gros levier de qualité individuel
	83	+- Gemma 2 : entraîné sur les logits du teacher pendant TOUT le pretraining (au lieu des cibles one-hot) — la raison principale pour laquelle Gemma-2-2B joue en classe ~7B.
	84	+- Llama 3.2 1B : pruning structuré de Llama-3.1-8B + logits 8B/70B comme cibles token par token.
	85	+- Qwen3-0.6B : distillation off-policy puis on-policy (l'élève génère, KL contre le teacher sur ses propres échantillons). MiniLLM : KL inverse (l'élève n'étale pas sa masse sur les queues du teacher).
	86	+- → Pour Forge : stocker offline les top-16 logits d'un teacher (Qwen2.5-1.5B via l'endpoint MLX du cluster MacLustr) et ajouter un terme KL au kernel cross-entropy fusionné. Un seul kernel à modifier.
	87	+
	88	+### 2.3 Hyperparamètres et optimiseur
	89	+- µP (µTransfer) : paramétrisation rendant le LR optimal invariant à la largeur — on règle sur un proxy minuscule, on transfère. Petits modèles = gros LR (~3e-3–1e-2 µP-scaled à 100–350M, vs le folklore 3e-4).
	90	+- Schedule WSD (warmup–stable–decay, MiniCPM) : plateau plat + cooldown court (10–20%) ≥ cosine, prolongeable sans fixer max_steps à l'avance, chaque checkpoint du plateau est reprenable — parfait pour le design « resumable » de Forge. Bonus : c'est pendant le cooldown qu'on injecte les données les plus précieuses.
	91	+- Muon : orthogonalisation Newton–Schulz du momentum sur les matrices 2D (AdamW pour embeddings/normes) — tous les records NanoGPT-speedrun depuis oct. 2024, ~35% plus rapide vers la loss cible. Cœur = ~15 lignes, 5 itérations Newton–Schulz = quelques matmuls → trivial à exprimer avec le kernel matmul Metal existant.
	92	+- Batch critique ≈ 4M tokens ; 0.5–1M tokens/step est sûr à 100–350M (grad accum).
	93	+
	94	+### 2.4 Stabilité et petits plus architecturaux
	95	+- QK-norm (RMSNorm sur Q et K avant RoPE) : tue les spikes de loss — adopté par Gemma 3. Peu coûteux, recommandé.
	96	+- z-loss (~1e-4·log²Z) : se compose trivialement avec le kernel CE fusionné de Forge.
	97	+- Attention sinks (biais appris dans le dénominateur du softmax) : une ligne dans le kernel online-softmax.
	98	+- Tokenizer : sweet spot 32k–64k pour 100–500M (embeddings liés) ; SuperBPE (tokens multi-mots) : +4% moyen / +8.2% MMLU à compute fixe et 27–33% de tokens en moins.
	99	+- SSM/hybrides (Mamba, Zamba) sous 1B : sous-performants seuls sur le raisonnement court-contexte, complexité d'implémentation élevée — basse priorité.
	100	+
	101	+---
	102	+
	103	+## 3. Post-training, raisonnement, inférence : la vraie « révolution »
	104	+
	105	+### 3.1 Distillation de raisonnement — ça marche, avec un plafond
	106	+- R1-Distill-Qwen-1.5B : SFT sur ~800K traces générées par R1 → 83.9% MATH-500, devant GPT-4o. Pas de RL sur le petit modèle : DeepSeek montre que distiller un gros raisonneur bat le RL direct à petite échelle.
	107	+- Distilling Step-by-Step (2023) : un T5 770M fine-tuné sur des rationales extraites d'un LLM bat PaLM 540B few-shot.
	108	+- Le plafond (Small Model Learnability Gap, ACL 2025) : ≤3B, les modèles apprennent mieux de chaînes courtes et simples que des méga-traces o1-style. Remède : Mix Distillation (mélanger CoT longs + courts, ou teachers gros + moyens).
	109	+
	110	+### 3.2 Test-time compute — le multiplicateur le plus spectaculaire
	111	+- HuggingFace (recette DeepMind) : Llama-3.2-1B + beam search guidé par un PRM = Llama-3.1-8B sur MATH-500 ; le 3B dépasse le 70B (22× sa taille).
	112	+- arXiv 2502.06703 : 3B > 405B, 0.5B > GPT-4o, R1-Distill-1.5B + TTS (91.6) > o1-preview (85.5) — à 100–1000× moins de FLOPs d'inférence que les géants.
	113	+- Baseline gratuite : self-consistency (échantillonner k réponses CoT, vote majoritaire).
	114	+- Limite honnête : il faut un domaine vérifiable (maths/code) et un bon vérificateur — le PRM est la contrainte réelle.
	115	+
	116	+### 3.3 Récursion / profondeur adaptative — la capacité sans les paramètres
	117	+- Huginn-3.5B : bloc récurrent bouclé en espace latent au test — la perf de raisonnement monte avec les itérations jusqu'à l'équivalent d'un ~50B. Architecture indépendante de la taille.
	118	+- HRM 27M → Samsung TRM 7M : un réseau de 2 couches récursant ≤16 fois atteint ~45% ARC-AGI-1 (au-dessus des scores publiés de R1/o3-mini). Caveat : solveurs spécialisés par tâche, pas des LM généraux — mais la preuve que l'itération peut remplacer les paramètres sur le raisonnement structuré.
	119	+
	120	+### 3.4 Post-training à petite échelle (le playbook SmolLM2)
	121	+1. SFT sur ~1.1M instructions curées (SmolTalk), 2 epochs.
	122	+2. DPO sur UltraFeedback (le plus fiable à cette échelle).
	123	+3. GRPO/RLVR bref : sur R1-Distill-1.5B, 7K exemples, 24h, 42$ → AMC23 63→80%. Les gains arrivent dans les 50–100 premiers steps ; au-delà, déstabilisation. Ordre gagnant : distiller d'abord, RL ensuite (le RL échoue sur une base sans raisonnement préexistant).
	124	+4. Leçon générale : la curation des données > la nouveauté algorithmique sous 1B.
	125	+
	126	+### 3.5 Décodage — de la cohérence gratuite à l'inférence
	127	+- min-p sampling (ICLR 2025) : troncature dynamique selon la confiance du top-token — cohérence préservée même à haute température. Défaut solide (une réplication 2025 conteste l'ampleur, pas la direction).
	128	+- DoLa : contraster logits couches tardives vs précoces — +12–17 pts TruthfulQA sans entraînement.
	129	+- Pénalités de répétition : disproportionnellement importantes sous 1B (les petits modèles bouclent plus). → À ajouter au CLI `forge generate` : min-p + repetition penalty, quelques lignes.
	130	+- RAG : transforme la tâche de « tout savoir » en « lire et synthétiser » — ce que les sous-1B font bien (RETRO : 7.5B + retrieval ≈ GPT-3 175B).
	131	+- Débouché économique : les sous-1B sont les drafters idéaux du speculative decoding (2.3–6.5× de speedup pour les gros modèles).
	132	+
	133	+### 3.6 Ce à quoi les sous-1B servent vraiment (2025–2026)
	134	+Gemma 3 270M, Qwen3-0.6B, LFM2-350M : function calling, extraction structurée, résumé sur contexte fourni, classification/guardrails, autocomplete, draft models. Des tâches étroites, vérifiables, fine-tunées — pas du chat ouvert. (Gemma 270M fine-tuné : 58% → 85% sur une tâche mobile en quelques minutes de QLoRA.)
	135	+
	136	+---
	137	+
	138	+## 4. La recette Forge — feuille de route par phases
	139	+
	140	+Phase A — quick wins (configs seulement)
	141	+1. Config deep-and-thin : à 50M, viser ~20 couches × d_model 448 (ratio ~22) au lieu de 10 × 640 ; comparer sur le même train.bin.
	142	+2. Ratio GQA 3:1 (déjà supporté via `n_kv_heads`).
	143	+3. Sur-entraîner : viser des milliers de tokens/param ; répéter le corpus curé ≤4 epochs.
	144	+
	145	+Phase B — entraînement (petits changements de code)
	146	+4. Schedule WSD en plus de cosine dans `scheduler.h` (+ injection de données premium pendant le cooldown).
	147	+5. QK-norm + z-loss (kernel CE fusionné) pour la stabilité.
	148	+6. Muon pour les matrices 2D (Newton–Schulz = 5 matmuls, kernel existant), AdamW ailleurs.
	149	+7. Partage de blocs immédiat (`share_pattern` en config) : profondeur effective ×2 à RAM constante — taillé pour la mémoire unifiée.
	150	+
	151	+Phase C — données (le plus gros gain)
	152	+8. Pipeline FineWeb-Edu (subset) + Cosmopedia-style dans `prepare_data.py` ; ou rester TinyStories-style mais élargir progressivement le domaine (TinyStories-Instruct pour l'instruction-following à 50–200M).
	153	+9. Tokenizer 32k (lié) quand on sort de TinyStories ; envisager SuperBPE.
	154	+
	155	+Phase D — capacité (moyen terme)
	156	+10. Distillation : top-k logits d'un teacher servi par le cluster MacLustr (endpoint MLX), terme KL dans le kernel CE.
	157	+11. SFT + DPO à la SmolTalk sur le meilleur checkpoint.
	158	+12. `forge generate` : min-p, repetition penalty ; puis self-consistency (vote majoritaire) comme premier pas de test-time compute.
	159	+
	160	+---
	161	+
	162	+## 5. Sources principales
	163	+
	164	+Données : [TinyStories](https://arxiv.org/abs/2305.07759) · [phi-1](https://arxiv.org/abs/2306.11644) · [phi-1.5](https://arxiv.org/pdf/2309.05463) · [SmolLM2](https://arxiv.org/abs/2502.02737) · [FineWeb/-Edu](https://arxiv.org/html/2406.17557v1) · [DCLM](https://arxiv.org/abs/2406.11794) · [Data-constrained scaling](https://arxiv.org/pdf/2305.16264) · [Cosmopedia](https://github.com/huggingface/blog/blob/main/cosmopedia.md) · [BabyLM](https://arxiv.org/abs/2504.08165) · [Beyond Chinchilla-Optimal](https://arxiv.org/abs/2401.00448)
	165	+
	166	+Architecture/training : [MobileLLM](https://arxiv.org/abs/2402.14905) · [Recursive Transformers](https://arxiv.org/abs/2410.20672) · [MiniLLM](https://arxiv.org/abs/2306.08543) · [GKD](https://arxiv.org/html/2306.13649v1) · [Llama 3.2](https://ai.meta.com/blog/llama-3-2-connect-2024-vision-edge-mobile-devices/) · [Qwen3](https://arxiv.org/pdf/2505.09388) · [Gemma 3](https://arxiv.org/pdf/2503.19786) · [MiniCPM/WSD](https://arxiv.org/pdf/2404.06395) · [µP](https://arxiv.org/abs/2203.03466) · [Muon](https://kellerjordan.github.io/posts/muon/) · [Vocab scaling](https://arxiv.org/pdf/2407.13623) · [SuperBPE](https://arxiv.org/abs/2503.13423) · [Critical batch size](https://allenai.org/blog/critical-batch-size)
	167	+
	168	+Post-training/inférence : [R1-Distill](https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B) · [Distilling Step-by-Step](https://arxiv.org/abs/2305.02301) · [Learnability Gap](https://arxiv.org/abs/2502.12143) · [HF test-time compute](https://huggingfaceh4-blogpost-scaling-test-time-compute.hf.space/) · [1B > 405B?](https://huggingface.co/papers/2502.06703) · [RL small LLMs](https://arxiv.org/abs/2503.16219) · [RETRO](https://arxiv.org/abs/2112.04426) · [min-p](https://arxiv.org/abs/2407.01082) · [DoLa](https://arxiv.org/abs/2309.03883) · [Quiet-STaR](https://arxiv.org/abs/2403.09629) · [Huginn](https://arxiv.org/abs/2502.05171) · [Samsung TRM](https://arxiv.org/abs/2510.04871) · [Gemma 3 270M](https://developers.googleblog.com/en/introducing-gemma-3-270m/) · [LFM2](https://arxiv.org/abs/2511.23404)

added tools/prepare_hf_data.py +248 −0

		`@@ -0,0 +1,248 @@`
	1	+# Author: Simon-Pierre Boucher — contact@spboucher.ai
	2	+"""Stream any mix of Hugging Face pretraining datasets into Forge's
	3	+train.bin / val.bin format (uint16 tokens, llm.c-style header — same as
	4	+prepare_data.py, whose tokenizer + parallel encoder this reuses).
	5	+
	6	+Built for the sub-1B recipes in SMALL_MODELS_RESEARCH.md: classifier-filtered
	7	+web (FineWeb-Edu, DCLM), synthetic textbooks (Cosmopedia), math
	8	+(FineMath, OpenWebMath), reference text (Wikipedia), instructions (SmolTalk),
	9	+and weighted mixtures of them (e.g. SmolLM2's 60/40 FineWeb-Edu/DCLM).
	10	+Everything is STREAMED — no full-corpus downloads; you pay only for the
	11	+megabytes you keep.
	12	+
	13	+Usage:
	14	+ # single source
	15	+ python3 tools/prepare_hf_data.py --source fineweb-edu --out data/fwe \
	16	+ --max-train-mb 500 --vocab-size 16384
	17	+
	18	+ # weighted mixture (name:weight, comma-separated)
	19	+ python3 tools/prepare_hf_data.py --mix fineweb-edu:0.6,dclm:0.4 \
	20	+ --out data/web-mix --max-train-mb 2000
	21	+
	22	+ # research-backed presets
	23	+ python3 tools/prepare_hf_data.py --preset smollm-web --out data/smollm-web
	24	+ python3 tools/prepare_hf_data.py --list # show sources & presets
	25	+
	26	+Requires: pip install datasets (numpy already required by prepare_data.py)
	27	+"""
	28	+import argparse
	29	+import os
	30	+import random
	31	+import sys
	32	+
	33	+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
	34	+import prepare_data as pd # tokenizer training + parallel BPE encoder + .bin writer
	35	+
	36	+SEP = "<\|endoftext\|>" # document separator; same convention as TinyStories
	37	+
	38	+
	39	+def messages_to_text(row):
	40	+ """Flatten a chat-format row (SmolTalk-style `messages`) to plain text."""
	41	+ parts = []
	42	+ for m in row.get("messages", []):
	43	+ role = m.get("role", "user")
	44	+ parts.append(f"<\|{role}\|>\n{m.get('content', '')}")
	45	+ return "\n".join(parts)
	46	+
	47	+
	48	+# name -> (repo, config, split, extractor). All public, all stream-capable.
	49	+SOURCES = {
	50	+ # -- narrow-domain / tiny-model corpora ------------------------------------
	51	+ "tinystories": ("roneneldan/TinyStories", None, "train",
	52	+ lambda r: r["text"]),
	53	+ # -- classifier-filtered web (the highest-leverage curation known) ---------
	54	+ "fineweb-edu": ("HuggingFaceFW/fineweb-edu", "sample-10BT", "train",
	55	+ lambda r: r["text"]),
	56	+ "fineweb": ("HuggingFaceFW/fineweb", "sample-10BT", "train",
	57	+ lambda r: r["text"]),
	58	+ "fineweb-edu-dedup":("HuggingFaceTB/smollm-corpus", "fineweb-edu-dedup", "train",
	59	+ lambda r: r["text"]),
	60	+ "dclm": ("mlfoundations/dclm-baseline-1.0", None, "train",
	61	+ lambda r: r["text"]),
	62	+ # -- synthetic textbooks (phi/Cosmopedia recipe) ----------------------------
	63	+ "cosmopedia": ("HuggingFaceTB/smollm-corpus", "cosmopedia-v2", "train",
	64	+ lambda r: r["text"]),
	65	+ # -- math (upsample during LR decay per SmolLM2/3) --------------------------
	66	+ "finemath": ("HuggingFaceTB/finemath", "finemath-4plus", "train",
	67	+ lambda r: r["text"]),
	68	+ "openwebmath": ("open-web-math/open-web-math", None, "train",
	69	+ lambda r: r["text"]),
	70	+ # -- reference / encyclopedic ----------------------------------------------
	71	+ "wikipedia-en": ("wikimedia/wikipedia", "20231101.en", "train",
	72	+ lambda r: r["text"]),
	73	+ "wikipedia-fr": ("wikimedia/wikipedia", "20231101.fr", "train",
	74	+ lambda r: r["text"]),
	75	+ # -- generic web baseline ---------------------------------------------------
	76	+ "c4": ("allenai/c4", "en", "train",
	77	+ lambda r: r["text"]),
	78	+ # -- instructions (SFT-style text; useful for a decay-phase blend) ----------
	79	+ "smoltalk": ("HuggingFaceTB/smoltalk", "all", "train",
	80	+ messages_to_text),
	81	+}
	82	+
	83	+# Research-backed mixtures (SMALL_MODELS_RESEARCH.md §1). Weights are document
	84	+# sampling probabilities over the interleaved stream.
	85	+PRESETS = {
	86	+ # SmolLM2's best web mix: 60% FineWeb-Edu + 40% DCLM
	87	+ "smollm-web": {"fineweb-edu": 0.6, "dclm": 0.4},
	88	+ # phi-style: filtered web + synthetic textbooks
	89	+ "textbooks": {"fineweb-edu": 0.6, "cosmopedia": 0.4},
	90	+ # single-stage high-quality blend for <500M models (uniform, no staging)
	91	+ "smol-full": {"fineweb-edu": 0.55, "cosmopedia": 0.25,
	92	+ "finemath": 0.10, "openwebmath": 0.10},
	93	+ # LR-decay-phase blend: premium math/instructions upsampled
	94	+ "decay-anneal": {"fineweb-edu": 0.4, "finemath": 0.3, "smoltalk": 0.3},
	95	+ # bilingual reference blend
	96	+ "wiki-bilingue":{"wikipedia-en": 0.5, "wikipedia-fr": 0.5},
	97	+}
	98	+
	99	+
	100	+def open_stream(name, seed, shuffle_buffer):
	101	+ try:
	102	+ from datasets import load_dataset
	103	+ except ImportError:
	104	+ sys.exit("prepare_hf_data: pip install datasets")
	105	+ repo, config, split, extract = SOURCES[name]
	106	+ ds = load_dataset(repo, config, split=split, streaming=True)
	107	+ if shuffle_buffer > 0:
	108	+ ds = ds.shuffle(seed=seed, buffer_size=shuffle_buffer)
	109	+ it = iter(ds)
	110	+
	111	+ def docs():
	112	+ for row in it:
	113	+ text = extract(row)
	114	+ if text and len(text) >= 64: # drop empty/near-empty docs
	115	+ yield text
	116	+
	117	+ return docs()
	118	+
	119	+
	120	+def stream_mixture(weights, seed, shuffle_buffer):
	121	+ """Yield documents, choosing the source of each by weight. A source that
	122	+ runs dry is dropped and the rest renormalize (streamed sets are huge —
	123	+ this mostly matters for small ones like tinystories)."""
	124	+ rng = random.Random(seed)
	125	+ streams = {n: open_stream(n, seed + i, shuffle_buffer)
	126	+ for i, n in enumerate(weights)}
	127	+ names = list(weights)
	128	+ while names:
	129	+ r = rng.random() * sum(weights[n] for n in names)
	130	+ acc = 0.0
	131	+ pick = names[-1]
	132	+ for n in names:
	133	+ acc += weights[n]
	134	+ if r <= acc:
	135	+ pick = n
	136	+ break
	137	+ try:
	138	+ yield next(streams[pick])
	139	+ except StopIteration:
	140	+ print(f" source exhausted: {pick}")
	141	+ names.remove(pick)
	142	+
	143	+
	144	+def write_corpus(docs, train_path, val_path, train_mb, val_mb, val_every):
	145	+ """Split the stream into train/val text files (every Nth doc to val until
	146	+ its budget fills), separated by <\|endoftext\|>."""
	147	+ train_budget = int(train_mb * 1024 * 1024)
	148	+ val_budget = int(val_mb * 1024 * 1024)
	149	+ train_n = val_n = 0
	150	+ report_every = max(1, train_budget // (20 * 1024 * 1024) or 1) * 10_000_000
	151	+ next_report = report_every
	152	+ with open(train_path, "w") as ftr, open(val_path, "w") as fva:
	153	+ for i, doc in enumerate(docs):
	154	+ piece = doc.rstrip() + "\n" + SEP + "\n"
	155	+ if val_n < val_budget and i % val_every == val_every - 1:
	156	+ fva.write(piece)
	157	+ val_n += len(piece)
	158	+ else:
	159	+ ftr.write(piece)
	160	+ train_n += len(piece)
	161	+ if train_n >= next_report:
	162	+ print(f" {train_n / 1e6:.0f} MB train / {val_n / 1e6:.1f} MB val")
	163	+ next_report += report_every
	164	+ if train_n >= train_budget and val_n >= val_budget:
	165	+ break
	166	+ print(f"corpus: {train_n / 1e6:.1f} MB train, {val_n / 1e6:.1f} MB val")
	167	+
	168	+
	169	+def main():
	170	+ ap = argparse.ArgumentParser(
	171	+ description="Stream HF datasets (single, mix, or preset) into Forge .bin files")
	172	+ g = ap.add_mutually_exclusive_group()
	173	+ g.add_argument("--source", choices=sorted(SOURCES), help="single dataset")
	174	+ g.add_argument("--mix", help="name:weight,name:weight,…")
	175	+ g.add_argument("--preset", choices=sorted(PRESETS), help="research-backed mixture")
	176	+ g.add_argument("--list", action="store_true", help="list sources & presets")
	177	+ ap.add_argument("--out", help="output dir (train.bin/val.bin/corpus)")
	178	+ ap.add_argument("--vocab-size", type=int, default=16384)
	179	+ ap.add_argument("--tokenizer", default=None,
	180	+ help="existing forgebpe .model to reuse (e.g. data/tinystories/"
	181	+ "tok4096.model); otherwise trained on this corpus")
	182	+ ap.add_argument("--tokenizer-mb", type=float, default=20.0)
	183	+ ap.add_argument("--max-train-mb", type=float, default=500.0)
	184	+ ap.add_argument("--max-val-mb", type=float, default=5.0)
	185	+ ap.add_argument("--val-every", type=int, default=200,
	186	+ help="route every Nth document to val until its budget fills")
	187	+ ap.add_argument("--seed", type=int, default=1337)
	188	+ ap.add_argument("--shuffle-buffer", type=int, default=10000,
	189	+ help="streaming shuffle buffer per source (0 disables)")
	190	+ ap.add_argument("--keep-text", action="store_true",
	191	+ help="keep the intermediate train.txt/val.txt corpus files")
	192	+ args = ap.parse_args()
	193	+
	194	+ if args.list or not (args.source or args.mix or args.preset):
	195	+ print("sources:")
	196	+ for n, (repo, cfg, _, _) in sorted(SOURCES.items()):
	197	+ print(f" {n:<18} {repo}" + (f" [{cfg}]" if cfg else ""))
	198	+ print("presets:")
	199	+ for n, w in PRESETS.items():
	200	+ print(f" {n:<18} " + ", ".join(f"{k}:{v}" for k, v in w.items()))
	201	+ return
	202	+ if not args.out:
	203	+ ap.error("--out is required")
	204	+
	205	+ if args.source:
	206	+ weights = {args.source: 1.0}
	207	+ elif args.preset:
	208	+ weights = PRESETS[args.preset]
	209	+ else:
	210	+ weights = {}
	211	+ for part in args.mix.split(","):
	212	+ name, _, w = part.partition(":")
	213	+ if name not in SOURCES:
	214	+ ap.error(f"unknown source '{name}' (see --list)")
	215	+ weights[name] = float(w) if w else 1.0
	216	+
	217	+ os.makedirs(args.out, exist_ok=True)
	218	+ print("mixture: " + ", ".join(f"{n}:{w}" for n, w in weights.items()))
	219	+
	220	+ train_txt = os.path.join(args.out, "train.txt")
	221	+ val_txt = os.path.join(args.out, "val.txt")
	222	+ docs = stream_mixture(weights, args.seed, args.shuffle_buffer)
	223	+ write_corpus(docs, train_txt, val_txt, args.max_train_mb, args.max_val_mb,
	224	+ args.val_every)
	225	+
	226	+ model = args.tokenizer or os.path.join(args.out, f"tok{args.vocab_size}.model")
	227	+ if not os.path.exists(model):
	228	+ import train_tokenizer as tt
	229	+ with open(train_txt, "rb") as f:
	230	+ sample = f.read(int(args.tokenizer_mb * 1024 * 1024))
	231	+ print(f"training {args.vocab_size}-vocab BPE on {len(sample) / 1e6:.1f} MB")
	232	+ tt.write_model(model, tt.train(sample, args.vocab_size))
	233	+ print(f"wrote {model}")
	234	+
	235	+ pd.tokenize_file(train_txt, model, os.path.join(args.out, "train.bin"),
	236	+ args.max_train_mb + 1)
	237	+ pd.tokenize_file(val_txt, model, os.path.join(args.out, "val.bin"),
	238	+ args.max_val_mb + 1)
	239	+
	240	+ if not args.keep_text:
	241	+ os.remove(train_txt)
	242	+ os.remove(val_txt)
	243	+ print("done — point `forge train --data` at", args.out,
	244	+ f"(set model vocab_size to match the tokenizer: {model})")
	245	+
	246	+
	247	+if __name__ == "__main__":
	248	+ main()
	249