diff --git a/docs/17-dropout.md b/docs/17-dropout.md new file mode 100644 index 0000000..d4844f3 --- /dev/null +++ b/docs/17-dropout.md @@ -0,0 +1,155 @@ +# Phase T18: Dropout(device RNG + mask)— Design Document + +## Goal + +在已有的 tape autograd 引擎(T4)+ tiny transformer(T5)之上,**手写一个 dropout 算子**: +训练时按 Bernoulli(keep = 1−p) 生成一个 0/1 mask,丢弃的元素置 0、保留的元素按 +**inverted dropout** 乘 `1/(1−p)`(让训练期望与推理一致);推理(eval)时 dropout 是**恒等**。 +新增一个 autodiff `dropout` 节点:**前向生成并施加 mask,反向施加同一个 mask**。 +接到模型的标准位置(residual 之前的 attention / MLP 子块输出;attention-probs dropout 不做,见下)。 +通过 `Config.dropout` / `--dropout` 暴露 `p`,**默认 `p=0`**。 + +明确范围(T18 只做这些): + +1. 一个 device 端 **counter-based RNG**(Philox 风格的 bit-mix),按 `(seed, 元素下标)` 无状态地产出 + 每元素的 Bernoulli 抽样 → 0/1 mask(保留=1,丢弃=0),同 seed **逐位可复现**。 +2. 一个 `dropout` autodiff 节点(fwd 生成 mask + 施加 inverted scaling;bwd 用**缓存的同一 mask**)。 +3. 模型里加 **training / eval 开关**:train 走 dropout、eval/采样/导出走恒等。 +4. `p` 经 `Config.dropout` 落地,`bin/train` 加 `--dropout` flag。 + +明确**不做**:attention-probs(softmax 后)dropout——本项目 attention 是**一个 fused batched SDPA 算子** +(`ops::attention`,softmax 在 kernel 内部不物化 probs 给外部施加 mask),在其上插 dropout 要么改 fused kernel、 +要么退回组合路径,**不值当**且偏离「标准 residual/ffn dropout」这条主线。文档明确记下「只做 residual-path dropout」。 + +## Module Layout + +``` +csrc/ops/dropout.cu # 新:counter-based RNG mask 生成 + 施加 (fwd) / 反向施加同 mask + # fp32 + bf16 两条(activation 流可能是 bf16,对齐 cast.cu 风格) + +crates/xtrain-cuda/ +├── build.rs # 新增 dropout.cu +└── src/ffi.rs # 新增 launch_dropout_{f32,bf16} 声明(no_cuda 门控) + +crates/xtrain-tensor/ +└── src/tensor.rs # 新增 Tensor::dropout_mask_apply(p, seed) -> (out, mask) + # Tensor::dropout_apply_mask(&mask) -> out(bwd 用) + +crates/xtrain-autodiff/ +├── src/ops.rs # 新增节点 dropout(x, p, seed)(p==0 提前返回 x.clone(),零节点) +└── tests/autograd.rs # 新增:固定 seed grad-check(mask 跨 ± 扰动固定)+ 期望保持数值检查 + +crates/xtrain-model/ +├── src/config.rs # Config 加 dropout: f32(默认 0) +├── src/model.rs # train/eval 开关(Cell)+ 在 attn/ffn 子块输出接 dropout; +│ # per-site 确定性 seed(与 checkpoint recompute 兼容) +└── tests/dropout.rs # 新增:p=0 逐位一致 / eval 恒等 / 期望保持 / p>0 小训练收敛 + +crates/xtrain-train/src/bin/train.rs # --dropout flag → Config.dropout;训练 model.train(),sample 前 model.eval() +``` + +为什么 RNG/mask 落在 `tensor.rs`(而非引擎):和 `scale`/`silu` 一样是一个 device kernel 的薄封装; +autodiff 层只负责把它包成带 backward 的 `Var` 节点(对齐 T4 既有分层)。 + +## Key Design Decisions + +### RNG:counter-based(Philox 风格),无状态、可复现、与重计算兼容 + +mask[i] 只由 `(seed, i)` 决定,**不读取任何可变 RNG 状态**: + +``` +key = seed XOR (i * 0x9E3779B97F4A7C15) // golden-ratio 常数打散下标 +h = splitmix64(key) // 几轮 bit-mix(xorshift+乘法) +u = (h >> 40) as f32 / 2^24 // [0,1) 均匀 +keep = u >= p // Bernoulli(keep = 1−p) +out[i] = keep ? x[i] * (1/(1−p)) : 0 +``` + +选 counter-based 而非「per-step 推进一个全局 LCG 状态」的关键原因 = **激活重计算(T13)**: +checkpoint 的 segment 在 backward 时会**重跑一遍 forward**(`segment_fn` 再执行)。 +若 dropout 用「调用时推进的可变状态」,重跑会拿到**不同的 mask** → 梯度与前向用的 mask 不一致 → 错。 +counter-based + **每个 dropout 站点一个确定性 seed**(见下)保证:重跑同 seed → **同 mask**, +重计算依旧逐位一致(T13 的硬闸门不被 dropout 破坏)。 + +> 复现性:同一 `(seed, p, shape)` 下 mask 逐位确定;fp32/bf16 mask 判定都在 fp32 里算 `u`(bf16 仅存/取 +> activation),所以两精度的 mask **同分布**(drop 与否由 fp32 `u` 决定,不受 bf16 舍入影响)。 + +### 每个 dropout 站点的确定性 seed(兼容 checkpoint 重算) + +模型持有一个 `base_seed`(`Cell`,每个训练 step 自增一次 → 每步换 mask)。`block_forward` +收到 `block_seed = base_seed XOR layer_index`,块内两处 dropout 再各 XOR 一个站点常量 +(attn=0xA77, ffn=0xF7N)派生出**该站点的 seed**。这些都是**纯函数**(只看 `base_seed + layer_index + +站点常量`,无可变推进),所以: + +- 同一 step 内不同站点 mask 不同(seed 不同); +- checkpoint 重算 `block_forward` 时,`block_seed` 由捕获的 `base_seed`/`layer_index` 重新算出 → **同 seed → 同 mask**; +- 跨 step mask 变化(`base_seed` 每步 +1)。 + +`base_seed` 的自增放在**训练入口**(`loss_batched` 训练态调用时 advance 一次)。eval/`forward`/采样 +**不 advance、不插 dropout**(恒等)。 + +### train / eval 开关 + +`TinyTransformer` 加一个 `Cell training`(默认 **false** = eval,安全:未显式开训练就不丢弃): + +- `model.train()` / `model.eval()` 切换(builder 风格 `with_training(bool)` 也提供,给测试)。 +- `forward_batched` 里:`p > 0 && training` 才在 attn/ffn 子块输出插 `ops::dropout`;否则**完全不建 dropout 节点**。 +- 因此 **`p == 0`** 或 **eval** → forward 图与改动前**逐字节相同**(`ops::dropout` 在 `p==0` 时也提前 + `return x.clone()`,双保险)→ 满足「p=0 与无 dropout 逐位一致」回归闸门。 + +训练 loop(`train`)开 `model.train()`;`eval_loss` / `generate` / 导出 `forward` 走 eval(恒等)—— +导出的模型权重不含任何 dropout,xserv 闭环不受影响。 + +### dropout 接在哪(wiring) + +接**两处 residual-path dropout**(标准 Pre-LN transformer 位置,对齐 GPT/LLaMA 训练实践): + +``` +h = h + dropout( attention(rms_norm(h)) ) # attn 子块输出,残差前 +h = h + dropout( swiglu_mlp(rms_norm(h)) ) # ffn 子块输出,残差前 +``` + +**不做** attention-probs dropout(理由见 Goal:fused SDPA 不物化 probs)。embedding dropout 也不做(非必需)。 + +### dropout 节点的 backward(为什么 grad-check 成立) + +``` +fwd: out = x ⊙ mask ⊙ (1/(1−p)) # mask 由 seed 生成,缓存进 backward 闭包 +bwd: dx = d ⊙ mask ⊙ (1/(1−p)) # 用同一个缓存 mask +``` + +dropout 在 **固定 mask** 下是一个逐元素线性映射 `out_i = c_i · x_i`(`c_i ∈ {0, 1/(1−p)}`), +其梯度就是 `dx_i = c_i · d_i`。finite-diff grad-check 之所以成立,关键是**前向缓存的 mask 在 ± 扰动两次 +forward 里保持不变**——本设计天然满足:mask 只由 `(seed, i)` 决定,与 `x` 的值无关,扰动 `x` 不改 mask。 +(grad-check 直接对 `ops::dropout` 节点跑:同一个 `seed` 调两次 forward 得到同一 mask,函数处处可微。) + +### 与既有特性的组合 + +- **bf16(T12)**:activation 流是 bf16 时,dropout kernel 走 bf16 分支(load→fp32 判 mask→store bf16), + mask 判定在 fp32,和 cast.cu 既有 bf16 elementwise 同风格;grad 也在 activation dtype(接回 bf16 链)。 +- **重计算(T13)**:见上「counter-based + 确定性 seed」——重算 mask 与前向逐位相同,T13 闸门不破。 +- **DDP(T8)**:每 rank 独立跑自己的 forward/backward,各自的 mask 由各 rank 的 `base_seed` 决定。 + 本任务的 DDP 闸门是「loss 对单卡 / 跨 rank 参数一致」,在 **dropout 关(默认 p=0)** 的回归配置下跑, + 不引入跨 rank mask 同步需求(p>0 时各 rank mask 本就该不同,属正常 DDP 语义)。 +- **梯度累积(T16)/ flash(T14)**:本分支独立于二者,不依赖其未合并改动。 + +## 验证方法 + +全部 `#![cfg(not(no_cuda))]` 门控;本地只 `cargo check`/`fmt`,构建 + 实跑在 dash5(8× RTX 5090, sm_120)。 + +**硬闸门(全绿,诚实正确性,不放宽容差)**: + +1. **固定 seed grad-check**(`autograd.rs::dropout_bwd`):对 `ops::dropout(x, p, seed)` 同一 seed + 跑 finite-diff(mask 跨 ± 扰动固定)→ `dx` 对中心差分通过(线性 op,用 `cfg_linear` 容差)。 +2. **train/eval + 期望保持**(`dropout.rs`): + - eval 恒等:`dropout` 关时 `out == x` **逐位**; + - 期望保持:大张量、训练态、对多组随机 mask 取均值,`E[out] ≈ x`(inverted scaling 正确),给数值; + - 实际 keep 比例 ≈ `1−p`(验证 RNG 分布)。 +3. **p=0 逐位一致**(`dropout.rs`):同 init 两个模型,一个不设 dropout、一个 `dropout=0`, + 同 batch forward+backward → **logits/loss/每参数 grad 逐位相同**(`|Δ| == 0`)。 +4. **p>0 小训练收敛**(`dropout.rs`,或 dash5 短跑):小模型开 `p=0.1` 训若干步,**loss 下降、无 NaN**。 +5. **全回归套绿**:autograd grad-checks、structural、batched==looped、bf16、recompute(逐位一致)、 + overfit 27/27、AdamW(GPU bit-exact + host vs torch)、DDP(loss-match + 跨 rank)、 + **xserv 闭环**(导出 md5 vs registry、token-identical;导出/推理 dropout **关**,导出模型不受影响)。 + +dash5 capture 每个闸门的 pass + 关键数字(max rel-err、期望 vs input、p=0 的 `|Δ|`、训练 loss 轨迹)。