#!/usr/bin/env bash

source ./env.sh

FILE=$1
CONFIG=$2
GPUS=$3
NODE=${NODE:-1}
#NODE_RANK=${NODE_RANK:-0}
NODE_RANK=$SLURM_NODEID
ADDR=${ADDR:-127.0.0.1}
PORT=${PORT:-12345}

export ONEFLOW_FUSE_OPTIMIZER_UPDATE_CAST=true 

echo "NODE_RANK=$NODE_RANK"

python3 -m oneflow.distributed.launch \
--nproc_per_node $GPUS --nnodes $NODE --node_rank $NODE_RANK --master_addr $ADDR --master_port $PORT \
$FILE --config-file $CONFIG ${@:4}

