Add customization point for init_process_group kwargs

examples/nlp_example.py

@@ -125,7 +125,7 @@ def collate_fn(examples):
     lr_scheduler = get_linear_schedule_with_warmup(
         optimizer=optimizer,
         num_warmup_steps=100,
-        num_training_steps=len(train_dataloader) * num_epochs,
+        num_training_steps=(len(train_dataloader) * num_epochs) // gradient_accumulation_steps,
     )
 
     # Now we train the model

src/accelerate/accelerator.py

@@ -22,7 +22,7 @@
 from packaging import version
 
 from .data_loader import prepare_data_loader
-from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, KwargsHandler
+from .kwargs_handlers import DistributedDataParallelKwargs, GradScalerKwargs, InitProcessGroupKwargs, KwargsHandler
 from .optimizer import AcceleratedOptimizer
 from .state import AcceleratorState, DistributedType, is_deepspeed_available
 from .utils import (
@@ -114,15 +114,10 @@ def __init__(
                 deepspeed_plugin, DeepSpeedPlugin
             ), "`deepspeed_plugin` must be a DeepSpeedPlugin object."
 
-        self.state = AcceleratorState(fp16=fp16, cpu=cpu, deepspeed_plugin=deepspeed_plugin, _from_accelerator=True)
-
-        self.device_placement = device_placement
-        self.split_batches = split_batches
-        self.dispatch_batches = dispatch_batches
-
         # Kwargs handlers
         self.ddp_handler = None
         self.scaler_handler = None
+        self.init_handler = None
         if kwargs_handlers is not None:
             for handler in kwargs_handlers:
                 assert isinstance(handler, KwargsHandler), f"Unsupported kwargs handler passed: {handler}."
@@ -136,6 +131,20 @@ def __init__(
                         raise ValueError("You can only pass one `GradScalerKwargs` in `kwargs_handler`.")
                     else:
                         self.scaler_handler = handler
+                elif isinstance(handler, InitProcessGroupKwargs):
+                    if self.init_handler is not None:
+                        raise ValueError("You can only pass one `InitProcessGroupKwargs` in `kwargs_handler`.")
+                    else:
+                        self.init_handler = handler
+
+        kwargs = self.init_handler.to_kwargs() if self.init_handler is not None else {}
+        self.state = AcceleratorState(
+            fp16=fp16, cpu=cpu, deepspeed_plugin=deepspeed_plugin, _from_accelerator=True, **kwargs
+        )
+
+        self.device_placement = device_placement
+        self.split_batches = split_batches
+        self.dispatch_batches = dispatch_batches
 
         # Mixed precision attributes
         self.scaler = None

src/accelerate/kwargs_handlers.py

@@ -14,6 +14,8 @@
 
 import copy
 from dataclasses import dataclass
+from datetime import timedelta
+from typing import Optional
 
 
 class KwargsHandler:
@@ -71,3 +73,16 @@ class GradScalerKwargs(KwargsHandler):
     backoff_factor: float = 0.5
     growth_interval: int = 2000
     enabled: bool = True
+
+
+@dataclass
+class InitProcessGroupKwargs(KwargsHandler):
+    """
+    Use this object in your :class:`~accelerate.Accelerator` to customize the initialization of the distributed
+    processes. Please refer to the documentation of this `method
+    <https://pytorch.org/docs/stable/distributed.html#torch.distributed.init_process_group>`__ for more information on
+    each argument.
+    """
+
+    init_method: Optional[str] = None
+    timeout: timedelta = timedelta(seconds=1800)

src/accelerate/state.py

@@ -138,7 +138,9 @@ class AcceleratorState:
 
     _shared_state = {}
 
-    def __init__(self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None, _from_accelerator: bool = False):
+    def __init__(
+        self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None, _from_accelerator: bool = False, **kwargs
+    ):
         self.__dict__ = self._shared_state
         if not getattr(self, "initialized", False):
             self.backend = None
@@ -161,7 +163,7 @@ def __init__(self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None,
                 ), "DeepSpeed is not available => install it using `pip3 install deepspeed` or build it from source"
                 self.distributed_type = DistributedType.DEEPSPEED
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl")
+                    torch.distributed.init_process_group(backend="nccl", **kwargs)
                     self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
@@ -175,7 +177,7 @@ def __init__(self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None,
             elif int(os.environ.get("LOCAL_RANK", -1)) != -1 and not cpu:
                 self.distributed_type = DistributedType.MULTI_GPU
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend="nccl")
+                    torch.distributed.init_process_group(backend="nccl", **kwargs)
                     self.backend = "nccl"
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()
@@ -213,7 +215,7 @@ def __init__(self, fp16: bool = None, cpu: bool = False, deepspeed_plugin=None,
                             "please try exporting rank 0's hostname as MASTER_ADDR"
                         )
                 if not torch.distributed.is_initialized():
-                    torch.distributed.init_process_group(backend, rank=rank, world_size=size)
+                    torch.distributed.init_process_group(backend, rank=rank, world_size=size, **kwargs)
                     self.backend = backend
                 self.num_processes = torch.distributed.get_world_size()
                 self.process_index = torch.distributed.get_rank()