| #!/bin/bash |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| set -euo pipefail |
|
|
| SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" |
| REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" |
|
|
| BENCHMARK="${1:-libero}" |
| NUM_GPUS="${2:-1}" |
| OUTPUT_DIR="${OUTPUT_DIR:-$REPO_ROOT/output/quickstart_${BENCHMARK}}" |
| PER_DEVICE_BS="${PER_DEVICE_BS:-1}" |
| MAX_STEPS="${MAX_STEPS:-5}" |
|
|
| echo "============================================" |
| echo " DreamZero Quick Start" |
| echo "============================================" |
| echo " Benchmark: $BENCHMARK" |
| echo " GPUs: $NUM_GPUS" |
| echo " Max steps: $MAX_STEPS" |
| echo " Output: $OUTPUT_DIR" |
| echo "============================================" |
|
|
| |
| CKPT_DIR="${CHECKPOINT_DIR:-$REPO_ROOT/checkpoints}" |
| WAN22_DIR="$CKPT_DIR/Wan2.2-TI2V-5B" |
| TOKENIZER_DIR="$CKPT_DIR/umt5-xxl" |
| CLIP_DIR="$CKPT_DIR/clip-encoder" |
|
|
| if [ ! -d "$WAN22_DIR" ]; then |
| echo "ERROR: 未找到 Wan2.2 权重: $WAN22_DIR" |
| echo "请先运行: bash scripts/data/download_checkpoints.sh" |
| exit 1 |
| fi |
| if [ ! -d "$TOKENIZER_DIR" ]; then |
| echo "ERROR: 未找到 tokenizer: $TOKENIZER_DIR" |
| echo "请先运行: bash scripts/data/download_checkpoints.sh" |
| exit 1 |
| fi |
| if [ ! -f "$CLIP_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" ]; then |
| echo "ERROR: 未找到 CLIP encoder: $CLIP_DIR" |
| echo "请先运行: bash scripts/data/download_checkpoints.sh" |
| exit 1 |
| fi |
|
|
| |
| case "$BENCHMARK" in |
| libero) |
| DATA_CFG="dreamzero/libero" |
| NUM_FRAMES=12 |
| ACTION_HORIZON=12 |
| NUM_VIEWS=1 |
| MAX_STATE_DIM=44 |
| MAX_ACTION_DIM=32 |
| NUM_FRAME_PER_BLOCK=2 |
| NUM_ACTION_PER_BLOCK=12 |
| LR=1e-5 |
| ;; |
| manifeel) |
| DATA_CFG="dreamzero/manifeel" |
| NUM_FRAMES=12 |
| ACTION_HORIZON=12 |
| NUM_VIEWS=3 |
| MAX_STATE_DIM=44 |
| MAX_ACTION_DIM=32 |
| NUM_FRAME_PER_BLOCK=2 |
| NUM_ACTION_PER_BLOCK=12 |
| LR=1e-5 |
| ;; |
| robotwin) |
| DATA_CFG="dreamzero/robotwin" |
| NUM_FRAMES=12 |
| ACTION_HORIZON=12 |
| NUM_VIEWS=1 |
| MAX_STATE_DIM=44 |
| MAX_ACTION_DIM=32 |
| NUM_FRAME_PER_BLOCK=2 |
| NUM_ACTION_PER_BLOCK=12 |
| LR=1e-5 |
| ;; |
| *) |
| echo "ERROR: 未知 benchmark: $BENCHMARK (可选: libero, manifeel, robotwin)" |
| exit 1 |
| ;; |
| esac |
|
|
| |
| if [ "$NUM_GPUS" -le 2 ]; then |
| DEEPSPEED_CFG="groot/vla/configs/deepspeed/zero2.json" |
| elif [ "$NUM_GPUS" -le 8 ]; then |
| DEEPSPEED_CFG="groot/vla/configs/deepspeed/zero2_offload.json" |
| else |
| DEEPSPEED_CFG="groot/vla/configs/deepspeed/zero3_multinode.json" |
| fi |
|
|
| cd "$REPO_ROOT" |
|
|
| torchrun --standalone --nproc_per_node "$NUM_GPUS" \ |
| groot/vla/experiment/experiment.py \ |
| report_to=none \ |
| data="$DATA_CFG" \ |
| train_architecture=full \ |
| model=dreamzero/vla \ |
| model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \ |
| model/dreamzero/transform=dreamzero_cotrain \ |
| num_frames="$NUM_FRAMES" \ |
| action_horizon="$ACTION_HORIZON" \ |
| num_views="$NUM_VIEWS" \ |
| max_state_dim="$MAX_STATE_DIM" \ |
| max_action_dim="$MAX_ACTION_DIM" \ |
| num_frame_per_block="$NUM_FRAME_PER_BLOCK" \ |
| num_action_per_block="$NUM_ACTION_PER_BLOCK" \ |
| num_state_per_block=1 \ |
| max_chunk_size=4 \ |
| frame_seqlen=50 \ |
| image_resolution_width=320 \ |
| image_resolution_height=160 \ |
| per_device_train_batch_size="$PER_DEVICE_BS" \ |
| max_steps="$MAX_STEPS" \ |
| save_strategy=no \ |
| optim=adamw_bnb_8bit \ |
| training_args.learning_rate="$LR" \ |
| training_args.deepspeed="$DEEPSPEED_CFG" \ |
| training_args.bf16=true \ |
| training_args.tf32=true \ |
| training_args.eval_bf16=true \ |
| output_dir="$OUTPUT_DIR" \ |
| dit_version="$WAN22_DIR" \ |
| text_encoder_pretrained_path="$WAN22_DIR/models_t5_umt5-xxl-enc-bf16.pth" \ |
| image_encoder_pretrained_path="$CLIP_DIR/models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \ |
| vae_pretrained_path="$WAN22_DIR/Wan2.2_VAE.pth" \ |
| tokenizer_path="$TOKENIZER_DIR" |
|
|
| echo "" |
| echo "============================================" |
| echo " Quick start 完成! (benchmark: $BENCHMARK)" |
| echo "============================================" |
|
|
