[Inductor] Run compiled model failed  on 2023_08_17 nightly
Labels: triaged, module: ddp, oncall: pt2, module: inductor, module: dynamo

### 🐛 Describe the bug

I attempted to use the latest Torch 2.1 compiler, but encountered errors. I believe this might be related to DistributedDataParallel (DDP) and the usage of torch.no_grad, but I'm unable to resolve it. Here are the detailed specifics.

By the way, the following code works well in Torch 2.0.

### Error logs

File "/opt/miniconda3/envs/python38/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 3905, in forward
    return compiled_fn(full_args)
  File "/opt/miniconda3/envs/python38/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 1482, in g
    return f(*args)
  File "/opt/miniconda3/envs/python38/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 2533, in runtime_wrapper
    all_outs = call_func_with_args(
  File "/opt/miniconda3/envs/python38/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 1506, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/opt/miniconda3/envs/python38/lib/python3.8/site-packages/torch/_functorch/aot_autograd.py", line 1594, in rng_functionalization_wrapper
    return compiled_fw(args)
  File "/opt/miniconda3/envs/python38/lib/python3.8/site-packages/torch/_inductor/codecache.py", line 373, in __call__
    return self.get_current_callable()(inputs)
  File "/opt/miniconda3/envs/python38/lib/python3.8/site-packages/torch/_inductor/compile_fx.py", line 631, in run
    return model(new_inputs)
  File "/opt/miniconda3/envs/python38/lib/python3.8/site-packages/torch/_inductor/codecache.py", line 400, in _run_from_cache
    return compiled_graph.compiled_artifact(inputs)
  File "/tmp/torchinductor_root/li/cliuv2iyw3ztzp2iy5xgusyzfr7qvm7e2nlflfn2pomd4xqsts65.py", line 461, in call
    assert_size_stride(arg11_1, (2, 64, 128, 128), (1048576, 16384, 128, 1))
AssertionError: expected size 64==64, stride 1==16384 at dim=1

### Minified repro

```
import os
import torch
import torch.distributed as dist
import torch.nn as nn

from torch.nn.parallel import DistributedDataParallel as DDP


class TestNet(nn.Module):
    def __init__(self):
        super(TestNet, self).__init__()

        channels = (3, 32, 64, 128)
        self.num_levels = len(channels)
        for i in range(1, self.num_levels):
            dla_layer = Tree(
                channels[i - 1],
                channels[i],
            )
            layer_name = f"level{i}"
            self.add_module(layer_name, dla_layer)

    def forward(self, x):
        outs = []
        for i in range(1, self.num_levels):
            x = getattr(self, "level{}".format(i))(x)
            outs.append(x)
                    
        return tuple(outs)

class Tree(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
    ):
        super(Tree, self).__init__()
        self.conv = nn.Conv2d(2 * out_channels, out_channels, 1, 1, bias=False)
        self.tree1 = nn.Conv2d(in_channels, out_channels, 3, stride=1, padding=1)
        self.tree2 = nn.Conv2d(out_channels, out_channels, 3, stride=1, padding=1)
        
        self.project = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, 1, stride=1, bias=False),
            nn.BatchNorm2d(out_channels),
        )

    def forward(self, x):
        x1 = self.tree1(x)
        identity = self.project(x) 
        x2 = self.tree2(identity)
        x = self.conv(torch.concat([x2, x1], 1))
        return x


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # initialize the process group
    dist.init_process_group(rank=rank, world_size=world_size)


def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    model = TestNet().cuda()
    ddp_model = DDP(model, device_ids=[rank]) # important
    model = torch.compile(ddp_model.module)
    setattr(ddp_model, 'module', model )

    with torch.no_grad():
        outputs = ddp_model(torch.randn(2, 3, 128, 128).cuda())

if __name__ == '__main__':
    demo_basic(0, 1)
```

### Versions

```
PyTorch version: 2.1.0.dev20230816+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 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.27.2
Libc version: glibc-2.31

Python version: 3.8.17 (default, Jul  5 2023, 21:04:15)  [GCC 11.2.0] (64-bit runtime)
Python platform: Linux-5.4.143-2-velinux1-amd64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Tesla V100-SXM2-32GB
Nvidia driver version: 470.129.06
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):                          8
On-line CPU(s) list:             0-7
Thread(s) per core:              2
Core(s) per socket:              4
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           85
Model name:                      Intel(R) Xeon(R) Platinum 8260 CPU @ 2.40GHz
Stepping:                        7
CPU MHz:                         2400.000
BogoMIPS:                        4800.00
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       128 KiB
L1i cache:                       128 KiB
L2 cache:                        4 MiB
L3 cache:                        35.8 MiB
NUMA node0 CPU(s):               0-7
Vulnerability Itlb multihit:     Not affected
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:   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 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology cpuid 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 abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat umip pku ospke avx512_vnni md_clear arch_capabilities

Versions of relevant libraries:
[pip3] flake8==6.1.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.4
[pip3] pytorch-triton==2.1.0+e6216047b8
[pip3] torch==2.1.0.dev20230816+cu118
[pip3] torchaudio==2.1.0.dev20230816+cu118
[pip3] torchvision==0.16.0.dev20230816+cu118
[pip3] triton==2.0.0
[conda] No relevant packages
```

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov