verify ddp training

Author: mingukkang

ddp_script.sh

@@ -0,0 +1,7 @@
+export CUDA_VISIBLE_DEVICES=0,...,N_GPUS
+export NCCL_DEBUG=INFO
+export NCCL_SOCKET_IFNAME=^docker0,lo
+export MASTER_ADDR="MASTER_IP"
+export MASTER_PORT="MASTER_PORT"
+
+python3 src/main.py -t -e -rm_API -c CONFIG_PATH -DDP -n NUM_NODES -nr CURRENT_NODES -eval_type EVAL_TYPE ...

src/main.py

@@ -123,8 +123,6 @@ def main():
 
     if train_config['distributed_data_parallel'] and world_size > 1:
         print("Train the models through DistributedDataParallel (DDP) mode.")
-        os.environ['MASTER_ADDR'] = 'localhost'
-        os.environ['MASTER_PORT'] = '2222'
         mp.spawn(prepare_train_eval, nprocs=gpus_per_node, args=(gpus_per_node, world_size, run_name,
                                                                  train_config, model_config, hdf5_path_train))
     else:

src/utils/misc.py

@@ -133,7 +133,10 @@ def setup(rank, world_size, backend="nccl"):
         )
     else:
         # initialize the process group
-        dist.init_process_group(backend, rank=rank, world_size=world_size)
+        dist.init_process_group(backend,
+                                init_method="tcp://%s:%s" % (os.environ['MASTER_ADDR'], os.environ['MASTER_PORT']),
+                                rank=rank,
+                                world_size=world_size)
 
 
 def cleanup():