Repository: pytorch
Issue: First all_to_all_single / barrier takes a large amount of time.
Labels: oncall: distributed,triaged

### 🐛 Describe the bug

I'm getting profiles like this:

<img width="1224" alt="Screenshot 2024-09-05 at 1 16 44 PM" src="https://github.com/user-attachments/assets/6785039e-c902-42b0-bc35-9196ce1d85e9">

The first `all_to_all_single` / `barrier` are taking large amounts of time (and invoking thousands of `cudaMemcpyAsync`, `cudaGetDeviceCount`, and more), but subsequent `torch.distributed` operations are near instantaneous.

Here's a reproducible example:
```
import os
import torch
import torch.distributed as dist

os.environ["CUDA_VISIBLE_DEVICES"] = os.environ["LOCAL_RANK"]
print(f"set cuda visible devices to {os.environ['CUDA_VISIBLE_DEVICES']}")

def main():
    # Initialize the distributed environment
    dist.init_process_group(backend='nccl')
    
    # Get the world size and rank of the current process
    world_size = dist.get_world_size()
    rank = dist.get_rank()

    # Create a new process group
    pg = dist.new_group(ranks=list(range(world_size)))

    # Create input tensor with unique values for each rank
    input_tensor = torch.tensor([rank * 10 + i for i in range(world_size)], dtype=torch.float32).cuda()

    # Create output tensor to store the result
    output_tensor = torch.zeros(world_size, dtype=torch.float32).cuda()

    # Synchronize before starting the timer
    torch.cuda.synchronize()

    # Create CUDA events for timing
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)

    # Perform timed all-to-all communication for 2 iterations
    print("Timing ALL_TO_ALL_SINGLE for 2 iterations...")
    
    for iteration in range(2):
        # Start timing
        start_event.record()

        # Perform all-to-all communication
        # dist.all_to_all_single(output_tensor, input_tensor, group=pg)
        dist.barrier(pg)

        # End timing
        end_event.record()

        # Synchronize to ensure the operation is complete
        torch.cuda.synchronize()

        # Calculate elapsed time
        elapsed_time = start_event.elapsed_time(end_event) / 1000  # Convert to seconds

        # Print timing results
        print(f"Rank {rank}: Iteration {iteration + 1} time: {elapsed_time:.6f} seconds")

    # Print the final result for each rank
    print(f"Rank {rank}: Input: {input_tensor}, Output: {output_tensor}")

    # Assert that the result is correct
    expected = torch.tensor([i * 10 + rank for i in range(world_size)], dtype=torch.float32).cuda()
    assert torch.all(output_tensor.eq(expected)), f"Rank {rank}: Result does not match expected values"

    # Clean up the distributed environment
    dist.destroy_process_group(pg)
    dist.destroy_process_group()

if __name__ == "__main__":
    main()
```

run it with:
```
torchrun --nproc_per_node=4 --standalone -m script_name
```

It should give output like:
```
set cuda visible devices to 0
set cuda visible devices to 2
set cuda visible devices to 3
set cuda visible devices to 1
Timing ALL_TO_ALL_SINGLE for 2 iterations...
Timing ALL_TO_ALL_SINGLE for 2 iterations...
Timing ALL_TO_ALL_SINGLE for 2 iterations...
Timing ALL_TO_ALL_SINGLE for 2 iterations...
Rank 2: Iteration 1 time: 6.860195 secondsRank 3: Iteration 1 time: 6.856339 seconds

Rank 0: Iteration 1 time: 6.771554 seconds
Rank 1: Iteration 1 time: 6.891210 seconds
Rank 2: Iteration 2 time: 0.000255 seconds
Rank 1: Iteration 2 time: 0.000166 seconds
Rank 0: Iteration 2 time: 0.000229 secondsRank 3: Iteration 2 time: 0.000307 seconds

Rank 0: Input: tensor([0., 1., 2., 3.], device='cuda:0'), Output: tensor([ 0., 10., 20., 30.], device='cuda:0')
Rank 3: Input: tensor([30., 31., 32., 33.], device='cuda:0'), Output: tensor([ 3., 13., 23., 33.], device='cuda:0')
Rank 1: Input: tensor([10., 11., 12., 13.], device='cuda:0'), Output: tensor([ 1., 11., 21., 31.], device='cuda:0')
Rank 2: Input: tensor([20., 21., 22., 23.], device='cuda:0'), Output: tensor([ 2., 12., 22., 32.], device='cuda:0')
```

### Versions

Collecting environment information...
PyTorch version: 2.4.0+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.4 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.35

Python version: 3.10.12 (main, Mar 22 2024, 16:50:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-112-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 H100 80GB HBM3
GPU 1: NVIDIA H100 80GB HBM3
GPU 2: NVIDIA H100 80GB HBM3
GPU 3: NVIDIA H100 80GB HBM3
GPU 4: NVIDIA H100 80GB HBM3
GPU 5: NVIDIA H100 80GB HBM3
GPU 6: NVIDIA H100 80GB HBM3
GPU 7: NVIDIA H100 80GB HBM3

Nvidia driver version: 535.183.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.3.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.3.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      46 bits physical, 57 bits virtual
Byte Order:                         Little Endian
CPU(s):                             176
On-line CPU(s) list:                0-175
Vendor ID:                          GenuineIntel
Model name:                         Intel(R) Xeon(R) Platinum 8468V
CPU family:                         6
Model:                              143
Thread(s) per core:                 2
Core(s) per socket:                 44
Socket(s):                          2
Stepping:                           8
BogoMIPS:                           4800.00
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves avx_vnni avx512_bf16 wbnoinvd arat avx512vbmi umip pku ospke waitpkg avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid bus_lock_detect cldemote movdiri movdir64b fsrm md_clear serialize tsxldtrk avx512_fp16 arch_capabilities
Virtualization:                     VT-x
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          4.1 MiB (88 instances)
L1i cache:                          2.8 MiB (88 instances)
L2 cache:                           176 MiB (88 instances)
L3 cache:                           195 MiB (2 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-87
NUMA node1 CPU(s):                  88-175
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI Syscall hardening, KVM SW loop
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Mitigation; TSX disabled

Versions of relevant libraries:
[pip3] lion-pytorch==0.1.4
[pip3] numpy==1.26.4
[pip3] onnx==1.15.0
[pip3] onnxruntime==1.17.0
[pip3] optree==0.11.0
[pip3] pytorch-ranger==0.1.1
[pip3] pytorch-triton==3.0.0+45fff310c8
[pip3] stlpips_pytorch==0.0.2
[pip3] torch==2.4.0+cu124
[pip3] torch-fidelity==0.3.0
[pip3] torch-optimizer==0.3.0
[pip3] torch-summary==1.4.5
[pip3] torchaudio==2.4.0+cu124
[pip3] torchdiffeq==0.2.3
[pip3] torchmetrics==1.3.0.post0
[pip3] torchvision==0.19.0+cu124
[pip3] triton==3.0.0
[conda] Could not collect

cc @XilunWu @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o