Reduce scale factor and val loss across ranks for DDP (#1461)

Fixes #1458 .

### Description
Reduce scale factor and val loss across ranks for DDP

### Checks
<!--- Put an `x` in all the boxes that apply, and remove the not
applicable items -->
- [ ] Avoid including large-size files in the PR.
- [ ] Clean up long text outputs from code cells in the notebook.
- [ ] For security purposes, please check the contents and remove any
sensitive info such as user names and private key.
- [ ] Ensure (1) hyperlinks and markdown anchors are working (2) use
relative paths for tutorial repo files (3) put figure and graphs in the
`./figure` folder
- [ ] Notebook runs automatically `./runner.sh -t <path to .ipynb file>`

---------

Signed-off-by: Can-Zhao <[email protected]>

Author: Can-Zhao

generative/2d_ldm/train_diffusion.py

@@ -126,6 +126,11 @@ def main():
                 )
                 print(f"Scaling factor set to {1/torch.std(z)}")
     scale_factor = 1 / torch.std(z)
+    print(f"Rank {rank}: local scale_factor: {scale_factor}")
+    if ddp_bool:
+        dist.barrier()
+        dist.all_reduce(scale_factor, op=torch.distributed.ReduceOp.AVG)
+    print(f"Rank {rank}: final scale_factor -> {scale_factor}")
 
     # Define Diffusion Model
     unet = define_instance(args, "diffusion_def").to(device)
@@ -261,9 +266,15 @@ def main():
                             timesteps=timesteps,
                         )
                         val_loss = F.mse_loss(noise_pred.float(), noise.float())
-                        val_recon_epoch_loss += val_loss.item()
+                        val_recon_epoch_loss += val_loss
                     val_recon_epoch_loss = val_recon_epoch_loss / (step + 1)
 
+                    if ddp_bool:
+                        dist.barrier()
+                        dist.all_reduce(val_recon_epoch_loss, op=torch.distributed.ReduceOp.AVG)
+
+                    val_recon_epoch_loss = val_recon_epoch_loss.item()
+
                     # write val loss and save best model
                     if rank == 0:
                         tensorboard_writer.add_scalar("val_diffusion_loss", val_recon_epoch_loss, epoch + 1)

generative/3d_ldm/train_diffusion.py

@@ -127,6 +127,11 @@ def main():
                     )
                 print(f"Scaling factor set to {1/torch.std(z)}")
     scale_factor = 1 / torch.std(z)
+    print(f"Rank {rank}: local scale_factor: {scale_factor}")
+    if ddp_bool:
+        dist.barrier()
+        dist.all_reduce(scale_factor, op=torch.distributed.ReduceOp.AVG)
+    print(f"Rank {rank}: final scale_factor -> {scale_factor}")
 
     # Define Diffusion Model
     unet = define_instance(args, "diffusion_def").to(device)
@@ -243,9 +248,15 @@ def main():
                             timesteps=timesteps,
                         )
                         val_loss = F.mse_loss(noise_pred.float(), noise.float())
-                        val_recon_epoch_loss += val_loss.item()
+                        val_recon_epoch_loss += val_loss
                     val_recon_epoch_loss = val_recon_epoch_loss / (step + 1)
 
+                    if ddp_bool:
+                        dist.barrier()
+                        dist.all_reduce(val_recon_epoch_loss, op=torch.distributed.ReduceOp.AVG)
+
+                    val_recon_epoch_loss = val_recon_epoch_loss.item()
+
                     # write val loss and save best model
                     if rank == 0:
                         tensorboard_writer.add_scalar("val_diffusion_loss", val_recon_epoch_loss, epoch)