Hugging Face's logo

poet70
/
vla-sft-code-dreamzero

Model card Files
xet
 Community
vla-sft-code-dreamzero
File size: 4,592 Bytes
ac29381
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
#!/bin/bash
# DreamZero 单节点快速启动脚本 (5 步验证)
#
# 用法:
#   bash scripts/train/quick_start.sh libero          # 单卡验证 LIBERO
#   bash scripts/train/quick_start.sh manifeel 4      # 4 卡 ManiFeel
#   bash scripts/train/quick_start.sh robotwin 2      # 2 卡 RoboTwin
#
# 前置条件: 权重文件已下载到 checkpoints/ 目录

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

# ============ 自动选择 DeepSpeed 配置 ============
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 "============================================"