#!/bin/bash
#SBATCH -p gpu                           # Specify partition [Compute/Memory/GPU]
#SBATCH -c 64                            # Specify number of 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 xxyyyyyy                      # Specify project name
#SBATCH -J nccl-test                     # Specify job name
#SBATCH -o ./logs/finetune-%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}"
: "${PTHREADS:=2}"
: "${BATCH_SIZE:=4}"
: "${DEEPSPEED_STAGE:=3}"
: "${MODEL_SIZE:=7b}"
: "${TASK:=finetune}"
: "${RUN_WITH:=conda}"
: "${ENV_PATH:=}"
: "${SCALING_TYPE:=}"
: "${WO_LORA:=NO}"
: "${PROJ_PATH:=}"
: "${SHARED_PATH:=}"
: "${CACHE_PATH:=}"
: "${ENV_PATH:=}"

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)

LOG_DIR="./logs/${NTHREADS}nth-${PTHREADS}pth-${SLURM_JOB_ID}"

mkdir -p $LOG_DIR/node_log

export LOG_DIR=$LOG_DIR

export NCCL_TIMEOUT=3600000
export NCCL_DEBUG=DEBUG
export NCCL_SOCKET_IFNAME=hsn
export NCCL_SOCKET_NTHREADS=$NTHREADS
export NCCL_NSOCKS_PERTHREAD=$PTHREADS
export NCCL_DEBUG_FILE=${LOG_DIR}/nccl-${SLURM_JOB_ID}.log
export NCCL_TOPO_DUMP_FILE=${LOG_DIR}/nccl-topo-${SLURM_JOB_ID}.log
export BATCH_SIZE=$BATCH_SIZE
export DEEPSPEED_STAGE=$DEEPSPEED_STAGE
export MODEL_SIZE=$MODEL_SIZE

export TORCH_NCCL_BLOCKING_WAIT=0
export TORCH_EXTENSIONS_DIR=$CACHE_PATH
export HF_HUB_CACHE="$CACHE_PATH/huggingface"
export HF_HOME="$CACHE_PATH/huggingface"
export HF_DATASETS_CACHE="$CACHE_PATH/huggingface"
export TORCH_HOME=$CACHE_PATH
export XDG_CACHE_HOME=$CACHE_PATH
export HF_DATASETS_OFFLINE=1 
export HF_HUB_OFFLINE=1

echo -------ENVIRONMENT-------
echo Python Path: $(which python)
echo Batch Size: $BATCH_SIZE
echo Deepspeed Stage: $DEEPSPEED_STAGE
echo Model Size: $MODEL_SIZE
echo Train with LoRA: $WO_LORA
echo -------------------------
echo NTHREADS: $NTHREADS
echo PTHREADS: $PTHREADS
echo NODES: $COUNT_NODE
echo HOSTNAMES: $HOSTNAMES
echo MASTER_ADDR: $MASTER_ADDR
echo MASTER_PORT: $MASTER_PORT
echo -------------------------

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

#!/bin/bash

echo "------INFO--------"
echo node_number: $SLURM_PROCID
echo hostname: $(hostname)
echo "------------------"
echo ""

echo "----LAUNCHING TRAINING----"

accelerate launch \
    --num_processes $((4 * $COUNT_NODE)) \
    --num_machines $COUNT_NODE \
    --multi_gpu \
    --mixed_precision bf16 \
    --machine_rank $SLURM_PROCID \
    --main_process_ip $MASTER_ADDR \
    --main_process_port $MASTER_PORT \
    --dynamo_backend inductor \
    $PROJ_PATH/scripts/train_wo_lora.py \
        --pretrained_model_name_or_path "meta-llama/Llama-2-$MODEL_SIZE-chat-hf" \
        --train_file $SHARED_PATH/datasets/alpaca_json/alpaca_all.json \
        --validation_file $SHARED_PATH/datasets/alpaca_json/alpaca_validation.json \
        --seed 42 \
        --max_seq_length 1300 \
        --output_dir $PROJ_PATH/checkpoint \
        --num_train_epochs 1 \
        --per_device_train_batch_size $BATCH_SIZE \
        --per_device_eval_batch_size $BATCH_SIZE \
        --save_steps 700 \
        --save_total_limit 5 \
        --learning_rate 8e-5 \
        --weight_decay 0.01 \
        --warmup_ratio 0.05 \
        --lr_scheduler_type cosine \
        --gradient_accumulation_steps 1 \
        --deepspeed "$PROJ_PATH/deepspeed_config/deepspeed_$DEEPSPEED_STAGE.json" \
        --gradient_checkpointing True \
        --tf32 True \
        --bf16 True \
        --max_grad_norm 1.0 \
        --logging_steps 10 \
        --dataloader_num_workers 16 \
        --ddp_find_unused_parameters False \
        --log_dir $LOG_DIR \
        --node_number $SLURM_PROCID