xtrain/docs/17-dropout.md

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<bool>）+ 在 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<u64>，每个训练 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<bool> 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 轨迹）。