join(divide_by_initial_world_size=True, enable=True) [source]
 
A context manager to be used in conjunction with an instance of torch.nn.parallel.DistributedDataParallel to be able to train with uneven inputs across participating processes. This context manager will keep track of already-joined DDP processes, and “shadow” the forward and backward passes by inserting collective communication operations to match with the ones created by non-joined DDP processes. This will ensure each collective call has a corresponding call by already-joined DDP processes, preventing hangs or errors that would otherwise happen when training with uneven inputs across processes. Once all DDP processes have joined, the context manager will broadcast the model corresponding to the last joined process to all processes to ensure the model is the same across all processes (which is guaranteed by DDP). To use this to enable training with uneven inputs across processes, simply wrap this context manager around your training loop. No further modifications to the model or data loading is required.  Warning This module works only with the multi-process, single-device usage of torch.nn.parallel.DistributedDataParallel, which means that a single process works on a single GPU.   Warning This module currently does not support custom distributed collective operations in the forward pass, such as SyncBatchNorm or other custom defined collectives in the model’s forward pass.   Parameters 
 
divide_by_initial_world_size (bool) – If True, will divide gradients by the initial world_size DDP training was launched with. If False, will compute the effective world size (number of ranks that have not depleted their inputs yet) and divide gradients by that during allreduce. Set divide_by_initial_world_size=True to ensure every input sample including the uneven inputs have equal weight in terms of how much they contribute to the global gradient. This is achieved by always dividing the gradient by the initial world_size even when we encounter uneven inputs. If you set this to False, we divide the gradient by the remaining number of nodes. This ensures parity with training on a smaller world_size although it also means the uneven inputs would contribute more towards the global gradient. Typically, you would want to set this to True for cases where the last few inputs of your training job are uneven. In extreme cases, where there is a large discrepancy in the number of inputs, setting this to False might provide better results. 
enable (bool) – Whether to enable uneven input detection or not. Pass in enable=False to disable in cases where you know that inputs are even across participating processes. Default is True.    Example: >>>  import torch
>>>  import torch.distributed as dist
>>>  import os
>>>  import torch.multiprocessing as mp
>>>  import torch.nn as nn
>>>  # On each spawned worker
>>>  def worker(rank):
>>>      dist.init_process_group("nccl", rank=rank, world_size=2)
>>>      torch.cuda.set_device(rank)
>>>      model = nn.Linear(1, 1, bias=False).to(rank)
>>>      model = torch.nn.parallel.DistributedDataParallel(
>>>          model, device_ids=[rank], output_device=rank
>>>      )
>>>      # Rank 1 gets one more input than rank 0.
>>>      inputs = [torch.tensor([1]).float() for _ in range(10 + rank)]
>>>      with model.join():
>>>          for _ in range(5):
>>>              for inp in inputs:
>>>                  loss = model(inp).sum()
>>>                  loss.backward()
>>>  # Without the join() API, the below synchronization will hang
>>>  # blocking for rank 1's allreduce to complete.
>>>  torch.cuda.synchronize(device=rank)