Gahow Wang
9a25616a30
docs/09-batched-forward.md: the launch-bound diagnosis recap, the [B*S,dim]-flatten + fused batched-attention design (RoPE per-seq position + causal masking inline in softmax), the attention forward/backward via strided-batched GEMM, autograd implications, the looped-split/merge dead-end post-mortem (1127 tok/s, host round-trips), verification methods + before→after throughput, and the v3 recommendation (per-rank batch 16-32, single/small world until KI-5 bucketed all-reduce lands). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
188 lines
12 KiB
Markdown
188 lines
12 KiB
Markdown
# Phase T10: Batched 多序列 Forward — Design Document
|
||
|
||
## Goal
|
||
|
||
修 **KI-1 的根因**。v2 暴露、v3 重诊断的结论（见 [docs/known-issues.md](known-issues.md)）：
|
||
吞吐瓶颈**不是** all-reduce，而是 **单序列模型设计 launch-bound**——T5 的模型一次只过**一条**
|
||
序列，每个 op 的 GEMM 都是 `[seq, dim]` 这种小矩阵，喂不饱 GPU；"batch" 是靠**循环 B 次 forward +
|
||
让 tape SUM 梯度**伪造的，于是 GPU util 只有 0–15%、显存占 ~8%，加大 global batch 也只是按比例
|
||
增加串行 kernel-launch（v3 实测 gbatch 32→256 仅 +1.2%）。
|
||
|
||
T10 的目标：给 model + autograd 加 **batch 维**，把一个 step 的 B 条序列**一次性**过模型，让线性层变成
|
||
**一个大 GEMM** 填满 GPU——这是 launch-bound 的根本解。**硬闸门是正确性**：所有既有 grad-check /
|
||
PyTorch 对拍 / overfit / DDP 跨 rank 一致**必须仍过**（PyTorch 对拍现在用 **B>1**），xserv 推理仍**逐
|
||
token 一致**；在此之上拿吞吐收益。
|
||
|
||
## 核心设计：linears 摊平为 `[B*S, dim]`，attention 批量化
|
||
|
||
一个关键观察：**transformer 里只有 attention 需要"序列感知"**，其余全是 per-token 的。
|
||
|
||
- **线性投影 / 逐元素 / norm / embedding / CE 天然是 2D `[rows, dim]` 上的运算**，不在乎 `rows` 是
|
||
`seq` 还是 `B*S`。所以把激活**摊平成 `[B*S, dim]`** 喂进去，这些 op **原样复用**：
|
||
q/k/v/o、gate/up/down、lm_head 全部变成**一个大 `[B*S, dim] × [dim, out]` GEMM**（填 GPU 的收益所在）；
|
||
embedding 对 `[B*S]` ids gather；rms_norm / qk_norm / swiglu / silu / 逐元素按行（per-token）；
|
||
cross_entropy 对 `[B*S, vocab]` vs `[B*S]` targets，**mean over B*S 行 == 各序列 loss 的均值**。
|
||
- **attention 是唯一序列感知的 op**：reshape 成 `[B, nh, S, hd]`，**每个序列内**做因果 SDPA（per-seq
|
||
因果 mask，**绝不跨序列 attend**），写回 `[B*S, dim]`。在 `(B, nh)` 维上批量。
|
||
- **RoPE 位置必须 per-sequence 复位**：位置 = 在**自己序列内**的下标（`row % S`），**不是**摊平后的全局
|
||
行号。这是正确性陷阱，显式处理（kernel 加 `period` 参数）。
|
||
|
||
```text
|
||
ids [B*S]
|
||
└ embedding → [B*S, dim]
|
||
每层 block：
|
||
rms_norm → [B*S, dim] (per-token)
|
||
attention:
|
||
x@Wq/Wk/Wv → [B*S, dim] (大 GEMM)
|
||
reshape → [B*S, nh, hd]
|
||
qk-norm + RoPE(period=S) → [B*S, nh, hd] (RoPE 位置 = row % S)
|
||
→ [B, nh, S, hd] → [B*nh, S, hd]
|
||
fused batched SDPA(因果) → [B*nh, S, hd] (2× 批量GEMM + 1× causal-softmax)
|
||
→ [B*S, dim]; @Wo → [B*S, dim] (大 GEMM)
|
||
+residual; rms_norm; SwiGLU MLP(大 GEMM); +residual
|
||
final rms_norm; @lm_head → [B*S, vocab]
|
||
cross_entropy([B*S, vocab], targets[B*S]) → 标量（B*S 行的 mean）
|
||
```
|
||
|
||
`forward(ids[seq])` 现在只是 `forward_batched(ids, batch=1)` 的特例 → **采样 / 推理路径 batch=1 不变**。
|
||
|
||
## Module Layout
|
||
|
||
```
|
||
csrc/ops/attention.cu # 新：causal 行 softmax（含 scale + per-seq 因果 mask）
|
||
csrc/ops/nn.cu # rope/rope_dx 加 period（位置 = tok % period）
|
||
csrc/ops/model.cu # 加 transpose_4d12（[B,S,nh,hd]<->[B,nh,S,hd]）
|
||
crates/xtrain-cuda/
|
||
├── src/cublas.rs # 加 sgemm_strided_batched（批量 GEMM，行主序同 sgemm 套路）
|
||
├── src/ffi.rs # +cublasSgemmStridedBatched / launch_softmax_causal / transpose_4d12；rope 加 period
|
||
└── build.rs # +attention.cu
|
||
crates/xtrain-tensor/src/tensor.rs # 加 attention/attention_backward（fused，2/4× 批量GEMM）、transpose_4d12；rope(+period)
|
||
crates/xtrain-autodiff/
|
||
├── src/ops.rs # 加 attention 节点、transpose_4d12 节点；rope(+period)
|
||
└── tests/autograd.rs # +rope_batched / transpose_4d12 / attention(batched) grad-check
|
||
crates/xtrain-model/
|
||
├── src/model.rs # forward_batched(ids,batch)/loss_batched；attention 走 fused 批量 op；删 causal_mask 叶
|
||
├── src/lib.rs # +batched_ids_tensor
|
||
├── tests/batched.rs # 新：batched == looped 单序列（logits + 梯度）
|
||
├── tests/parity_dump.rs + parity.py # PyTorch 对拍改 B=2,S=4（per-seq RoPE + per-seq mask）
|
||
crates/xtrain-train/src/train_loop.rs # 真 batch：一次 batched forward/backward 替代 loop+SUM；clip pre-scale 1.0
|
||
crates/xtrain-distributed/
|
||
├── src/ddp.rs # 每 rank 一次 batched forward（b_local 条）；clip pre-scale 1.0
|
||
└── tests/ddp_correctness.rs # 单卡基线也改 batched（对齐）
|
||
```
|
||
|
||
## Key Design Decisions
|
||
|
||
### ① 为什么"摊平 linears"是主要收益
|
||
|
||
GPU 填不满是因为 GEMM 太小（`[256, 384]` 这种）。摊平成 `[B*S, dim]` 后，B=16/S=256 时左矩阵是
|
||
`[4096, 384]`——同一个 cuBLAS GEMM kernel，但 M 大了 16×，一次 launch 干 16× 的活，启动开销被摊薄，
|
||
SM 占用率上去。**embedding / rms_norm / silu / CE 等本就按行**，摊平后只是行数变多，**零改动复用**。
|
||
|
||
### ② RoPE 位置 per-sequence 复位（正确性陷阱）
|
||
|
||
摊平后第 `r` 行是序列 `r/S` 的第 `r%S` 个 token。RoPE 角度 = `pos * freq`，**`pos` 必须是 `r%S`**，
|
||
否则第 2 条及之后的序列会用错位置、与单序列结果不一致、PyTorch 对拍直接挂。kernel 加 `period`
|
||
（= 序列长度 S）参数，`pos = tok % period`；`period == tokens`（单序列）时退化为原"位置=行号"。
|
||
反向同样按 `period` 取角度（RoPE 是正交映射，反向 = 逆旋转，不需缓存 forward）。
|
||
|
||
### ③ Attention：先摊平到 `[B*nh, S, hd]`，fused 批量 SDPA
|
||
|
||
q/k/v 投影后是 `[B*S, nh, hd]`（顺序 b,s,head,hd）。要喂批量 GEMM，需排成 `[B*nh, S, hd]`（每个 batch
|
||
元素 = 一个 (b, head) 的整条序列）：
|
||
|
||
```text
|
||
[B*S, nh, hd] → reshape [B, S, nh, hd] → transpose_4d12 → [B, nh, S, hd] → reshape [B*nh, S, hd]
|
||
```
|
||
|
||
`transpose_4d12`（轴 1,2 互换，新加的结构 op，自反式 backward）是关键。然后 **fused attention op**：
|
||
|
||
```text
|
||
scores[B*nh,S,S] = Q · Kᵀ (cublasSgemmStridedBatched)
|
||
P[B*nh,S,S] = softmax(causal(scores / √hd)) (launch_softmax_causal，行内做 mask+scale)
|
||
out[B*nh,S,hd] = P · V (cublasSgemmStridedBatched)
|
||
```
|
||
|
||
**整个 attention 不论 B*nh 多大，都是 3 次 kernel launch**（2 批量 GEMM + 1 softmax），没有 per-head /
|
||
per-seq 的 Python 循环，也没有 host 往返。
|
||
|
||
**因果 mask 内联在 softmax kernel**：第 `r` 行的 query 位置 = `r % S`，列 `j > r%S` 是未来 → 概率置 0，
|
||
不需要 `[S,S]` 的 `-1e9` 加性 mask 张量（连 `causal_mask` 叶子都删了）。`1/√hd` 的 scale 也折进
|
||
softmax（取 max/exp 前乘），省一个 scale pass。
|
||
|
||
> **踩坑记录（重要）**：T10 的第一版 attention 用"per-(batch,head) 循环 + `split_heads`/`merge_heads`"，
|
||
> 而 split/merge_heads 内部走 **host 往返**（`to_device(Cpu)` 再传回）。linears 确实摊平成大 GEMM 了，
|
||
> 但 B=16 时 attention 路径里成千上万次小 kernel + 几 MB 的 host memcpy **反而把吞吐拖到 1127 tok/s
|
||
> （比单序列还慢）**。教训：摊平 linears 是必要非充分——attention 的 host 往返 / launch 风暴必须一起干掉。
|
||
> fused 批量 SDPA 之后才到 **25.6K tok/s**。
|
||
|
||
### ④ Attention backward（手写，grad-check 兜底）
|
||
|
||
fused op 的反向是标准 attention 反传，全用批量 GEMM + 复用既有 `softmax_backward`：
|
||
|
||
```text
|
||
dP = dOut · Vᵀ ; dV = Pᵀ · dOut
|
||
dScores = softmax_jacobian(P, dP) · (1/√hd) (复用行 softmax 反向；mask 处 P=0 → 自动零梯度)
|
||
dQ = dScores · K ; dK = dScoresᵀ · Q
|
||
```
|
||
|
||
新加的 `attention` / `transpose_4d12` 都有 finite-diff grad-check（见 ⑥）。
|
||
|
||
### ⑤ 训练 loop：真 batch 替代 loop+SUM
|
||
|
||
旧：循环 B 次 `model.loss(单序列)` + `backward()`，靠 tape 把 B 份梯度 SUM，clip 时 ×`1/B` 还原 batch-mean。
|
||
新：`batched_ids_tensor` 把 B 条序列摊平成 `[B*S]`，**一次** `loss_batched` + **一次** `backward()`。
|
||
CE 已是 B*S 行的 mean = batch-mean loss，**backward 直接给出 batch-mean 梯度**，所以 **clip pre-scale = 1.0**
|
||
（不再有 loop+SUM+×1/B）。
|
||
|
||
**DDP 同理且保持等价**：每 rank 跑 `b_local = B_global/world` 条的**一次** batched forward → backward
|
||
梯度 = 本地 mean `Σ_local/b_local`；`all_reduce_average`（跨 rank 求和 /world）得
|
||
`Σ_global/(world·b_local) = Σ_global/B_global` = 全局 batch-mean → clip pre-scale 也 = 1.0。
|
||
单卡基线（`ddp_correctness` 里）同步改成对全局 batch 的一次 batched forward，二者对齐。
|
||
|
||
## 验证方法
|
||
|
||
**双闸门，都必须过。**
|
||
|
||
### 正确性（无回归）
|
||
|
||
- **算子级 finite-diff**（`xtrain-autodiff`，15 个）：新增 `rope_batched`（per-seq 位置）、`transpose_4d12`、
|
||
`attention(batched)` 的 dQ/dK/dV，连同既有 12 个全过。
|
||
实测：`attn(batched) dQ 7.5e-3 / dK 1.5e-2 / dV 2.9e-4`，`transpose_4d12 8.2e-5`，`rope_batched 4.5e-4`。
|
||
- **batched == looped 单序列**（`xtrain-model/tests/batched.rs`，新）：同一组权重，batched forward 对
|
||
"逐序列单独 forward 再拼接" 的 logits + 每参数梯度逐一对比。实测 **logits 完全一致(0.0)**，
|
||
梯度 max rel **6.4e-4**，loss 完全一致。
|
||
- **PyTorch 对拍 B>1**（`parity.py`，改 B=2/S=4）：等价 PyTorch 模型（per-seq RoPE `pos=row%S`、per-seq
|
||
因果 mask、QK-norm、SwiGLU）对拍 forward logits + 全部 25 个参数梯度。实测 **loss relerr 5e-8、logits
|
||
6.9e-6、梯度全在 rtol 2e-2 内**。
|
||
- **overfit**（27/27 token）、**checkpoint 逐位**、**AdamW 对 torch**、**DDP 跨 rank 参数 bit-identical
|
||
(0.0) + DDP loss 对单卡 5.7e-7**：全过。
|
||
- **xserv 闭环**：短训→导出→xserv serve，对 xtrain 贪心**仍逐 token 一致**；采样路径 batch=1 仍工作。
|
||
|
||
### 吞吐（收益）
|
||
|
||
单卡 dim384/12L/12h、batch 16、seq 256，back-to-back：
|
||
|
||
| 路径 | tok/s | GPU util | 显存 |
|
||
|---|---|---|---|
|
||
| **before**（单序列 launch-bound，KI-1 基线）| ~1653 | 0–15% | ~3 GB |
|
||
| T10 第一版（looped split/merge，host 往返）| 1127 | ~11% | 14.8 GB |
|
||
| **after**（fused 批量 attention，batch 16）| **25627** | **37% 均值 / 54% 峰** | 10.2 GB |
|
||
| **after**（batch 32）| **40263** | — | — |
|
||
|
||
→ 单卡相对 KI-1 基线 **~15.5×**（batch 16）/ **~24×**（batch 32），GPU util 0–15% → 37–54%。
|
||
|
||
**DDP 4 卡（dim384, per-rank batch 32, global 128）**：1 卡 40.3K → 4 卡 47.2K tok/s（global），
|
||
仅 **~1.17×**。这是 batched forward **新暴露**的下一层瓶颈：单卡 compute 快了 15–24×后，每步**对全部
|
||
67M 参数的 eager all-reduce + host 侧 optimizer/clip 同步**成了 DDP 的主导开销（不随卡数缩小）。
|
||
注意：**单卡 batch 32 = 40K tok/s 已经把 KI-1 时代的 4 卡 3163 tok/s 干翻 ~12×**——根因（单卡
|
||
launch-bound）已修。DDP 近线性需要 **bucketed / 与 backward overlap 的 all-reduce**（KI-1 修复项 2），
|
||
此前 all-reduce 非瓶颈做了没用，现在才有意义——列为 v3+ 的 follow-up（本 Phase 范围外）。
|
||
|
||
## 给 v3 的 note（已解锁）
|
||
|
||
batched forward 就位后，v3（dim512/16L、200–300M tok）建议：**per-rank batch 16–32、seq 256，4 卡
|
||
global batch 64–128**。按单卡 ~25K tok/s（dim384，dim512 略低）、4 卡放大，200–300M tok 的训练时间从
|
||
KI-1 时代估的 17–26h 压到**数小时级**。bucketed/overlapped all-reduce（KI-1 修复项 2）现在才有意义，
|
||
作为 v3 之后的进一步优化项。
|