set -x

if [ "$#" -lt 2 ]; then
    echo "Usage: run_deepseek_6b7.sh <nproc_per_node> <save_path> [other_configs...]"
    exit 1
fi

nproc_per_node=$1
save_path=$2

source /opt/dtk/env.sh
export NCCL_P2P_LEVEL=PXB # SYS

# Runs Mixtral 8x7B model
export HIP_DIRECT_DISPATCH=0
export HSA_FORCE_FINE_GRAIN_PCIE=1
export OMP_NUM_THREADS=1
export GPU_MAX_HW_QUEUES=10
#export NVTE_FLASH_ATTN_TRITON=1
export NCCL_ALGO=Ring
export NCCL_SOCKET_IFNAME=enp33s0f3u1

export NCCL_NCHANNELS_PER_PEER=16
export NCCL_MIN_NCHANNELS=32 # 20
export NCCL_MAX_NCHANNELS=32 # 20
export NCCL_IB_TIMEOUT=22
export CUDA_DEVICE_MAX_CONNECTIONS=1

export NCCL_IB_HCA=mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1,mlx5_8:1,mlx5_9:1
export NCCL_NET_GDR_LEVEL=7
export NCCL_NET_GDR_READ=1
export RCCL_SDMA_COPY_ENABLE=0
export NCCL_TOPO_FILE="/public/home/fugx1/datasets/rccl-test/topo-input.xml"
# export NCCL_TOPO_FILE="/workspace/rccl-test/rccl-tests-0204/topo-input.xml"
export GLOG_minloglevel=3 # 打印error级别的nccl日志

export PATH=/opt/hpc/software/mpi/hpcx/2.12.0/gcc-8.3.1/bin/:$PATH
export LD_LIBRARY_PATH=/opt/hpc/software/mpi/hpcx/2.12.0/gcc-8.3.1/lib/:$LD_LIBRARY_PATH

# 导入hipblaslt库
export LD_LIBRARY_PATH=/public/home/fugx1/tests1/test03/whl/hipblaslt-install-dtk-25.04-0212/lib:$LD_LIBRARY_PATH

# 更新rocblas
export LD_LIBRARY_PATH=/public/home/fugx1/tests1/test03/whl/rocblas-install-0224/lib:$LD_LIBRARY_PATH


RANK=$OMPI_COMM_WORLD_RANK
LOCAL_RANK=$OMPI_COMM_WORLD_LOCAL_RANK
WORLD_SIZE=$OMPI_COMM_WORLD_SIZE


# Shift the arguments so $@ refers to the rest
shift 2

torchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \
     -m verl.trainer.fsdp_sft_trainer \
    data.train_files=/public/home/fugx1/ds/gsm8k/gsm8k-train.parquet \
    data.val_files=/public/home/fugx1/ds/gsm8k/gsm8k-test.parquet \
    data.prompt_key='question' \
    data.response_key='answer' \
    +data.prompt_dict_keys=['question'] \
    +data.response_dict_keys=['answer'] \
    data.micro_batch_size_per_gpu=1 \
    model.partial_pretrain=/public/opendas/DL_DATA/llm-models/DeepSeek-V3-bf16 \
    trainer.default_local_dir=$save_path \
    trainer.project_name=gsm8k-sft  \
    trainer.experiment_name=gsm8k-sft-deepseek-v3-671b \
    trainer.total_epochs=1 \
    trainer.logger=['console'] \
    trainer.default_hdfs_dir=null $@
