#!/bin/bash
export OMP_NUM_THREADS=1
# export CUDA_LAUNCH_BLOCKING=0
# export NCCL_DEBUG=WARN
# export TORCH_DISTRIBUTED_DEBUG=INFO

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


export NCCL_DEBUG=WARN
# export NCCL_CROSS_NIC=2

export TRITON_CACHE_DIR=/mnt/localdisk/.triton_cache_victor

cd /home/victor/neon/diffusion-infer
echo "Starting training..."

/home/victor/anaconda3/envs/stable-audio-tools-env/bin/torchrun \
    --master_addr=$MASTER_ADDR \
    --master_port=$MASTER_PORT \
    --nnodes=$SLURM_JOB_NUM_NODES \
    --nproc_per_node=$SLURM_GPUS_ON_NODE \
    --rdzv_id $SLURM_JOB_ID \
    --rdzv_backend c10d \
    --rdzv_endpoint "$MASTER_ADDR:$MASTER_PORT" \
    train.py