#!/bin/bash
GPU_NUM=4
PER_DEVICE_BATCH_SIZE=24
GRADIENT_ACCUMULATION=8
CUDA_VISIBLE_DEVICES=3,5,6,7 torchrun --nproc_per_node ${GPU_NUM} --nnodes=1 --node_rank=0 llm_train.py \
    --model_name llama_1_7b_hf \
    --model_path ./result/llama_1_7b_hf_wikitext2/weight_matching/model_folded \
    --dataset wikipedia_en \
    --per_device_train_batch_size ${PER_DEVICE_BATCH_SIZE} \
    --gradient_accumulation_steps ${GRADIENT_ACCUMULATION} \
    --ddp_find_unused_parameters True \
    --max_steps 500 \
    --learning_rate 1e-3 --weight_decay 5e-4 \
    --lr_scheduler_type constant \
    --save_strategy steps --save_steps 500 \
    --prediction_loss_only true \
    --eval_steps 500 \
    --output_dir ./result/llama_1_7b_hf_wikitext2/weight_matching/finetuned --overwrite_output_dir \
    --report_to wandb --seed 42 \
    --bf16 \
    --logging_steps 1 \
    --logging_first_step true \
    --deepspeed ds_config.json \
    --ppl_eval true \
    --zero_shot_eval true
# CUDA_VISIBLE_DEVICES=2,3,4,5,6 torchrun --nproc_per_node ${GPU_NUM} --nnodes=1 --node_rank=0 llm_train.py \
#     --model_name llama_1_7b_hf \
#     --model_path ./result/llama_1_7b_hf_wikitext2/weight_matching/model_folded \
#     --dataset wikitext2 \
#     --per_device_train_batch_size ${PER_DEVICE_BATCH_SIZE} \
#     --gradient_accumulation_steps ${GRADIENT_ACCUMULATION} \
#     --ddp_find_unused_parameters True \
#     --max_steps 100000 --warmup_steps 5000 \
#     --learning_rate 1e-4 --weight_decay 0.01 \
#     --save_strategy steps --save_steps 10000 \
#     --output_dir ./result/llama_1_7b_hf_wikitext2/weight_matching/finetuned --overwrite_output_dir \
#     --report_to wandb --seed 42 \
#     --bf16 \
#     --logging_steps 10 \
#     --logging_first_step true


