#!/bin/bash
# DreamZero 多节点 torchrun 启动脚本（不依赖 SLURM）
#
# 用法:
#   bash scripts/cluster/torchrun_multinode.sh <NODE_LIST> <BENCHMARK> [额外参数]
#
# 示例:
#   bash scripts/cluster/torchrun_multinode.sh "node1,node2,node3,node4" libero
#   bash scripts/cluster/torchrun_multinode.sh "10.0.0.1,10.0.0.2" manifeel --max_steps=100000
#
# 前置条件:
#   - 所有节点可互相 SSH 免密访问
#   - 代码和数据在所有节点上路径一致
#   - 所有节点已安装所需依赖

set -euo pipefail

if [ $# -lt 2 ]; then
    echo "用法: $0 <NODE_LIST> <BENCHMARK> [额外参数]"
    echo "示例: $0 node1,node2,node3,node4 libero"
    exit 1
fi

NODE_LIST="$1"
BENCHMARK="$2"
shift 2

# ============ 解析节点列表 ============
IFS=',' read -ra NODES <<< "$NODE_LIST"
NNODES=${#NODES[@]}
FIRST_NODE="${NODES[0]}"
MASTER_PORT=${MASTER_PORT:-29500}

echo "============================================"
echo " DreamZero Multi-Node Training"
echo "============================================"
echo " Nodes:      $NNODES ($NODE_LIST)"
echo " Master:     $FIRST_NODE:$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_${NNODES}nodes}"

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

# ============ 运行 ============
torchrun \
    --nnodes="$NNODES" \
    --nproc_per_node=8 \
    --rdzv_id="dz_${BENCHMARK}" \
    --rdzv_backend=c10d \
    --rdzv_endpoint="$FIRST_NODE:$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 \
    save_steps=2000 \
    save_total_limit=4 \
    training_args.learning_rate=1e-5 \
    training_args.deepspeed="groot/vla/configs/deepspeed/zero3_multinode.json" \
    training_args.bf16=true \
    training_args.tf32=true \
    output_dir="$OUTPUT_DIR" \
    dataloader_num_workers=8 \
    optim=adamw_bnb_8bit \
    "$@"