#!/bin/bash

# If any command fails, the script will exit immediately.
set -e

# Set the environment variables.
export HF_ALLOW_CODE_EVAL=1
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
export TOKENIZERS_PARALLELISM=true

## Reference: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-debug
# export NCCL_DEBUG=INFO

## Reference: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-p2p-level
export NCCL_P2P_LEVEL=LOC

## Reference: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-p2p-disable
# export NCCL_P2P_DISABLE=1

## Reference: https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-shm-disable
# export NCCL_SHM_DISABLE=1

# GPU_IDS=""
GPU_IDS="0"

# If you want to run more than one job on the same machine, you need to specify different MASTER_PORT for each job.
MASTER_PORT=48763

# The config file path.
CONFIG_PATH="configs/testings/vanilla/task_config.yaml"

# The DeepSpeed config file path.
DS_CONFIG_PATH="configs/testings/common/ds_config.json"

echo "GCC Version: $(gcc --version | head -n 1)"
echo "G++ Version: $(g++ --version | head -n 1)"
echo "CUDA Toolkit Version: $(nvcc --version | grep release | sed 's/.*release //; s/,//')"

deepspeed \
    --include=localhost:${GPU_IDS} \
    --master_port ${MASTER_PORT} \
    ds_main.py \
    -c ${CONFIG_PATH} \
    -dsc ${DS_CONFIG_PATH} \
    -m test \
    -mt vanilla
