Repository: pytorch
Issue: Communications slowdown when using multiple hosts
Labels: oncall: distributed

### 🐛 Describe the bug

Hi,
I'm doing experiments with distributed training with torch (related to [this](https://github.com/pytorch/pytorch/issues/120428)). I found that when I'm training my model on two hosts (on any number of hosts > 1) DtoH operations have significant impact on collective communications (in particular, AG).
Below is the overall trace from 16 GPU (2 hosts) which shows what I'm talking about:
<img width="1369" alt="image" src="https://github.com/pytorch/pytorch/assets/36795440/dbe09d75-4b17-44f2-b759-335835460189">

The trace was obtained with the following code (After ranks and other torch.distributed setup):
```  from torch.profiler import profile, record_function, ProfilerActivity
    from torch.profiler import schedule
    my_schedule = schedule( # just setup scheduler 
    skip_first=1,
    wait=1,
    warmup=1,
    active=58,
    repeat=1)

    n_gpus = torch.distributed.get_world_size()
    print(f"Running on {n_gpus} gpus")
    W = 8192*4
    H = 8192*4
    # large cuda tensor
    activation_tensor = torch.rand(W, H, dtype=torch.bfloat16, device=torch.device('cuda'))
    # large cuda tensor 
    weight_tensor = torch.rand(W, H//n_gpus, dtype=torch.bfloat16, device=torch.device('cuda'))
    cpu_buffer = torch.empty(W, H, dtype=torch.bfloat16, device='cpu', pin_memory=True)
    to_tensor = torch.zeros_like(activation_tensor)
    N_times = 30 # how many times to run each kernel (loops)
    comm_stream = torch.cuda.Stream(priority=-1)
    offload_stream = torch.cuda.Stream()
    torch.distributed.barrier() # just to start from here simultaneously
    torch.cuda.synchronize() # wait all

    with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True,schedule=my_schedule) as prof:
        for op_idx in range(N_times):
            with torch.cuda.stream(comm_stream):
                torch.distributed.all_gather_into_tensor(to_tensor, weight_tensor)
            with torch.cuda.stream(offload_stream):
                cpu_buffer.copy_(activation_tensor, non_blocking=True) # Here is DtoH

            prof.step()
        torch.cuda.synchronize() # wait all kernels in all streams

        for op_idx_1 in range(N_times):
            with torch.cuda.stream(comm_stream):
                torch.distributed.reduce_scatter_tensor(weight_tensor, to_tensor)
                torch.distributed.all_gather_into_tensor(to_tensor, weight_tensor)
            with torch.cuda.stream(offload_stream):
                activation_tensor.copy_(cpu_buffer, non_blocking=True) # Here is HtoD

            prof.step()
        torch.cuda.synchronize() # wait all kernels in all streams
```
Here I noticed that the first loop (AG + DtoH)  looks bizarre:

<img width="676" alt="image" src="https://github.com/pytorch/pytorch/assets/36795440/93af6672-e4eb-4948-a6a6-e73e740311f0">

Here is about 606ms per AG, but for the second loop with the **same** AG, but with HtoD (btw, thus, the load on a bus is approximately the same during that second loop) looks like this:

<img width="1003" alt="image" src="https://github.com/pytorch/pytorch/assets/36795440/6c373e24-b2c1-4153-86df-02a47120a22b">
i.e. 25ms per AG, which is about 40 times slowdown for AG with DtoH.

So the question of this issue is **why there is so significant slowdown when using DtoH in another stream?**


P.S.
For the same experiment with one host only there is no such unusual effect.
I use IB-connection between hosts.

### Versions

torch==2.1.0a0+32f93b1
CUDA used to build PyTorch: 12.2
CUDA runtime version: 12.2.128
Nvidia driver version: 525.60.13
Python version: 3.10.12
NCCL 2.18.5

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