# DFlash LoRA OOM 修复记录

## OOM 根因分析

1. **SHARD_GRAD_OP (ZeRO-2)** — 每卡持有完整 Qwen3-8B 参数 (~16GB bf16)，参数未分片
2. **SDPA + 4D additive mask** — FlashAttention 不支持 4D additive mask，fallback 到 math backend，每层 materialize 完整 attention scores (`bsz × 32heads × 2048 × 2048`)
3. **大 vocab logits** — `[bsz, 2048, 151936]` bf16 ≈ 1.18GB，加上梯度和 boolean indexing 拷贝，峰值 ~3-4GB
4. **机器只有 2 张 H100**，脚本默认 `NUM_GPUS=4`

## 已完成的改动

### 1. FSDP sharding 改为 FULL_SHARD (ZeRO-3)
- 文件: `SpecForge/scripts/train_dflash_lora.py:347`
- `ShardingStrategy.SHARD_GRAD_OP` → `ShardingStrategy.FULL_SHARD`
- 效果: 参数跨卡分片，每卡省 ~8-12GB

### 2. 降 batch-size，提高 accumulation-steps
- 文件: `SpecForge/scripts/run_train_dflash_lora.sh`
- `--batch-size 2` → `1`，`--accumulation-steps 4` → `8`
- 效果: 等效 global batch size 不变，峰值显存减半

## 待验证 / 后续优化

- [ ] 运行时传 `bash run_train_dflash_lora.sh 2` 确保用 2 卡
- [x] 如仍 OOM，考虑 chunked cross-entropy loss 避免大 vocab logits 全量 materialize
- [x] 长期可探索自定义 attention kernel 支持 block-sparse mask，绕过 SDPA math fallback

### 3. flex_attention + BlockMask 替换 4D additive mask
- 文件: `SpecForge/specforge/core/dflash_lora.py`, `specforge/modeling/draft/dflash_lora.py`, `scripts/train_dflash_lora.py`
- 从非 LoRA 版 `dflash.py` 移植 `_get_or_create_block_mask()` 方法，适配 LoRA 场景 (Q_LEN == KV_LEN == seq_len)
- LoRA 版 mask: context causal + block bidirectional (非 LoRA 版是 [context, noise] concat KV)
- 用 `--attention-backend flex_attention` 启用 (默认)，退回 `--attention-backend additive` 走原有 4D mask
- HuggingFace model 用 `attn_implementation="flex_attention"` 加载
- 效果: 不再 fallback 到 SDPA math backend，省去 `[bsz, heads, seq, seq]` attention scores 的显存

### 4. chunked cross-entropy loss
- 文件: `SpecForge/specforge/core/dflash_lora.py`, `specforge/modeling/draft/dflash_lora.py`, `scripts/train_dflash_lora.py`
- 从非 LoRA 版 `dflash.py` 移植 `_chunked_lm_loss()` 方法
- 分 chunk 过 lm_head + CE loss + gradient checkpointing，避免 materialize 完整 `[bsz, seq, vocab]` logits
- 用 `--lm-head-chunk-size 256` 启用 (默认 0 = 不启用)
- `DFlashLoRADraftModel.forward()` 新增 `output_hidden_states` 参数，chunked 时返回 hidden states
- 效果: logits 峰值显存从 O(seq_len × vocab_size) 降至 O(chunk_size × vocab_size)