0728-1322

diffusion.py

@@ -75,15 +75,22 @@ import socket
 import sys
 # %%
 def ddp_setup(rank: int, world_size: int, local_world_size, master_addr, master_port):
-  """
-  Args:
-      rank: Unique identifier of each process
-     world_size: Total number of processes
-  """
-  os.environ["MASTER_ADDR"] = master_addr
-  os.environ["MASTER_PORT"] = master_port
-#   print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!ddp_setup, rank =", rank)
-  init_process_group(backend="nccl", rank=rank, world_size=world_size*local_world_size)
+    """
+    Args:
+       rank: Unique identifier of each process
+       world_size: Total number of processes
+    """
+
+    print("inside ddp_setup")
+    os.environ["MASTER_ADDR"] = master_addr
+    os.environ["MASTER_PORT"] = master_port
+    print("ddp_setup, rank =", rank)
+    init_process_group(
+            backend="nccl", 
+            init_method=f"tcp://{master_addr}:{master_port}", 
+            rank=rank, 
+            world_size=world_size*local_world_size
+            )
 
 # %%
 # notebook_login()
@@ -736,7 +743,7 @@ if __name__ == "__main__":
     #master_addr = os.environ["SLURM_NODELIST"].split(",")[0]
     master_addr = os.environ.get("MASTER_ADDR", "localhost")
     master_port = "12355"
-    world_size = int(os.environ["SLURM_NTASKS"])
+    world_size = 1#int(os.environ["SLURM_NTASKS"])
     local_world_size = torch.cuda.device_count()
 
     ############################ training ################################

phoenix_diffusion.sbatch

@@ -2,7 +2,7 @@
 #SBATCH -J diffusion # Job name
 #SBATCH -A gts-jw254-coda20 
 #SBATCH -qembers
-#SBATCH -N1 --gpus-per-node=RTX_6000:3 # -C A100-80GB              # Number of nodes and cores per node required
+#SBATCH -N2 --gpus-per-node=RTX_6000:3 # -C A100-80GB              # Number of nodes and cores per node required
 #SBATCH --ntasks-per-node=1
 #SBATCH --mem-per-gpu=32G                        # Memory per core
 #SBATCH -t 10:00                                    # Duration of the job (Ex: 15 mins)