# DreamZero 训练指南

---

## 1. 超参调优

### 全量 SFT 推荐参数

| 参数 | 推荐值 | 说明 |
|------|--------|------|
| `learning_rate` | 1e-5 | 全量微调，比 LoRA 小 10x |
| `weight_decay` | 1e-5 | 轻微权重衰减 |
| `warmup_ratio` | 0.05 | 5% 步数线性预热 |
| `lr_scheduler_type` | cosine | 余弦退火 |
| `adam_beta1` | 0.95 | Adam beta1 |
| `adam_beta2` | 0.999 | Adam beta2 |

### LoRA 推荐参数

| 参数 | 推荐值 |
|------|--------|
| `learning_rate` | 1e-4 |
| LoRA rank | 64 |
| LoRA alpha | 128 |
| LoRA target modules | q_proj, k_proj, v_proj, o_proj |

### 学习率调整策略

- **从零开始训练**: `learning_rate=1e-5`, `warmup_ratio=0.05`
- **从 checkpoint resume**: 学习率通常重置为初始值或减半
- **Loss 震荡**: 减小学习率至 5e-6
- **Loss 不降**: 检查数据预处理是否正确

---

## 2. Batch Size 和 Global Batch Size

```bash
# global_batch_size = per_device_train_batch_size * num_gpus * gradient_accumulation_steps
# 32 GPU 训练: 推荐 global_batch_size=128
global_batch_size=128
per_device_train_batch_size=1
# gradient_accumulation_steps = 128 / (1 * 32) = 4
```

---

## 3. 显存优化决策树

```
训练时 OOM?
├── 使用 8-bit Adam? → 未使用 → 加 optim=adamw_bnb_8bit
├── 开 gradient_checkpointing? → 未开 → 加 model.gradient_checkpointing=true
├── 可降低视频分辨率? → 是 → 降低 image_resolution_width/height
├── 可减少帧数? → 是 → 减少 num_frames
└── 以上都不行 → 使用 ZeRO-3 + CPU offload
```

### 各优化节约的显存

| 优化手段 | 显存节约 | 计算开销 |
|---------|---------|---------|
| 8-bit Adam | ~50% 优化器显存 | 无（更慢的更新） |
| Gradient Checkpointing | ~30% 激活显存 | ~15% 计算时间 |
| ZeRO-3 (vs ZeRO-2) | ~20% 总显存 | ~5% 通信开销 |
| CPU Offload | 额外 ~30% | ~20% 时间 |

---

## 4. Checkpoint 管理

### 保存

```bash
save_steps=2000          # 每 2000 步保存
save_total_limit=4       # 保留最近 4 个 + 最优
save_strategy=steps      # 按步数保存
```

### Resume

```bash
# 从指定 checkpoint 恢复
torchrun ... experiment.py \
  ++training_args.resume_from_checkpoint=./output/libero_full/checkpoint-50000
```

**注意**: ZeRO-3 checkpoint 按 rank 分片保存，恢复时 GPU 数量必须一致。

### ZeRO-3 Checkpoint 合并

如果需要在不同 GPU 数之间迁移 checkpoint：

```bash
# 使用 DeepSpeed 的 zero_to_fp32.py
python zero_to_fp32.py \
  --checkpoint_dir ./checkpoint-50000 \
  --output_file ./pytorch_model.bin
```

---

## 5. 多节点训练注意事项

### NCCL 配置

```bash
export NCCL_DEBUG=WARN           # 减少日志量
export NCCL_IB_DISABLE=0         # 启用 InfiniBand
export NCCL_IB_TIMEOUT=22        # IB 超时
export NCCL_IB_RETRY_CNT=4       # IB 重试次数
export NCCL_SOCKET_IFNAME=^docker0,lo  # 排除 docker/loopback
```

### 常见问题

| 问题 | 原因 | 解决 |
|------|------|------|
| NCCL timeout | 网络配置错误 | 检查 NCCL_IB_*、NCCL_SOCKET_IFNAME |
| 训练挂起 | DataLoader worker 死锁 | 设置 `dataloader_pin_memory=false` |
| Loss 不一致 | 数据加载不统一 | 检查 `seed` 和 DataLoader shuffle |
| OOM | 显存不足 | 参考显存优化决策树 |

---

## 6. DataLoader 优化

| 场景 | num_workers | pin_memory | prefetch_factor |
|------|-------------|------------|-----------------|
| 单 GPU | 4 | false | 2 |
| 单节点 8 GPU | 8 | true | 4 |
| 多节点 | 8 | true | 4 |

**注意**: `dataloader_pin_memory=true` 在 decord 视频后端可能导致问题，如遇训练挂起可设为 false。

---

## 7. 实验跟踪（Wandb）

```bash
# 配置 Wandb
export WANDB_API_KEY=your_key
export WANDB_ENTITY=your_entity

# 启动训练
torchrun ... experiment.py \
  report_to=wandb \
  wandb_project=dreamzero-sft \
  ...
```

Wandb 记录内容：
- `loss`: 总损失
- `dynamics_loss_avg`: 视频预测损失
- `action_loss_avg`: 动作预测损失
- `learning_rate`: 学习率
- `grad_norm`: 梯度范数
- `training_step_time`: 每步时间
- `model_forward_time`: 前向时间

---

## 8. 从 LoRA 切换到全量 SFT

```bash
# LoRA
train_architecture=lora
optim=adamw_torch
learning_rate=1e-4

# 全量 SFT
train_architecture=full
optim=adamw_bnb_8bit   # 必须！否则 5B 优化器状态 ~40GB
learning_rate=1e-5
```

---

## 9. 基准 Loss 参考

| Benchmark | Step 1 | Step 5 | Step 100 | Step 1k | Step 10k | Step 100k |
|-----------|--------|--------|----------|---------|----------|-----------|
| LIBERO | ~2-3 | ~1.0-1.5 | ~0.8-1.0 | ~0.6-0.8 | ~0.5-0.7 | ~0.35-0.5 |
| ManiFeel | ~1.5-2.5 | ~0.8-1.2 | ~0.6-0.8 | ~0.5-0.6 | ~0.35-0.5 | ~0.25-0.35 |
| RoboTwin | ~3-5 | ~1.5-2.0 | ~1.0-1.5 | ~0.8-1.2 | ~0.6-0.9 | ~0.4-0.6 |

Loss 持续高于参考范围的可能原因：
- 学习率不合适
- 数据预处理不一致（state/action 未正确 padding）
- 视频帧率与 state 频率不对齐

---

## 10. 故障排除

### "CUDA out of memory"

```bash
# 添加这些参数
per_device_train_batch_size=1
model.gradient_checkpointing=true
optim=adamw_bnb_8bit
training_args.deepspeed="groot/vla/configs/deepspeed/zero2_offload.json"
```

### "RuntimeError: Expected all tensors to be on the same device"

可能是 VAE 或模型设备不匹配。检查 `experiment.py` 中的设备分配。

### 训练突然变慢

- 检查 CPU 内存（dataloader_workers 可能导致内存交换）
- 检查磁盘 IO（数据是否在 SSD 上）
- 减少 dataloader_num_workers

### NCCL 初始化失败

```bash
export NCCL_DEBUG=INFO        # 查看详细 NCCL 日志
export NCCL_IB_DISABLE=1      # 如果 IB 不可用
export NCCL_SOCKET_IFNAME=eth0  # 指定网络接口
```