#!/usr/bin/env bash
# smoke：4 卡 ZeRO-2，64 个 clip，12 步，第 6/12 步存权重并验证 4 个 val clip。
# 用法：bash scripts/smoke_train.sh [arope|plain] [额外参数...]
#   ARM=plain  或第一个位置参数选臂；GPUS 覆盖卡号（默认 0,1,2,3，必须与 accelerate_zero2_4gpu.yaml 的 4 进程一致）
set -euo pipefail
ROOT=/opt/dlami/nvme/zhiyangdeng/ActionRoPE
export DIFFSYNTH_SKIP_DOWNLOAD=True
export PYTHONPATH="$ROOT"
export CUDA_VISIBLE_DEVICES="${GPUS:-0,1,2,3}"
ARM="${1:-${ARM:-arope}}"; [ $# -gt 0 ] && shift
RUN="${RUN:-smoke_${ARM}_4gpu}"
mkdir -p "$ROOT/outputs/$RUN"
"$ROOT/.venv/bin/accelerate" launch --config_file "$ROOT/configs/accelerate_zero2_4gpu.yaml" \
  "$ROOT/actionrope/train.py" \
  --arm "$ARM" --limit 64 --max_steps 12 --save_every 6 --val_every 6 --val_n 4 \
  --warmup_steps 4 --num_workers 2 --seed 0 \
  --output "$ROOT/outputs/$RUN" "$@" 2>&1 | tee "$ROOT/outputs/$RUN/train.log"