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

[junukwon7]

🐛 Describe the bug

Full error message (no traceback):

AppleInternal/Library/BuildRoots/20d6c351-ee94-11ec-bcaf-7247572f23b4/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Types/MPSNDArray.mm:705: failed assertion '[MPSNDArray initWithDevice:descriptor:] Error: product of dimension sizes > 2**31 '

How to reproduce

1. Install stable-diffusion using instructions for macOS
2. run python scripts/txt2img.py --prompt "a horse" --plms --n_samples 1 --n_rows 1 --n_iter 1: runs well.
3. But if you add --W 1024 --H 1024 flag, which means width and height respectively, then it'll return the error.

• default width and height is 512, so no flag means 512x512

I'm finding a way to reproduce it without installing the whole procedure, so I'll update the procedure soon.

Edit: repro by @Birch-san

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')

It fails when n_samples is 2 or over 7. Which looks pretty weird.

About vram?

As you would all expect, the error seems to be something about VRAM. However, there remains question.

1. The error seems to be the size exceeding INT_MAX(2**31)
   The error doesn't occur at --W 512 --H 512 or lower resolution.
2. The error is a software issue
   Unlike errors like CUDA out of memory, this error isn't about the real memory limit.
   If the error was due to lack of VRAM, the code above (--W 1024 --H 1024) should run on M1 Max 64GB since --W 512 --H 512 runs well on my M1 8G macbook. Also, the limit 2**31 is a fixed number, which would not change from the current memory usage.

So, my expectation is that something is being computed in 32-bit, which shouldn't be.

This might not be torch's problem - maybe(surely) metal.

However, all helps will be accepted gracefully.

Thanks.

Versions

PyTorch version: 1.13.0.dev20220824
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.5.1 (arm64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2.5)
CMake version: version 3.24.1
Libc version: N/A

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:14) [Clang 12.0.1 ] (64-bit runtime)
Python platform: macOS-12.5.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
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

Versions of relevant libraries:
[pip3] numpy==1.23.2
[pip3] pytorch-lightning==1.7.2
[pip3] torch==1.13.0.dev20220824
[pip3] torch-fidelity==0.3.0
[pip3] torchaudio==0.13.0.dev20220824
[pip3] torchmetrics==0.9.3
[pip3] torchvision==0.14.0.dev20220824
[conda] numpy 1.23.2 py38h579d673_0 conda-forge
[conda] pytorch 1.13.0.dev20220824 py3.8_0 pytorch-nightly
[conda] pytorch-lightning 1.7.2 pypi_0 pypi
[conda] torch-fidelity 0.3.0 pypi_0 pypi
[conda] torchaudio 0.13.0.dev20220824 py38_cpu pytorch-nightly
[conda] torchmetrics 0.9.3 pypi_0 pypi
[conda] torchvision 0.14.0.dev20220824 py38_cpu pytorch-nightly

cc @kulinseth @albanD

[Karric]

I had the same issue, but with M1 mac + tensorflow (for mac), no pytorch but I am using numpy. Not sure if that's helpful.

[junukwon7]

I had the same issue, but with M1 mac + tensorflow (for mac), no pytorch but I am using numpy. Not sure if that's helpful.

@Karric Thanks for your reply. It seems like the error is something about metal then. Could you provide more information about the situation with m1 mac + tensorflow?

Thanks.

[i3oc9i]

I have the same issue on my Mac Studio Ultra with 128G :(

[pcuenca]

We are having the same problem in 🤗 diffusers, but I fail to see where the offending dimensions are being generated or used. The error arises when attempting to run the unet module of the model on input latents with shape (4, 4, 64, 64). This is what torchinfo.summary has to say about the output shapes of that model's layers:

===================================================================================================================
Layer (type:depth-idx)                                            Output Shape              Param #
===================================================================================================================
UNet2DConditionModel                                              [4, 4, 64, 64]            510,807,680
├─Conv2d: 1-4                                                     [4, 320, 64, 64]          (recursive)
├─Timesteps: 1-2                                                  [4, 320]                  --
├─TimestepEmbedding: 1-3                                          [4, 1280]                 --
│    └─Linear: 2-1                                                [4, 1280]                 410,880
│    └─SiLU: 2-2                                                  [4, 1280]                 --
│    └─Linear: 2-3                                                [4, 1280]                 1,639,680
├─Conv2d: 1-4                                                     [4, 320, 64, 64]          (recursive)
├─ModuleList: 1                                                   --                        --
│    └─CrossAttnDownBlock2D: 2-4                                  [4, 320, 32, 32]          --
│    └─CrossAttnDownBlock2D: 2-5                                  [4, 640, 16, 16]          --
│    └─CrossAttnDownBlock2D: 2-6                                  [4, 1280, 8, 8]           --
│    └─DownBlock2D: 2-7                                           [4, 1280, 8, 8]           --
├─UNetMidBlock2DCrossAttn: 1-5                                    [4, 1280, 8, 8]           --
│    └─ModuleList: 2                                              --                        --
│    │    └─ResnetBlock2D: 3-1                                    [4, 1280, 8, 8]           31,138,560
│    └─ModuleList: 2                                              --                        --
│    │    └─SpatialTransformer: 3-2                               [4, 1280, 8, 8]           34,760,960
│    └─ModuleList: 2                                              --                        --
│    │    └─ResnetBlock2D: 3-3                                    [4, 1280, 8, 8]           31,138,560
├─ModuleList: 1                                                   --                        --
│    └─UpBlock2D: 2-8                                             [4, 1280, 16, 16]         --
│    └─CrossAttnUpBlock2D: 2-9                                    [4, 1280, 32, 32]         --
│    └─CrossAttnUpBlock2D: 2-10                                   [4, 640, 64, 64]          --
│    └─CrossAttnUpBlock2D: 2-11                                   [4, 320, 64, 64]          --
├─GroupNorm: 1-6                                                  [4, 320, 64, 64]          640
├─SiLU: 1-7                                                       [4, 320, 64, 64]          --
├─Conv2d: 1-8                                                     [4, 4, 64, 64]            11,524
===================================================================================================================

It looks to me that those sizes shouldn't trigger any internal Metal limitation about 32-bit sizes being exceeded. Is there any hint as what might be causing this? Thank you!

[pcuenca]

In fact, as @FahimF reports, the crash does not occur for larger batches such as those that would result in latent shapes of (6, 4, 64, 64) or (8, 4, 64, 64). Is there any way to debug where the problem might be triggering?

[Birch-san]

I have a repro (from stable-diffusion's attention forward-pass).

https://github.com/CompVis/stable-diffusion/blob/69ae4b35e0a0f6ee1af8bb9a5d0016ccb27e36dc/ldm/modules/attention.py#L180

from torch import einsum, ones
# crashes with "product of dimension sizes > 2**31"
# this is equivalent to invoking stable-diffusion with --n_samples 2
einsum('b i d, b j d -> b i j', ones(32, 4096, 40, device='mps'), ones(32, 4096, 40, device='mps')).shape

# doesn't crash, even though it's bigger
# this is equivalent to invoking stable-diffusion with --n_samples 3
einsum('b i d, b j d -> b i j', ones(48, 4096, 40, device='mps'), ones(48, 4096, 40, device='mps')).shape

Perhaps related to this?

#80808 (comment)

Depending on the size/number of dimensions, different algorithm might get selected leading to small differences.

Perhaps at the larger size of 48: we hit tensor size which persuades einsum() to use a different algorithm, which doesn't crash?

I'm not sure whether this is necessarily the same line of stable-diffusion code on which @junukwon's is crashing (i.e. due to image dimensions). but if we're lucky it's the same mechanism.

[junukwon7]

@Birch-san

Thanks for your repro!
I've tried running a bash script for that, and found out that 16*i fails when i == 2 or i >= 8.

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')

made this as a python file and ran

for i in {1..80}; do python3 mpsndarray.py --n_samples ${i}; done

The result was:

1 pass
2 ERR
3 ~ 7 pass
8 ~ 35 ERR
36 ~ RuntimeError: Invalid buffer size: 36.00 GB

M1 Max MBP with 64GB RAM

[pcuenca]

Regarding the reproducibility code, @patil-suraj replaced einsum with matmul in the diffusers codebase and the problem still occurs in exactly the same way.

In fact, following up on @Birch-san's example, this behaves exactly like they describe:

# Crashes
t1 = torch.rand((32, 4096, 40))
t2 = torch.rand((32, 4096, 40))
torch.matmul(t1.to("mps"), t2.to("mps").transpose(1, 2)).shape

# Doesn't crash, even though it's bigger
t1 = torch.rand((48, 4096, 40))
t2 = torch.rand((48, 4096, 40))
torch.matmul(t1.to("mps"), t2.to("mps").transpose(1, 2)).shape

In addition, this is something different but very weird:

t1 = torch.rand((48, 4096, 40))
t2 = torch.rand((48, 4096, 40))

x_mps = einsum('b i d, b j d -> b i j', t1.to('mps'), t2.to('mps'))
x_cpu = einsum('b i d, b j d -> b i j', t1, t2)
print((x_mps.to("cpu") - x_cpu).abs().max())
# tensor(9.4567) !?

As you can see, the output seems to be wrong. However, if we do other operations first, the einsum is now correct:

t1 = torch.rand((48, 4096, 40))
t2 = torch.rand((48, 4096, 40))

x_mm_mps = torch.matmul(t1.to("mps"), t2.to("mps").transpose(1, 2))
x_mm_cpu = torch.matmul(t1, t2.transpose(1, 2))
print((x_mm_mps.to("cpu") - x_mm_cpu).abs().max())
# tensor(0.)

x_mps = einsum('b i d, b j d -> b i j', t1.to('mps'), t2.to('mps'))
x_cpu = einsum('b i d, b j d -> b i j', t1, t2)
print((x_mps.to("cpu") - x_cpu).abs().max())
# tensor(0.)

Is there any guidance or hint as to what could be going on here?

TL:DR:

• torch.matmul crashes for some sizes when using mps.
• torch.einsum produces inconsistent results "sometimes". Doing some previous operations with similar memory requirements seems to prevent the issue, maybe?