Device agnostic gradient reduction

[pietern]

There are separate implementations for the distributed data parallelism for CUDA and CPU tensors.

The CUDA one has evolved quite a bit and includes:

• Dynamic bucketing of gradient tensors to reduce the number of reduction calls
• Start reduction when a bucket is full (i.e. all gradients participating have been computed)
• Handling of fp32 and fp16 gradients
• Handling of multiple model copies across devices
• Buffer broadcast

The CPU one has not evolved and does the following:

1. Wait for all gradients to be computed
2. Flatten all gradients into one big tensor
3. Allreduce
4. Unflatten

Ideally, we turn the CUDA approach into a device agnostic one, such that we can remove the CPU specific implementation. Parts of the CUDA approach are written in C++ (see #12852, #12954, #13496, #13607) and these can be made device agnostic as well. Both the bucketing and autograd hooks are now defined in Python land but can be turned into C++ code as well.

Separately, @mcarilli and team made really nice improvements to standard DDP in Apex. The ideal bucketing order is discovered dynamically such that it reflects the execution order instead of definition order, which can cause significant speedups if these orderings are mismatched (e.g. layers that are used last are defined first).

Goals:

• Use the same DDP implementation regardless of backend
• Enable use of DDP for models defined with the C++ frontend
• Discover ideal bucketing order dynamically

Let's use this as a starting point for discussion.

cc @mcarilli @mrshenli

[mcarilli]

stashing random thought for visibility: On paper, the logic to discover/establish a bucket structure is totally orthogonal to the logic that then uses that bucket structure to kick off allreduces, so that logic from Apex could possibly be tacked onto torch DDP as an additional separate method, for example, without breaking anything else.

[pietern]

Thanks @mcarilli, you're right. We can have a first version of this that uses the order in which the parameters are defined and bolt on the reordering. Then we could have a separate mechanism that discovers the ordering by either:

1. Installing a separate set of autograd hooks, which then propagate to the reducer after 1 iteration (and like you did in Apex it would need all workers to agree on the same ordering). This is at parity with your Apex implementation.

2. Walking the autograd graph before executing it (or in parallel with executing it) to discover the order, agree on it between workers, and apply it to the reducer. This would allow for really dynamic stuff where the unused paths are very large and the ordering changes all the time. I don't know how practical it is to walk the autograd graph though... Any ideas?

[pietern]

@mrshenli We can also use this to be explicit about the multi-device awareness in the forward pass for model parallelism, if it's not already taken care of today.

[mrshenli]

@pietern could you please elaborate on how this could help the forward pass for model parallelism. Do you mean users no long need to manually move intermediate outputs across devices?

[pietern]

Not for the forward pass, but the backwards pass needs to be aware that gradients may be produced on different devices. I meant the mode of execution where you use multiple devices in the forward pass, as something to be aware of when hooking into the backwards pass. I.e. we keep a list of gradients per device, assuming that we see the same gradient on every device. This is not the case when doing combined model parallel and data parallel. Between this and model replication/broadcast that is device aware (so that you can run 2 or more models in the same process, like DDP supports today), there is some work needed to support combined model parallelism within process and data parallelism across processes.

For example, let's say you use devices 0 and 1 and wrap your model with DDP, you'd want to execute:

model = DDP(model, device_ids=[[0,1], [2,3]])

Currently, model replication only works if the model uses a single GPU.

Though you could argue that this mode of execution is not preferred and you should be using different processes anyway, I think that we should keep it in mind when designing this new reducer.

[pietern]

This new reducer should also solve #13273. I.e. if there are tensors that won't get their autograd hooks called, we should somehow be notified of this upfront (or be notified of all tensors that will see their gradients computed), such that we can zero their gradients and unblock their reduction.

[pietern]

Example of mixing model parallelism and data parallelism in this post.

The new gradient reducer should deal with this out of the box.

[pietern]

With #18953 merged, the only item remaining to merge DistributedDataParallelCPU into DistributedDataParallel proper is to have a non-CUDA specific broadcast_coalesced.