Repository: pytorch
Issue: torch.compile uses more memory when using less parameters
Labels: module: memory usage,module: convolution,triaged,oncall: pt2

### 🐛 Describe the bug

When running a residual block, I noticed that compiled networks that have a 1x1 skip convolution use less memory when applied to all layers rather than just layers where the number of channels change. Up to a certain input size, the program OOMs only when 1x1 convs happen in the layers with changing channel counts. I would expect the opposite because applying 1x1 convs to all layers would add parameters thus increasing memory usage. 


Also, this is not the case if there is no normalization and activation happens after the convolution. Is this just a quirk of the compile function? I know that pre-activation is not the norm but it seems to have produced better results here: https://towardsdatascience.com/resnet-with-identity-mapping-over-1000-layers-reached-image-classification-bb50a42af03e

Code to reproduce:
```python
import torch.nn.functional as F
import torch.nn as nn
import torch._dynamo
import torch, gc

class ResBlock(nn.Module):
    def __init__(self, in_c: int, nc: int, conv_shortcut: bool = False):
        '''
        in_c: number of input channels
        nc: number of output channels
        '''
        super().__init__()
        self.norm = nn.GroupNorm(8, in_c)
        self.conv = nn.Conv2d(in_c, nc, 3, padding=1)
        self.skip = nn.Conv2d(in_c, nc, 1) if (in_c != nc or conv_shortcut) else None
    
    def forward(self, x):
        h = self.conv(F.relu(self.norm(x)))
        #h = F.relu(self.conv(self.norm(x))) # bug also happens in this case
        #h = F.relu(self.norm(self.conv(x))) # bug also happens in this case
        #h = self.conv(F.relu(x)) # bug also happens in this case
        #h = F.relu(self.conv(x)) # bug does NOT happens in this case
        x = x if self.skip is None else self.skip(x)
        return h + x

device = 'cuda'
for ENABLE_CONV_SHORTCUT in [True, False]:
    print('conv shortcut:', ENABLE_CONV_SHORTCUT)
    with torch.no_grad():
        with torch.amp.autocast('cuda'):
            model = nn.Sequential(
                    ResBlock(256, 256, conv_shortcut=ENABLE_CONV_SHORTCUT),
                    ResBlock(256, 512, conv_shortcut=ENABLE_CONV_SHORTCUT),
                    ResBlock(512, 512, conv_shortcut=ENABLE_CONV_SHORTCUT),
                    ResBlock(512, 1024, conv_shortcut=ENABLE_CONV_SHORTCUT),
                    ResBlock(1024, 1024, conv_shortcut=ENABLE_CONV_SHORTCUT),
                    ResBlock(1024, 2048, conv_shortcut=ENABLE_CONV_SHORTCUT),
                    ResBlock(2048, 2048, conv_shortcut=ENABLE_CONV_SHORTCUT),
            ).to(device)
            model = torch.compile(model)
            print('N params:', sum([p.numel() for p in model.parameters()]))
            x = torch.zeros((8, 256, 64, 64)).to(device)
            for _ in range(2):
                __ = model(x)
                torch.cuda.synchronize()
            print('max mem alloc (GiB):', torch.cuda.max_memory_allocated() / (1024**3))
            print('max mem reser (GiB):', torch.cuda.max_memory_reserved() / (1024**3))

    # clear gpu mem
    del model, x, __
    torch._C._cuda_clearCublasWorkspaces()
    torch._dynamo.reset()
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    print()
```

Output (OS: Ubuntu 22.04, CPU: i7-12700H, GPU: RTX 3060 Mobile Max-Q):
```
1x1 conv shortcut in all layers: True
N params: 83245568
max mem alloc (GiB): 1.231997013092041
max mem reser (GiB): 1.888671875

1x1 conv shortcut in all layers: False
N params: 77671168
max mem alloc (GiB): 1.3362345695495605
max mem reser (GiB): 2.07421875
```

### Error logs

_No response_

### Minified repro

_No response_

### Versions

PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04.1) 11.3.0
Clang version: Could not collect
CMake version: version 3.26.4
Libc version: glibc-2.35

Python version: 3.10.9 | packaged by conda-forge | (main, Feb  2 2023, 20:20:04) [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-5.19.0-46-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.2.91
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 3060 Laptop GPU
Nvidia driver version: 535.54.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.7.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:                   39 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          20
On-line CPU(s) list:             0-19
Vendor ID:                       GenuineIntel
Model name:                      12th Gen Intel(R) Core(TM) i7-12700H
CPU family:                      6
Model:                           154
Thread(s) per core:              2
Core(s) per socket:              14
Socket(s):                       1
Stepping:                        3
CPU max MHz:                     4700.0000
CPU min MHz:                     400.0000
BogoMIPS:                        5376.00
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 art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid rdseed adx smap clflushopt clwb intel_pt sha_ni xsaveopt xsavec xgetbv1 xsaves split_lock_detect avx_vnni dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp hwp_pkg_req hfi umip pku ospke waitpkg gfni vaes vpclmulqdq rdpid movdiri movdir64b fsrm md_clear serialize arch_lbr ibt flush_l1d arch_capabilities
Virtualization:                  VT-x
L1d cache:                       544 KiB (14 instances)
L1i cache:                       704 KiB (14 instances)
L2 cache:                        11.5 MiB (8 instances)
L3 cache:                        24 MiB (1 instance)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-19
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
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
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.25.0
[pip3] torch==2.0.1
[pip3] torch-tb-profiler==0.4.1
[pip3] torchaudio==2.0.2
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.8.0               h6a678d5_0  
[conda] libmagma                  2.7.1                hc72dce7_3    conda-forge
[conda] libmagma_sparse           2.7.1                hc72dce7_4    conda-forge
[conda] magma                     2.7.1                h2c23e93_0  
[conda] mkl                       2023.1.0         h6d00ec8_46342  
[conda] mkl-service               2.4.0           py310h5eee18b_1  
[conda] mkl_fft                   1.3.6           py310h1128e8f_1  
[conda] mkl_random                1.2.2           py310h1128e8f_1  
[conda] numpy                     1.25.0          py310h5f9d8c6_0  
[conda] numpy-base                1.25.0          py310hb5e798b_0  
[conda] pytorch                   2.0.1           py3.10_cuda11.7_cudnn8.5.0_0    pytorch
[conda] pytorch-cuda              11.7                 h778d358_5    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-tb-profiler         0.4.1                    pypi_0    pypi
[conda] torchaudio                2.0.2               py310_cu117    pytorch
[conda] torchtriton               2.0.0                     py310    pytorch
[conda] torchvision               0.15.2              py310_cu117    pytorch
[conda] triton                    2.0.0                    pypi_0    pypi


cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305