Gahow Wang
6b8c1e4e0f
Counter-based (stateless) RNG → Bernoulli(keep=1-p) mask, inverted 1/(1-p) scaling at train, identity at eval. New autodiff `dropout` op (fwd generates + applies mask, bwd applies the SAME cached mask). Wired at the two residual-path sites (attn / ffn outputs); attention-probs dropout deliberately skipped (fused SDPA doesn't materialise probs). Documents the RNG choice, per-site deterministic seed (so T13 recompute reproduces the same mask), train/eval switch, p=0 bit-identity, and the acceptance gates. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
156 lines
9.9 KiB
Markdown
156 lines
9.9 KiB
Markdown
# 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 轨迹）。
|