# 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 轨迹）。
