diff --git a/docs/11-bf16-mixed-precision.md b/docs/11-bf16-mixed-precision.md index 53ff44d..7e2e43c 100644 --- a/docs/11-bf16-mixed-precision.md +++ b/docs/11-bf16-mixed-precision.md @@ -76,8 +76,23 @@ - **dim768 bf16 能跑 per-rank batch 32**(v4 OOM 的触发点)。 - 测 dim768 **bf16 vs fp32** 的峰值显存 + steady-state tok/s(预期:显存↓、tok/s↑)。 -## 实测结果(dash5, 8× RTX 5090, sm_120) +## 实测结果(dash5, 1× RTX 5090 32GB, sm_120) -> 见 `docs/known-issues.md` KI-2 的 before→after 表(fp32 batch32 OOM → bf16 batch32 fit;显存 A→B;tok/s A→B)。 +**闸门 ① fp32 不回归**:全套测试在原紧容差绿(autograd 15 / structural 5 / GEMM 5 / batched==looped / overfit 27/27 / AdamW GPU bit-exact + host 对 torch / checkpoint 逐位 / DDP 2)。**xserv 闭环**:v3 ckpt 用 T12 代码重导 `model.safetensors` 与 registry **md5 逐位一致**(`b04fc9f9a0c9af04c47d9ca649aea12e`)——export/fp32 数值零漂移。 + +**闸门 ② bf16 正确性 + 收敛**: +- **looser-tol(tests/bf16.rs)**:同 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(master 未动)。 +- **收敛**:dim768 短训 150 步,bf16-b16 loss 轨迹对住 fp32-b16(step50 `4.40` vs `4.40`、step149 `3.984` vs `3.988`),单调降、无发散。 + +**闸门 ③ 显存 + 吞吐(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:bf16 显存 −29% / tok/s +13%;**bf16 解 fp32-batch32 OOM**,batch32 达 40.8K tok/s(+29% vs fp32-b16)。KI-2 标 **FIXED**。 diff --git a/docs/known-issues.md b/docs/known-issues.md index e2a9f23..a99c9f6 100644 --- a/docs/known-issues.md +++ b/docs/known-issues.md @@ -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 后显存成约束。