Merge pull request #387 from allenai/epwalsh/dist-init

initialize distributed group before logging
allenai · Dec 1, 2023 · d1c185b · d1c185b
2 parents e30d29f + 7f54682
commit d1c185b
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/scripts/train.py b/scripts/train.py
@@ -46,9 +46,9 @@ def main(cfg: TrainConfig) -> None:
             "setting has no effect."
         )
 
-    # Initialize process group and set device.
-    dist.init_process_group(backend="nccl")
     barrier()
+
+    # Set CUDA device.
     torch.cuda.set_device(f"cuda:{get_local_rank()}")
     device = torch.device("cuda")
 
@@ -239,6 +239,9 @@ def dummy_init_fn(module: torch.nn.Module) -> None:
 
 
 if __name__ == "__main__":
+    # Initialize process group.
+    dist.init_process_group(backend="nccl")
+
     prepare_cli_environment()
 
     try: