Spaces:
Running on Zero
Running on Zero
File size: 1,426 Bytes
3e936b2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | #!/bin/bash
NNODES="${NNODES:-4}"
NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
MASTER_ADDR="${MASTER_ADDR:-localhost:29551}"
RDZV_ID=20047
USE_WANDB="${USE_WANDB:-0}" # set USE_WANDB=1 to enable W&B logging (fill in wandb_key/wandb_entity in the config)
NODE_IPS=("${NODE_IP_0:-127.0.0.1}" "${NODE_IP_1:-your_node_ip_1}" "${NODE_IP_2:-your_node_ip_2}" "${NODE_IP_3:-your_node_ip_3}")
LOCAL_IPS=$(hostname -I 2>/dev/null)
LOCAL_ADDR=""
for node_ip in "${NODE_IPS[@]}"; do
if echo "$LOCAL_IPS" | grep -qw "$node_ip"; then
LOCAL_ADDR="$node_ip"
break
fi
done
if [ -z "$LOCAL_ADDR" ]; then
echo "ERROR: failed to match local IP against NODE_IPS; check NODE_IPS configuration"
echo "local IPs: $LOCAL_IPS"
exit 1
fi
echo "LOCAL_ADDR=$LOCAL_ADDR"
CONFIG=configs/echo_infinity-long.yaml
CONFIG_NAME=$(basename "$CONFIG" .yaml)
LOGDIR="logs/${CONFIG_NAME}"
WANDB_SAVE_DIR=wandb
echo "CONFIG=$CONFIG"
echo "LOGDIR=$LOGDIR"
mkdir -p "$LOGDIR"
LOGFILE="${LOGDIR}/train_$(date +%Y%m%d_%H%M%S).log"
echo "LOGFILE=$LOGFILE"
WANDB_ARGS="--disable-wandb"
[ "${USE_WANDB}" != "0" ] && WANDB_ARGS=""
torchrun --nnodes=$NNODES --nproc_per_node=$NPROC_PER_NODE \
--rdzv_id=$RDZV_ID --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR \
--local-addr=$LOCAL_ADDR \
train.py \
--config_path $CONFIG \
--logdir $LOGDIR \
--wandb-save-dir $WANDB_SAVE_DIR \
${WANDB_ARGS} \
2>&1 | tee "$LOGFILE"
|