Repository: pytorch
Issue: mmap fails on 64k page aarch64 systems for AOTI model loading
Labels: module: crash,module: arm,oncall: pt2,export-triaged,oncall: export,module: aotinductor

### 🐛 Describe the bug

The AOTI loader mmap's a file with an offset `weights_offset`. An asssertion ensures that `weights_offset` is a multiple of 16k. The offset in à `mmap` syscall needs to be a multiple of the page size though, causing this mmap to fail on kernels with 64k pages but not on 4k pages.

https://github.com/pytorch/pytorch/blob/629840e038ee623911bedc8fef1ab84acce5ba39/torch/csrc/inductor/aoti_runtime/model.h#L601

**Failing syscall with 64k page kernel:**
```
newfstatat(AT_FDCWD, "/pytorch-llama/torchchat/exportedModels/llama3.1.json", 0xffffec3d0838, 0) = -1 ENOENT (No such file or directory)
openat(AT_FDCWD, "/pytorch-llama/torchchat/exportedModels/llama3.1.so", O_RDONLY) = 3
mmap(NULL, 6144375816, PROT_READ|PROT_WRITE, MAP_PRIVATE, 3, 0xd8000) = -1 EINVAL (Invalid argument) # 0xd8000 is no multiple of 65536
```

**Same syscalls with 4k page kernel:**
```
newfstatat(AT_FDCWD, "/pytorch-llama/torchchat/exportedModels/llama3.1.json", 0xfffff418a778, 0) = -1 ENOENT (No such file or directory)
openat(AT_FDCWD, "/pytorch-llama/torchchat/exportedModels/llama3.1.so", O_RDONLY) = 3
mmap(NULL, 6144375816, PROT_READ|PROT_WRITE, MAP_PRIVATE, 3, 0xd8000) = 0xfff9d10a7000 # 0xd8000 is a multiple of 4096
```

**How to reproduce:**
Follow these steps https://learn.arm.com/learning-paths/servers-and-cloud-computing/pytorch-llama/pytorch-llama/ on an aarch64 system with 64k page size kernel. The error occurs during model loading when running `torchchat.py` (last step).

**Error message:**
```
PyTorch version 2.5.0.dev20240828+cpu available.
Warning: checkpoint path ignored because an exported DSO or PTE path specified
Warning: checkpoint path ignored because an exported DSO or PTE path specified
Using device=cpu
Loading model...
Time to load model: 0.05 seconds
Error: mmap() failed
Traceback (most recent call last):
  File "/pytorch-llama/torchchat/build/builder.py", line 480, in _initialize_model
    model.forward = torch._export.aot_load(
  File "/usr/local/lib/python3.10/dist-packages/torch/_export/__init__.py", line 300, in aot_load
    runner = torch._C._aoti.AOTIModelContainerRunnerCpu(so_path, 1)  # type: ignore[call-arg]
RuntimeError: create_func_( &container_handle_, num_models, device_str.c_str(), cubin_dir.empty() ? nullptr : cubin_dir.c_str()) API call failed at /pytorch/torch/csrc/inductor/aoti_runner/model_container_runner.cpp, line 70

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/pytorch-llama/torchchat/torchchat.py", line 88, in <module>
    generate_main(args)
  File "/pytorch-llama/torchchat/generate.py", line 901, in main
    gen = Generator(
  File "/pytorch-llama/torchchat/generate.py", line 253, in __init__
    self.model = _initialize_model(self.builder_args, self.quantize, self.tokenizer)
  File "/pytorch-llama/torchchat/build/builder.py", line 484, in _initialize_model
    raise RuntimeError(f"Failed to load AOTI compiled {builder_args.dso_path}")
RuntimeError: Failed to load AOTI compiled exportedModels/llama3.1.so
```

**GDB backtrace:**
```
Catchpoint 6 (call to syscall mmap), __GI___mmap64 (offset=884736, fd=3, flags=2, prot=3, len=6144375816, addr=<optimized out>) at ../sysdeps/unix/sysv/linux/mmap64.c:58
58      in ../sysdeps/unix/sysv/linux/mmap64.c
(gdb) bt
#0  __GI___mmap64 (offset=884736, fd=3, flags=2, prot=3, len=6144375816, addr=<optimized out>)
    at ../sysdeps/unix/sysv/linux/mmap64.c:58
#1  __GI___mmap64 (addr=<optimized out>, len=6144375816, prot=3, flags=2, fd=3, offset=884736)
    at ../sysdeps/unix/sysv/linux/mmap64.c:46
#2  0x0000e286e96d2fe8 in torch::aot_inductor::AOTInductorModelBase<torch::aot_inductor::AOTInductorModel>::load_constants() () from /pytorch-llama/torchchat/exportedModels/llama3.1.so
#3  0x0000e286e96ee31c in torch::aot_inductor::AOTInductorModelContainer::AOTInductorModelContainer(unsigned long, std::string const&, std::optional<std::string> const&) ()
   from /pytorch-llama/torchchat/exportedModels/llama3.1.so
#4  0x0000e286e96cd0e8 in AOTInductorModelContainerCreateWithDevice ()
   from /pytorch-llama/torchchat/exportedModels/llama3.1.so
#5  0x0000e287174d8888 in torch::inductor::AOTIModelContainerRunner::AOTIModelContainerRunner(std::string const&, unsigned long, std::string const&, std::string const&) ()
   from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#6  0x0000e287174d97bc in torch::inductor::AOTIModelContainerRunnerCpu::AOTIModelContainerRunnerCpu(std::string const&, unsigned long) () from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_cpu.so
#7  0x0000e2871c841618 in pybind11::cpp_function::initialize<pybind11::detail::initimpl::constructor<std::string const&, int>::execute<pybind11::class_<torch::inductor::AOTIModelContainerRunnerCpu>, , 0>(pybind11::class_<torch::inductor::AOTIModelContainerRunnerCpu>&)::{lambda(pybind11::detail::value_and_holder&, std::string const&, int)#1}, void, pybind11::detail::value_and_holder&, std::string const&, int, pybind11::name, pybind11::is_method, pybind11::sibling, pybind11::detail::is_new_style_constructor>(pybind11::class_<torch::inductor::AOTIModelContainerRunnerCpu>&&, void (*)(pybind11::detail::value_and_holder&, std::string const&, int), pybind11::name const&, pybind11::is_method const&, pybind11::sibling const&, pybind11::detail::is_new_style_constructor const&)::{lambda(pybind11::detail::function_call&)#3}::_FUN(pybind11::detail::function_call&) ()
   from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_python.so
#8  0x0000e2871c2c0814 in pybind11::cpp_function::dispatcher(_object*, _object*, _object*) ()
   from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_python.so
#9  0x0000b02a971d4b54 in ?? ()
#10 0x0000b02a971cb100 in _PyObject_MakeTpCall ()
#11 0x0000b02a971e58c4 in ?? ()
#12 0x0000b02a971e1a28 in ?? ()
#13 0x0000b02a971cb568 in ?? ()
#14 0x0000e2871c2be5dc in pybind11_meta_call ()
   from /usr/local/lib/python3.10/dist-packages/torch/lib/libtorch_python.so
#15 0x0000b02a971cb100 in _PyObject_MakeTpCall ()
#16 0x0000b02a971c2334 in _PyEval_EvalFrameDefault ()
#17 0x0000b02a971d57e8 in _PyFunction_Vectorcall ()
#18 0x0000b02a971c1bcc in _PyEval_EvalFrameDefault ()
#19 0x0000b02a971d57e8 in _PyFunction_Vectorcall ()
#20 0x0000b02a971bd764 in _PyEval_EvalFrameDefault ()
#21 0x0000b02a971ca164 in _PyObject_FastCallDictTstate ()
#22 0x0000b02a971e162c in ?? ()
#23 0x0000b02a971cb078 in _PyObject_MakeTpCall ()
#24 0x0000b02a971c1f68 in _PyEval_EvalFrameDefault ()
#25 0x0000b02a971d57e8 in _PyFunction_Vectorcall ()
#26 0x0000b02a971bd764 in _PyEval_EvalFrameDefault ()
#27 0x0000b02a972be070 in ?? ()
#28 0x0000b02a972bdef4 in PyEval_EvalCode ()
#29 0x0000b02a972f151c in ?? ()
#30 0x0000b02a972e9c38 in ?? ()
#31 0x0000b02a972f11cc in ?? ()
```

### Versions

PyTorch version: 2.5.0.dev20240820+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.5 LTS (aarch64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.31.4
Libc version: glibc-2.35

Python version: 3.10.12 (main, Jan 17 2025, 14:35:34) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.8.0-<....>
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                         aarch64
CPU op-mode(s):                       64-bit
Byte Order:                           Little Endian
CPU(s):                               72
On-line CPU(s) list:                  0-71
Vendor ID:                            ARM
Model name:                           Neoverse-V2
Model:                                0
Thread(s) per core:                   1
Core(s) per cluster:                  72
Socket(s):                            -
Cluster(s):                           1
Stepping:                             r0p0
Frequency boost:                      disabled
CPU max MHz:                          3447.0000
CPU min MHz:                          81.0000
BogoMIPS:                             2000.00
Flags:                                fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm ssbs sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh bti
L1d cache:                            4.5 MiB (72 instances)
L1i cache:                            4.5 MiB (72 instances)
L2 cache:                             72 MiB (72 instances)
L3 cache:                             114 MiB (1 instance)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-71
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:   Not affected
Vulnerability Spec store bypass:      Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:             Mitigation; __user pointer sanitization
Vulnerability Spectre v2:             Not affected
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.5.0.dev20240828+cpu
[pip3] torchao==0.4.0+git174e630a
[conda] Could not collect

cc @malfet @snadampal @milpuz01 @chauhang @penguinwu @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @desertfire @chenyang78