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

perf: KI-3 fixed — dim1024 batch32 fits, mem 31.1→14.6GB, tok/s 39.7K→31.5K

Browse Source 
Per-block activation recompute (T13) measured on dash5 (1× RTX 5090 32GB, bf16,
batch32 seq256, steady-state):

- Correctness (exact, hard gate): recompute on-vs-off grads are BIT-IDENTICAL —
  fp32 AND bf16: loss / logits / every param grad max rel = 0.00e0 (not "within
  tol", exactly equal). Full suite green with recompute on/off; DDP loss-match
  5.67e-7; DDP+recompute 2-rank descends 11.079→6.010.
- dim768 (18L/24h ffn2048, core 127M): peak mem 31144→14562 MiB (−53%), tok/s
  39.7K→31.5K (−20%, the extra-forward tradeoff, in the predicted 20–35% band).
- dim1024 (18L/32h ffn2730, core 226M): recompute OFF OOMs (hits 32100/32607
  MiB → OutOfMemory); recompute ON fits at 16596 MiB, ~23K tok/s, converges.
  → KI-3 payoff achieved: dim1024 batch32 unblocked, v8 can proceed.

Fill docs/12 bench table; mark KI-3 FIXED in docs/known-issues.md.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-06-17 09:50:29 +08:00
parent 69c5f07359
commit 0150263055
2 changed files with 29 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
docs/12-activation-recompute.md
View File

@@ -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 解锁**。