deepspeed: resolve error for multinode (fix for FA3)

bghira · Oct 10, 2024 · 390d43a · 390d43a
1 parent cfa6c6e
commit 390d43a
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/helpers/training/diffusion_model.py b/helpers/training/diffusion_model.py
@@ -56,6 +56,11 @@ def load_diffusion_model(args, weight_dtype):
         import torch
 
         if torch.cuda.is_available():
+            rank = (
+                torch.distributed.get_rank()
+                if torch.distributed.is_initialized()
+                else 0
+            )
             primary_device = torch.cuda.get_device_properties(0)
             if primary_device.major >= 9:
                 try: