Repository: pytorch
Issue: Performance degradation for `aten::fill_` and `aten::add_` in backward pass
Labels: module: multiprocessing,triaged

### 🐛 Describe the bug

I am using Pytorch's distributed framework to train a simple model whose parameters are organized into a `num_rows`-by-`num_cols` matrix. I've found that, for a fixed number of training points and `num_cols`, the backward pass slows down dramatically after `num_rows` crosses a threshold. I profiled my code and realized that when `num_rows` reaches this threshold, the `CPU time avg` for backward pass calls to `aten::fill_` and `aten::add_` spike and remain high as `num_rows` grows. This slowdown is more pronounced for larger world sizes. Below is an MWE demonstrating this behavior.

```python
# mwe.py
import argparse
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from torch.nn.functional import mse_loss
from typing import Callable

    
def init_process(
        rank: int, 
        world_size: int, 
        fcn: Callable, 
        path_shared: str,
        num_rows: int, 
        num_cols: int, 
        batch_size: int,
        seed: int
    ):
    dist.init_process_group(
        'gloo', init_method=f'file://{path_shared}',
        rank=rank, world_size=world_size
    )
    fcn(rank, num_rows, num_cols, batch_size, seed)


def train(
        rank: int, 
        num_rows: int, 
        num_cols: int, 
        batch_size: int,
        seed: int
    ) -> None:

    gen = torch.Generator().manual_seed(seed)

    # Generate fake data
    sz = 100000
    x = torch.randint(0, num_rows, (sz,), generator=gen)
    y = torch.randn(sz, generator=gen)

    # Set "model" parameters
    param = torch.randn(num_rows, num_cols, requires_grad=True, generator=gen)

    with torch.autograd.profiler.profile(enabled=True) as prof:

        cnt = 0
        while cnt < sz: 

            # Get batch of data
            last_idx = min(cnt, cnt + batch_size)
            x_batch = x[cnt:last_idx]
            y_batch = y[cnt:last_idx]
            cnt += batch_size

            # Forward pass (a minimum working forward pass)
            preds = (param[x_batch] * param[x_batch]).sum(dim=1)
            loss = mse_loss(preds, y_batch)

            # Backward pass
            with torch.autograd.profiler.record_function('backward_pass'):
                loss.backward()
            param.grad.zero_()

    if rank == 0:
        print(prof.key_averages().table(sort_by="self_cpu_time_total"))

    dist.destroy_process_group()


if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('--num_rows', type=int)
    parser.add_argument('--world_size', type=int)
    args = parser.parse_args()

    num_cols = 4
    batch_size = 512
    seed = 12345
    
    # ---------- DISTRIBUTED TRAINING ---------- #

    # Refresh file for shared file-system initialization
    script_dir = os.path.dirname(os.path.realpath(__file__))
    path_shared = os.path.join(script_dir, 'shared')
    if os.path.exists(path_shared):
        os.remove(path_shared)

    mp.set_start_method('spawn')
    processes = []
    for rank in range(args.world_size):
        p = mp.Process(
            target=init_process, 
            args=(
                rank, args.world_size, train, path_shared,
                args.num_rows, num_cols, batch_size, seed + rank
            )
        )
        p.start()
        processes.append(p)

    for p in processes:
        p.join()
```

If you run the script for various values of `num_rows` (see commands below), you should see that the average CPU times for `aten::fill_` and `aten::add_` are low for 8190, 8191 and 8192 but much higher for 8193, 8194, and 8195. 

```bash
$ python mwe.py --num_rows 8190 --world_size 2
$ python mwe.py --num_rows 8191 --world_size 2
$ python mwe.py --num_rows 8192 --world_size 2
$ python mwe.py --num_rows 8193 --world_size 2
$ python mwe.py --num_rows 8194 --world_size 2
$ python mwe.py --num_rows 8195 --world_size 2
```

Below are two screenshots of profile traces for `num_rows = 8191` and `num_rows = 8194`. 

<img width="1489" alt="Screenshot 2024-02-21 at 3 35 19 PM" src="https://github.com/pytorch/pytorch/assets/44176344/1dca94e7-db54-4262-9f31-b4c0fddbf99e">
<img width="1486" alt="Screenshot 2024-02-21 at 3 35 51 PM" src="https://github.com/pytorch/pytorch/assets/44176344/891faa38-ae00-44ad-a526-3a28c13c44ab">


Notice that in the `num_rows = 8194` trace, there are extended calls to `aten::fill_` and `aten::add_` that are not present in the `num_rows = 8191` trace. 

Does anyone know what causes this slowdown and how I can avoid it?


### Versions

```
Collecting environment information...
PyTorch version: 2.1.2
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Red Hat Enterprise Linux release 8.9 (Ootpa) (x86_64)
GCC version: (GCC) 8.5.0 20210514 (Red Hat 8.5.0-20)
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.28

Python version: 3.10.12 | packaged by conda-forge | (main, Jun 23 2023, 22:40:32) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-4.18.0-513.9.1.el8_9.x86_64-x86_64-with-glibc2.28
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
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
CPU(s):              24
On-line CPU(s) list: 0-23
Thread(s) per core:  1
Core(s) per socket:  12
Socket(s):           2
NUMA node(s):        2
Vendor ID:           GenuineIntel
CPU family:          6
Model:               79
Model name:          Intel(R) Xeon(R) CPU E5-2650 v4 @ 2.20GHz
Stepping:            1
CPU MHz:             2575.204
CPU max MHz:         2900.0000
CPU min MHz:         1200.0000
BogoMIPS:            4390.23
L1d cache:           32K
L1i cache:           32K
L2 cache:            256K
L3 cache:            30720K
NUMA node0 CPU(s):   0,2,4,6,8,10,12,14,16,18,20,22
NUMA node1 CPU(s):   1,3,5,7,9,11,13,15,17,19,21,23
Flags:               fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 invpcid_single pti ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm cqm rdt_a rdseed adx smap intel_pt xsaveopt cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts md_clear flush_l1d

Versions of relevant libraries:
[pip3] numpy==1.25.2
[pip3] torch==2.1.2
[pip3] torchvision==0.15.2a0
[conda] blas                      1.0                         mkl  
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2023.1.0         h213fc3f_46344  
[conda] mkl-fft                   1.3.6                    pypi_0    pypi
[conda] mkl-random                1.2.2                    pypi_0    pypi
[conda] mkl-service               2.4.0                    pypi_0    pypi
[conda] mkl_fft                   1.3.6           py310h1128e8f_1  
[conda] mkl_random                1.2.2           py310h1128e8f_1  
[conda] numpy                     1.25.2                   pypi_0    pypi
[conda] numpy-base                1.25.2          py310hb5e798b_0  
[conda] pytorch                   2.1.2              py3.10_cpu_0    pytorch
[conda] pytorch-mutex             1.0                         cpu    pytorch
[conda] torch                     2.1.2                    pypi_0    pypi
[conda] torchvision               0.15.2a0                 pypi_0    pypi
```

cc @VitalyFedyunin