0728-1727

diffusion.py

@@ -73,8 +73,10 @@ import torch.distributed as dist
 import argparse
 import socket
 import sys
+from datetime import timedelta
+
 # %%
-def ddp_setup(rank: int, world_size: int, local_world_size, master_addr, master_port):
+def ddp_setup(rank: int, world_size: int, master_addr, master_port):
     """
     Args:
        rank: Unique identifier of each process
@@ -89,7 +91,8 @@ def ddp_setup(rank: int, world_size: int, local_world_size, master_addr, master_
             backend="nccl", 
             init_method=f"tcp://{master_addr}:{master_port}", 
             rank=rank, 
-            world_size=world_size*local_world_size
+            world_size=world_size,
+            timeout=timedelta(minutes=4)
             )
 
 # %%
@@ -659,15 +662,17 @@ class DDPM21CM:
 #num_train_image_list = [6000]#[60]#[8000]#[1000]#[100]#
 def train(rank, world_size, local_world_size, master_addr, master_port):
     #print("before ddp_setup")
-    ddp_setup(rank, world_size, local_world_size, master_addr, master_port)
+    global_rank = rank + local_world_size * int(os.environ["SLURM_NODEID"])
+
+    ddp_setup(global_rank, world_size, master_addr, master_port)
     #print("after ddp_setup")
-    local_rank = rank % local_world_size
-    torch.cuda.set_device(local_rank)
+    #local_rank = rank % local_world_size
+    torch.cuda.set_device(rank)
     #print("after set device")
-    print(f"rank = {rank}, local_rank = {local_rank}, world_size = {world_size}, local_world_size = {local_world_size}")
+    print(f"rank = {rank}, global_rank = {global_rank}, world_size = {world_size}, local_world_size = {local_world_size}")
 
     config = TrainConfig()
-    config.device = f"cuda:{local_rank}"
+    config.device = f"cuda:{rank}"
     config.world_size = local_world_size
     
     #[3200]#[200]#[1600,3200,6400,12800,25600]
@@ -741,10 +746,13 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     #master_addr = os.environ["SLURM_NODELIST"].split(",")[0]
-    master_addr = os.environ.get("MASTER_ADDR", "localhost")
-    master_port = "12355"
-    world_size = 1#int(os.environ["SLURM_NTASKS"])
+    #master_addr = os.environ.get("MASTER_ADDR", "localhost")
+    #master_port = "12355"
+    master_addr = os.environ["MASTER_ADDR"]
+    master_port = os.environ["MASTER_PORT"]
     local_world_size = torch.cuda.device_count()
+    total_nodes = int(os.environ["SLURM_NNODES"])
+    world_size = local_world_size * total_nodes #6#int(os.environ["SLURM_NTASKS"])
 
     ############################ training ################################
     #world_size = torch.cuda.device_count()

phoenix_diffusion.sbatch

@@ -22,4 +22,10 @@ conda activate diffusers
 conda env list
 cat $0
 
+MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
+MASTER_PORT=$((10000 + RANDOM % 10000)) #12355
+
+export MASTER_ADDR=$MASTER_ADDR
+export MASTER_PORT=$MASTER_PORT
+
 srun python diffusion.py --train 1 --sample 0