gahow/xtrain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: KI-2 FIXED — dim768 bf16 fits batch 32, tok/s 31.5K→40.8K

Browse Source 
bf16 mixed precision (fp32 master) solves the v4 dim768 fp32 batch-32
OOM and speeds up the now-compute-bound dim768 GEMMs (dash5 1× RTX
5090 32GB, dim768/18L/24h×32 ffn2048 seq256, steady-state):

  config          batch  peak mem   tok/s   fits 32GB
  fp32            16      27.2 GB    31.5K   yes
  bf16            16      19.3 GB    35.5K   yes   (-29% mem / +13% tok/s)
  fp32            32      —          —       OOM
  bf16            32      31.1 GB    40.8K   yes   (+29% vs fp32-b16)

Verified on dash5: fp32 suite green at tight tol + xserv export md5
bit-identical to registry; bf16 looser-tol (loss 1.2e-4, logits p99
6.8e-3, grad 1.0e-2) + 150-step convergence tracks fp32 (3.984 vs
3.988); 2-GPU bf16 DDP at per-rank batch 32 trains cleanly.

Mark KI-2 FIXED; fill docs/11 results.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-16 14:28:20 +08:00
parent 48922cb628
commit 320c1ae4fb
2 changed files with 39 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
docs/known-issues.md
View File

@@ -7,12 +7,33 @@
## 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。)_
_(KI-1 fixed in T10. KI-5 fixed in T11. **KI-2（bf16 混合精度）FIXED in T12**——fp32 master + bf16 linears/激活，dim768 fp32-batch32 OOM 解除（bf16 fits, 31.1/32.6GB），bf16-vs-fp32 同 batch 显存 −29% / tok/s +13%，收敛对住 fp32。见下方 Fixed。)_
---
## Fixed
### KI-2 · bf16 混合精度（fp32 master）— `FIXED` (T12)
- **触发点（v4 surfaced）**：dim768 fp32 在单卡 32GB 里 per-rank batch 32（global 256）OOM，被迫降到 per-rank 16。bf16（激活减半）找回 batch-32 甜点区，并加速已 compute-bound 的 dim768 GEMM；附带：xserv 推理 BF16-only，bf16 训练更贴闭环。
- **设计（标准 AMP，opt-in，[docs/11-bf16-mixed-precision.md](11-bf16-mixed-precision.md)）**：**fp32 master weights** + AdamW/clip/DDP all-reduce 全程 fp32；**bf16 compute**=linears(q/k/v/o, gate/up/down, lm_head)走 `cublasGemmEx`（CUDA_R_16BF in/out, CUBLAS_COMPUTE_32F 累加）+ 激活流 bf16（含 attention probs / logits）；**fp32 稳定**=RMSNorm/QK-norm、softmax、RoPE、cross-entropy 内部 upcast→fp32→downcast。**无 loss scaling**（bf16 8-bit 指数=fp32 动态范围）。关键钩子=autodiff `cast` 算子：前向把 fp32 master leaf 降到 bf16 喂 matmul，**反向把 bf16 grad 升回 fp32** → fp32 leaf 累加 fp32 grad，优化器一行不改。fp32 路径按 dtype 分派、逐字节不变（hard gate）。
- **正确性（双闸门全绿）**：
- **fp32 不回归**：全套在原紧容差绿——autograd 15、structural 5、GEMM 对 cuBLAS 5、batched==looped、overfit 27/27、AdamW GPU bit-exact + host 对 torch、checkpoint 逐位、DDP loss 对单卡 + 跨 rank、**xserv 闭环**（v3 ckpt 用 T12 代码重导 safetensors 与 registry **md5 逐位一致** `b04fc9f9…`）。
- **bf16 looser-tol**：同 fp32 master 跑 fp32 vs bf16——loss rel **1.2e-4**、logits mean rel **2.0e-3** / p99 **6.8e-3**、grad worst scaled-mean **1.0e-2**，无 NaN，grad 仍 fp32。
- **收敛**：dim768 短训 150 步，bf16-b16 loss 轨迹对住 fp32-b16（step50 4.40 vs 4.40、step149 **3.984 vs 3.988**），单调降、无发散。
- **显存 + 吞吐（payoff，dash5 1× RTX 5090 32GB, dim768/18L/24h×32 ffn2048 seq256, steady-state）**：
| config | per-rank batch | 峰值显存 | tok/s | fits 32GB? |
|---|---|---|---|---|
| fp32 | 16 (v4 fallback) | 27.2 GB | 31.5K | ✅ |
| **bf16** | 16 | **19.3 GB（−29%）** | **35.5K（+13%）** | ✅ |
| fp32 | 32 | — | — | ❌ **OOM** |
| **bf16** | **32（甜点区）** | **31.1 GB** | **40.8K** | ✅ **解 OOM** |
→ **同 batch 16：bf16 显存 27.2→19.3GB（−29%）、tok/s 31.5K→35.5K（+13%）**；**bf16 解 fp32-batch32 OOM**（31.1/32.6GB fit），batch32 达 **40.8K tok/s（+29% vs fp32-b16）**。残留：norm/softmax/CE 的 fp32 upcast 是 transient，但仍占峰值——若 v5 要更大 batch，下一杠杆是 KI-3 激活重计算。
- **commit**：见 T12 提交链（`cuda: bf16 cuBLAS GemmEx` / `autodiff: bf16 mixed-precision path` / `train: --bf16 flag` / `perf: keep bf16 logits` / 本条）。
---
### 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 卡几乎无差（详见下方历史诊断）。
@@ -91,11 +112,6 @@ _(KI-1 fixed in T10. KI-5 **FIXED** in T11——device caching/pool allocator
## 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 后显存成约束。