diff --git a/docs/runs/01-v1-tinystories-dim256.md b/docs/runs/01-v1-tinystories-dim256.md new file mode 100644 index 0000000..d189f7e --- /dev/null +++ b/docs/runs/01-v1-tinystories-dim256.md @@ -0,0 +1,179 @@ +# Scaling Run v1: TinyStories 全量 + dim256/8L — Design Document + +## Goal + +在 v0-baseline（dim32/4L、core ~41K 参、只喂 TinyStories valid 的 3MB 切片）之上，做第一次**有意义 +的放大**： + +1. **数据放大**：从 3MB 切片 → **TinyStories 全量 train**（468.3M tokens），并解决「2GB 语料每次重新 + 跑 from-scratch BPE 太慢」——tokenize **一次**、把 token-id 流缓存到盘，后续直接读缓存。 +2. **模型放大**：dim 32→256、层 4→8、头 2→8，把 **transformer core 做到 ~8M 参**（embedding+lm_head + 因 gpt2 50257 vocab 固定再加 ~25.7M，属预期，单列出来）。 +3. **参数化阶梯**：把模型尺寸从硬编码改成 CLI 可调（`--heads/--head-dim/--layers/--ffn`），让 v2/v3 + 只改 flag 即可，不再动代码。 +4. 训完存 registry（`~/projects/tiny-models/v1-tinystories-dim256/`）+ 导出 xserv 格式验证可服务，并给出 + **相比 v0 的具体提升**（同一保留集 val loss + 同 prompt 并排采样）。 + +> 范围（escape hatch 已评估）：全量 2GB 下载 + 全量 tokenize 实测都很快（下载即时、tokenize ~75s 出 +> 468M token 缓存），所以 v1 **用了全量语料**。训练在 bounded 预算内做 2500 步（≈5.1M token），不追求把 +> 全量跑满一个 epoch——v1 的目的是「相对 v0 的清晰、可量化提升 + 参数化阶梯 + 设计文档」，不是榨满模型。 + +## 数据 + +| 项 | v0-baseline | v1 | +|----|-------------|----| +| 来源 | TinyStories **valid** 的 3MB 字节切片 | TinyStories **全量 train**（hf-mirror.com）| +| 原始大小 | 3 MB | 1.92 GB（`TinyStories-train.txt`，Content-Length 1924281556）| +| token 数 | ~72 万 | **468,260,367**（≈468M）| +| tokenizer | 复用 xserv from-scratch GPT-2 BPE（vocab 50257）| 同 | +| 缓存 | 无（每次重 tokenize）| **`.u16.bin`**：468M token 的 u16 流（936MB），首跑 tokenize 一次写盘，后续直接读 | +| 验证集 | 无独立 val 切片 | 全量末尾保留 **1,000,000 token** 作 held-out val（训练不触及）| + +**下载**：`curl -sL https://hf-mirror.com/datasets/roneneldan/TinyStories/resolve/main/TinyStories-train.txt` +（hf-mirror 302 跳到 xethub CDN，直连可下；HF 直连被墙）。 + +**缓存设计（`crates/xtrain-train/src/data.rs`）**：gpt2 vocab=50257 < 65536，token id 用 **u16** 无损存储。 +`Corpus::load_cached` 首跑 tokenize 整个语料并写 `.u16.bin`（纯 little-endian `[u16]`，无 header，按 +路径为 key），后续 run 直接读缓存跳过 BPE。实测：全量 1.92GB tokenize 一次 ~75s；之后每次 run 读缓存 +**即时**。这把「2GB × 每次 BPE」的反复开销摊成一次。 + +**为什么比 v0 更高质/更大**：v0 只喂了 valid 集的一个 3MB 字节切片（~72 万 token，且是 byte-range 抓的、 +首尾残story），覆盖的故事极少、词汇/句式重复度高 → 模型只能记住极少数模板（采样里反复 "mommy's mommy's +mommy"）。v1 用全量 train（468M token、数百万个完整小故事），故事/句式/词汇覆盖面大几个数量级，随机窗口 +采样能见到远更丰富的语言结构。 + +## 架构 + +v1 = 一个更大的、同构的 tiny Qwen3（RoPE + RMSNorm + per-head QK-norm + SwiGLU + 独立 lm_head，MHA）， +forward 图与 v0 完全同构，只是 dims 变大。 + +| 维度 | v0-baseline | v1 | +|------|-------------|----| +| dim（= heads·head_dim）| 32 | **256** | +| n_layers | 4 | **8** | +| n_heads | 2 | **8** | +| head_dim | 16 | **32** | +| ffn_hidden（SwiGLU）| 64 | **1024** | +| vocab | 50257 | 50257 | +| **core 参数**（除 embed+lm_head）| **41,376** | **8,393,472**（≈8.39M）| +| embed + lm_head（2×vocab×dim）| 3,216,448 | 25,731,584（≈25.7M）| +| **总参数** | 3,257,824 | **34,125,056**（≈34.13M）| + +**core 的量法**：`Config::core_params() = num_params() − 2·vocab·dim`。gpt2 50257 vocab 在 dim256 下让 +embedding + lm_head 固定占 ~25.7M——这两张表是**词表大小**的函数，不是模型容量，所以阶梯按 **core** 量 +（v1 core 8.39M 命中 ~8M 目标）。这也是为什么 v1 总参 34M「看着大」但有效容量是 8.39M core。 + +**相比 v0 的架构变化**：纯放大，无结构改动（QK-norm/RoPE/SwiGLU/MHA 都在 v0 就有，T9 已对齐 xserv）。 +唯一工程改动是把尺寸**参数化**（见下）。 + +### 参数化阶梯（实现） + +`Config::from_arch(vocab, n_heads, head_dim, n_layers, ffn)` 派生 `dim = n_heads·head_dim`； +`bin/train` 与 `bin/export_safetensors` 都从 CLI flag 读架构（`--heads/--head-dim/--layers/--ffn`）， +默认值复刻 v0 tiny config。v2/v3 只改 flag： + +```sh +# v1 +./train tokenizer.json data/tinystories-train.txt \ + --heads 8 --head-dim 32 --layers 8 --ffn 1024 \ + --steps 2500 --batch 16 --seq 128 --max-lr 6e-4 --min-lr 6e-5 \ + --val-tokens 1000000 --eval-every 250 --ckpt /tmp/xtrain_v1.ckpt +``` + +## 超参 + +| 项 | 值 | 备注 | +|----|----|----| +| optimizer | 手写 AdamW（GPU 端 step）| wd=0.1，β/eps 用 xtrain-optim 默认 | +| LR schedule | 线性 warmup → cosine decay | max_lr **6e-4** → min_lr **6e-5** | +| warmup | steps/20 = 125 步 | | +| grad clip | global-norm 1.0 | | +| steps | **2500** | bounded（≈25 min 单卡）| +| batch | **16** | 单序列模型，靠多次 forward 让 tape SUM 梯度，clip 时 ×1/batch 取均值 | +| seq_len | **128** | v0 是 64 | +| tokens/step | 16×128 = 2048 | 总训练 token ≈ 5.12M | +| 精度 | f32（训练）| 导出 xserv 时转 BF16（见 T9）| + +**算力**：dash5 单卡 RTX 5090（GPU 1，sm_120），吞吐 ≈ **3.3K tok/s**（单序列设计 GPU 利用率 ~25-29%， +是已知瓶颈，见 docs/06）；wall-clock ≈ **25.9 min**（1551s, EXIT=0）。DDP 多卡路径存在（T8，~1.87x@2）， +v1 单卡已足以清晰超过 v0，未启用——留作 v2 提速杠杆。 + +## 结果 + +- **train loss**：start 10.8590 → end 2.6247 +- **best val loss（held-out 1M token）**：**2.5847**（step 2499） +- val loss 曲线（每 250 步）： + +| step | 249 | 499 | 749 | 999 | 1249 | 1499 | 1749 | 1999 | 2249 | 2499 | +|------|----|----|----|----|------|------|------|------|------|------| +| val | 3.8609 | 3.3534 | 3.1114 | 2.9702 | 2.8498 | 2.7643 | 2.7046 | 2.6496 | 2.6124 | **2.5847** | + +单调下降、未见过拟合（val 一路降到末步），说明 2500 步仍欠拟合——更多步数/数据还能继续降（v2 杠杆）。 + +### 采样（greedy，xtrain 直采，同 prompt） + +``` +[Once upon a time] → Once upon a time, there was a little girl named Lily. She loved to play + outside in the sunshine. One day, she saw a big, scary dog. The dog was + scared and didn't know what to do +[The little] → The little girl was so happy that she had been able to help. + <|endoftext|> Once upon a time there was a little girl named Lucy. She was + three years old and loved to explore. One day, +[One day] → One day, she saw a big, shiny ball in the park. She wanted to play with it, + but she was too scared to go. She went to the park and saw a big, scary dog +``` + +温度 0.8 采样同样连贯（多角色、完整情节），见 `RUN.md`。 + +## 相比 v0 的提升 + +**同一保留集（v1 train 末尾 1M token）上的 val loss**——`bin/train --eval-ckpt` 加载各自 checkpoint、 +在**同一 held-out 1M token** 上算 cross-entropy，把两个模型放到同一指标（公平对比）： + +| 模型 | core 参数 | 训练数据 | **val loss（同一 1M held-out）** | +|------|-----------|----------|------------------------------| +| v0-baseline | 41K | 3MB 切片（~72万 tok）| **3.8050** | +| v1 | 8.39M（**×203**）| 全量 468M（**×650**）| **2.5847**（**低 1.22**）| + +### 并排采样（greedy 40 tok，xserv 服务，同 prompt） + +| prompt | v0-baseline | v1 | +|--------|-------------|----| +| `Once upon a time` | …a little girl named Lily. **Timmy** loved to play with her mommy. One day, **Timmy's mommy's mommy's mommy**. "I'm sorry, I | …a little girl named Lily. She loved to play outside in the sunshine. One day, she saw a big, scary dog. The dog was scared and didn't know what to do | +| `The little` | The little girl named Lily. She loved to play with her mommy. One day, **Timmy's mommy's mommy's mommy**. "I'm sorry, I can't have a good time | The little girl was so happy that she had been able to help. | +| `One day` | One day, **Timmy's mommy's mommy's mommy**. "I'm sorry, I can't be careful and be careful. I'm sorry, I can't have a good time. | One day, she saw a big, shiny ball in the park. She wanted to play with it, but she was too scared to go. | + +**结论**：v0（41K core / 3MB 数据）只学到极少的模板，主语/指代崩坏（Lily↔Timmy 混用）、立刻陷入 +**"mommy's mommy's mommy"** 退化循环——它记住的是少数 n-gram，没有连贯的故事建模。v1（8.39M core / +全量数据）能稳定保持单一主角、写出有场景（sunshine/park）、有情节（saw a dog → scared）、跨句一致的 +完整小故事，并正确生成 `<|endoftext|>` 分隔下一篇——这正是 TinyStories 想要 tiny 模型学会的东西。 +**val loss 低 1.22 + 采样从"循环复读"到"连贯叙事"**，v1 是相对 v0 的清晰、可量化提升。 + +## xserv 验证 + +导出 HF Qwen3 safetensors（命名映射 + 2D 权重转置 [in,out]→[out,in] + BF16，见 T9 `docs/08`，91 tensors）， +存入 registry 后用 `xserv-cli` 加载并贪心生成——**逐 token 对住 xtrain 自身的贪心生成**（闭环在 v1 规模仍成立）： + +``` +$ xserv-cli ~/projects/tiny-models/v1-tinystories-dim256 --max-tokens 40 +Model: qwen3, layers=8, hidden=256, heads=8/8 kv, vocab=50257 +Loaded 91 tensors +xserv> Once upon a time, there was a little girl named Lily. She loved to play outside in the + sunshine. One day, she saw a big, scary dog. The dog was scared and didn't know what to do +xserv> The little girl was so happy that she had been able to help. +xserv> One day, she saw a big, shiny ball in the park. She wanted to play with it, but she was + too scared to go. +``` + +## v2 提案 + +v1 的 val 曲线一路单调下到末步（无过拟合）= **欠拟合**，说明同规模再多喂步数/数据还能降。建议 v2 沿两个轴同时拉： + +- **数据/步数**：把训练 token 从 ~5M 拉到 ~50-100M（DDP 2-4 卡把 wall-clock 压回 ~30min；T8 路径已就绪， + 只需把 `train_ddp` 也接上参数化 config + cache + best-val checkpoint）。 +- **模型**：dim 384 / 12 heads·32 / 12 layers / ffn 1536 → core ≈ **27M**（仍是 tiny，但容量翻 ~3x）。 + 词表不变 → embed+lm_head 仍 ~38.6M，总 ~66M。 + +阶梯已参数化，v2 只改 `--dim/--heads/--layers/--ffn/--steps` flag + 加 DDP 启动，不动模型代码。 +预期 val loss 进一步明显下降（目标 < 2.2），采样在更长上下文/更复杂情节上更稳。 + diff --git a/docs/runs/README.md b/docs/runs/README.md new file mode 100644 index 0000000..d8ba7ef --- /dev/null +++ b/docs/runs/README.md @@ -0,0 +1,25 @@ +# Scaling Runs + +xtrain 的 scaling 阶段：在 v0-baseline 之上逐版放大**数据 + 参数**，每版一份 +`docs/runs/NN-.md` 设计文档（数据来源 / 架构 + 参数 / 超参 / 结果 val-loss + 采样 / +相比上一版的提升），训练完存入 dash5 模型 registry（`~/projects/tiny-models//`）并导出 +xserv 格式验证可服务。 + +模型核心参数（`core params`）= `Config::core_params()` = 总参数减去两张 `vocab×dim` 表 +（token embedding + lm_head）。gpt2 vocab=50257 使这两张表固定占 ~25.7M（dim256 时），它**不反映 +模型容量**，所以阶梯按 core 来量。 + +## 对比表 + +val loss 一栏给的是**同一 held-out 1M token**（v1 train 末尾切片）上、用 `bin/train --eval-ckpt` +对两个 checkpoint 各自评出来的——同一指标、公平对比。 + +| 版本 | 数据 | 架构 (dim/L/heads·hd/ffn) | core 参数 | 总参数 | val loss | 备注 | +|---|---|---|---|---|---|---| +| [v0-baseline](../../docs/05-training-loop.md) | TinyStories valid 3MB 切片 (~72 万 tok) | 32 / 4 / 2·16 / 64 | ~41K | 3.26M | **3.8050** | 太小不可用；采样陷入 "mommy's mommy's mommy" 循环 | +| [v1-tinystories-dim256](01-v1-tinystories-dim256.md) | TinyStories **全量 train** (468.3M tok, u16 缓存) | 256 / 8 / 8·32 / 1024 | 8.39M | 34.13M | **2.5847** | 全量数据 + dim256/8L；val 低 1.22，采样连贯成篇；~25.9min/单卡 | + +## 下一档（提案） + +- **v2**（待派发）：见 `01-v1-*.md` 末尾 "v2 提案"。 +