#!/bin/bash
# Please run this script under ${project_id} in project directory of
#   https://github.com/shizhediao/llm-ft
#     COMMIT: d5fecf30ba8011067b10cf51fede53a5ab6574e4

# Parses arguments
model_name_or_path=models/qwen2.5-1.5B-bd-stage1
dataset_path=data/nemotron_no_thinking_subset
output_dir=output_models/finetune_qwen2.5-1.5B-bd-stage2-test
deepspeed_args="--master_port=11000"
conversation_template=qwen2_5_no_reasoning
export CUDA_HOME=/lustre/fsw/portfolios/nvr/users/chengyuew/anaconda3/envs/lmflow
# Safety related arguments
trust_remote_code=1

# 动态检测最新的checkpoint
latest_checkpoint=""
if [ -d "${output_dir}" ]; then
    # 查找所有checkpoint目录，按数字排序，取最新的一个
    latest_checkpoint=$(find "${output_dir}" -name "checkpoint-*" -type d | sort -V | tail -1)
    if [ -n "${latest_checkpoint}" ]; then
        echo "找到最新的checkpoint: ${latest_checkpoint}"
    else
        echo "在 ${output_dir} 中没有找到checkpoint"
        latest_checkpoint=""
    fi
else
    echo "输出目录 ${output_dir} 不存在，将从头开始训练"
    latest_checkpoint=""
fi

# Finetune
exp_id=finetune
project_dir=$(cd "$(dirname $0)"/../..; pwd)
log_dir=${project_dir}/log/${exp_id}
mkdir -p ${output_dir} ${log_dir}

# 构建resume_from_checkpoint参数
resume_arg=""
if [ -n "${latest_checkpoint}" ]; then
    resume_arg="--resume_from_checkpoint ${latest_checkpoint}"
fi

cmd="deepspeed ${deepspeed_args} \
  examples/finetune.py \
    --model_name_or_path ${model_name_or_path} \
    --trust_remote_code ${trust_remote_code} \
    --dataset_path ${dataset_path} \
    --output_dir ${output_dir} \
    ${resume_arg} \
    --conversation_template ${conversation_template} \
    --num_train_epochs 3 \
    --learning_rate 2e-6 \
    --lr_scheduler_type constant_with_warmup \
    --warmup_ratio 0.03 \
    --disable_group_texts 1 \
    --block_size 2048 \
    --per_device_train_batch_size 4 \
    --gradient_accumulation_steps 1 \
    --deepspeed configs/ds_config_zero2_no_offload.json \
    --bf16 \
    --run_name finetune \
    --validation_split_percentage 0 \
    --logging_steps 1 \
    --do_train \
    --ddp_timeout 72000 \
    --save_steps 1000 \
    --dataloader_num_workers 8 \
    --preprocessing_num_workers 32 \
    --learn_padding 1 \
    --bd_size 128 \
    --mdm 1 \
    --conplemenrary_mask 1 \
    --save_total_limit 10 \
    --gradient_checkpointing 1 \
    --tokenizer_name models/Qwen2.5-1.5B-Instruct
    > >(tee ${log_dir}/train.log) \
    2> >(tee ${log_dir}/train.err >&2)"

echo $cmd
eval $cmd
