Repository: pytorch
Issue: [AOTInductor] AOT compiled models do not consistently check input shapes
Labels: oncall: pt2,oncall: export

### 🐛 Describe the bug

AOT compiled models do not consistently provide the same runtime input Tensor checks and safety as do either plain PyTorch models or PyTorch exported graph models (via `torch.export`). Oddly, whether these checks happen with AOT models seems  to depend on the shape of the data used to export the PyTorch model (???).

In the repro below, I export a toy model as a graph and then as an AOT compiled model and run inference on both. In the first case, I seed the "happy path" model with 2D input and both the torch.export'ed and AOT compiled models expectedly raise `RuntimeException`s and complain about input Tensor shape mismatch. However, in the second case, the "sad path", the model is exported with 1D example data and the AOT compiled model **does not** error on invalid input data. Inference completes successfully which is **not** the expected behavior -- I'd expected the AOT compiled model to raise a RuntimeException as it did in the "happy path".


**repro**
```python
import torch

class Model(torch.nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self._layer = torch.nn.Linear(2, 1)

    def forward(self, X: torch.Tensor) -> torch.Tensor:
        return self._layer(X)

model = Model()


# HAPPY PATH: 2D example data
print("-" * 50 + "\nHAPPY PATH")
with torch.inference_mode():
    device = "cuda:0"
    # device = "cpu"

    X = torch.rand((1, 2), device=device)
    model = model.to(device)

    # AOTI compile on CPU
    example_inputs = X
    exported_program = torch.export.export(model, (example_inputs,),)
    exported_model = exported_program.module().to(device)
    so_path = torch._inductor.aot_compile(exported_model, (example_inputs,),)
    aot_model = torch._export.aot_load(so_path, device)

    # Inference: OK
    print(">> OK: Exported Graph Model: ", exported_model(X))
    print(">> OK: AOT Compiled Model: ", aot_model(X))

    # Inference: FAILS
    # NOTE: shape of X_t is invalid
    X_t = torch.rand((3, 1), device=device)
    try:
        exported_model(X_t)
    except RuntimeError as exc:
        print(">> FAIL: Exported Graph Model: ", exc)

    try:
        aot_model(X_t)
    except RuntimeError as exc:
        print(">> FAIL: AOT Compiled Model: ", exc)
print("-" * 50)


print("-" * 50 + "\nSAD PATH")
# SAD PATH: 1D example data
with torch.inference_mode():
    device = "cuda:0"

    # make example data 1D
    X = torch.rand((2,), device=device)
    model = model.to(device)

    # AOTI compile on CPU
    example_inputs = X
    exported_program = torch.export.export(model, (example_inputs,),)
    exported_model = exported_program.module().to(device)
    so_path = torch._inductor.aot_compile(exported_model, (example_inputs,),)
    aot_model = torch._export.aot_load(so_path, device)

    # Inference: OK
    print(">> OK: Exported Graph Model: ", exported_model(X))
    print(">> OK: AOT Compiled Model: ", aot_model(X))

    # Inference: FAILS
    # NOTE: shape of X_t is invalid
    X_t = torch.rand((3, 1), device=device)
    try:
        exported_model(X_t)
    except RuntimeError as exc:
        print(">> FAIL: Exported Graph Model: ", exc)

    print(">> THIS SHOULD FAIL BUT IS OK: AOT Compiled Model: ", aot_model(X_t))
print("-" * 50)
```

**output**
```bash
--------------------------------------------------
HAPPY PATH
>> OK: Exported Graph Model:  (tensor([[-0.0415]], device='cuda:0'),)

>> OK: AOT Compiled Model:  tensor([[-0.0415]], device='cuda:0')

>> FAIL: Exported Graph Model:  Expected input at *args[0].shape[0] to be equal to 1, but got 3

Error: mat1 and mat2 shapes cannot be multiplied (3x1 and 2x1)
Exception raised from meta at ../aten/src/ATen/native/LinearAlgebra.cpp:189 (most recent call first):
...
>> FAIL: AOT Compiled Model:  run_func_( container_handle_, input_handles.data(), input_handles.size(), output_handles.data(), output_handles.size(), cuda_stream_handle, proxy_executor_handle_) API call failed at ../torch/csrc/inductor/aoti_runner/model_container_runner.cpp, line 75
--------------------------------------------------

--------------------------------------------------
SAD PATH
>> OK: Exported Graph Model:  (tensor([-0.0204], device='cuda:0'),)

>> OK: AOT Compiled Model:  tensor([-0.0204], device='cuda:0')

>> FAIL: Exported Graph Model:  Unexpected number of dimensions in input at *args[0].shape (expected torch.Size([2]), got torch.Size([3, 1]))

>> THIS SHOULD FAIL BUT IS OK: AOT Compiled Model:  tensor([-0.0676], device='cuda:0')
--------------------------------------------------
```

This was observed on CUDA + `torch==2.4.1,2.5.1,nightly`.

### Versions

```bash
PyTorch version: 2.4.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04 LTS (x86_64)
GCC version: (Ubuntu 13.2.0-23ubuntu4) 13.2.0
Clang version: Could not collect
CMake version: version 3.28.3
Libc version: glibc-2.39

Python version: 3.12.3 (main, Sep 11 2024, 14:17:37) [GCC 13.2.0] (64-bit runtime)
Python platform: Linux-6.8.0-1018-aws-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 12.0.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA L40S
Nvidia driver version: 560.35.03
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
Address sizes:                        48 bits physical, 48 bits virtual
Byte Order:                           Little Endian
CPU(s):                               4
On-line CPU(s) list:                  0-3
Vendor ID:                            AuthenticAMD
Model name:                           AMD EPYC 7R13 Processor
CPU family:                           25
Model:                                1
Thread(s) per core:                   2
Core(s) per socket:                   2
Socket(s):                            1
Stepping:                             1
BogoMIPS:                             5300.00
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 constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr rdpru wbnoinvd arat npt nrip_save vaes vpclmulqdq rdpid
Hypervisor vendor:                    KVM
Virtualization type:                  full
L1d cache:                            64 KiB (2 instances)
L1i cache:                            64 KiB (2 instances)
L2 cache:                             1 MiB (2 instances)
L3 cache:                             8 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-3
Vulnerability Gather data sampling:   Not affected
Vulnerability Itlb multihit:          Not affected
Vulnerability L1tf:                   Not affected
Vulnerability Mds:                    Not affected
Vulnerability Meltdown:               Not affected
Vulnerability Mmio stale data:        Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Not affected
Vulnerability Spec rstack overflow:   Vulnerable: Safe RET, no microcode
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; Retpolines; IBPB conditional; IBRS_FW; STIBP always-on; RSB filling; PBRSB-eIBRS Not affected; BHI Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==2.1.2
[pip3] nvidia-cublas-cu12==12.1.3.1
[pip3] nvidia-cuda-cupti-cu12==12.1.105
[pip3] nvidia-cuda-nvrtc-cu12==12.1.105
[pip3] nvidia-cuda-runtime-cu12==12.1.105
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.0.2.54
[pip3] nvidia-curand-cu12==10.3.2.106
[pip3] nvidia-cusolver-cu12==11.4.5.107
[pip3] nvidia-cusparse-cu12==12.1.0.106
[pip3] nvidia-nccl-cu12==2.20.5
[pip3] nvidia-nvjitlink-cu12==12.6.68
[pip3] nvidia-nvtx-cu12==12.1.105
[pip3] onnxruntime-gpu==1.20.0
[pip3] torch==2.4.1
[pip3] triton==3.0.0
[conda] Could not collect
```

cc @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4