xtrain/docs/known-issues.md

xtrain — Known Issues & Perf Backlog

已知问题（性能 / 正确性 / 建模）与延后项的活文档：记录现象、复现、根因、拟修复、优先级、状态。 发现即记，修复即标 FIXED（附 commit）。

---

Open

(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）：新增 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）：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 卡几乎无差（详见下方历史诊断）。

• 修复：xtrain-cuda 加 device caching/pool allocator（docs/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。

• 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、docs/06-performance.md。

---

Deferred（来自 T7，放大后重启）

(KI-3 已在 T13 FIXED，见上方 Fixed。本节暂无待启项。)

---

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 即可缓解占比）。