Training iresnet with torch.compile is slower than eager mode for torch 2.1.0
Labels: triaged, oncall: pt2

### 🐛 Describe the bug

When I compiled "iresnet200" using "torch.compile", I noticed that iresnet's speedup was abnormal. 
The following table shows the speedup of  inductor mode comapred to eager mode between  networks.

| networks      | resnet152 | iresnet200 |
| ------------- | --------- | ---------- |
| training mode | 1.38x     | **0.7x**   |
| eval mode     | 2.86x     | 1.4x       |

* resnet152 is imported from torchvision.models

Can you help me with the following questions?
1. The inducer can give iresnet in eval mode a delightful acceleration. But why does it actually reduce the speedup of iresnet in training mode？
2.  Is there a way to improve the speedup of iresnet in training mode ？

### Error logs

## logs of resnet152  
```
eager train time 0: 1.672048583984375
eager train time 29: 0.05983129501342774
~~~~~~~~~~
compile train time 0: 129.01534375
compile train time 29: 0.04168806457519531
~~~~~~~~~~
(train) eager median: 0.05907711982727051, compile median: 0.04255897521972656, speedup: 1.3881236454182198x
~~~~~~~~~~

eager eval time 0: 1.4442802734375
eager eval time 29: 0.022561792373657227
~~~~~~~~~~
compile eval time 0: 58.9843515625
compile eval time 29: 0.00785203218460083
~~~~~~~~~~
(eval) eager median: 0.022537727355957028, compile median: 0.007875584125518798, speedup: 2.8617213652672873x
~~~~~~~~~~
```

##  logs of iresnet200
```
eager train time 0: 1.701213134765625
eager train time 29: 0.09520025634765625
~~~~~~~~~~
compile train time 0: 258.081875
compile train time 29: 0.12313497924804688
~~~~~~~~~~
(train) eager median: 0.0953922576904297, compile median: 0.12323174285888672, speedup: 0.774088359682325x


eager eval time 0: 1.4410618896484375
eager eval time 29: 0.03419750213623047
~~~~~~~~~~
compile eval time 0: 107.7888359375
compile eval time 29: 0.024010751724243166
~~~~~~~~~~
(eval) eager median: 0.03415500831604004, compile median: 0.02414489555358887, speedup: 1.4145850513303746x
```


### Minified repro

```python
import numpy as np
import torch
import torch._dynamo as dynamo
import torch.nn as nn
import torch.optim as optim
from torch import nn
from torch.utils.checkpoint import checkpoint

__all__ = ['iresnet18', 'iresnet34', 'iresnet50', 'iresnet100', 'iresnet200']
using_ckpt = False


def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes,
                     out_planes,
                     kernel_size=3,
                     stride=stride,
                     padding=dilation,
                     groups=groups,
                     bias=False,
                     dilation=dilation)


def conv1x1(in_planes, out_planes, stride=1):
    """1x1 convolution"""
    return nn.Conv2d(in_planes,
                     out_planes,
                     kernel_size=1,
                     stride=stride,
                     bias=False)


class IBasicBlock(nn.Module):
    expansion = 1

    def __init__(self,
                 inplanes,
                 planes,
                 stride=1,
                 downsample=None,
                 groups=1,
                 base_width=64,
                 dilation=1):
        super(IBasicBlock, self).__init__()
        if groups != 1 or base_width != 64:
            raise ValueError(
                'BasicBlock only supports groups=1 and base_width=64')
        if dilation > 1:
            raise NotImplementedError(
                "Dilation > 1 not supported in BasicBlock")
        self.bn1 = nn.BatchNorm2d(
            inplanes,
            eps=1e-05,
        )
        self.conv1 = conv3x3(inplanes, planes)
        self.bn2 = nn.BatchNorm2d(
            planes,
            eps=1e-05,
        )
        self.prelu = nn.PReLU(planes)
        self.conv2 = conv3x3(planes, planes, stride)
        self.bn3 = nn.BatchNorm2d(
            planes,
            eps=1e-05,
        )
        self.downsample = downsample
        self.stride = stride

    def forward_impl(self, x):
        identity = x
        out = self.bn1(x)
        out = self.conv1(out)
        out = self.bn2(out)
        out = self.prelu(out)
        out = self.conv2(out)
        out = self.bn3(out)
        if self.downsample is not None:
            identity = self.downsample(x)
        out += identity
        return out

    def forward(self, x):
        if self.training and using_ckpt:
            return checkpoint(self.forward_impl, x)
        else:
            return self.forward_impl(x)


class IResNet(nn.Module):
    fc_scale = 7 * 7

    def __init__(self,
                 block,
                 layers,
                 dropout=0,
                 num_features=512,
                 zero_init_residual=False,
                 groups=1,
                 width_per_group=64,
                 replace_stride_with_dilation=None,
                 fp16=False):
        super(IResNet, self).__init__()
        self.extra_gflops = 0.0
        self.fp16 = fp16
        self.inplanes = 64
        self.dilation = 1
        if replace_stride_with_dilation is None:
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError("replace_stride_with_dilation should be None "
                             "or a 3-element tuple, got {}".format(
                                 replace_stride_with_dilation))
        self.groups = groups
        self.base_width = width_per_group
        self.conv1 = nn.Conv2d(3,
                               self.inplanes,
                               kernel_size=3,
                               stride=1,
                               padding=1,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(self.inplanes, eps=1e-05)
        self.prelu = nn.PReLU(self.inplanes)
        self.layer1 = self._make_layer(block, 64, layers[0], stride=2)
        self.layer2 = self._make_layer(block,
                                       128,
                                       layers[1],
                                       stride=2,
                                       dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(block,
                                       256,
                                       layers[2],
                                       stride=2,
                                       dilate=replace_stride_with_dilation[1])
        self.layer4 = self._make_layer(block,
                                       512,
                                       layers[3],
                                       stride=2,
                                       dilate=replace_stride_with_dilation[2])
        self.bn2 = nn.BatchNorm2d(
            512 * block.expansion,
            eps=1e-05,
        )
        self.dropout = nn.Dropout(p=dropout, inplace=True)
        self.fc = nn.Linear(512 * block.expansion * self.fc_scale, num_features)
        self.features = nn.BatchNorm1d(num_features, eps=1e-05)
        nn.init.constant_(self.features.weight, 1.0)
        self.features.weight.requires_grad = False

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, 0, 0.1)
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, IBasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)

    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                nn.BatchNorm2d(
                    planes * block.expansion,
                    eps=1e-05,
                ),
            )
        layers = []
        layers.append(
            block(self.inplanes, planes, stride, downsample, self.groups,
                  self.base_width, previous_dilation))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(
                block(self.inplanes,
                      planes,
                      groups=self.groups,
                      base_width=self.base_width,
                      dilation=self.dilation))

        return nn.Sequential(*layers)

    def forward(self, x):
        # with torch.cuda.amp.autocast(self.fp16):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.prelu(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.bn2(x)
        x = torch.flatten(x, 1)
        x = self.dropout(x)
        x = self.fc(x)
        x = self.features(x)
        return x


def _iresnet(arch, block, layers, pretrained, progress, **kwargs):
    model = IResNet(block, layers, **kwargs)
    if pretrained:
        raise ValueError()
    return model


def iresnet18(pretrained=False, progress=True, **kwargs):
    return _iresnet('iresnet18', IBasicBlock, [2, 2, 2, 2], pretrained,
                    progress, **kwargs)


def iresnet34(pretrained=False, progress=True, **kwargs):
    return _iresnet('iresnet34', IBasicBlock, [3, 4, 6, 3], pretrained,
                    progress, **kwargs)


def iresnet50(pretrained=False, progress=True, **kwargs):
    return _iresnet('iresnet50', IBasicBlock, [3, 4, 14, 3], pretrained,
                    progress, **kwargs)


def iresnet100(pretrained=False, progress=True, **kwargs):
    return _iresnet('iresnet100', IBasicBlock, [3, 13, 30, 3], pretrained,
                    progress, **kwargs)


def iresnet200(pretrained=False, progress=True, **kwargs):
    return _iresnet('iresnet200', IBasicBlock, [6, 26, 60, 6], pretrained,
                    progress, **kwargs)


def generate_data(b, device_id=0):
    channel = 3
    heigh = 112
    width = 112
    return (
        torch.randn(b, channel, heigh, width).to(torch.float32).cuda(device_id),
        torch.randint(256, (b,)).cuda(device_id),
    )


def init_model(device_id=0):
    # return resnet152(num_classes=256).cuda(device_id)
    return iresnet200(num_features=256).cuda(device_id)


def train(model, data, opt, grad_scaler):
    opt.zero_grad(True)
    with torch.cuda.amp.autocast(enabled=True):
        predict = model(data[0])
        loss = torch.nn.CrossEntropyLoss()(predict, data[1])
    loss = grad_scaler.scale(loss)
    loss.backward()
    opt.step()
    return loss


def eval(model, data):
    with torch.cuda.amp.autocast(enabled=True):
        predict = model(data[0])
        loss = torch.nn.CrossEntropyLoss()(predict, data[1])
    return loss


def timed(fn):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    result = fn()
    end.record()
    torch.cuda.synchronize()
    return result, start.elapsed_time(end) / 1000


def demo_train():
    N_ITERS = 30

    device_id = 0

    dynamo.reset()
    model = init_model(device_id=device_id)

    opt = optim.SGD(model.parameters(), lr=0.001)
    grad_scaler = torch.cuda.amp.GradScaler(init_scale=2.0)

    eager_times = []
    for i in range(N_ITERS):
        inp = generate_data(16, device_id=device_id)
        _, eager_time = timed(lambda: train(model, inp, opt, grad_scaler))
        eager_times.append(eager_time)
        print(f"eager train time {i}: {eager_time}")
    print("~" * 10)

    dynamo.reset()
    model = init_model(device_id=device_id)
    opt = optim.SGD(model.parameters(), lr=0.001)

    grad_scaler = torch.cuda.amp.GradScaler(init_scale=2.0)

    train_opt = torch.compile(train,
                              options={
                                  "triton.cudagraphs": True,
                              },
                              backend="inductor")

    compile_times = []
    for i in range(N_ITERS):
        inp = generate_data(16, device_id=device_id)
        _, compile_time = timed(lambda: train_opt(model, inp, opt, grad_scaler))
        compile_times.append(compile_time)
        print(f"compile train time {i}: {compile_time}")
    print("~" * 10)
    eager_med = np.median(eager_times)
    compile_med = np.median(compile_times)
    speedup = eager_med / compile_med
    print(
        f"(train) eager median: {eager_med}, compile median: {compile_med}, speedup: {speedup}x"
    )
    print("~" * 10)


def demo_eval():
    N_ITERS = 30

    device_id = 0

    dynamo.reset()
    model = init_model(device_id=device_id)

    eager_times = []
    for i in range(N_ITERS):
        inp = generate_data(16, device_id=device_id)
        with torch.no_grad():
            _, eager_time = timed(lambda: eval(model, inp))
        eager_times.append(eager_time)
        print(f"eager eval time {i}: {eager_time}")
    print("~" * 10)

    dynamo.reset()
    model = init_model(device_id=device_id)

    eval_opt = torch.compile(eval,
                             options={
                                 "triton.cudagraphs": True,
                             },
                             backend="inductor")

    compile_times = []
    for i in range(N_ITERS):
        inp = generate_data(16, device_id=device_id)
        with torch.no_grad():
            _, compile_time = timed(lambda: eval_opt(model, inp))
        compile_times.append(compile_time)
        print(f"compile eval time {i}: {compile_time}")
    print("~" * 10)
    eager_med = np.median(eager_times)
    compile_med = np.median(compile_times)
    speedup = eager_med / compile_med
    print(
        f"(eval) eager median: {eager_med}, compile median: {compile_med}, speedup: {speedup}x"
    )
    print("~" * 10)


if __name__ == "__main__":
    demo_train()
    # demo_eval()

```

### Versions

```
PyTorch version: 2.1.0+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.22.1
Libc version: glibc-2.31

Python version: 3.8.12 (default, Oct 12 2021, 13:49:34)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.4.52-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.6.124
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Tesla V100-SXM2-32GB
Nvidia driver version: 510.47.03
cuDNN version: Probably one of the following:
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn.so.8.4.0
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.4.0
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.4.0
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.4.0
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.4.0
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.4.0
/usr/local/cuda-11.6/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.4.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
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 48 bits virtual
CPU(s):                          96
On-line CPU(s) list:             0-95
Thread(s) per core:              2
Core(s) per socket:              24
Socket(s):                       2
NUMA node(s):                    2
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           85
Model name:                      Intel(R) Xeon(R) Gold 6240R CPU @ 2.40GHz
Stepping:                        7
CPU MHz:                         3200.000
CPU max MHz:                     4000.0000
CPU min MHz:                     1000.0000
BogoMIPS:                        4800.00
Virtualization:                  VT-x
L1d cache:                       1.5 MiB
L1i cache:                       1.5 MiB
L2 cache:                        48 MiB
L3 cache:                        71.5 MiB
NUMA node0 CPU(s):               0-23,48-71
NUMA node1 CPU(s):               24-47,72-95
Vulnerability Itlb multihit:     KVM: Mitigation: Split huge pages
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:   Mitigation; TSX disabled
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 pni pclmulqdq dtes64 ds_cpl vmx 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 intel_ppin ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm mpx rdt_a avx512f avx512dq 
rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req pku ospke avx512_vnni md_clear flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] numpy==1.22.4
[pip3] torch==2.1.0+cu118
[pip3] torchaudio==2.1.0+cu118
[pip3] torchdata==0.7.0
[pip3] torchtext==0.16.0+cpu
[pip3] torchvision==0.16.0+cu118
[pip3] triton==2.1.0
[conda] numpy                     1.22.4                   pypi_0    pypi
[conda] torch                     2.1.0+cu118              pypi_0    pypi
[conda] torchaudio                2.1.0+cu118              pypi_0    pypi
[conda] torchdata                 0.7.0                    pypi_0    pypi
[conda] torchtext                 0.16.0+cpu               pypi_0    pypi
[conda] torchvision               0.16.0+cu118             pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi
```

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305