#!/bin/bash

#SBATCH --job-name=JACK
#SBATCH --partition=kshdnormal01
#SBATCH --nodes=24
#SBATCH --cpus-per-task=32
#SBATCH --ntasks-per-node=1
#SBATCH --gres=dcu:4
#SBATCH --mem=100G
#SBATCH --wait-all-nodes=1
#SBATCH --exclusive
#SBATCH --output log/%j.out.log
#SBATCH --error log/%j.err.log
#SBATCH -x c01r3n06,c01r1n04,c02r2n04,c02r2n19

M_NODE=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
M_ADDR=$(scontrol show node=$M_NODE | grep NodeAddr | awk -F' ' '{print $1}' | awk -F'=' '{print $2}')
M_PORT=12345

#echo "SLURM_NODEID=$SLURM_NODEID"
echo "SLURMD_NODENAME=$SLURMD_NODENAME"

source ./env.sh

export NODE=$SLURM_NNODES
#export NODE_RANK=$SLURM_NODEID
export ADDR=$M_ADDR
export PORT=$M_PORT

echo "NODE=$SLURM_NNODES"
#echo "NODE_RANK=$SLURM_NODEID"
echo "ADDR=$M_ADDR"
echo "PORT=$M_PORT"

#srun bash tools/train-13B.sh tools/train_net.py configs/bert_large_pretrain.py 4 
srun bash tools/train-13B.sh tools/train_net.py configs/gpt2-13B_pretrain.py 4


