#!/bin/bash
source /mnt/shared-storage-user/caipengxiang/workspace/ChemBOMAS/train_regression/env_volcengine.sh

# 读取命令行参数
DATA_NAME="suzuki_60" # 数据集名称
if [ -z "./train_regression/${DATA_NAME}" ]; then
    echo "Error: DATA_NAME environment variable is required"
    exit 1
fi

# 检查所需路径是否都存在
PROJECT_DIR="/mnt/shared-storage-user/caipengxiang/workspace/ChemBOMAS"

required_paths=(
    "/mnt/shared-storage-user/caipengxiang/H200-share/models/share/step1_llama3_8b_0916_yearly_pistachio_ep3"
    "${PROJECT_DIR}/train_regression/data4regression"
    "yield_ft_ds_config.json"
)

for path in "${required_paths[@]}"; do
    if [ ! -e "$path" ]; then
        echo "Error: Required path not found: $path"
        exit 1
    fi
done

export TOKENIZERS_PARALLELISM=false #Disable parallelism to avoid deadlocks
# export NVIDIA_TF32_OVERRIDE=0 #Disable the auto-conversion from NVIDIA PF32 to TF32

# torchrun --nnodes=${NNODES} \
#             --node_rank=${RANK} \
#             --nproc_per_node=${GPUS_PER_NODE} \
#             --master_addr=${MASTER_ADDR} \
#             --master_port=${MASTER_PORT} \

deepspeed --num_nodes=${NNODES} \
            --num_gpus=${GPUS} \
            --master_addr=${MASTER_ADDR} \
            --master_port=${MASTER_PORT} \
            --node_rank=${RANK} \
            yield_ft_ds.py \
            --pretrained_model_path "/mnt/shared-storage-user/caipengxiang/H200-share/models/share/step1_llama3_8b_0916_yearly_pistachio_ep3" \
            --lora_adapter_path "None" \
            --yield_predictor_path "${PROJECT_DIR}/share/jianpeng/llama_ep3_1115-18/predictor.pt" \
            --num_epoch 100 \
            --lr 1e-4 \
            --data_path "${PROJECT_DIR}/train_regression/data4regression" \
            --data_name "${DATA_NAME}" \
            --per_device_train_batch_size 16 \
            --save_root "${PROJECT_DIR}/train_regression/saved_models/${DATA_NAME}" \
            --gradient_accumulation_steps 1 \
            --use_lora 1 \
            --log_file "training_ds.log" \
            --deepspeed_config "yield_ft_ds_config.json" \
            --mlp_lr_multiplier 1 \
            --save_interval 30 \
            --save_lora_adapter 1 \
            --save_predictor 1 \
            --lora_adapter_save_path "${PROJECT_DIR}/share/llama-3.1-8B/pretrained_lora/${DATA_NAME}"