Unexpected High PCIe traffic in Distributed Training since PT 2
Labels: oncall: distributed, triaged, module: fsdp

### 🐛 Describe the bug

Since PT 2, we have noticed significant amount of PCIe traffic between host and device, which is something we didn't expect to happen and not observed in PT 1.x version. This applies to both DDP and FSDP, and we observed as much as 2G/s to 4G/s traffic throughout of the whole training stage (for our jobs with 3B+ models) which eventually harm the training due to continuous high consumption of the PCIe and the host.

To reproduce, I am using the same code from the official DDP tutorial: https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#initialize-ddp-with-torch-distributed-run-torchrun
except:
1. modify model size to test different model sizes
2. add a training loop to observe the continuous traffic
```python
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim

from torch.nn.parallel import DistributedDataParallel as DDP

from tqdm import tqdm

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 100000000)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(100000000, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic():
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    print(f"Start running basic DDP example on rank {rank}.")

    # create model and move it to GPU with id rank
    device_id = rank % torch.cuda.device_count()
    model = ToyModel().to(device_id)
    print(f"\n--> model has {sum(p.numel() for p in model.parameters() if p.requires_grad)/1e6} Million params\n")
    ddp_model = DDP(model, device_ids=[device_id])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    for _ in tqdm(range(1000000000)):
        optimizer.zero_grad()
        outputs = ddp_model(torch.randn(20, 10))
        labels = torch.randn(20, 5).to(device_id)
        loss_fn(outputs, labels).backward()
        optimizer.step()
```

PCIe traffic:

  | Prior 2.0 (1.13.1+cu117) | Prior 2.0 (1.13.1+cu117) | 2.x (2.1.0.dev20230603+cu118) | 2.x (2.1.0.dev20230603+cu118)
-- | -- | -- | -- | --
model size | device to host (DtoH) | host to device (HtoD) | device to host (DtoH) | host to device (HtoD)
1.6B | 8.8M/s | 21M/s | 122M/s | 950M/s
160M | 10.8M/s | 24.5M/s | 122M/s | 950M/s
16M | 21M/s | 48M/s | 73M/s | 480M/s
1.6M | 9M/s | 45M/s | 25M/s | 170M/s
0.16M | 8.5M/s | 41.5M/s | 7M/s | 36M/s

Some screenshots (for 1.6B model):
PT 1.x
<img width="1692" alt="image" src="https://github.com/pytorch/pytorch/assets/20955448/681a33fe-eb13-4619-9656-27a93ce08f16">
PT 2.x
<img width="1697" alt="image" src="https://github.com/pytorch/pytorch/assets/20955448/bb73ab7a-4e79-4668-909b-b58124ceb7c2">





Some other notes from our experiences:
1. we can confirm the traffic came only from back path (i.e. `loss.backward()`)
2. for DDP, we observed very little to zero PCIe traffic in PT 1.x (as expected), but very high PCIe traffic in PT 2.x
3. for FSDP, we observed high PCIe traffic in PT 1.x as well (which isn't as expected as we didn't turn on any offloading in FSDP), and in PT 2.x, the traffic is even higher (roughly 2X for the same training)
4. We were able to bring PCIe traffic down to almost zero in all these experiments by turning on CUDA_LAUNCH_BLOCKING. However, that is only for debugging purpose as turning on blocking will slow down the training.


### Versions

```
PyTorch version: 2.1.0.dev20230603+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Debian GNU/Linux 11 (bullseye) (x86_64)
GCC version: (Debian 10.2.1-6) 10.2.1 20210110
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.10.10 (main, Mar 21 2023, 18:45:11) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-4.18.0-372.19.1.el8_6.x86_64-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 515.48.07
cuDNN version: Could not collect
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
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 48 bits virtual
CPU(s):                          80
On-line CPU(s) list:             0-79
Thread(s) per core:              2
Core(s) per socket:              20
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           85
Model name:                      Intel Xeon Processor (Cascadelake)
Stepping:                        5
CPU MHz:                         2400.000
BogoMIPS:                        4800.00
Virtualization:                  VT-x
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       1.3 MiB
L1i cache:                       1.3 MiB
L2 cache:                        160 MiB
L3 cache:                        32 MiB
NUMA node0 CPU(s):               0-39
NUMA node1 CPU(s):               40-79
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          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
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
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 rep_good nopl xtopology cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 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 mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat pku ospke avx512_vnni md_clear arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.25.0rc1
[pip3] pytorch-triton==2.1.0+9820899b38
[pip3] torch==2.1.0.dev20230603+cu118
[pip3] torchvision==0.16.0.dev20230603+cu118
[conda] numpy                     1.25.0rc1                pypi_0    pypi
[conda] pytorch-triton            2.1.0+9820899b38          pypi_0    pypi
[conda] torch                     2.1.0.dev20230603+cu118          pypi_0    pypi
[conda] torchvision               0.16.0.dev20230603+cu118          pypi_0    pypi
```


cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @lessw2020 @HamidShojanazeri