PyTorch Runs On the GPU of Apple M1 Macs Now! - Announcement With Code Samples

🧐 A backend is a highly optimized low-level library that enables take full use of the specific instructions that the GPU has. This makes computations way faster and in parallel.

😎 Early Benchmark Tensorflow vs PyTorch

Results of PyTorch on Apple Hardware (and some Nvidia)

A vision ResNet50
A Huggingface Bert model

import torch
torch.tensor([1,2,3], device="mps") #that's it!

Relative to Nvidia

😎 The Nvidia GPU supports Mixed Precision Training so you get even more out of the hardware!

For NLP let's try BERT training

The Max GPU is considerably faster for this Attention-based model, probably the extra ram helps here more than the extra GPU cores.

Run this benchmark and contribute to this table 🚀

Installation

# The version of Anaconda may be different depending on when you are installing`
$ curl -L -O https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-MacOSX-arm64.sh
$ sh Miniforge3-MacOSX-arm64.sh
# and follow the prompts. The defaults are generally good.`

$ conda create --name="pt" "python<3.11"

# activate the environment
$ conda activate pt

# install Pytorch
$ conda install pytorch torchvision -c pytorch

# install dependencies of this training script 😎
$ pip install wandb tqdm transformers datasets

Verifying the installation

import torch

torch.__version__
>>> '1.13.1'

torch.tensor([1,2,3], device="mps")

Training a Model

You can verify that the GPU is being used by going to Activity Monitor and check the GPU history window, you should get something like this while training:

As soon as you launch training, the bars fill up!

How to run this benchmark on your machine?

python train_pets.py --device="mps" --gpu_name="M1Pro GPU 16 Cores"

Pass the --gpu_name flag to group the runs, I am not able to detect this automatically on Apple.
To run on cpu pass --device="cpu" or for CUDA --device="cuda" (you need a linux PC with an Nvidia GPU)
You can also pass other params, and play with different batch_size and model_name.

If you need any help, please contact me or reply to this:

See the training script

Conclusions