Adding alltoall_single collective to collective quantization API (#63154)
authorMarjan Fariborz <marjanf@fb.com>
Fri, 27 Aug 2021 19:45:01 +0000 (12:45 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Fri, 27 Aug 2021 19:46:31 +0000 (12:46 -0700)
commit6a76ee04de5f10b76cc8f97cc254da43905d170b
treec62b77bebcf9d8d0b9e90eb198cc877d8301d7fe
parent04108592a362848a5d3af4332f7628a14e312174
Adding alltoall_single collective to collective quantization API (#63154)

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

The collective quantization API now supports alltoall, alltoall_single, and allscatter. The test is also included.
ghstack-source-id: 136856877

Test Plan: buck test mode/dev-nosan //caffe2/test/distributed/algorithms/quantization:DistQuantizationTests_nccl -- test_all_to_all_single_bfp16

Reviewed By: wanchaol

Differential Revision: D30255251

fbshipit-source-id: 856f4fa12de104689a03a0c8dc9e3ecfd41cad29
test/distributed/algorithms/quantization/test_quantization.py
torch/distributed/algorithms/quantization/quantization.py