Repository: pytorch
Issue: [DDP + Dynamo] Tracing DDP AllReduce (Compiled DDP)
Labels: oncall: distributed,feature,triaged,module: ddp,oncall: pt2

### 🚀 The feature, motivation and pitch

**Background**
DistributedDataParallel (DDP) uses `Reducer` to bucket and issue  `allreduce` calls. The main entry point of `Reducer` is through the gradient hook of `AccumuldateGrad`. However, both `Reducer` and the gradient hook are not not traceable by Dynamo.  As a result, when using `torch.compile` to trace a DDP module, all the `allreduce` buckets will only be called at the end of the backend graph, resulting in zero overlapping between communication and backward computation. The current solution is to use `DDPOptimizer` to analyze the bucketing process of `Reducer` and breaks the forward graph into multiple ones so that each corresponding backward graph issues one `allreduce`.

With the recent development of CompiledAutograd, we would like to revisit tracing the whole DDP into a graph and capture all the `allreduce` in the backward graph.

**Approach**
The initial concept involves utilizing CompiledAutograd to trace the gradient hook registered by `Reducer` and properly record the `allreduce` calls. We don't need to trace the entire `Reducer` but only the bucketing process and `allreduce` calls. While CompiledAutograd could trace the gradient hook as a node, the gradient hook function must be traceable by Dynamo to be inlined in the graph challenge arises because the gradient hook o`Reducer` is written in C++ and not a torch op, Dynamo could not trace it. While it is possible to address this issue, the solutions can be difficult to maintain and handling rebucketing and graph breaks remains uncertain.

As an alternative approach, we propose the complete deactivation of `Reducer` and the registration of gradient hooks at the Python layer during DDP compilation. Unlike the gradient hook within `Reducer`, which has to manages bucketing, the Python gradient hook solely invokes the functional `allreduce` for the relevant gradient.  The Python gradient hook is simple to maintain. This tracing process will result in very bad communication performance but the `allreduce` calls are completely captured, allowing the subsequent optimization. We will implement the `allreduce` bucketing in Inductor.

**Pros/Cons**
The advantages of the proposed solution inherently support graph breaks. Since the graph does not contain any `allreduce` buckets, any graph break won't disturb the allreduce process, and the bucketing function in Inductor can continue functioning smoothly. Additionally, rebucketing is seamlessly managed by Dynamo/CompiledAutograd. When there are changes to the used parameters, Dynamo and CompiledAutograd will automatically recompile the model and trigger the rebucketing process.

A potential concern with the proposed solution lies in the compilation speed, particularly it can be slow for large models. The first iteration has to invoke unbucketing `allreduce`, which can be quite slow.

**Milestones**
1. Trace the gradient hook registered in the Python layer with `CompiledAutograd`
2. Trace DDP communication hook (the basic ones, like bf16_compress).
3. Disable `Reducer` when compiling.
4. Implement bucketing in Inductor.

We have validated milestones 1 and 2 and are working on milestone 3 and 4. After all the 4 milestones are done. We will have a MVP of fully traceable DDP.

cc., @wconstab, @jansel, @yf225, @voznesenskym, @ezyang, @animesh

### Alternatives

The proposed approach may also work with https://github.com/pytorch/pytorch/pull/109537/. We will verify with the PR after it is landed.

### Additional context

_No response_

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @XilunWu @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @zou3519 @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @wanchaol @fduwjj @wz337 @kiukchung @d4l3k @lucasllc @tianyu-l @Xia-Weiwen