Repository: pytorch
Issue: Memory leak in `all_reduce` on CPU tensors with Gloo
Labels: oncall: distributed

### 🐛 Describe the bug

I want to synchronize the parameters by averaging them across all workers in distributed training.

In this specific setup, for some reason, doing that `all_reduce` on GPU (NCCL or Gloo) causes a CUDA OOM (see https://github.com/rwth-i6/returnn/issues/1482, maybe https://github.com/pytorch/pytorch/issues/116177, maybe https://github.com/NVIDIA/nccl/issues/962). So my workaround for this is to move the parameters to CPU first, then do the average (via Gloo), then move it back to GPU. The code is as follows:

```python
for param in module.parameters():
    # Separately move each param to CPU (instead of the whole module), to safe CPU memory.
    param_cpu = param.to(torch.device("cpu"))
    # On CPU, we are likely using Gloo, and Gloo does not support AVG
    dist.all_reduce(param_cpu.data, op=dist.ReduceOp.SUM)
    param_cpu.data /= dist.get_world_size()
    param.data = param_cpu.to(param.device)
```

Is there anything wrong with this code?

This seem to cause a memory leak of CPU memory. I added some logs which print the used memory, and I see that exactly during this sync, the CPU memory always increases a bit. It's not much, so it trains for quite a while, but then at some point I get a CPU OOM.

More logs here: https://github.com/rwth-i6/returnn/issues/1490

### Versions

```
PyTorch version: 2.1.0+cu121      
Is debug build: False               
CUDA used to build PyTorch: 12.1                                                
ROCM used to build PyTorch: N/A                                                 
                                                                                
OS: Ubuntu 22.04.3 LTS (x86_64)                                                 
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0                          
Clang version: 15.0.7                                                           
CMake version: Could not collect                                                                                                                                
Libc version: glibc-2.35                                                                                                                                        
                                                                                                                                                                
Python version: 3.11.2 (main, Feb  7 2023, 13:52:42) [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-46-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce GTX 980
Nvidia driver version: 530.41.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True  

...

Versions of relevant libraries:
[pip3] lovely-numpy==0.2.10
[pip3] numpy==1.26.1
[pip3] torch==2.1.0
[pip3] torch-complex==0.4.3
[pip3] torchaudio==2.1.0
[pip3] triton==2.1.0
[conda] No relevant packages
```

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu @fegin @XilunWu @wanchaol @fduwjj @wz337 @tianyu-l @wconstab @yf225