Accelerate PyTorch using `MPS` on Apple silicon

accelarate_pytorch_mps.md

Accelerate PyTorch using MPS on Apple silicon

1. Download Xcode from the App Store
2. Opan a new Terminal tab and run sudo xcode-select -switch /Applications/Xcode.app/Contents/Developer
3. Accept the Xcode licence by running: sudo xcodebuild -license
4. Try running /usr/bin/xcodebuild -version to ensure Xcode is correctly installed and configured
5. Create a new conda env and install pytorch-nightly:

conda create -n nanogpt python=3.10
conda activate nanogpt
conda install pytorch -c pytorch-nightly
conda install numpy

6. Start Python and run the following code:

import torch

device = "cpu"
if torch.cuda.is_available():
    device = "cuda"
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
    device = "mps"

print(f"using device: {device}")

using device: mps

Resources: https://developer.apple.com/metal/pytorch/, Let's reproduce GPT-2 (124M), ChatGPT4o