| #!/bin/bash |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| 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 \ |
| "$@" |
|
|
