0726-1631

diffusion.py

@@ -732,7 +732,8 @@ if __name__ == "__main__":
     parser.add_argument("--sample", type=int, required=False, help="whether to sample", default=0)
     args = parser.parse_args()
 
-    master_addr = os.environ["SLURM_NODELIST"].split(",")[0]
+    #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"])
     local_world_size = torch.cuda.device_count()
@@ -740,7 +741,7 @@ if __name__ == "__main__":
     ############################ training ################################
     #world_size = torch.cuda.device_count()
     if args.train:
-        print(f" training, ip_addr = {socket.gethostbyname(socket.gethostname())}, local_world_size = {local_world_size}, world_size = {world_size} ".center(120,'-'))
+        print(f" training, ip_addr = {socket.gethostbyname(socket.gethostname())}, master_addr = {master_addr}, local_world_size = {local_world_size}, world_size = {world_size} ".center(120,'-'))
         mp.spawn(
                 train, 
                 args=(world_size, local_world_size, master_addr, master_port), 

phoenix_diffusion.sbatch

@@ -3,6 +3,7 @@
 #SBATCH -A gts-jw254-coda20 
 #SBATCH -qembers
 #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)
 #SBATCH -oReport-%j                         # Combined output and error messages file