Repository: pytorch
Issue: Why the cuda peak memory usage increase after dist.reuduce_scatter operation ? 
Labels: oncall: distributed

### 🐛 Describe the bug
``` Python
import torch
import torch.distributed as dist
import gc

def see_memory_usage(rank,message, force=False, ram=False, ranks=[0], group=None):
    device = torch.device(f"cuda:{rank}")
    if not force:
        return
    gc.collect()
    global peak_memory
    max_ma = round(torch.cuda.max_memory_allocated(device) / (1024 * 1024),2)
    peak_memory = max(peak_memory, max_ma)
    print(
        '____rank___',rank,
        message,
        f"MA {round(torch.cuda.memory_allocated(device) / (1024 * 1024),2 )} MB \
        Max_MA {max_ma} MB \
        CA {round(torch.cuda.memory_reserved(device) / (1024 * 1024),2)} MB \
        Max_CA {round(torch.cuda.max_memory_reserved(device) / (1024 * 1024))} MB \
        PEAK_MA {peak_memory} MB    ")

dist.init_process_group(backend="nccl")

rank = dist.get_rank()
size = dist.get_world_size()
dp_groups = [[0,1,2,3]]
for g in dp_groups:
    proc_group = dist.new_group(ranks=g)
    if rank in g:
        dp_group = proc_group
local_rank = dist.get_rank(dp_group)


e = []
if rank == 0 or rank == 1:
    for i in range(0, 4):
        e.append(torch.rand(size=(1,50000000), device=f"cuda:{rank}"))
elif rank == 2 or rank == 3:
    for i in range(0, 4):
        e.append(torch.rand(size=(1,50000000), device=f"cuda:{rank}"))
see_memory_usage(rank,'**** \n memory consumption before dist.reduce_scatter ', force=True)
dist.reduce_scatter(output=e[rank],
                input_list=e,
                group=dp_group)
see_memory_usage(rank,'**** \n memory consumption after dist.reduce_scatter ', force=True)
```


## Output
``` Shell
____rank___ 3 **** 
 memory consumption before dist.reduce_scatter  MA 762.94 MB         Max_MA 762.94 MB         CA 768.0 MB         Max_CA 768 MB         PEAK_MA 762.94 MB    
____rank___ 1 **** 
 memory consumption before dist.reduce_scatter  MA 762.94 MB         Max_MA 762.94 MB         CA 768.0 MB         Max_CA 768 MB         PEAK_MA 762.94 MB    
____rank___ 2 **** 
 memory consumption before dist.reduce_scatter  MA 762.94 MB         Max_MA 762.94 MB         CA 768.0 MB         Max_CA 768 MB         PEAK_MA 762.94 MB    
____rank___ 0 **** 
 memory consumption before dist.reduce_scatter  MA 762.94 MB         Max_MA 762.94 MB         CA 768.0 MB         Max_CA 768 MB         PEAK_MA 762.94 MB    
____rank___ 1 **** 
 memory consumption after dist.reduce_scatter  MA 762.94 MB         Max_MA 1525.88 MB         CA 1532.0 MB         Max_CA 1532 MB         PEAK_MA 1525.88 MB    
____rank___ 2 **** 
 memory consumption after dist.reduce_scatter  MA 762.94 MB         Max_MA 1525.88 MB         CA 1532.0 MB         Max_CA 1532 MB         PEAK_MA 1525.88 MB    
____rank___ 3 **** 
 memory consumption after dist.reduce_scatter  MA 762.94 MB         Max_MA 1525.88 MB         CA 1532.0 MB         Max_CA 1532 MB         PEAK_MA 1525.88 MB    
____rank___ 0 **** 
 memory consumption after dist.reduce_scatter  MA 762.94 MB         Max_MA 1525.88 MB         CA 1532.0 MB         Max_CA 1532 MB         PEAK_MA 1525.88 MB 
```
### Versions

pytorch 1.7.1 backend nccl