register_comm_hook(state, hook) [source]
 
Registers a communication hook which is an enhancement that provides a flexible hook to users where they can specify how DDP aggregates gradients across multiple workers. This hook would be very useful for researchers to try out new ideas. For example, this hook can be used to implement several algorithms like GossipGrad and gradient compression which involve different communication strategies for parameter syncs while running Distributed DataParallel training.  Parameters 
 
state (object) – 
Passed to the hook to maintain any state information during the training process. Examples include error feedback in gradient compression, peers to communicate with next in GossipGrad, etc. It is locally stored by each worker and shared by all the gradient tensors on the worker.  
hook (callable) – 
Averages gradient tensors across workers and defined as: hook(state: object, bucket: dist._GradBucket) -> torch.futures.Future: This function is called once the bucket is ready. The hook can perform whatever processing is needed and return a Future indicating completion of any async work (ex: allreduce). If the hook doesn’t perform any communication, it can also just return a completed Future. The Future should hold the new value of grad bucket’s tensors. Once a bucket is ready, c10d reducer would call this hook and use the tensors returned by the Future and copy grads to individual parameters. We also provide an API called get_future to retrieve a Future associated with the completion of c10d.ProcessGroup.work.      Warning Grad bucket’s tensors will not be predivided by world_size. User is responsible to divide by the world_size in case of operations like allreduce.   Warning DDP communication hook can only be registered once and should be registered before calling backward.   Warning The Future object that hook returns should contain a result that has the same shape with the tensors inside grad bucket.   Warning DDP communication hook does not support single-process multiple-device mode. Gradbucket tensors should consist of only a single tensor.   Warning get_future API supports only NCCL backend and will return a torch._C.Future which is an internal type and should be used with caution. It can still be used by register_comm_hook API, but it is subject to some subtle differences compared to torch.futures.Future.   Warning DDP communication hook is experimental and subject to change.   Example::

Below is an example of a noop hook that returns the same tensors. >>> def noop(state: object, bucket: dist._GradBucket): -> torch.futures.Future
>>>     fut = torch.futures.Future()
>>>     fut.set_result(bucket.get_tensors())
>>>     return fut
 >>> ddp.register_comm_hook(state = None, hook = noop)
  Example::

Below is an example of a Parallel SGD algorithm where gradients are encoded before allreduce, and then decoded after allreduce. >>> def encode_and_decode(state: object, bucket: dist._GradBucket): -> torch.futures.Future
>>>     tensors = [t / process_group.world_size for t in bucket.get_tensors()]
>>>     encoded_tensors = encode(tensors) # encode gradients
>>>     fut = process_group.allreduce(encoded_tensors).get_future()
>>>     # Define the then callback to decode.
>>>     def decode(fut):
>>>         decoded_tensors = decode(fut.value()) # decode gradients
>>>         return decoded_tensors
>>>     return fut.then(decode)
 >>> ddp.register_comm_hook(state = None, hook = encode_and_decode)