Repository: pytorch
Issue: [RFC] Modularize DistributedDataParallel
Labels: oncall: distributed,feature,triaged

## Updated on 12/09/2020

Added more details on implementation plan in [this comment below](https://github.com/pytorch/pytorch/issues/37002#issuecomment-741958753).

---------


## Summary

This project aims at decomposing existing `DistributedDataParallel` (DDP) implementation into multiple smaller pluggable and customizable building blocks. So that applications can customize DDP to best serve specific applications. More specifically, each type of individual building block exposes a fixed set of APIs but can implement different algorithms. The new DDP constructor takes building blocks to assemble a DDP instance. Existing DDP API can stay intact and internally calls into the new DDP implementation.

This project is motivated by the following requirements:

* Enable new features without polluting existing DDP API
    * Many of the new features would require adding more configurations at initialization. Instead of adding all of them into DDP constructor, these configurations should be self-contained inside the implementation of a building block.
* Make DDP/Reducer code more readable
    * Most of DDP core lives in `reducer.cpp`. It starts from a nicely organized implementation, but as we are gradually adding more features, the growing complexity starts to detrimental to readability.
* Allow applications to customize their DDP by picking and extending sub-components
    * Different applications have vastly different requirements, which we will elaborate in the *Targeted Features* section. Packing everything in one implementation would inevitably leads to unnecessary slowdowns. Instead, applications should be able to just pick the algorithms they need and get rid of those unnecessary complexities. 
* Add C++ DDP API
* Gradually remove legacy features, e.g., `device_ids`, `dim`, `check_reduction`, etc. 
    * `device_ids` is necessary for single-process multi-GPU mode, which is not the recommended mode as GIL contention and the per-iteration scatter and gather could introduce additional overhead. The recommended one process per module replica/device use case does not need `device_ids` or `dim`.
* Gradually move DDP from `torch.nn.parallel` to `torch.distributed` 
    * Currently DDP API lives in `torch.nn.parallel` but the core `Reducer` implementation is in the distributed package, which creates a weird dependency that the single-machine package depends on the distributed package. We can resolve this by moving DDP into the `torch.distributed` package.



## Compositional Building Blocks

We propose to decompose the main logic of `DistributedDataParallel` into the following four building blocks


* **Grad Reader**: responsible for reading grad from the local model, which is done by manipulating DDP hooks, including
    * Where to read the grad from: e.g., `parameter.grad` or from distributed autograd context.
    * When to read grad: e.g., `AccumlateGrad` , forward output tensor, etc.
    * Which grad to expect: e.g., `find_unused_parameters` should skip hooks on unused parameters
* **Grad Bucketer**: responsible for organizing grads into buckets, including
    * How do grads maps to buckets: e.g., static mapping determined at construction time, dynamic mapping based on grad-ready order, etc.
    * Determine when a bucket is ready for communication
    * Grad/bucket compression
* **Comm Scheduler**: responsible for how to communicate buckets, including
    * What comm strategy to use: e.g., `AllReduce` or gossip
    * How to prioritize communication
* **Grad Writer**: responsible for writing the communicated grad back to the local model, including
    * Where to read the grad from: e.g., `parameter.grad` or from distributed autograd context.
    * Trigger post DDP-grad hooks if there are any: e.g., this is useful to overlap optimizer `step` with DDP communication. 



## Targeted New Features

* **Mix DDP with RPC**
    * Use DDP in conjunction with RPC functions.  Currently, DDP hook is only registered to the `AccumulateGrad` function and looks for grad from `param.grad`. To support RPC, DDP need to read grad from and write grad to the distributed autograd context.
    * Compositional building blocks: 
        * Grad Reader
        * Grad Writer
* **Delay Allreduce**
    * There are certain scenarios where we cannot make assumptions on how many times a `AccumulateGrad` hook will be called or we cannot traverse the real autograd graph from the outputs (e.g., with checkpoints + reused layers). For these cases, we can have a slower version of DDP which does not overlap comm with comp, so that DDP should be able to support all use cases that would work with local autograd. (see https://github.com/pytorch/pytorch/pull/35083)
    * Compositional building blocks: 
        * Grad Reader
        * Grad Writer
* **Tune parameter to bucket mapping**
    * This tunes the order of `AllReduce` parameters. Ideally, the `AllReduce` order should be the same as gradient ready order. So that DDP can launch communications as soon as possible. We might have to do some profiling to detect the gradient ready order in the backward pass, and might need to change it over iterations. All processes would need to agree on the mapping, but as the mapping adjustment should occur infrequently, we should be able to afford introducing one more `AllReduce` to achieve consensus. (see https://github.com/pytorch/pytorch/pull/35137)
    * Compositional building blocks: 
        * Grad Bucketer
* **Layer Drop/Skip Comm for Unused Params**
    * Currently, DDP always `AllReduce` all buckets (i.e., all grads) even if some parameters are not used in the forward pass. Some applications intentionally skip layers to speed up training and prevent overfitting. In some other applications, different inputs might reach different part of the model.  Allowing skip unused parameters in the communication as well will help to speed up training. One challenge is that all processes need to have a global consensus on which gradients to skip. (see [paper](https://arxiv.org/pdf/1909.11556.pdf) as one example use case)
    * Compositional building blocks:
        * Grad Reader
        * Grad Bucketer
        * Grad Writer
* **Overlap DDP with optimizer**
    * Existing DDP enforces a hard barrier between backward and optimizer step. However, this is not necessary, as optimizers should be able to start working as soon as some grad is ready. Therefore, we should allow optimizer step to overlap with DDP backward. 
    * Compositional building blocks:
        * Grad Writer
* **Overlap DDP backward with next forward**
    * Description: the reason for computing grad is to use the updated parameter in the next iteration. The next forward does not have to be block until all parameter are updated. Instead, it can start as long as the first layer parameters are ready. So we should allow overlapping DDP backward with next iteration’s forward, and only block forward on a layer if its parameter hasn’t been updated yet.
    * Challenge: API design
    * Compositional building blocks:
        * per sub-module pre-forward hooks (to block forward until parameter is updated, this is already available)
* **Prioritize and preempt allreduce**
    * Grads becomes ready in the backward order, but parameters are consumed by the next forward in the opposite order. If we want to overlap backward with next forward well, we need to assign higher priorities to `AllReduce` operations for first layers. However, as buckets for last layers will be ready first, they will start early. Hence, we need some way to implement *preemption* in allreduce operations.  (see [paper1](https://dl.acm.org/doi/10.1145/3341301.3359642) and [paper2](https://arxiv.org/abs/1905.03960))
    * Compositional building blocks:
        * Comm Scheduler
* **Gossip DDP**
    * Instead of using `AllReduce` to globally sync gradients, it’s proven that gossip and async param sync also works and can potentially lead to considerable training speedups. 
    * Compositional building blocks:
        * Comm Scheduler
* **Gradient Compression**

    * Mixed precision training has been proven to help accelerate training considerably. DDP should support this use case as well by allowing use less bits to store and communicate grads.
    * Compositional building blocks:
        * Grad Bucketer



cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @xush6528 @osalpekar