M1 GPU Acceleration

GPU acceleration for PyTorch is now available on Apple Silicon. I wanted to document how to use GPU acceleration across the frameworks I use.

Pytorch

As of 2022-05-20.

Install PyTorch 1.12. Only works with the nightly build.

1
import torch
2
import torchvision.models as models
3
from torchsummary import summary
4

5
print(torch.__version__)
6
mps_device = torch.device("mps")
7

8
print(mps_device)
9

10
# Create a Tensor directly on the mps device
11
x = torch.ones((1, 3, 224, 224), device=mps_device)
12
print(x.shape)
13

14
# Move your model to mps just like any other device
15
model = models.resnet18()
16
summary(model, (3, 244, 244))
17
model.to(mps_device)
18

19
# Now every call runs on the GPU
20
pred = model(x)
21

22
print(pred, pred.shape)

HuggingFace

It can’t be installed via pip or conda, so you need to build from source. I used the rust tokenizer.

# install rust on arm terminal
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh


# intsall tokenizer
git clone https://github.com/huggingface/tokenizers
cd tokenizers/bindings/python
pip install setuptools_rust
python setup.py install

# install transformers
pip install git+https://github.com/huggingface/transformers

# install datasets
pip install git+https://github.com/huggingface/datasets

1
from transformers import AutoTokenizer, BertModel
2

3
device = "mps"
4
sentence  = 'Hello World!'
5
tokenizer = AutoTokenizer.from_pretrained('bert-large-uncased', use_fast=True)
6
model     = BertModel.from_pretrained('bert-large-uncased')
7

8
inputs    = tokenizer(sentence, return_tensors="pt").to(device)
9
model     = model.to(device)
10
outputs   = model(**inputs)
11
print(outputs)

M1 GPU Acceleration

Pytorch

HuggingFace

Ref