
export WANDB_API_KEY='xxxxxxxxxxxxxxxxxx'
export TOKENIZERS_PARALLELISM=true
export NCCL_IB_DISABLE=1
export CHIEF_PORT=23333


run_name=$(date +'%Y%m%d-%H%M%S')

accelerate launch \
    --config_file configs/accelerate_config.yaml \
    --num_processes $((HOST_NUM*HOST_GPU_NUM)) \
    --num_machines ${HOST_NUM} \
    --machine_rank ${INDEX} \
    --main_process_ip ${CHIEF_IP} \
    --main_process_port ${CHIEF_PORT} \
    train_grpo_flux_kontext.py \
    --seed 42 \
    --pretrained_model_name_or_path prefix=/path/to/kontext \
    --data_json_path data/preprocessd_data.json \
    --gradient_checkpointing \
    --train_batch_size 1 \
    --dataloader_num_workers 8 \
    --gradient_accumulation_steps 4 \
    --max_train_steps 10000 \
    --learning_rate 1e-5 \
    --mixed_precision bf16 \
    --checkpointing_steps 100 \
    --wandb_project GRPO_OminiKontext \
    --output_dir runs/ \
    --run_name $run_name \
    --h 512 \
    --w 512 \
    --t 1 \
    --sampling_steps 16 \
    --eta 0.3 \
    --lr_warmup_steps 0 \
    --sampler_seed 1223627 \
    --max_grad_norm 1.0 \
    --weight_decay 0.0001 \
    --num_generations 12 \
    --shift 3 \
    --use_group \
    --timestep_fraction 0.6 \
    --init_same_noise \
    --clip_range 1e-4 \
    --adv_clip_max 5.0 \
    --max_checkpoints_to_keep 3 \
    --reward_model vlm