Repository: pytorch
Issue: `torch.distributed.all_reduce` allocates excess GPU memory when using NCCL backend
Labels: oncall: distributed,module: nccl,has workaround

### 🐛 Describe the bug

### Problem

If I perform the following steps in a distributed setting (NCCL backend):

1.  Create two tensors on two workers, one tensor on GPU0 and the other on GPU1.
2. Run `torch.distributed.all_reduce` on these tensors across the two processes.

Then this will result in the GPU1 worker allocating a non-trivial amount of memory (~800mb) on GPU0.

This problem seems to occur when using torch 1.12.1 with CUDA 11.3 but does not occur when I use torch 1.9.0 with CUDA 11.1.

### Minimal reproducing example

The following script will reproduce this behavior, e.g. if you run this script and watch `nvidia-smi` you will see that GPU0 allocates roughly 2x the memory of GPU1.

```python
import time

import torch
import torch.multiprocessing as mp

mp = mp.get_context("forkserver") # Also tried with "spawn"

import torch.distributed as dist

import socket
from contextlib import closing

from torch import multiprocessing as mp


def find_free_port(address: str = "127.0.0.1") -> int:
    """Helper function to find a free port for distributed tests below"""
    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
        s.bind((address, 0))
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        port = s.getsockname()[1]
    return port


def init_and_run_all_reduce(worker_ind: int, port: int):
    print(f"Worker {worker_ind}: started")

    print(f"Worker {worker_ind}: creating TCP store")
    store = torch.distributed.TCPStore(
        host_name="127.0.0.1", port=port, world_size=2, is_master=worker_ind == 0,
    )

    print(f"Worker {worker_ind}: Starting process group")
    dist.init_process_group(
        backend="nccl", store=store, rank=worker_ind, world_size=2,
    )

    device = torch.device(worker_ind)

    t = torch.tensor([worker_ind], device=device)

    print(f"Worker {worker_ind}: Creating tensor {t} and running all_reduce on it")

    dist.all_reduce(t)

    print(f"Worker {worker_ind}: Done, reduced tensor: {t}")

    print(f"Worker {worker_ind} entering 10 second sleep")
    time.sleep(10)
    dist.barrier()


if __name__ == "__main__":
    processes = []
    port = find_free_port()
    for i in range(2):
        processes.append(
            mp.Process(
                target=init_and_run_all_reduce,
                kwargs={"worker_ind": i, "port": port},
            )
        )

        processes[-1].start()
        time.sleep(1)

    for p in processes:
        p.join()
```


### Versions

```
PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.9.13 (main, Aug 25 2022, 23:26:10)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.0-91-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: 
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000
GPU 2: NVIDIA RTX A6000
GPU 3: NVIDIA RTX A6000
GPU 4: NVIDIA RTX A6000
GPU 5: NVIDIA RTX A6000
GPU 6: NVIDIA RTX A6000
GPU 7: NVIDIA RTX A6000

Nvidia driver version: 495.29.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.1
[pip3] torch==1.12.1
[pip3] torchvision==0.13.1
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.3.1               h2bc3f7f_2  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py39h7f8727e_0  
[conda] mkl_fft                   1.3.1            py39hd3c417c_0  
[conda] mkl_random                1.2.2            py39h51133e4_0  
[conda] numpy                     1.23.1           py39h6c91a56_0  
[conda] numpy-base                1.23.1           py39ha15fc14_0  
[conda] pytorch                   1.12.1          py3.9_cuda11.3_cudnn8.3.2_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchvision               0.13.1               py39_cu113    pytorch
```

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu