Only initialize when necessary

pomonam · Jul 7, 2024 · c918176 · c918176
1 parent 2df875b
commit c918176
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/kronfluence/utils/model.py b/kronfluence/utils/model.py
@@ -96,7 +96,8 @@ def apply_fsdp(
         RuntimeError:
             If the distributed initialization fails.
     """
-    dist.init_process_group("nccl", rank=rank, world_size=world_size)
+    if not dist.is_initialized():
+        dist.init_process_group("nccl", rank=rank, world_size=world_size)
     device = torch.device(f"cuda:{local_rank}")
     torch.cuda.set_device(local_rank)