#!/bin/bash
set -e

# 激活 conda 环境
source /mnt/shared-storage-user/yuanfei/miniconda3/bin/activate /mnt/shared-storage-user/yuanfei/feyuan/envs/llamax2

# 设置环境变量
export PYTORCH_ALLOC_CONF=expandable_segments:False
export HF_HUB_OFFLINE=1
export WANDB_DISABLED=True
export DISABLE_VERSION_CHECK=1
export NCCL_DEBUG=WARN

echo "============================================================"
echo "步骤 1: 准备数据"
echo "============================================================"

# 创建数据目录
mkdir -p "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_scripts"

# 初始化 dataset_info.json
# 如果文件不存在或为空，创建一个空的 JSON 对象
if [ ! -f "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_scripts/dataset_info.json" ] || [ ! -s "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_scripts/dataset_info.json" ]; then
    echo "{}" > "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_scripts/dataset_info.json"
    echo "已初始化 dataset_info.json"
fi

# 注册数据集
python "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_scripts/data_prepare.py" \
    --data_file "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/smol/lst_nothink_data/upsampled_50k_all_subtasks_lst_nothink_data.jsonl" \
    --dataset_name "upsampled_50k_all_subtasks_lst_nothink_data" \
    --dataset_info_path "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_scripts/dataset_info.json"

echo ""
echo "============================================================"
echo "步骤 2: 生成训练配置"
echo "============================================================"

# 构建 generate_yaml 参数
YAML_ARGS=(
    --model_path "/mnt/shared-storage-user/yuanfei/hf_hub/Qwen3-8B"
    --dataset "upsampled_50k_all_subtasks_lst_nothink_data"
    --dataset_dir "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_scripts"
    --output_dir "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_outputs/Qwen3-8B_template-llama3/upsampled_50k_all_subtasks_lst_nothink_data_cutoff16384_20260625_015012"
    --yaml_output_path "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_outputs/Qwen3-8B_template-llama3/upsampled_50k_all_subtasks_lst_nothink_data_cutoff16384_20260625_015012/training.yaml"
    --template "llama3"
    --num_train_epochs 1
    --learning_rate 1e-5
    --per_device_train_batch_size 2
    --gradient_accumulation_steps 8
    --cutoff_len 16384
    --warmup_ratio 0.03
    --weight_decay 0.01
    --lr_scheduler_type "cosine_with_min_lr"
    --lr_scheduler_kwargs '{"min_lr": 2.0e-6}'
    --flash_attn "fa2"
    --deepspeed "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_scripts/ds_z3_config.json"
    --logging_steps 10
    --save_steps 50000
    --save_total_limit 0
)

# 添加条件参数
YAML_ARGS+=(--gradient_checkpointing)
# packing disabled
YAML_ARGS+=(--plot_loss)
# no validation set

# Layer-selective tuning 参数
# no trainable_layers
# no top_trainable_layers
# no specific_layers
# no freeze_extra_modules

python "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_scripts/generate_training_yaml.py" "${YAML_ARGS[@]}"

echo ""
echo "============================================================"
echo "步骤 3: 开始训练"
echo "============================================================"

llamafactory-cli train "/mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_outputs/Qwen3-8B_template-llama3/upsampled_50k_all_subtasks_lst_nothink_data_cutoff16384_20260625_015012/training.yaml"

echo ""
echo "============================================================"
echo "训练完成!"
echo "============================================================"
echo "模型保存在: /mnt/shared-storage-user/yuanfei/yileitu/multidomain_rcot/ft_outputs/Qwen3-8B_template-llama3/upsampled_50k_all_subtasks_lst_nothink_data_cutoff16384_20260625_015012"