MPS (Apple Silicon) Support

[dsanmart]

Will there be MPS support for the Gemma models? It would enable access to a larger community.

[lamroger]

Took a look, a few things.

Linux / mps support looks to be in progress still https://github.com/pytorch/pytorch/issues/81224 so running in a container isn't ready yet.

MPS has some limitations around complex tensors atm. Since gemma uses RoPE, it uses complex tensors and errors out if you run it locally.

https://github.com/pytorch/pytorch/pull/116764/files#diff-fe061f10677283971d77576718d3a04a00b2225d72c043fd59222a882b92c64bR654

https://github.com/google/gemma_pytorch/blob/01062c9ef4cf89ac0c985b25a734164ede017d0b/gemma/model.py#L426

Running locally with python scripts/run.py --ckpt gemma-2b-it.ckpt --variant 2b --device mps

[bghira]

pytorch 2.3 has bf16 and complex tensor support, and dockerised containers now work @lamroger