#!/bin/bash
# GUI-Shift GRPO Training Script
# Based on arxiv 2505.12493
# Hardware: 8x H100 recommended. For smaller GPUs, reduce per_device_train_batch_size.

# Step 1: Build the K-step GUI Transition dataset
python build_data.py

# Step 2: Run GRPO training
torchrun --nproc_per_node="8" \
    train_gui_shift.py \
    --deepspeed zero3.json \
    --model_name_or_path Qwen/Qwen2.5-VL-7B-Instruct \
    --data_file_paths data/ui_transition_training_2000_k_1_no_reasoning.jsonl \
    --image_folders data \
    --freeze_vision_modules true \
    --max_prompt_length 1024 \
    --max_completion_length 256 \
    --num_generations 8 \
    --per_device_train_batch_size 2 \
    --gradient_accumulation_steps 8 \
    --num_train_epochs 4 \
    --learning_rate 1e-6 \
    --lr_scheduler_type cosine \
    --warmup_ratio 0.1 \
    --bf16 \
    --attn_implementation flash_attention_2 \
    --beta 0.04 \
    --logging_steps 1 \
    --save_strategy epoch \
    --output_dir outputs/gui-shift-qwen2.5-vl-7b \
    --push_to_hub true \
    --hub_model_id xxilu/gui-shift-qwen2.5-vl-7b \
    --report_to trackio \
    --run_name gui-shift-k1-2k