Repository: pytorch
Issue: Unstable DtoH execution time
Labels: oncall: distributed

### 🐛 Describe the bug

Hi, 
I faced with unexpected behaviour in DtoH operations. The time spent for operation execution varies along GPUs. 
And the problem is that the difference reach about 50%.
Bellow is the code for reproducing the results (same as in my previous issue [here](https://github.com/pytorch/pytorch/issues/120428)).


```
import torch
import torch.distributed as dist
from torch.profiler import profile, record_function, ProfilerActivity
from torch.profiler import schedule

my_schedule = schedule(
    skip_first=1,
    wait=1,
    warmup=1,
    active=20,
    repeat=1)

dist.init_process_group("nccl")
n_gpus = torch.cuda.device_count()
print(f"Running on {n_gpus} gpus")
W = 2048
H = 8192
torch.cuda.set_device(torch.distributed.get_rank()) # I initialised device this time
activation_tensor = torch.rand(W, H, 128, dtype=torch.bfloat16, device='cuda') # large tensor activation
weight_tensor = torch.rand(W, H, 128//n_gpus, dtype=torch.bfloat16, device='cuda') # large tensor
cpu_buffer = torch.empty(W, H, 128, dtype=torch.bfloat16, device='cpu', pin_memory=True)
to_tensor = torch.zeros_like(activation_tensor)
N_times = 30
comm_stream = torch.cuda.Stream()
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) # AG
        with torch.cuda.stream(offload_stream):
            cpu_buffer.copy_(activation_tensor, non_blocking=True) # DtoH
        prof.step()
    torch.cuda.synchronize() # wait all kernels in all streams
prof.export_chrome_trace("my_profile.json")
dist.destroy_process_group()
```

I did experiments with multiple devices (3 GPUs/ 4 GPUs / 5 GPUs). Bellow are traces (combined from all GPUs into single trace)
* 3 GPUs

![3GPU](https://github.com/pytorch/pytorch/assets/36795440/ca90ff53-d8be-4872-93cf-ee283c13a093)

* 4 GPUs
 
![4GPU](https://github.com/pytorch/pytorch/assets/36795440/0e70ddcd-0c08-4427-ae42-20b636fed06f)


* 5 GPUs

![5GPU](https://github.com/pytorch/pytorch/assets/36795440/2692cdc3-a849-472e-83e1-dde96b25649b)


Here can be seen that time for DtoH varies a lot across the GPUs.
The same shown in my another issue [here](https://github.com/pytorch/pytorch/issues/120709)

The below is screenshot form trace (DtoH part)
<img width="1189" alt="image" src="https://github.com/pytorch/pytorch/assets/36795440/c9c30a39-7353-4611-afdd-5b9eb1ca5c04">

<img width="1184" alt="image" src="https://github.com/pytorch/pytorch/assets/36795440/dda5d36f-2ef3-476d-9ac8-773e6de90e07">

As can be seen, GPU0 is about 2 times slower, than GPU4. The same effect in 1-host experiments as well.
And for HtoD operations the same pattern. Im surprised that the slowdowns are permanent (i.e. all DtoH have approximately the same execution time during the whole experiment)

### 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