export http_proxy=http://star-proxy.oa.com:3128 
export https_proxy=http://star-proxy.oa.com:3128
wandb login --relogin 9f00c4cc90ef7ea81eabc2df4b2306f49cb73bf8
wandb online

mkdir -p ift_log

# data=continue_pretrain if input c, SFT_new if input s
choise=$1


if [ "$choise" == "c" ]; then
    data="continue_pretrain"
elif [ "$choise" == "s" ]; then
    data="SFT_new"
else
    echo "Invalid input"
    exit 1
fi

torchrun --standalone --nproc_per_node=8 train_gpt2.py \
    --load_checkpoint datalaw_log/power_135_0_0/state_step021672.pt \
    --input_folder ift/uniform/$data \
    --save_every 200 \
    --val_loss_every 200 \
    --run_name $data \
    --warmup_ratio 0.05 \
    --warmdown_ratio 0.9 \
    --sequence_length 512 \
    --device_batch_size 16 \
    --num_epochs 5 \
    --weight_decay 0.1 \
    --learning_rate 0.0003 \
    --batch_size 128 \
    --bf16 \
    --model_size base\
    --output_dir ift_log/ \
    --wandb_project ift \
    --grad_max_norm 1.0