TORCHRUN 515 Bytes
Newer Older
chenzk's avatar
v1.0  
chenzk committed
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
#!/bin/bash
set -x
start_time="$(date +%FT%T)"

if [[ "$master_addr" == "" && "$master_port" == "" ]]; then
  nnodes=1
  node_rank=0
  master_port=12345
  nproc_per_node=4
  torchrun \
  --master_port=$master_port \
  --node_rank=$node_rank \
  --nproc_per_node=$nproc_per_node \
  --nnodes=$nnodes \
  "$@"
else
  torchrun \
  --master_addr=$master_addr \
  --master_port=$master_port \
  --node_rank=$node_rank \
  --nproc_per_node=$nproc_per_node \
  --nnodes=$nnodes \
  "$@"
fi

echo "start_time: $start_time"