#!/bin/bash
# DreamZero SLURM 多节点训练启动脚本
#
# 用法:
#   sbatch scripts/cluster/slurm_train.sh libero          # 默认 4 节点 × 8 GPU
#   sbatch --nodes=8 scripts/cluster/slurm_train.sh manifeel  # 8 节点
#
# 或者覆盖参数:
#   sbatch scripts/cluster/slurm_train.sh robotwin \
#     --max_steps=100000 --training_args.learning_rate=5e-6

#SBATCH --job-name=dreamzero-sft
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=8
#SBATCH --cpus-per-task=64
#SBATCH --mem-per-cpu=8G
#SBATCH --time=72:00:00
#SBATCH --output=logs/%x-%j.out
#SBATCH --error=logs/%x-%j.err
#SBATCH --partition=gpu

set -euo pipefail

# ============ 参数解析 ============
BENCHMARK=${1:-libero}
shift || true

# ============ 环境变量 ============
export NCCL_DEBUG=WARN
export NCCL_IB_DISABLE=0
export NCCL_IB_TIMEOUT=22
export NCCL_IB_RETRY_CNT=4
export NCCL_SOCKET_IFNAME=^docker0,lo
export HF_ENDPOINT=https://hf-mirror.com

NODES=${SLURM_NNODES:-4}
GPUS_PER_NODE=${SLURM_GPUS_PER_NODE:-8}
TOTAL_GPUS=$((NODES * GPUS_PER_NODE))

# 自动获取 MASTER_ADDR
if [ -n "${SLURM_NODELIST:-}" ]; then
    MASTER_ADDR=$(scontrol show hostname "$SLURM_NODELIST" | head -n1)
else
    MASTER_ADDR="localhost"
fi
MASTER_PORT=${MASTER_PORT:-29500}

echo "============================================"
echo " DreamZero SLURM Training"
echo "============================================"
echo " Nodes:       $NODES"
echo " GPUs/node:   $GPUS_PER_NODE"
echo " Total GPUs:  $TOTAL_GPUS"
echo " Master:      $MASTER_ADDR:$MASTER_PORT"
echo " Benchmark:   $BENCHMARK"
echo "============================================"

# ============ 路径 ============
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/output/${BENCHMARK}_full_${NODES}nodes}"

# ============ 运行训练 ============
torchrun \
    --nnodes="$NODES" \
    --nproc_per_node="$GPUS_PER_NODE" \
    --rdzv_id="dz_sft_${BENCHMARK}" \
    --rdzv_backend=c10d \
    --rdzv_endpoint="$MASTER_ADDR:$MASTER_PORT" \
    "$REPO_ROOT/groot/vla/experiment/experiment.py" \
    report_to=wandb \
    data="dreamzero/${BENCHMARK}" \
    wandb_project=dreamzero-sft \
    train_architecture=full \
    model=dreamzero/vla \
    model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \
    model/dreamzero/transform=dreamzero_cotrain \
    per_device_train_batch_size=1 \
    global_batch_size=$((TOTAL_GPUS * 4)) \
    save_steps=2000 \
    save_total_limit=4 \
    training_args.learning_rate=1e-5 \
    training_args.warmup_ratio=0.05 \
    training_args.deepspeed="groot/vla/configs/deepspeed/zero3_multinode.json" \
    training_args.bf16=true \
    training_args.tf32=true \
    training_args.eval_bf16=true \
    output_dir="$OUTPUT_DIR" \
    model.gradient_checkpointing=true \
    dataloader_num_workers=8 \
    dataloader_pin_memory=true \
    optim=adamw_bnb_8bit \
    bf16=true tf32=true eval_bf16=true \
    "$@"