# 显存与内存估算（bf16 AMP + GradNorm + PCGrad）

由 `scripts/estimate_memory.py` 生成。模型规模：18 层主干（9 Dense + 9 MoE，每 MoE 层 7 路由 + 1 共享专家）+ DINOv3 ViT-B/16 + 6 层校准 + 1024 检测 token + 24 控制 token。

| 项目 | 数值 |
|---|---|
| 总参数 | **725.62 M** |
| 可训练 (Stage1, DINOv3 冻结) | 639.96 M |
| 可训练 (Stage2, DINOv3 解冻) | 725.62 M |
| 序列长度（拼接后） | 2848 |

## 显存（含 15% 余量）

| Batch Size | Stage2 峰值 | 推荐单卡 GPU | HF Sandbox 选项 |
|---:|---:|---|---|
| 1 | ~16 GB | T4 16GB（紧）/ L4 24GB | `t4-small` |
| 2 | ~18 GB | L4 24GB | `l4x1` |
| 4 | ~22 GB | L4 24GB / A10G 24GB | `a10g-small` |
| **8 (目标)** | **~30 GB** | **A10G Large 48GB / A100 40GB** | **`a10g-large`** |
| 16 | ~46 GB | A100 80GB / H100 80GB | `a100-large` |

显存细分（BS=8 Stage2）：
- 权重 (bf16): 1.35 GB
- 优化器 (AdamW fp32 m+v + 主副本): 8.11 GB
- 主激活 (bf16, 18 + 6 层): ~12.6 GB
- PCGrad retain_graph 开销: ~6.3 GB
- 缓冲 / cuDNN workspace / 碎片: ~2 GB

如显存不足：
- 开 `gradient_checkpointing`（激活降至 ~1/3，可把 BS=8 塞进 A10G 24GB 大约 28GB）
- BS=4 + `grad_accum_steps=2` 等价 BS=8 训练
- 关 PCGrad（节省 ~6 GB），但牺牲多任务收敛质量

## 主机内存 / 磁盘

| 项目 | 数值（BS=8） |
|---|---|
| 主机 RAM 推荐 | ≥ 32 GB（DataLoader 4 workers × prefetch 2 + 模型 CPU 副本） |
| 磁盘（一个 weather 子集，sandbox 验证） | ~50 GB |
| 磁盘（synthetic 全量 121 帧 × 7 weather × 5843 clip） | ~700 GB |
| 磁盘（synthetic + lidar + hdmap 全量） | ~3 TB |

## 设备选择建议

- **本地烟囱（CPU/小卡）**：`scripts/smoke_test.py` 用极小张量验证 forward+backward，不需要 GPU。
- **HF Sandbox**：`a10g-large` (48 GB)，BS=8 + bf16 + PCGrad 一次成功；约 $1.05/小时（HF 价格随时调整请以官方为准）。
- **HF Jobs 全量训练**：`a100x1` (80 GB) 或 `h100x1`，BS=8~16。

## 复现命令

```bash
# 升级依赖到最新（写入 requirements.lock.txt）
python scripts/update_deps.py --torch-index https://download.pytorch.org/whl/cu124

# 估算
python scripts/estimate_memory.py

# Sandbox 推送
python scripts/push_to_sandbox.py --repo your-username/wjad-sandbox --gpu a10g-large

# Jobs 全量
python scripts/push_to_jobs.py --repo your-username/wjad --flavor a100x1
```