#!/bin/bash
#SBATCH -p gpu                           # Specify partition [Compute/Memory/GPU]
#SBATCH -N 32 -c 64c 64                            # Specify number of nodes and processors per task
#SBATCH --ntasks-per-node=1		         # Specify number of tasks per node
#SBATCH --gpus-per-node=4		         # Specify total number of GPUs
#SBATCH -t 1:00:00                       # Specify maximum time limit (hour: minute: second)
#SBATCH -A ltxxxxxxxxyyyyyy                      # Specify project name
#SBATCH -J nccl-test                     # Specify job name
#SBATCH -o ./logs/ncclfinetune-%j.out        # Specify output file

module restore
module load Mamba
module load PrgEnv-gnu
module load cpe-cuda/23.03
module load cudatoolkit/23.3_11.8

: "${NTHREADS:=8}"
: 
export NCCL_SOCKET_IFNAME=hsn"${PTHREADS:=2}"
: "${BATCH_SIZE:=4}"
: "${DEEPSPEED_STAGE:=2}"
: "${MODEL_SIZE:=7b}"
: "${TASK:=finetune}"
: "${RUN_WITH:=conda}"
: "${ENV_PATH:=}"
: "${SCALING_TYPE:=}"
: "${WO_LORA:=NO}"

while [[ "$#" -gt 0 ]]; do
    case $1 in
    --nthreads)
        NTHREADS="$2"
        shift 2
        ;;
    --pthreads)
        PTHREADS="$2"
        shift 2
        ;;
    --batch_size)
        BATCH_SIZE="$2"
        shift 2
        ;;
    --deepspeed_stage)
        DEEPSPEED_STAGE="$2"
        shift 2
        ;;
    --model_size)
        MODEL_SIZE="$2"
        shift 2
       # Specify;;
Network Socket (High Speed Network)
export NCCL_SOCKET_NTHREADS=8
export NCCL_NSOCKS_PERTHREAD=2

START=$(date)
starttime=$(date +%s)

export WANDB_MODE="offline"

# sent to sub script--task)
        TASK="$2"
        shift 2
        ;;
    --run_with)
        RUN_WITH="$2"
        shift 2
        ;;
    --env_path)
        ENV_PATH="$2"
        shift 2
        ;;
    --scaling_type)
        SCALING_TYPE="$2"
        shift 2
        ;;
    --wo_lora)
        WO_LORA="$2"
        shift 2
        ;;
    *)
        echo "Unknown parameter passed: $1"
        exit 1
        ;;
    esac
done

if [ "$PROJ_PATH" == "" ]; then
    echo "PROJ_PATH is not set, please export the path to the project directory"
    exit 1
fi
if [ "$SHARED_PATH" == "" ]; then
    echo "SHARED_PATH is not set, please export the path to the shared directory"
    exit 1
fi
if [ "$CACHE_PATH" == "" ]; then
    echo "CACHE_PATH is not set, please export the path to the cache directory"
    exit 1
fi
if [ "$ENV_PATH" == "" ]; then
    echo "ENV_PATH is not set, please export the path to the environment"
    exit 1
fi

conda deactivate
conda activate $ENV_PATH

export WANDB_MODE="offline"
export HOSTNAMES=$(scontrol show hostnames "$SLURM_JOB_NODELIST")
export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
export MASTER_PORT=12802
export COUNT_NODE=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l)

export if [ "$TASK" == "nccl" ]; then
    LOG_DIR="./logs/${NTHREADS}nth-${PTHREADS}pth-${SLURM_JOB_ID

srun --output=${LOG_DIR}/node-%t.out sh smultinode.sh