Gahow Wang
2ff4573a31
Backfill docs/14-gqa.md gate table (dash5 numbers); add T15 evolution row + cumulative 模型架构 line; README build-journey T15 row + Phase 2 prose + doc index range (00..14). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
71 lines
14 KiB
Markdown
71 lines
14 KiB
Markdown
# xtrain 演进总览 — 按维度记录每次变化
|
||
|
||
每个里程碑（**T# 基建 phase** 或 **v# 训练 run**）在四个维度上分别改了什么、结果如何。
|
||
这是活文档：**每次新 run 收尾时追加一行/一段**。细节见各 `docs/runs/0N-*.md`、各 phase 设计文档、`docs/known-issues.md`。
|
||
|
||
四个维度：**算法**（autograd/优化器/精度/反向数学）· **模型架构**（dim/层/头/算子）· **Infra**（构建/显存/并行/吞吐）· **数据集**（语料/token/epoch/tokenizer）。
|
||
|
||
---
|
||
|
||
## 一、基建 phase（T1–T13 + Phase 2 systems-depth）—— 主要动「算法」与「Infra」
|
||
|
||
| Phase | 维度 | 变化 | 结果 / 验证 |
|
||
|---|---|---|---|
|
||
| T1 | Infra | Rust↔CUDA FFI 构建链（build.rs+nvcc, `no_cuda` cfg），gitea↔dash5 流 | vector-add 跑通 |
|
||
| T2 | Infra | Tensor 抽象（dtype/shape/Storage, H↔D 拷贝）+ elementwise kernel | roundtrip 保真 |
|
||
| T3 | 算法 | 手写 tiled GEMM fwd/bwd + finite-diff 梯度检查 harness | fwd vs cuBLAS 1e-7；bwd vs finite-diff |
|
||
| T4 | 算法 | tape autograd 引擎 + 11 算子 backward（含梯度扇出累加）；attention 由 matmul+softmax 组合 | 每算子 finite-diff ≤2e-2 |
|
||
| T5 | 模型架构 | 组装 tiny decoder（RoPE+RMSNorm+SwiGLU）+ embedding/reshape/transpose 算子 | overfit 27/27 + PyTorch 对拍 B>1 |
|
||
| T6 | 算法 + 数据集 | 手写 AdamW + 训练 loop + LR sched + grad clip + checkpoint；gpt2 BPE + TinyStories | 真训出连贯英文 |
|
||
| T7 | Infra | cuBLAS matmul + GPU 端 AdamW/grad-norm + 去 per-op sync | **~3×**：2.7K→8.5K tok/s，零回归 |
|
||
| T8 | Infra | NCCL DDP（单进程 thread-per-GPU）+ 梯度 all-reduce | 多卡（当时弱扩展 ~1.4×） |
|
||
| T9 | 算法/模型架构 + Infra | **加 per-head QK-norm**（Qwen3 兼容）；safetensors 导出 | xserv 闭环：贪心**逐 token 一致** |
|
||
| T10 | 算法/Infra | **batched 多序列 forward**（linears flatten `[B·S,dim]` + fused batched SDPA + 每序列 RoPE） | **单卡 15–24×**；MFU 0.4%→14%（修 KI-1） |
|
||
| T11 | Infra | **device caching/pool allocator**（复用 op 输出显存，消 per-step cudaMalloc） | 单卡 2.3×；**8卡 461K tok/s** 近线性（修 KI-5） |
|
||
| T12 | 算法/Infra | **bf16 混合精度**（fp32 master，cuBLAS GemmEx，norm/softmax/CE 保 fp32） | dim768 OOM 解除，−29% 显存/+13% tok/s（修 KI-2） |
|
||
| T13 | 算法/Infra | **激活重计算**（per-block gradient checkpointing：前向 no-tape + 反向重算，`backward_seeded`） | 梯度对非重计算版**逐位一致**(0.00)；dim768 31.1→14.6GB；**dim1024 batch32 OOM→16.6GB 装下**（修 KI-3，解锁 v8） |
|
||
| T14 | 算法/Infra | **融合 flash-attention kernel**（手写单 kernel：online softmax、tiled over KV、**不物化 N×N scores**；flash 式 bwd：重算 scores + `D=ΣdO·O` 化简雅可比 + dQ/dK/dV）；opt-in `--flash`，默认保 composed（Phase 2） | fwd 对 composed 6.7e-5、bwd 对 composed dQ 1.7e-5、PyTorch B>1 7.9e-6、flash==composed loss rel 0.0；**峰值显存 −16%@seq1024 / −23%@seq2048**（不物化 N×N，收益随 seq 增长）；tok/s ~2.3–2.8× 慢（hd=64 小头维干不过 cuBLAS tensor-core，flash 已知权衡=胜场在显存）；md5 闭环逐位一致 |
|
||
| T15 | 模型架构 | **真 GQA**（`num_kv_heads<num_heads`：wk/wv 投影到 `kv_dim`，新 `repeat_kv` broadcast 算子把 K/V 复制 `group=nh/num_kv` 份喂给**未改动**的 composed/flash 两条 SDPA；分组约定对齐 xserv repeat_kv `dst=kvh·group+r`）；`repeat_kv` 反向=组内 group 行**确定性求和**（无 atomic）→ 多组 q 头梯度汇一个 kv 头；`num_kv_heads` 进 Config(默认=nh→MHA)、`--kv-heads` flag、导出写真 `num_key_value_heads`（Phase 2） | repeat_kv grad-check 2.1e-4(group3)+group1 identity 逐位；GQA flash==composed fp32 grad 4.1e-5/bf16 在带；**group1 对 MHA 逐位一致**(回归保护)；PyTorch GQA B>1 对拍 composed/flash 各 loss 1.7e-8/logits 2.3e-5/25 grad 进 rtol；小 GQA(8h/2kv) 训 600 步 10.9→3.15 连贯；**xserv 闭环真 GQA**(num_kv 2<8)：2/3 prompt token-identical、1 在 BF16 漂移处晚分叉；MHA 默认 export md5 逐位一致(b04fc9f9) |
|
||
| T16 | 算法/Infra | **梯度累积**（N 个 micro-step：每个 micro-loss `×1/N` 再 backward，tape SUM 累加 → 一次 AdamW step+zero；`--accum-steps`）；**DDP 只在累积边界 all-reduce**（中间 micro-step 不发 NCCL，`/world` 与 `1/N` 正交）；显存随 micro 不随有效 batch | 等效大 batch**逐位贴合**（loss rel 8.5e-8、grad rel 3.8e-5）；`accum=1` 逐位回归(0.00)；DDP+accum 对单卡 loss 5.7e-7/跨 rank 一致；**显存平**：同有效 batch 64，big-batch 27.7GB→accum(4×16) **7.2GB(−74%)**（big-batch OOM 而 accum 装下）；全回归+xserv 闭环 md5 一致 |
|
||
| T18 | 算法 | **dropout**（手写 counter-based 设备 RNG → Bernoulli mask，训练 inverted 1/(1-p) scaling、eval 恒等）；新 autodiff `dropout` 算子（fwd 生成+施加 mask，bwd 用同 mask），接 residual/ffn 两处；`--dropout` flag 默认 0 | 固定 seed grad-check 过；E[out]≈input + keep≈1-p；**p=0 与无 dropout 逐位一致**；recompute(T13) 组合下梯度仍逐位一致（counter-based seed 重算复现同 mask）；全回归 + xserv 闭环绿（导出/推理 dropout 关） |
|
||
|
||
---
|
||
|
||
## 二、Scaling runs（v0–v8）—— 主要动「模型架构」与「数据集」
|
||
|
||
架构始终是 **Qwen3-style**（RoPE + RMSNorm + QK-norm + SwiGLU，gpt2 50257 词表），逐版放大 dim/层/头（v8 起首次拨容量轴到 dim1024）；其余维度逐版变化如下：
|
||
|
||
| ver | 模型架构（dim/层/头·hd · 核心/总参） | 数据集（语料 · 实训 token · epoch） | 算法/精度 | Infra（GPU · 吞吐） | 结果（val · 备注） |
|
||
|---|---|---|---|---|---|
|
||
| v0 | dim32/4L/2h · 41K/3.26M | TinyStories 3MB 切片 · ~0.72M · — | fp32 单序列 | 1 GPU | val 3.80（toy，不可用） |
|
||
| v1 | dim256/8L/8h · 8.4M/34M | TinyStories 全量 · 5.1M · 0.01ep | fp32 单序列 | 1 GPU · 3.3K | val 2.58 |
|
||
| v2 | dim384/12L/12h · 28M/67M | TinyStories · 37M · 0.08ep | fp32 单序列 | 4 GPU DDP · 3.6K | val 1.71（暴露 KI-1/弱扩展） |
|
||
| v3 | dim512/16L/16h · 67M/119M | TinyStories · 246M · 0.53ep | fp32 **batched(T10)** | 1 GPU · 26K | val 1.30 |
|
||
| v4 | dim768/18L/24h · 127M/205M | TinyStories · 721M · 1.54ep | fp32 batched | **8 GPU(T11)** · 145K | val 1.17（仍欠拟合） |
|
||
| v5 | dim768/18L（**同 v4**） | TinyStories · 2.49B · **5.33ep** | **bf16(T12)** | 8 GPU · 217K | val **1.11**：⚠️**TinyStories 饱和**（3.5×数据仅↓5%） |
|
||
| v6 | dim768/18L（同 v4/v5） | **FineWeb-edu** 真实网页 · 2.29B · 1.02ep | bf16 | 8 GPU · 218K | val **3.07**：⚠️**FineWeb 留出集，与 v0–v5 不可比**（真实网页熵高，~3.0 是预期）；判据=采样质量+transfer。第一版脱离 TinyStories，**语言种类质变**（小故事→真实说明文）；transfer→TinyStories val 2.75(v5 native 1.11)，纯通用数据对窄分布有代价；val 末步仍单调降=未饱和 |
|
||
| v7 | dim768/18L（同 v4/v5/v6） | **同 v6 的 FineWeb-edu 子集**（非新数据）· 3.28B · **1.45ep** | bf16 | 8 GPU · 218K | val **3.01**（与 v6 可比）：⚠️**同子集多 epoch 近天花板**——唯一变量=epoch(1.02→1.45)，多喂 ~1B token val 仅 ↓0.05 且 ~step44000 后走平、采样无质变。与 v5 的 TinyStories 数据量饱和同类（重复老数据边际薄）；真·更多数据要**新 shards** |
|
||
| v8 | **dim1024**/18L/**32h** · **226M/329M**（+78% 容量，ffn 2730） | **同 v6/v7 的 FineWeb-edu 子集**（非新数据）· 2.36B · **1.05ep** | bf16 **+ 激活重计算(T13)** | 8 GPU · 129K（重算税） | val **2.98**（与 v6/v7 可比）：⭐**容量轴 A/B——容量有用**：唯一变量=dim768→dim1024，同 ~1ep v6 3.07→**2.98**（↓0.085），且 v8(1.05ep) < v7(1.45ep 更多老数据) 3.01 ⇒ 放大容量 > 重复老数据 ⇒ v6/v7 部分 capacity-limited。⚠️但增益仅 ~3%、val 末步**仍在降未饱和** ⇒ **单轴(数据/容量)单步都已 ~3%/lever = 全面边际递减，要双轴一起 scale(Chinchilla)** |
|
||
|
||
> 实训 token = steps×batch×seq（非数据集大小）。val 同一 1M-token TinyStories 留出集（v0–v5 可比；v6 起换 FineWeb-edu 留出集，分布不同、与 v0–v5 不可比；v6/v7/v8 同一 FineWeb 留出集、三版彼此可比 3.07/3.01/2.98）。
|
||
|
||
---
|
||
|
||
## 三、各维度的累积演进（轴向看一条线怎么走的）
|
||
|
||
- **算法**：手写 autograd(tape)+扇出累加 → AdamW/LR-sched/grad-clip → +QK-norm(Qwen3) → batched forward → bf16 混合精度(fp32 master) → 激活重计算(T13) → 融合 flash-attention(T14，online softmax + flash 式 bwd) → 梯度累积(T16，复用 tape SUM，等效大 batch 而显存随 micro) → dropout(T18，counter-based 设备 RNG + inverted scaling，train/eval 切换)。
|
||
- **模型架构**：固定 Qwen3-style；dim **32→256→384→512→768→1024**（v8 首拨容量轴，头数 24→32）；核心参数 **41K→226M**（总 3.26M→329M）。+QK-norm(T9，Qwen3 兼容) → **真 GQA(T15，`num_kv_heads<num_heads`，repeat_kv broadcast + 组内梯度求和；默认=nh→MHA 逐位回归)**——架构补齐到现代 LLM 标配（MHA/GQA/MQA 一条 `num_kv_heads` 轴），两条 SDPA(composed/flash) 共用同一 broadcast，导出真 `num_key_value_heads` 且 xserv 闭环。
|
||
- **Infra**：单卡 fp32 → cuBLAS/GPU-optim(T7) → NCCL DDP(T8) → batched forward(T10) → caching allocator(T11) → bf16(T12) → 激活重计算(T13，解锁 dim1024) → flash-attention(T14，不物化 N×N，attention 显存收益随 seq 增长) → 梯度累积(T16，DDP 只在累积边界通信，显存随 micro 不随有效 batch)。吞吐 **3.3K→217K tok/s**（dim768 bf16），dim1024+重算 ~129K（重算税）；MFU **0.4%→17%**（每次提升都对应一块 perf 基建，详见 known-issues + MFU 分析）。T13/T14/T16 是三条**显存杠杆**（重计算压激活峰值、flash 不物化 N×N attention scores、梯度累积解耦有效 batch 与激活显存），可叠加放大有效 batch。
|
||
- **数据集**：TinyStories 3MB 切片 → 全量 TinyStories（epoch 0.01→5.33，**至饱和**）→ **v6 毕业到 FineWeb-edu 真实网页**（2.255B 语料，1.02ep）→ **v7 同子集多 epoch（1.45ep，近顶）→ v8 同子集换大模型**（dim1024，1.05ep）。tokenizer 全程 gpt2 BPE（复用 xserv-tokenizer；v6 刻意不换 tokenizer 以隔离「数据来源」变量，KI-4 留后续版本）。
|
||
- **v5→v6 数据轴的质变**：v0–v5 都吃合成幼儿故事（TinyStories，低熵、词汇受控），v5 证明同尺寸模型在它上面已饱和；v6 第一版换成**真实教育类网页文本**（FineWeb-edu），语言种类发生质变——采样从「只会写小故事」变成「能写历史/科学/说明文」。
|
||
- ⚠️ **同子集多 epoch 也有天花板（v6→v7）**：v6 的 FineWeb val 才训 1.02ep、末步仍单调降，曾被读作「还没喂够」；v7 把**同一 2.255B 子集**喂到 1.45ep（多 ~1B token），FineWeb val 仅 ↓0.05（3.07→3.01）且 ~step44000 后走平、采样无质变 ⇒ **该子集在 dim768 已近天花板**。这与 v5 的 TinyStories 数据量饱和是**同一类现象**：**「重复喂老数据」边际都薄，无论是 v5 的同语料多 epoch 还是 v7 的同子集多 epoch**。真正抬天花板的是 v6「换更广的新语料」那一步——**杠杆在「更多样的新 token」，不在「同数据多读几遍」**。后续要继续降 val，必须补**新 FineWeb shards**（更多样、不重复），不是同子集加 epoch。
|
||
- ⚠️ **val 可比性**：v0–v5 的 val 是同一 TinyStories 1M 留出集（彼此可比）；**v6 起换 FineWeb-edu 留出集，分布不同、val 不能和 v0–v5（~1.1）比大小**——真实网页熵高，~3.0 是预期而非回退；**v6/v7/v8 同一 FineWeb 留出集、三版彼此可比**（3.07→3.01→2.98）。v6 的判据还有采样质量 + **transfer eval**（v6→TinyStories val 2.75 vs v5 native 1.11，量化「纯通用数据对窄分布的代价」）。
|
||
- ⭐ **容量轴有用，但也只有 ~3%（v8）**：v6/v7 在 dim768 上「吃不动更多数据」，v8 用最干净的 A/B 回答了「是数据见够还是容量不够」——**冻结数据子集、纯把 dim768→dim1024（core 127M→226M，+78%）**，同 ~1 epoch 下 FineWeb val **3.07→2.98（↓0.085）**，且 v8（1.05ep）还低于 v7（1.45ep 更多老数据）的 3.01。⇒ **容量有用，v6/v7 部分是 capacity-limited（不全是数据见够）**；放大容量比「给小模型多喂老数据」更值。**但增益只有 ~3%**，与数据轴单步杠杆同量级。
|
||
- 🧭 **元结论：单轴单步都已 ~3%/lever = 全面边际递减，要双轴一起 scale（Chinchilla 小尺度复现）**：把三条轴并起来看——数据量轴（v5/v7 同子集多 epoch，饱和，~1.6–5%/步）、数据广度轴（v6 换语料，是一次性换分布红利）、容量轴（v8，有用但 ~3%）——**到 v8，任何单轴的单步杠杆都收敛到 ~3%/lever**。而 v8 容量 +78% 却只配同样的 2.36B token、val 末步仍在降 ⇒ 数据立刻成新瓶颈。⇒ **要继续进步，容量与数据必须匹配地一起 scale，而不是单独猛拨一根轴**——这正是 Chinchilla 在这个 toy 尺度上的复现。
|
||
|
||
## 四、perf 杠杆台账（详见 [known-issues.md](known-issues.md)）
|
||
|
||
- **已修**：KI-1 单序列 launch-bound（T10）· KI-5 per-op cudaMalloc 串行（T11）· KI-2 bf16/OOM（T12）· KI-3 激活重计算（T13，解锁 dim1024，v8 用上）。
|
||
- **待办**：KI-4 大词表小 vocab · process-per-GPU（要更高多卡线性时）。
|
||
- 两次「先 profile 再动手」证伪了错误的拟修复（KI-1「加大batch」、KI-5「分桶all-reduce」），避免了无效大改——profile-first。
|