Let all_reduce_coalesced and all_gather_coalesced return Future objects (#64722)
authorShen Li <cs.shenli@gmail.com>
Fri, 10 Sep 2021 14:44:09 +0000 (07:44 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Fri, 10 Sep 2021 14:45:25 +0000 (07:45 -0700)
commit2a81e8b8f1526b375d1e78402f91bf8fd82d2b68
tree42ad1f682f0af9fc5ae628083b0e91c141d76879
parent88fff22023b201ee237ab0856d53a154cc1784bb
Let all_reduce_coalesced and all_gather_coalesced return Future objects (#64722)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/64722

`all_reduce_coalesced` and `all_gather_coalesced` are never publicly
released in our API docs. So, I would assume the blast radius to be small.

The motivation for this change to allow implementing
`all_reduce_coalesced` and `all_gather_coalesced` by re-using `allreduce`
and `allgather` C++ cores and perform flatten and copy only on the Python
side. With that, we can then remove `all_reduce_coalesced` and
`all_gather_coalesced` from C++ ProcessGroup APIs. For the async mode,
the copy-back logic after the communication will need to be chained
as a callback on the returned Future and use the chained child Future
as the return value (otherwise, we will need to wrap the child Future
into another work handle). This PR tries to test if we can directly
return a Future without breaking tests and internal use cases. If yes,
it will make the consolidation a lot easier.

cc pietern mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse agolynski SciPioneer H-Huang mrzzd cbalioglu gcramer23

Test Plan: Imported from OSS

Reviewed By: rohan-varma

Differential Revision: D30830994

Pulled By: mrshenli

fbshipit-source-id: dcde0ed9245e9e8fee357b3588b07d540a4b6318
test/distributed/test_c10d_gloo.py
torch/distributed/distributed_c10d.py