Support distributed training

[jansel]

This is a placeholer task to make distributed training work with TorchDynamo + AOT Autograd. The main work seems to be making sure the relevant ops can be traced with AOT Autograd and are properly added to the FX graph by TorchDynamo.

I expect most of the issues will be at the AOT Autograd level, because TorchDynamo treats most torch.* ops as black box. We should test and verify this though.

@alanwaketan can fill in details.

[Jack47]

any updates on this? I'm wondering DDP is supported yet

[jansel]

@wconstab has this prototype: #628

So it should be possible to get things working. Though we still have work to do to make the experience more smooth and performant.

[carmocca]

Does the limitation only apply to training? Or is distributed inference supported?

[wconstab]

DDP should be functional with dynamo today. What we're working on is profiling the performance and optimizing it as needed. The concern is that DDP's optimization which overlaps allreduce communication with backward computation is undone by dynamo+aot-autograd fusing the whole graph into one compiled forward and one compiled backward function. Depending on your cluster size, model and interconnect bandwidth, you may still see performance gains by using dynamo+ddp over eager+ddp.

[anijain2305]

Closing this as there is a separate workstream for PT2 Distributed.