# 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）| **`<corpus>.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 整个语料并写 `<path>.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），采样在更长上下文/更复杂情节上更稳。
</content>