[MPS] MPSNDArray error: product of dimension sizes > 2**32

[Tim-Roy]

🐛 Describe the bug

This is related to #84039 and #86279. The error is slightly different that the first (232 vs 231) and the second is specific to matmul. The machine has 64 GB of unified RAM and is not a physical limitation.

I discovered this bug when attempting to perform LoRa on Llama 8B. I can replicate the error by using a similar script as shown with the first bug report.

from torch import einsum, ones
import argparse

parser = argparse.ArgumentParser(description='mpsndarray test')
parser.add_argument('--n_samples', type=int, default=2)
args = parser.parse_args()
n_samples = args.n_samples

einsum('b i d, b j d -> b i j', ones(16 * n_samples, 4096, 40, device='mps'), ones(16 * n_samples, 4096, 40, device='mps')).shape

print(n_samples, 'passed')

I receive the error when running python repro.py --n_samples 17 or higher values.

Versions

PyTorch version: 2.3.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 14.5 (arm64)
GCC version: Could not collect
Clang version: 15.0.0 (clang-1500.3.9.4)
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.14 (main, Aug 19 2024, 18:36:18) [Clang 15.0.0 (clang-1500.3.9.4)] (64-bit runtime)
Python platform: macOS-14.5-arm64-arm-64bit
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:
Apple M2 Ultra

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] onnx==1.16.2
[pip3] torch==2.3.1
[pip3] torchaudio==2.3.1
[pip3] torchvision==0.18.1
[conda] Could not collect

cc @kulinseth @albanD @malfet @DenisVieriu97 @jhavukainen

[jhavukainen]

Hi @Tim-Roy! I think this issue got fixed with a tiling implementation for batch matmul that got just merged (#133430). I am able to reproduce the error with the version of torch you listed but when building locally with the above PR I'm seeing the test pass with values n_samples > 17.

Would you be able to try out either locally building on top of the current main branch or wait for a day and try with the nightly torch build to verify if this also solves your issue?

[Tim-Roy]

Thanks @jhavukainen - I'll wait for the nightly build to test and update you shortly.

[hvaara]

Note that @Tim-Roy is using macOS < 15, so the expected behavior is

from torch import einsum, ones
n_samples = 17
einsum('b i d, b j d -> b i j', ones(16 * n_samples, 4096, 40, device='mps'), ones(16 * n_samples, 4096, 40, device='mps')).shape
# Output:
# RuntimeError: Tiling of batch matmul for larger than 2**32 entries only available from MacOS15 onwards

[jhavukainen]

Thanks @hvaara that's a good comment here, I missed that detail. Indeed @Tim-Roy there is the unfortunate aspect to the fix that the necessary API in the MPSNDArray that allows me to tile the computation is only available for MacOS15 and later.

[Tim-Roy]

Thanks @hvaara that's a good comment here, I missed that detail. Indeed @Tim-Roy there is the unfortunate aspect to the fix that the necessary API in the MPSNDArray that allows me to tile the computation is only available for MacOS15 and later.

Oh I do see that now in the merged code. Well thanks for the update and good eye @hvaara.

[hvaara]

macOS 15.0 has been released, and with it I get:

from torch import einsum, ones
n_samples = 17
einsum('b i d, b j d -> b i j', ones(16 * n_samples, 4096, 40, device='mps'), ones(16 * n_samples, 4096, 40, device='mps')).shape
# torch.Size([272, 4096, 4096])

@Tim-Roy have you had a chance to update to latest macOS yet? If so, do you still see an error?

[Tim-Roy]

@hvaara Sorry I have not yet responded. This is for a work device and we have not yet upgraded. Should be shortly - I will respond soon.

[jhavukainen]

This issue has been addressed initially in commit 92f282ca520377f0 mentioned above and further expanding the range in commit afa313e669. Note that it will require update to MacOS15, otherwise user only gets the error and recommendation to update. This is due to tiling the operation requiring an API that was only released on MacOS 15.0.

I'll close the issue as addressed. Please reopen if you feel this is inaccurate.