Repository: pytorch
Issue: Slowdown in torch.distributed.new_group when scaling to large clusters. 
Labels: oncall: distributed,triaged,module: c10d

### 🐛 Describe the bug

When initializing torch.distributed comms with `new_group`, the time takes much longer, growing linearly with the number of GPUs that we use. We believe the reason for this is that we need to call `torch.distributed.new_group` on all ranks, not just the ones that are part of it. One example layout would be a model with op sharding and pipelining using 8 op shards and 100 pipe pieces. In this case, there are 100 op-shard groups, but each rank is only part of 1 of them. However, we currently need to call new_group all 100 times. Similarly, if we have comms between adjacent pipe pieces, there would be 99 of them, but each rank is only part of 2. That's just a hypothesis, the data we are seeing is that the time to create the comms increases as we increase the number of GPUs. 

<img width="1223" alt="Screen Shot 2022-01-10 at 3 37 56 PM" src="https://user-images.githubusercontent.com/4275604/148855625-76088261-4810-45fc-92c2-be26c4a3e4e4.png">


### Versions

Collecting environment information...
PyTorch version: 1.10.0a0+git302ee7b
Is debug build: False
CUDA used to build PyTorch: 11.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.27

Python version: 3.8.10 (default, Jun  4 2021, 15:09:15)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.129-63.229.amzn2.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.4.120
GPU models and configuration:
GPU 0: A100-SXM4-40GB
GPU 1: A100-SXM4-40GB
GPU 2: A100-SXM4-40GB
GPU 3: A100-SXM4-40GB
GPU 4: A100-SXM4-40GB
GPU 5: A100-SXM4-40GB
GPU 6: A100-SXM4-40GB
GPU 7: A100-SXM4-40GB

Nvidia driver version: 460.73.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.4
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.930
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.2
[pip3] pytorch-memlab==0.2.4
[pip3] torch==1.10.0a0+git302ee7b
[pip3] torchtyping==0.1.4
[conda] blas                      1.0                         mkl
[conda] mkl                       2021.4.0           h06a4308_640
[conda] mkl-include               2021.4.0           h06a4308_640
[conda] mkl-service               2.4.0            py38h7f8727e_0
[conda] mkl_fft                   1.3.1            py38hd3c417c_0
[conda] mkl_random                1.2.2            py38h51133e4_0
[conda] mypy                      0.930                    pypi_0    pypi
[conda] mypy-extensions           0.4.3                    pypi_0    pypi
[conda] numpy                     1.21.2           py38h20f2e39_0
[conda] numpy-base                1.21.2           py38h79a1101_0
[conda] pytorch-memlab            0.2.4                    pypi_0    pypi
[conda] torch                     1.10.0a0+git302ee7b          pypi_0    pypi
[conda] torchtyping               0.1.4                    pypi_0    pypi

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @SciPioneer @H-Huang