Setting Up PyTorch with MPS

Duration: 18 min

Installation Requirements

Before setting up PyTorch with MPS, ensure you have:

macOS 12.3 or later (MPS support added in PyTorch 1.12)
Apple Silicon Mac (M1, M2, M3, or later)
Python 3.8 or later
Xcode Command Line Tools

Installing Xcode Command Line Tools

xcode-select --install

This installs essential development tools including compilers and headers needed for PyTorch.

Creating a Virtual Environment

Always use a virtual environment to isolate dependencies:

# Create a new virtual environment
python3 -m venv pytorch-env

# Activate it
source pytorch-env/bin/activate

Installing PyTorch with MPS Support

Install PyTorch optimized for Apple Silicon:

# Install PyTorch with MPS support
pip install torch torchvision torchaudio

# Verify installation
python -c "import torch; print(torch.__version__)"

PyTorch automatically detects MPS availability on Apple Silicon. No additional configuration needed.

Verifying MPS Availability

Check if MPS is available on your system:

import torch

print(f"PyTorch version: {torch.__version__}")
print(f"MPS available: {torch.backends.mps.is_available()}")
print(f"MPS built: {torch.backends.mps.is_built()}")

if torch.backends.mps.is_available():
    print("✓ MPS is ready to use!")
    device = torch.device("mps")
    print(f"Device: {device}")
else:
    print("✗ MPS is not available")

Expected Output:

PyTorch version: 2.0.0
MPS available: True
MPS built: True
✓ MPS is ready to use!
Device: mps

Moving Tensors to MPS Device

Once MPS is available, move tensors and models to the device:

import torch

# Create a tensor on CPU
x_cpu = torch.randn(1000, 1000)

# Move to MPS device
device = torch.device("mps")
x_mps = x_cpu.to(device)

print(f"CPU tensor device: {x_cpu.device}")
print(f"MPS tensor device: {x_mps.device}")

# Or create directly on MPS
y_mps = torch.randn(1000, 1000, device=device)
print(f"Direct MPS tensor: {y_mps.device}")

Setting Up a Training Loop with MPS

Here's a minimal training loop using MPS:

import torch
import torch.nn as nn
import torch.optim as optim

# Set device
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
print(f"Using device: {device}")

# Define a simple model
model = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
).to(device)

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Dummy data
X = torch.randn(100, 784, device=device)
y = torch.randint(0, 10, (100,), device=device)

# Training step
for epoch in range(5):
    optimizer.zero_grad()
    output = model(X)
    loss = criterion(output, y)
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

Output:

Using device: mps
Epoch 1, Loss: 2.3012
Epoch 2, Loss: 2.1456
Epoch 3, Loss: 1.9823
Epoch 4, Loss: 1.8234
Epoch 5, Loss: 1.6789

Fallback to CPU

For compatibility, always include a fallback to CPU:

import torch

# Automatic fallback
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

# Or explicit check
if torch.backends.mps.is_available():
    device = torch.device("mps")
    print("Using MPS")
else:
    device = torch.device("cpu")
    print("MPS not available, using CPU")

Common Issues and Solutions

Issue: RuntimeError: MPS is not available

Solution: Ensure macOS 12.3+, PyTorch 1.12+, and Apple Silicon Mac

Issue: OutOfMemoryError on MPS

Solution: Reduce batch size or model size (MPS shares system memory)

Issue: Slow performance on MPS

Solution: Ensure model is on MPS device; check with tensor.device