File size: 578 Bytes
13c5606 | 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 | #!/bin/bash
set -x
export TOKENIZERS_PARALLELISM=false
export TORCH_NCCL_AVOID_RECORD_STREAMS=1
NNODES=${NNODES:=1}
NPROC_PER_NODE=${NPROC_PER_NODE:=$(nvidia-smi --list-gpus | wc -l)}
NODE_RANK=${NODE_RANK:=0}
MASTER_ADDR=${MASTER_ADDR:=0.0.0.0}
MASTER_PORT=${MASTER_PORT:=12345}
if [[ "$NNODES" == "1" ]]; then
additional_args="$additional_args --standalone"
else
additional_args="--rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT}"
fi
torchrun \
--nnodes=$NNODES \
--nproc-per-node=$NPROC_PER_NODE \
--node-rank=$NODE_RANK \
$additional_args $@ 2>&1 | tee log.txt
|