diff --git a/docs/12-activation-recompute.md b/docs/12-activation-recompute.md index a286b7a..dc472ff 100644 --- a/docs/12-activation-recompute.md +++ b/docs/12-activation-recompute.md @@ -74,15 +74,17 @@ - 同 config 测 steady-state tok/s recompute on vs off,报告慢多少(预计 ~20–35%,多一次前向)。 -## 实测结果(dash5 1× RTX 5090 32GB, dim768/18L/24h×32 ffn2048 seq256, bf16, steady-state) +## 实测结果(dash5 1× RTX 5090 32GB, bf16, batch32 seq256, steady-state) -> 待 dash5 实跑回填。 +**正确性(exact,硬闸门)**:on-vs-off 梯度对拍 —— **fp32 与 bf16 双路都逐位一致**:loss rel `0.00e0`、logits max rel `0.00e0`、**每个参数梯度 max rel `0.00e0`**(不是「在容差内」,是逐位相同——证实重算确实精确)。全套回归开/关重计算全绿:T4 15 算子 grad-check、5 结构、batched、bf16、overfit、AdamW(GPU+host)、GEMM、checkpoint roundtrip、**T8 DDP loss 对单卡 5.67e-7 + 跨 rank 0.0**;DDP+recompute 2 卡短训 loss 单调降(11.079→6.010)。 + +**显存 + 吞吐**(dim768 = 18L/24h×32/ffn2048 core 127M;dim1024 = 18L/32h×32/ffn2730 core 226M): | config | per-rank batch | 峰值显存 | tok/s | fits 32GB? | |---|---|---|---|---| -| dim768 recompute **off** | 32 | TBD | TBD | ✅ | -| dim768 recompute **on** | 32 | **TBD(↓)** | TBD(↓~xx%) | ✅ | -| **dim1024** recompute **off** | 32 | — | — | ❌ **OOM** | -| **dim1024** recompute **on** | 32 | **TBD** | TBD | ✅ **解 OOM** | +| dim768 recompute **off** | 32 | 31144 MiB | 39.7K | ✅ | +| **dim768 recompute on** | 32 | **14562 MiB(−53%)** | **31.5K(−20%)** | ✅ | +| **dim1024** recompute **off** | 32 | 32100 → **OOM** | — | ❌ **OOM** | +| **dim1024 recompute on** | 32 | **16596 MiB** | 23.1K | ✅ **解 OOM** | -**正确性**:on-vs-off 梯度 max rel = TBD(fp32)/ TBD(bf16);全套回归 + xserv 闭环全绿。 +→ dim768:重计算把峰值显存 **31.1→14.6GB(−53%,~砍半激活)**,代价 tok/s **−20%**(多一次前向,落在预测 20–35% 区间)。dim1024 batch32:**不开 OOM(撞 32100/32607MiB 上限)→ 开了 16.6GB 稳稳装下**,~23K tok/s 训练正常收敛 —— **KI-3 的目标达成,dim1024 解锁**。 diff --git a/docs/known-issues.md b/docs/known-issues.md index a99c9f6..bb3f120 100644 --- a/docs/known-issues.md +++ b/docs/known-issues.md @@ -7,12 +7,30 @@ ## Open -_(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。)_ +_(KI-1 fixed in T10. KI-5 fixed in T11. KI-2 fixed in T12. **KI-3(激活重计算 / gradient checkpointing)FIXED in T13**——per-block checkpoint(no-tape forward + recompute-on-backward),梯度对非重计算版**逐位一致**(fp32/bf16 max rel 0.00e0);dim768 bf16 batch32 峰值显存 31.1→14.6GB(−53%)/ tok/s −20%;**dim1024 batch32 不开 OOM → 开了 16.6GB 装得下**,解锁 v8。见下方 Fixed。)_ --- ## Fixed +### KI-3 · 激活重计算(gradient checkpointing)— `FIXED` (T13) +- **触发点(v8 surfaced)**:容量轴放大到 dim1024(核心 ~210M+)测是否 capacity-limited。autograd tape 为反向保存所有中间激活,激活显存随 dim 线性增长——dim768 bf16 batch32 已 31.1GB(T12 甜点区),**dim1024 batch32 再次 OOM**(实测撞 32100/32607MiB → `OutOfMemory`)。 +- **设计(per-block gradient checkpointing,opt-in,[docs/12-activation-recompute.md](12-activation-recompute.md))**:新增 `xtrain_autodiff::checkpoint(segment_fn, input, params)` 高阶原语(类比 `torch.utils.checkpoint`)。**前向**:把 input/params detach 成局部 leaf 跑 `segment_fn`,只取输出值,局部 tape 立即 drop → 段内激活释放(不留在外层 tape);checkpoint 节点 parents=[input, ..params]。**反向**:从保存的 input + 未变的 param 值重跑 `segment_fn` 重建局部 tape,用上游 grad seed(`Var::backward_seeded`,新增——段输出非标量)回传,恢复的 input/param 梯度 push 给真 parents,局部 tape drop → 重算激活释放。模型每个 transformer block 前向用它包裹(`--recompute` flag,默认关)。切粒度 = 每 block。 +- **正确性(exact,硬闸门全绿)**:重计算数学精确(同 `segment_fn`、同输入、同参数值、确定性 kernel → 重算激活逐位等于原激活)。**on-vs-off 梯度对拍 fp32/bf16 双路逐位一致**:loss rel `0.00e0`、logits max rel `0.00e0`、**每个参数梯度 max rel `0.00e0`**(不是容差内,是逐位)。全套回归开/关重计算全绿:autograd 15、structural 5、batched、bf16、overfit 27/27、AdamW(GPU bit-exact + host 对 torch)、checkpoint roundtrip、**DDP loss 对单卡 5.67e-7 + 跨 rank 0.0**;DDP+recompute 2 卡短训单调降(11.079→6.010)。非重计算路径图不变(默认关)→ T10/T11/T12 数值不回归。 +- **显存 + 吞吐(payoff,dash5 1× RTX 5090 32GB, bf16, batch32 seq256, steady-state)**: + + | config | per-rank batch | 峰值显存 | tok/s | fits 32GB? | + |---|---|---|---|---| + | dim768 (18L/24h ffn2048, core 127M) off | 32 | 31144 MiB | 39.7K | ✅ | + | **dim768 on** | 32 | **14562 MiB(−53%)** | **31.5K(−20%)** | ✅ | + | **dim1024** (18L/32h ffn2730, core 226M) off | 32 | 32100 → **OOM** | — | ❌ **OOM** | + | **dim1024 on** | 32 | **16596 MiB** | 23.1K | ✅ **解 OOM** | + + → dim768:重计算砍 ~半激活(**31.1→14.6GB,−53%**),代价 tok/s **−20%**(多一次前向,落在预测 20–35%)。**dim1024 batch32:不开 OOM → 开了 16.6GB 稳稳装下**,~23K tok/s 正常收敛 → **dim1024 解锁,v8 可展开**。 +- **commit**:T13 提交链(`autodiff: checkpoint primitive (recompute-on-backward)` / `model: per-block activation recompute (--recompute)` / `perf: KI-3 fixed …` 本条带 before/after / 文档 `docs: Phase T13 — activation recompute`)。 + +--- + ### 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)。 @@ -112,9 +130,7 @@ _(KI-1 fixed in T10. KI-5 fixed in T11. **KI-2(bf16 混合精度)FIXED in T1 ## Deferred(来自 T7,放大后重启) -### KI-3 · 激活重计算(gradient checkpointing)— `deferred` -- T7 延后理由:单序列、显存不紧。 -- **重启条件**:更大模型 / 更长 seq / 更大 batch 后显存成约束。 +_(KI-3 已在 T13 FIXED,见上方 Fixed。本节暂无待启项。)_ ---