File size: 3,451 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 | #!/bin/bash
# DreamZero RoboTwin Full Fine-Tuning Script (Wan2.2-TI2V-5B, custom RobotWinDataset)
export HYDRA_FULL_ERROR=1
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
# Activate DreamZero virtual environment
VENV_DIR="/root/autodl-tmp/venvs/dreamzero"
if [ -f "$VENV_DIR/bin/activate" ]; then
source "$VENV_DIR/bin/activate"
fi
# ============ CONFIGURATION ============
ROBOTWIN_DATA_DIR=${ROBOTWIN_DATA_DIR:-"/root/autodl-tmp/data/robotwin_gear"}
OUTPUT_DIR=${OUTPUT_DIR:-"$REPO_ROOT/checkpoints/dreamzero_robotwin_full"}
if [ -z "${NUM_GPUS:-}" ]; then
NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
fi
PER_DEVICE_BS=${PER_DEVICE_BS:-1}
CKPT_DIR="${CHECKPOINT_DIR:-/root/autodl-tmp/checkpoints}"
WAN22_DIR="$CKPT_DIR/Wan2.2-TI2V-5B"
TOKENIZER_DIR="$CKPT_DIR/umt5-xxl"
CLIP_DIR="$CKPT_DIR/DreamZero-DROID"
# ========================================
# Auto-download weights if missing
source /etc/network_turbo 2>/dev/null
if [ ! -d "$WAN22_DIR" ] || [ -z "$(ls -A "$WAN22_DIR" 2>/dev/null)" ]; then
echo "Downloading Wan2.2-TI2V-5B..."
huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir "$WAN22_DIR"
fi
if [ ! -d "$TOKENIZER_DIR" ] || [ -z "$(ls -A "$TOKENIZER_DIR" 2>/dev/null)" ]; then
echo "Downloading umt5-xxl..."
huggingface-cli download google/umt5-xxl --local-dir "$TOKENIZER_DIR"
fi
if [ ! -d "$ROBOTWIN_DATA_DIR" ]; then
echo "ERROR: RoboTwin data not found at $ROBOTWIN_DATA_DIR"; exit 1
fi
cd "$REPO_ROOT"
# ZeRO config
if [ "$NUM_GPUS" -le 2 ]; then
DEEPSPEED_CFG=${DEEPSPEED_CFG:-"groot/vla/configs/deepspeed/zero2_offload.json"}
else
DEEPSPEED_CFG=${DEEPSPEED_CFG:-"groot/vla/configs/deepspeed/zero2.json"}
fi
torchrun --standalone --nproc_per_node "$NUM_GPUS" \
groot/vla/experiment/experiment.py \
report_to="${REPORT_TO:-none}" \
data=dreamzero/robotwin \
wandb_project="${WANDB_PROJECT:-dreamzero-robotwin-sft}" \
train_architecture=full \
model=dreamzero/vla \
model/dreamzero/action_head=wan_flow_matching_action_tf_wan22 \
model/dreamzero/transform=dreamzero_cotrain \
num_frames=12 action_horizon=12 state_horizon=1 num_views=1 \
num_frame_per_block=2 num_action_per_block=12 \
num_state_per_block=1 max_chunk_size=4 frame_seqlen=50 \
image_resolution_width=320 image_resolution_height=160 \
max_state_dim=44 max_action_dim=32 \
seed=42 \
training_args.learning_rate="${LR:-1e-5}" \
training_args.deepspeed="$DEEPSPEED_CFG" \
save_steps="${SAVE_STEPS:-2000}" \
training_args.warmup_ratio=0.05 \
output_dir="$OUTPUT_DIR" \
per_device_train_batch_size="$PER_DEVICE_BS" \
max_steps="${MAX_STEPS:-200000}" \
weight_decay=1e-5 save_total_limit=2 \
upload_checkpoints=false bf16=true tf32=true eval_bf16=true \
dataloader_pin_memory=true dataloader_num_workers=4 \
save_lora_only=false save_strategy=steps \
robotwin_dataset_dir="$ROBOTWIN_DATA_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 " RoboTwin training finished! Output: $OUTPUT_DIR"
echo "============================================"
|