Apple Silicon M1 (MPS) compatibility

[aljungberg]

I think the Apple M1 is an interesting device for running this model. It's still unusual to find consumer GPUs with enough VRAM to run it. Even in the server class, a V100 won't do it, you'd need two. But the M1 with its unified memory architecture means it reasonably within reach for a regular end-user to get a high memory GPU device.

Unfortunately I couldn't get it to work yet. Posting anyhow just to share with the community.

The good news is that if you switch device from cuda:0 to mps:0, the code doesn't crash and seems to load the model reasonably quickly (testing on an M1 Max with 64 GB). Inference also doesn't crash.

The bad news is that the output is nonsense. Following the prompt from the readme it just repeats the word "in":

GPTNeoX20B is a 20B-parameter autoregressive Transformer model developed by EleutherAI. in in in in in in in in in in in in in in in in in in in in in in in in in in in in in in in in in in in in

For reference, by commenting out model = model.half().to_empty(device=device) it's possible to run the model on the CPU (specifying cpu instead of mps as the device). This is basically unusably slow but does prove that nothing is fundamentally broken with test:

... 36/36 [1:31:17<00:00, 152.16s/it]
GPTNeoX20B is a 20B-parameter autoregressive Transformer model developed by EleutherAI. It is a state-of-the-art language model that is trained on a large corpus of text data.

The model is trained on a corpus of over 1.

(This run took 1.5h, which isn't fast. I think is mostly because to run on the CPU we need FP32, doubling the memory requirements. So we'd expect to need 80 GB of RAM just for the model and I'm running this on a 64 GB device. So there's the performance overhead of memory compression, maybe even swapping.)

Possibly the reason it doesn't work on the GPU is that PyTorch MPS support is still in early stages. So there's a chance this will all just magically start to work with future updates. For this test I used the most recent nightly build, 1.13.0.dev20220803.