Gahow Wang
511ceebbb3
v4 surfaced the concrete bf16 trigger: dim768 fp32 OOMs at per-rank batch 32 (global 256) in 32GB, forcing per-rank 16 (global 128). bf16 (halve activation mem) would restore the batch-256 sweet spot. Record it on KI-2; keep KI-2 as the backlog item it is (still deferred). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
110 lines
12 KiB
Markdown
110 lines
12 KiB
Markdown
# xtrain — Known Issues & Perf Backlog
|
||
|
||
已知问题（性能 / 正确性 / 建模）与延后项的活文档：记录现象、复现、根因、拟修复、优先级、状态。
|
||
发现即记，修复即标 `FIXED`（附 commit）。
|
||
|
||
---
|
||
|
||
## Open
|
||
|
||
_(KI-1 fixed in T10. KI-5 **FIXED** in T11——device caching/pool allocator 消掉 per-op cudaMalloc 串行，单卡 ~2.3×、DDP scaling 从 ~1.3× 封顶恢复到 ~5×@8。见下方 Fixed。)_
|
||
|
||
---
|
||
|
||
## Fixed
|
||
|
||
### KI-5 · DDP 弱扩展性 — `FIXED` (T11, device caching/pool allocator)
|
||
- **根因（T11 重诊断，all-reduce **不是**瓶颈）**：每个 tape op 输出走 `Tensor::zeros`→`GpuBuffer::alloc`→`cudaMalloc`（同步、进程级串行的 driver 调用）。单进程 thread-per-GPU 下 N rank 每步几百次 alloc 在单 CUDA context 排队串行（`NOCOMM=1` 完全不通信时 fwd+bwd 仍 136→780ms 膨胀 ~6×，`nvidia-smi` 抽样 8 卡只 1–2 张在忙轮流跑）；单卡也吃这笔 per-op alloc。
|
||
- **原拟修复「分桶 all-reduce」经 T11 实测证伪并 revert**：grad all-reduce 每步只占 ~6–7%，融成一发对 1/2/4/8 卡几乎无差（详见下方历史诊断）。
|
||
- **修复**：`xtrain-cuda` 加 **device caching/pool allocator**（[docs/10-caching-allocator.md](10-caching-allocator.md)）——`GpuBuffer::alloc` 从 per-device、size-classed free-list 取，miss 才 `cudaMalloc`；`Drop` 归还 free-list（不 `cudaFree`）。训练定形状→命中率极高，warm-up 后每步 `cudaMalloc`≈0。线程安全：global registry 按 device id 分桶，每 device 的 free-list 各自 `Mutex`（registry 锁只在 clone 出 `Arc<Mutex<_>>` 时极短持有→跨 device 真并发）；buffer 记 alloc 时的 device，Drop 归还对应 pool。**透明**：物理 cap 可向上取整但 `len()`/memset/copy bounds 都用请求 `len`，尾部字节永不读到→数值逐位不变。memset 保留（复用 buffer 有陈旧字节）；skip-memset uninit 本次不做（malloc 已是瓶颈，memset async 开销小，逐 op 证明全覆盖风险大）。
|
||
- **before → after**（dash5, 8× RTX 5090, dim384/12L per-rank batch 32 seq 256, steady-state tok/s; before=`d422c68` after=pooled）：
|
||
|
||
| world | before tok/s | before speedup | after tok/s | after speedup |
|
||
|---|---|---|---|---|
|
||
| 1 | 39801 | 1.00× | **92385** | 1.00× |
|
||
| 2 | 47229 | 1.19× | 146821 | 1.59× |
|
||
| 4 | 52854 | 1.33× | 269867 | 2.92× |
|
||
| 8 | 48996 | 1.23× | **461270** | **4.99×** |
|
||
|
||
→ **单卡 40226→92638 tok/s (~2.3×)**；**8 卡 49K→461K tok/s (9.4×)**，scaling 从 ~1.3× 封顶恢复到 **~5×@8**；8 卡 `nvidia-smi` 抽样 **全 8 卡 95–99% util**（KI-5 时只 1–2/8 忙）。loss 轨迹逐位对住（单卡 10.9026→4.8453 before/after 一致）。
|
||
- **正确性（全绿，无回归）**：15 算子 grad-check、5 结构、GEMM 对 cuBLAS、batched==looped、overfit 27/27、AdamW GPU bit-exact + host 对 torch、checkpoint 逐位、DDP loss 对单卡 **5.67e-7** + 跨 rank diff 0.0（loosened `<1e-6`）、**xserv 闭环**（v3 ckpt 重导 safetensors 与 registry md5 逐位一致 + xserv 加载服务贪心 "Once upon a time," 对住）。
|
||
- **顺手**：DDP `ddp_correctness` 的 cross-rank `==0.0` → `<1e-6`（本机 PCIe-only NCCL run-to-run 跨 rank 非逐位可复现，diff≤1.2e-7 几 ULP 无害，承重闸门是 loss-match 5.67e-7）；`ddp_throughput_scaling` 扩到 world=8。
|
||
- **残留**：~5×@8 非完美线性（grad all-reduce ~7% + 8 卡 PCIe/launch 余量），但弱扩展悬崖已消。v4 若要更高线性度，下一步是 **process-per-GPU**（每 rank 独立 CUDA context，torchrun 式）。
|
||
- **commit**：见 T11 提交链（`cuda: device caching allocator` / `perf: KI-5 …` 那条带 before/after）。
|
||
- **历史诊断保留如下**（证伪「分桶 all-reduce」的过程）：
|
||
|
||
---
|
||
|
||
### KI-5 历史诊断 · DDP 弱扩展性 — T10 暴露，T11 重诊断（all-reduce **不是**瓶颈）
|
||
- **现象**：batched forward 修掉单卡 launch-bound 后，dim384/per-rank batch 32：1 卡 40.3K → 4 卡 47.2K tok/s（global），仅 ~1.17×。
|
||
- **T11 实测（dash5, 8× RTX 5090, dim384/12L, per-rank batch 32, seq 256, 原 ungrouped all-reduce, 50 步均, ms/step）**：
|
||
|
||
| world | fwd+bwd | grad all-reduce | clip+opt+zero | TOTAL | tok/s(global) | speedup |
|
||
|---|---|---|---|---|---|---|
|
||
| 1 | 136 | 0 | 8.6 | 145 | 36582 | 1.00× |
|
||
| 2 | 202 | 21 | 15 | 238 | 47267 | 1.29× |
|
||
| 4 | 342 | 29 | 21 | 392 | 51466 | 1.41× |
|
||
| 8 | 780 | 54 | 47 | 882 | 47719 | 1.30× |
|
||
|
||
→ grad all-reduce 每步只占 **~6–7%**；真正爆炸的是**逐 rank 的 fwd+bwd 时间随 world 线性膨胀**（同一 per-rank workload，136→780ms，~6×）。
|
||
- **「分桶 all-reduce」拟修复经 T11 实测证伪（无收益）**：把 ~150 个 per-tensor `ncclAllReduce` 用 `ncclGroupStart/End` 融成一发 → 1/2/4/8 卡 = 1.00/1.30/1.42/1.34×，**与不分桶几乎无差**（all-reduce 本就只占 7%）。flat-buffer 分桶同理。故回退（revert b8b5821），保留原 ungrouped 路径。
|
||
- **附带发现：T8 correctness 测试的 `max|p0−p1| == 0.0` 在本机 flaky**（与 T11 无关）。原 ungrouped 代码同一 GPU 重跑 6 次 cross-rank diff = {0.0, 0.0, 5.96e-8, 5.96e-8, 1.19e-7, 1.19e-7}，只 ~1/3 命中 `0.0`。即本机/本版 NCCL 的 all-reduce **run-to-run 跨 rank 不是逐位可复现**（PCIe-only 拓扑下 algorithm/chunk 选择不稳）。diff 都 ≤1.19e-7（几 ULP，数值无害，loss-match 仍 ~6e-7），但 `== 0.0` 断言过严 → 建议改为 `< 1e-6` 紧容差（**留作 follow-up，本次未改测试**）。
|
||
- **重新定位的根因**：**单进程 thread-per-GPU 模型下，N 个 rank 线程各自跑独立训练却互相串行**——`NOCOMM=1`（完全不做任何跨 rank 通信/barrier）时 fwd+bwd 仍 136→378→800ms 膨胀；`nvidia-smi` 抽样显示 8 卡同一时刻只有 1–2 张在忙、轮流跑。排除项：CPU 不缺（187 核, load 2.5）；`nvcc --default-stream per-thread` 不解决。**剩余怀疑：每个 op 输出走 `Tensor::zeros`→`cudaMalloc`+`cudaMemset`，而 `cudaMalloc` 是同步、进程级串行的 driver 调用；单 CUDA context 下 N rank 每步几百次 alloc 互相排队**——即 DDP 真瓶颈是 **per-op 显存分配 / driver 调用在单进程内串行**，不是梯度通信。
|
||
- **真正的修复方向（待定，非 T11 范围）**：① **caching/pool allocator**（op 输出复用显存，消掉每步几百次 `cudaMalloc`，单卡也受益）；或 ② **process-per-GPU**（每 rank 独立 CUDA context，torchrun 式，彻底解串行，但要改 launcher + 跨进程 UniqueId 分发）。先做 ① 再实测是否解 DDP 串行。
|
||
- **重启条件**：多卡 v4 需要扩展性时做。**单卡 batched 已 40K tok/s（v3 即单卡训完）**，多卡当前只有 ~1.4× 上限，v4 若要多卡须先修上面的真瓶颈。
|
||
|
||
---
|
||
|
||
### KI-1 · 单序列 launch-bound（"DDP 弱扩展性"的根因）— `FIXED` (T10, batched forward)
|
||
- **修复**：T10 给 model + autograd 加 batch 维——linears 摊平成 `[B*S, dim]` 一个大 GEMM 填满 GPU；attention 走 fused 批量 SDPA（`cublasSgemmStridedBatched` ×2 + 一个 causal-softmax kernel），RoPE 位置 per-sequence 复位（`row % S`）；训练 loop 用真 batch 一次 forward/backward 替代 "loop B 次 + SUM"。详见 [docs/09-batched-forward.md](09-batched-forward.md)。
|
||
- **before → after**（dim384/12L/12h, batch 16, seq 256, 1 卡, back-to-back A/B）：
|
||
| | tok/s | GPU util | 显存 |
|
||
|---|---|---|---|
|
||
| before（单序列 launch-bound）| ~1653 | 0–15% | ~3 GB |
|
||
| after（batched）| **25627**（batch16）/ **40263**（batch32）| **37% 均值 / 54% 峰** | ~10 GB |
|
||
|
||
→ 单卡 **~15.5×（batch16）/ ~24×（batch32）**，util 0–15% → 37–54%。
|
||
- **正确性（全绿，无回归）**：15 算子 grad-check（新增 batched-rope / transpose_4d12 / batched-attention dQ/dK/dV）、batched==looped 单序列（logits 0.0、grad 6.4e-4）、**PyTorch 对拍 B>1**（loss 5e-8 / logits 6.9e-6 / 全参数 grad 在 rtol 2e-2）、overfit 27/27、checkpoint 逐位、AdamW 对 torch、DDP loss 对单卡 5.7e-7 + 跨 rank 参数 bit-identical(0.0)、**xserv 加载导出权重对 xtrain 贪心仍逐 token 一致**（top token 同序、BF16 漂移 ~0.03）。
|
||
- **commit**：见 T10 提交链（`perf: KI-1 fixed — GPU util / tok/s` 那条带 before/after）。
|
||
- **DDP 残留弱扩展性 → KI-5**（这是 batching 后新暴露的 all-reduce 瓶颈，不是 KI-1 的单序列根因）。
|
||
- **历史诊断保留如下**（v2 暴露 → v3 重诊断的过程，证明根因不是 all-reduce）：
|
||
|
||
---
|
||
|
||
### KI-1 历史诊断 · DDP 弱扩展性（吞吐受单序列 launch-bound 限制）— v2 暴露，v3 重新诊断
|
||
- **现象**：4 卡 DDP 仅 ~3.2K tok/s，几乎不快于单卡（≈2× over 单卡，远低于近线性；T8 在 tiny micro-bench 为 3.0×@4）。
|
||
- **复现**：`dim384/12L, world=4, seq 256`。
|
||
- **v3 实测（dash5, 4× RTX 5090, dim384, 隔离 back-to-back A/B）**：
|
||
| global_batch | 每卡 | tok/s（4卡）| GPU util | 显存 |
|
||
|---|---|---|---|---|
|
||
| 32 | 8 | **3163** | 5–69%（spiky）| ~2–3 GB / 32 GB |
|
||
| 256 | 64 | **3200** | 0–15% | ~2–3 GB / 32 GB |
|
||
→ **加大 8× batch 仅 +1.2% 吞吐（噪声内）**。1 卡 dim384 ≈ 1653 tok/s，4 卡 3163 ≈ 2.1×。
|
||
- **原"拟修复"（加大 global batch）经 v3 实测 falsified**：gbatch256 时每 token 的 all-reduce 次数只有 gbatch32 的 1/8，若瓶颈是 all-reduce 应大幅提速——实际没有 → **all-reduce / 通信不是瓶颈**。
|
||
- **重新诊断的根因**：瓶颈是**单序列模型设计**（T5：每个 sequence 各跑一次独立 forward/backward，逐 op kernel-launch 开销，见 docs/06 延迟瓶颈）。GPU util 仅 0–15%、显存仅占 ~8% → 严重 **launch-bound / under-utilized**；GEMM 太小喂不饱 GPU。加大 batch 只是按比例增加串行 launch 次数，无法摊薄。4 卡相对单卡 ~2× 的固定天花板来自跨 rank 同步税，但**不是**靠调 batch 能修的。
|
||
- **真正的修复（需实作，非调参）**：
|
||
1. **batched（多序列）forward**——把一个 step 的多条序列在 batch 维一次性过模型，让 GEMM 大到能填满 GPU（这是 launch-bound 的根本解，但要改 T4/T5 的 single-sequence autograd/model，工作量大、有正确性风险）；
|
||
2. 在 (1) 之后，梯度 all-reduce **分桶 + 与 backward 重叠**（bucketed / overlapped all-reduce）才会有意义（当前 all-reduce 已非瓶颈，做了也无收益）。
|
||
- **参考**：[docs/07-distributed.md](07-distributed.md)、[docs/06-performance.md](06-performance.md)。
|
||
|
||
---
|
||
|
||
## Deferred（来自 T7，放大后重启）
|
||
|
||
### KI-2 · bf16 混合精度（fp32 master）— `deferred`
|
||
- T7 延后理由：tiny 规模延迟瓶颈、bf16 改变数值会威胁 fp32 正确性闸门。
|
||
- **重启条件**：模型放大（v2+ `dim≥384`）后 GEMM 渐成 compute-bound，tensor-core 收益显现。需 fp32 master weights + 单独 looser-tol 测试 + 收敛对比。
|
||
- **具体触发点（v4 surfaced）**：dim768 fp32 在单卡 32GB 显存里 per-rank batch 32（global 256）OOM，被迫降到 per-rank 16（global 128）训练。bf16（激活减半）能把 batch-256 的甜点区找回来。这是 v0–v3 tiny 规模延后 bf16 后第一次有 fp32 放不下的硬约束——v5 该先拉的杠杆。
|
||
|
||
### KI-3 · 激活重计算（gradient checkpointing）— `deferred`
|
||
- T7 延后理由：单序列、显存不紧。
|
||
- **重启条件**：更大模型 / 更长 seq / 更大 batch 后显存成约束。
|
||
|
||
---
|
||
|
||
## Modeling notes
|
||
|
||
### KI-4 · 大词表 embedding 占比过高
|
||
- gpt2 `vocab=50257` 在 dim 小时让 embed+lm_head 主导参数：v1 25.7M/34M、v2 38.6M/66.9M；core transformer 才是学习主体。
|
||
- 后续可考虑更贴合 TinyStories 的小 vocab（会牺牲 xserv gpt2-tokenizer 复用）；或在更大 dim 下让 core 自然成为主体（继续 scaling 即可缓解占比）。
|