Isolate distrib_run (#828)

src/accelerate/commands/launch.py

@@ -56,9 +56,6 @@
     logging.basicConfig(format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
 
 
-if is_torch_version(">=", "1.9.0"):
-    import torch.distributed.run as distrib_run
-
 logger = logging.getLogger(__name__)
 
 options_to_group = {
@@ -555,6 +552,8 @@ def simple_launcher(args):
 
 
 def multi_gpu_launcher(args):
+    if is_torch_version(">=", "1.9.0"):
+        import torch.distributed.run as distrib_run
     num_processes = getattr(args, "num_processes")
     num_machines = getattr(args, "num_machines")
     main_process_ip = getattr(args, "main_process_ip")
@@ -644,6 +643,8 @@ def multi_gpu_launcher(args):
 
 
 def deepspeed_launcher(args):
+    if is_torch_version(">=", "1.9.0"):
+        import torch.distributed.run as distrib_run
     if not is_deepspeed_available():
         raise ImportError("DeepSpeed is not installed => run `pip3 install deepspeed` or build it from source.")
     num_processes = getattr(args, "num_processes")

src/accelerate/utils/launch.py

@@ -21,10 +21,6 @@
 from .dataclasses import DistributedType
 
 
-if is_torch_version(">=", "1.9.0"):
-    import torch.distributed.run as distrib_run
-
-
 def get_launch_prefix():
     """
     Grabs the correct launcher for starting a distributed command, such as either `torchrun`, `python -m
@@ -43,6 +39,8 @@ def _filter_args(args):
     """
     Filters out all `accelerate` specific args
     """
+    if is_torch_version(">=", "1.9.0"):
+        import torch.distributed.run as distrib_run
     distrib_args = distrib_run.get_args_parser()
     new_args, _ = distrib_args.parse_known_args()