############################################
#! 2. Train Difficulty Classifer
############################################
LLM_path="" #TODO: LLM path
LLM_state_dict_dir="" #TODO: fine-tuned LLM save dir
LLM_state_dict_tag="" #TODO: fine-tuned LLM save tag
binary_data_path="" #TODO: binary train dataset path
save_dir="" #TODO: save dir


TIME_STR=`date +%Y-%m-%d_%H:%M:%S`
export NCCL_DEBUG=WARN         # 仅显示 NCCL 的警告和错误（默认是 INFO）
export TOKENIZERS_PARALLELISM=False
PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:256 \
deepspeed  \
    --include localhost:4,5,6,7 \
    --master_port 29502 train_DiffClf.py \
    --deepspeed_config ./config/ds_z1_offload_config.conf \
    --Qwen_model_path "$LLM_path" \
    --LLM_state_dict_dir "$LLM_state_dict_dir" \
    --LLM_state_dict_tag "$LLM_state_dict_tag" \
    --binary_path "$binary_data_path" \
    --epoch 3 \
    --save_step 512 \
    --time_str $TIME_STR \
    --micro_batch_size 1 \
    --n_accumulation_steps 16 \
    --use_optimizer True \
    --use_scheduler True \
    --save_dir "$save_dir" \
    --cot_length 512 \
    --time_step 4
