Repository: pytorch
Issue: Memory Usage of Conv2d Rapidly Increases when Params >1G
Labels: module: cudnn,module: cuda,module: memory usage,triaged

### 🐛 Describe the bug

We recently try to generate high resolution contents using diffusion models but notice unusual memory usage when resolution becoms large

<img width="1161" alt="image" src="https://github.com/pytorch/pytorch/assets/7783214/225c03a5-26d2-4f82-b476-002e48d0e4eb">

Shown in above, the memory usage increasing rapidly when resolution scales from `2048x1792` to `2048x2048` (channel 512x512).  This is not unexpected as the inputs only increases by <10%. 

A potential guess is that there are some hard-coded rules in the CUDNN library which switch the conv implemetation backend, e.g.,  im2col, as `512 * 2048^2 * 3^2 * 2 = 36G` and very close the increased memory usage. Similiar thoughts are also mentioned in [pytorch forum](https://discuss.pytorch.org/t/memory-usage-suddenly-increase-with-specific-input-shape-on-torch-nn-conv2d/99681/2)  and https://github.com/pytorch/pytorch/issues/26950. 

We have figured out a workaround to compute results patch-by-patch (@lmxyy will soon share the code), but utlimatley, we hope this issue can be solved in the PyTorch upstream so further high-res image and video generation works will not be affected.

The code and env to reproduce is attached below:


```python
import wandb

import torch
from torch import nn
from torch.nn import functional as F


def main(
    in_c = 512,
    out_c = 512,
    input_x = 2048,
    input_y = 2048,
    kernel_size = 3,
    stride = 1,
):
    padding = kernel_size // 2,
    model = nn.Conv2d(
        in_channels=in_c,
        out_channels=out_c,
        bias=False,
        kernel_size=kernel_size,
        padding=padding,
        stride=stride,
        dtype=torch.float16,
        device="cuda",
    )
    
    config = {
        "in_c": in_c,
        "out_c": out_c,
        "input_x": input_x,
        "input_y": input_y,
        "kernel_size": kernel_size,
        "stride": stride,
    }

    x = torch.zeros([1, in_c, input_x, input_y], dtype=torch.float16, device="cuda")
    
    print(x.shape)
    print(f"Before max used memory: {torch.cuda.max_memory_allocated()/ 1024**3} G")
    prev_g = torch.cuda.max_memory_allocated()/ 1024**3
    with torch.no_grad():
        y = model(x)
        print(y.shape)
        print(f"After max used memory: {torch.cuda.max_memory_allocated()/ 1024**3} G")
        curr_g = torch.cuda.max_memory_allocated()/ 1024**3
    
if __name__ == "__main__":
    import fire
    fire.Fire(main)
```

### Versions

```
Collecting environment information...
PyTorch version: 2.1.2+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.15.0-1032-oracle-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: 12.3.103
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
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:                   40 bits physical, 48 bits virtual
CPU(s):                          96
On-line CPU(s) list:             0-95
Thread(s) per core:              2
Core(s) per socket:              48
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       AuthenticAMD
CPU family:                      25
Model:                           1
Model name:                      AMD EPYC 7J13 64-Core Processor
Stepping:                        1
CPU MHz:                         2445.404
BogoMIPS:                        4890.80
Virtualization:                  AMD-V
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       3 MiB
L1i cache:                       3 MiB
L2 cache:                        24 MiB
L3 cache:                        192 MiB
NUMA node0 CPU(s):               0-95
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 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; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
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 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm rep_good nopl cpuid extd_apicid amd_dcm tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy svm cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw topoext perfctr_core invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves clzero xsaveerptr wbnoinvd arat npt nrip_save umip pku ospke vaes vpclmulqdq rdpid arch_capabilities

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.2
[pip3] pytorchvideo==0.1.5
[pip3] torch==2.1.2
[pip3] torchvision==0.16.2
[pip3] triton==2.1.0
[conda] numpy                     1.26.2                   pypi_0    pypi
[conda] pytorchvideo              0.1.5                    pypi_0    pypi
[conda] torch                     2.1.2                    pypi_0    pypi
[conda] torchvision               0.16.2                   pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi
```

cc @csarofeen @ptrblck @xwang233