Repository: pytorch
Issue: The first epoch is very slow when using torch.compile
Labels: triaged,oncall: pt2,module: dynamic shapes

### 🐛 Describe the bug

We tried using torch.compile on our training code and the first epoch was very slow.

### Error logs

```bash
# uncompiled mode: out = model(x)
$ python torch2_compile.py
```
2.0.0+cu117
/mnt/nfs/env/py-venv/py39_cuda117_torch200/lib/python3.9/site-packages/torch/autograd/__init__.py:200: UserWarning: CUDA initialization: Unexpected error from cudaGetDeviceCount(). Did you run some cuda functions before calling NumCudaDevices() that might have already set an error? Error 803: system has unsupported display driver / cuda driver combination (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:109.)
  Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
**Epoch 0/10 time: 0.09859752655029297**
Epoch 1/10 time: 0.09461498260498047
Epoch 2/10 time: 0.09749722480773926
Epoch 3/10 time: 0.09946417808532715
Epoch 4/10 time: 0.09831523895263672
Epoch 5/10 time: 0.0978553295135498
Epoch 6/10 time: 0.09766864776611328
Epoch 7/10 time: 0.0983278751373291
Epoch 8/10 time: 0.09094047546386719
Epoch 9/10 time: 0.09506630897521973
Epoch avg time: 0.09683477878570557



```bash
# compiled mode: out = compiled_model(x)
$ python torch2_compile.py
```
2.0.0+cu117
/mnt/nfs/env/py-venv/py39_cuda117_torch200/lib/python3.9/site-packages/torch/cuda/__init__.py:107: UserWarning: CUDA initialization: Unexpected error from cudaGetDeviceCount(). Did you run some cuda functions before calling NumCudaDevices() that might have already set an error? Error 803: system has unsupported display driver / cuda driver combination (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:109.)
  return torch._C._cuda_getDeviceCount() > 0
No CUDA runtime is found, using CUDA_HOME='/mnt/nfs/env/cuda/cuda-11.7'
[2023-03-28 22:20:57,063] torch._inductor.utils: [WARNING] using triton random, expect difference from eager
**Epoch 0/10 time: 18.54069495201111**
Epoch 1/10 time: 0.10959434509277344
Epoch 2/10 time: 0.10520339012145996
Epoch 3/10 time: 0.10396909713745117
Epoch 4/10 time: 0.10267758369445801
Epoch 5/10 time: 0.10393261909484863
Epoch 6/10 time: 0.10507917404174805
Epoch 7/10 time: 0.10506296157836914
Epoch 8/10 time: 0.10470414161682129
Epoch 9/10 time: 0.10574936866760254
Epoch avg time: 1.948666763305664





### Minified repro

```bash
$ CUDA_VISIBLE_DEVICES=0
$ python torch2_compile.py
```

```bash
$ cat torch2_compile.py
#!/usr/bin/env python3

import time
import torch
import torch._dynamo as dynamo
import torchvision.models as models

model = models.alexnet()

optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
compiled_model = torch.compile(model)

x = torch.randn(16, 3, 224, 224)
optimizer.zero_grad()

epoches=10
count = []
for epoch in range(epoches):
    start = time.time()

    #out = model(x)
    out = compiled_model(x)
    out.sum().backward()
    optimizer.step()

    end = time.time()
    count.append(end - start)
    print(f"Epoch {epoch}/{epoches} time: {end - start}")

print(f"Epoch avg time: {sum(count)/len(count)}")

```

### Versions

$ python ~/env/bin/collect_env.py
Collecting environment information...
/mnt/nfs/env/py-venv/py39_cuda117_torch200/lib/python3.9/site-packages/torch/cuda/__init__.py:107: UserWarning: CUDA initialization: Unexpected error from cudaGetDeviceCount(). Did you run some cuda functions before calling NumCudaDevices() that might have already set an error? Error 803: system has unsupported display driver / cuda driver combination (Triggered internally at ../c10/cuda/CUDAFunctions.cpp:109.)
  return torch._C._cuda_getDeviceCount() > 0
PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 7.3.1611 (Core)  (x86_64)
GCC version: (GCC) 7.5.0
Clang version: Could not collect
CMake version: version 3.26.1
Libc version: glibc-2.17

Python version: 3.9.15 (main, Nov 12 2022, 16:52:40)  [GCC 5.5.0] (64-bit runtime)
Python platform: Linux-3.10.0-693.21.1.std7a.el7.0.x86_64-x86_64-with-glibc2.17
Is CUDA available: False
CUDA runtime version: 11.7.64
CUDA_MODULE_LOADING set to: N/A
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: 530.30.02
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
CPU(s):                128
On-line CPU(s) list:   0-127
Thread(s) per core:    2
Core(s) per socket:    32
Socket(s):             2
NUMA node(s):          2
Vendor ID:             GenuineIntel
CPU family:            6
Model:                 106
Model name:            Intel(R) Xeon(R) Platinum 8350C CPU @ 2.60GHz
Stepping:              6
CPU MHz:               3500.000
BogoMIPS:              5206.54
Virtualization:        VT-x
L1d cache:             48K
L1i cache:             32K
L2 cache:              1280K
L3 cache:              49152K
NUMA node0 CPU(s):     0-31,64-95
NUMA node1 CPU(s):     32-63,96-127

Versions of relevant libraries:
[pip3] audio2numpy==0.1.2
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.23.5
[pip3] torch==2.0.0
[pip3] torchaudio==2.0.1
[pip3] torchvision==0.15.1
[pip3] triton==2.0.0
[conda] Could not collect


cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh