@Ruiyang_Chen
That’s a great question! This is actually one item on our roadmap.

I’m wondering is there a way to get the dependency graph for computation/communication operations?


It is possible. One option is to traverse the autograd graph from the output. This code 23 shows how DDP reducer implements this currently.
Another option is to do sth similar to APEX 1 by logging the grad ready order in the first iteration, and use that to prioritize comms.
Comm (async allreduce) operations are kicked off whenever a gradient bucket is ready.

This 6 is a brief description of how DDP works. You can start from there. This paper 4 seems did sth similar for PyTorch.
Do you mind if I ask the scope of the project you are working on? The PyTorch Distributed team hasn’t started working on this yet, but we have some ideas regarding how to do it. If you plan to modify PyTorch code or open source your solution, we can collaborate on this effort.