File size: 2,838 Bytes
0ec5e1f | 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 | #!/bin/bash
# GUI-Shift GRPO Training Script
set -e
# Default values
MODEL_NAME="Qwen/Qwen2.5-VL-7B-Instruct"
DATA_DIR="./data/gui_transition/filtered"
OUTPUT_DIR="./checkpoints/gui-shift"
K=1
NUM_GPUS=8
# Parse arguments
while [[ $# -gt 0 ]]; do
case $1 in
--model_name_or_path)
MODEL_NAME="$2"
shift 2
;;
--data_dir)
DATA_DIR="$2"
shift 2
;;
--output_dir)
OUTPUT_DIR="$2"
shift 2
;;
--k)
K="$2"
shift 2
;;
--num_gpus)
NUM_GPUS="$2"
shift 2
;;
*)
echo "Unknown option: $1"
exit 1
;;
esac
done
echo "=== GUI-Shift GRPO Training ==="
echo "Model: $MODEL_NAME"
echo "Data: $DATA_DIR"
echo "Output: $OUTPUT_DIR"
echo "K value: $K"
echo "GPUs: $NUM_GPUS"
echo ""
# Find data file for the specified K
DATA_FILE="${DATA_DIR}/k${K}_transition_filtered.jsonl"
if [ ! -f "$DATA_FILE" ]; then
DATA_FILE="${DATA_DIR}/k${K}_transition.jsonl"
if [ ! -f "$DATA_FILE" ]; then
echo "Error: Could not find data file for k=$K in $DATA_DIR"
exit 1
fi
fi
# Image folder (relative to trajectory data)
IMAGE_FOLDER="${DATA_DIR}/images"
echo "Using data file: $DATA_FILE"
echo "Using image folder: $IMAGE_FOLDER"
echo ""
# GRPO training hyperparameters (from paper Appendix A)
# learning_rate: 1e-6, temperature: 0.9, num_generations: 8
# num_train_epochs: 4, max_prompt_length: 1024, max_completion_length: 256
# per_device_train_batch_size: 2, gradient_accumulation_steps: 8
# epsilon: 0.2, beta: 0.04
export EXP_NAME="gui-shift-k${K}"
mkdir -p "runs/${EXP_NAME}/log"
# Note: This requires the VLM-R1 framework to be installed
# See: https://github.com/om-ai-lab/VLM-R1
torchrun --nproc_per_node="$NUM_GPUS" \
--nnodes="1" \
--node_rank="0" \
--master_addr="127.0.0.1" \
--master_port="12349" \
src/training/gui_grpo_trainer.py \
--model_name_or_path "$MODEL_NAME" \
--data_file_paths "$DATA_FILE" \
--image_folders "$IMAGE_FOLDER" \
--output_dir "$OUTPUT_DIR" \
--per_device_train_batch_size 2 \
--gradient_accumulation_steps 8 \
--gradient_checkpointing true \
--logging_steps 1 \
--num_train_epochs 4 \
--max_steps -1 \
--bf16 \
--attn_implementation flash_attention_2 \
--run_name "$EXP_NAME" \
--save_steps 400 \
--num_generations 8 \
--max_completion_length 256 \
--reward_funcs format accuracy \
--beta 0.04 \
--learning_rate 1.0e-6 \
--lr_scheduler_type cosine \
--warmup_ratio 0.1 \
--report_to wandb \
--dataset-name not_used \
--freeze_vision_modules true \
echo ""
echo "Training complete! Model saved to: $OUTPUT_DIR"
|