How to Use PyTorch with Apple MPS Backend

by

Apple Silicon (M1, M2, M3, M4) brings a capable GPU integrated into the same die as the CPU, with shared memory that eliminates the PCIe data transfer bottleneck of discrete GPUs. PyTorch supports it via the Metal Performance Shaders (MPS) backend, available since PyTorch 1.12. For ML workloads that fit in the Mac’s unified memory, MPS training is genuinely fast — not A100-fast, but fast enough to make a MacBook Pro a productive development machine without needing cloud GPU access for every experiment.

Requirements

  • macOS 12.3+ (Monterey or later)
  • Apple Silicon Mac (M1 or newer) or AMD GPU Mac (limited support)
  • PyTorch 1.12+ (PyTorch 2.x recommended for best MPS support)
pip install torch torchvision torchaudio

Enabling MPS

import torch

# Check availability
print(torch.backends.mps.is_available())   # True on supported hardware
print(torch.backends.mps.is_built())       # True if PyTorch was built with MPS

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

Moving Models and Data to MPS

import torch
import torch.nn as nn

device = torch.device("mps")

# Model to MPS
model = nn.Sequential(
    nn.Linear(784, 256),
    nn.ReLU(),
    nn.Linear(256, 10)
).to(device)

# Tensors to MPS
x = torch.randn(32, 784).to(device)
y = torch.randint(0, 10, (32,)).to(device)

# Forward pass on GPU
logits = model(x)
loss = nn.CrossEntropyLoss()(logits, y)
print(f"Loss: {loss.item():.4f} on {loss.device}")

Complete Training Loop with MPS

import torch, torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

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

# Data
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])
train_data = datasets.MNIST('./data', train=True, download=True, transform=transform)
loader = DataLoader(train_data, batch_size=256, shuffle=True, num_workers=0)
# Note: num_workers=0 avoids multiprocessing issues on macOS

model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 10)).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()

for epoch in range(5):
    total_loss = 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        optimizer.zero_grad()
        loss = criterion(model(x), y)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    print(f"Epoch {epoch+1}: avg loss {total_loss/len(loader):.4f}")

Writing Device-Agnostic Code

Write code that works on MPS, CUDA, and CPU without changes by using a device detection function:

def get_device():
    if torch.cuda.is_available():
        return torch.device("cuda")
    elif torch.backends.mps.is_available():
        return torch.device("mps")
    else:
        return torch.device("cpu")

device = get_device()
print(f"Training on: {device}")

MPS Performance: What to Expect

MPS GPU speedup varies significantly by workload. Typical ranges on M2/M3 Pro compared to CPU-only training: CNNs see 3–8x speedup; transformers see 2–5x; simple MLPs sometimes see less than 2x where memory bandwidth rather than compute is the bottleneck. The M1 Ultra and M2/M3 Max with their larger GPU core counts and higher memory bandwidth deliver the best results. Compared to a cloud NVIDIA A100, an M3 Max GPU is roughly 3–5x slower on typical deep learning workloads — but it is always on, always available, and costs nothing per experiment beyond electricity.

The unified memory architecture is MPS’s main structural advantage: data does not need to be explicitly transferred from CPU RAM to GPU VRAM. On a 96GB M2 Ultra, you can train models that would not fit in any consumer discrete GPU’s VRAM. This is practically significant for large language model fine-tuning where VRAM is the binding constraint.

Mixed Precision on MPS

Automatic Mixed Precision (AMP) is supported on MPS from PyTorch 2.0+, using torch.float16. Note that MPS uses torch.float16 rather than torch.bfloat16 — bfloat16 is not natively supported on most Apple Silicon GPUs (M3 Ultra and some configurations support it, but float16 is the safe default):

from torch.cuda.amp import autocast, GradScaler

# MPS autocast uses float16
scaler = GradScaler()  # gradient scaling for float16 stability

for x, y in loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad()
    with autocast(device_type='mps', dtype=torch.float16):
        output = model(x)
        loss = criterion(output, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Mixed precision on MPS typically provides 20–40% additional speedup on top of MPS baseline, plus reduced memory usage allowing larger batch sizes.

Benchmarking MPS vs CPU

import torch, time

def benchmark(device_str, n_iters=200):
    device = torch.device(device_str)
    model = torch.nn.Sequential(
        torch.nn.Conv2d(3, 64, 3, padding=1),
        torch.nn.ReLU(),
        torch.nn.Conv2d(64, 128, 3, padding=1),
        torch.nn.ReLU(),
        torch.nn.AdaptiveAvgPool2d(1),
        torch.nn.Flatten(),
        torch.nn.Linear(128, 10)
    ).to(device)
    x = torch.randn(32, 3, 64, 64).to(device)

    # Warmup
    for _ in range(10):
        _ = model(x)
    if device_str == 'mps':
        torch.mps.synchronize()

    t0 = time.perf_counter()
    for _ in range(n_iters):
        _ = model(x)
    if device_str == 'mps':
        torch.mps.synchronize()
    elapsed = (time.perf_counter() - t0) * 1000 / n_iters
    print(f"{device_str}: {elapsed:.2f}ms/iter")

benchmark('cpu')
benchmark('mps')

Always call torch.mps.synchronize() before stopping the timer — MPS operations are asynchronous and timing without synchronisation measures dispatch latency, not compute time.

Figure 1 — Apple Silicon ML performance: MPS vs CPU vs cloud A100

Approximate Speedup Relative to M3 Pro CPU (CNN training) CPU (M3 Pro) 1x MPS M2 Pro ~4x MPS M3 Max ~6x A100 (cloud) ~20-30x Approximate — varies by model architecture, batch size, and operation mix

MPS Limitations and Workarounds

MPS does not support every PyTorch operation — some fall back to CPU silently. Check if an op is supported:

# Enable fallback to CPU for unsupported MPS ops (slower but safe)
import os
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

# Then import torch and use MPS normally
# Unsupported ops will run on CPU with a warning

Key known limitations as of PyTorch 2.x: some sparse tensor operations are unsupported; certain custom CUDA extensions will not work (they require CUDA, not Metal); complex-valued tensors have limited support; and torch.bfloat16 is unsupported on most models (use float16 instead). The list shrinks with each PyTorch release. For production training at scale, a cloud GPU is still the right choice — MPS is optimal for development, quick experiments, and training smaller models where the Mac’s always-on convenience outweighs the raw speed difference.

Saving Models Trained on MPS

# Move to CPU before saving — more portable
model.cpu()
torch.save(model.state_dict(), 'model.pt')

# Or save directly (file will contain MPS tensors — less portable)
torch.save(model.state_dict(), 'model_mps.pt')

# Loading on any device
state_dict = torch.load('model.pt', map_location='cpu')
model.load_state_dict(state_dict)
model = model.to(device)  # move to target device after loading

Using Hugging Face Transformers on MPS

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

device = torch.device("mps")
model_name = "distilbert-base-uncased"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name).to(device)

inputs = tokenizer("Hello, MPS!", return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)
print(outputs.logits)

Most Hugging Face models run on MPS without modification. The Trainer API detects MPS automatically when you pass no_cuda=True and MPS is available. For fine-tuning smaller models (up to ~7B parameters depending on Mac RAM), MPS provides a viable local alternative to cloud training.

MPS makes Apple Silicon Macs genuinely useful for ML development. The workflow is: develop and debug on MPS where iteration is fast and free, then scale final training runs to cloud GPU when you need maximum speed. The device-agnostic code pattern above ensures your training script runs on both with no changes. For models up to a few hundred million parameters, MPS on a high-end Mac is fast enough to run full experiments without cloud access — which changes the economics of ML experimentation significantly for solo practitioners and small teams.

Figure 1 — Approximate training speedup relative to Apple Silicon CPU (CNN workload)

CPU (M3 Pro) 1x MPS M2 Pro ~4x MPS M3 Max ~6x A100 (cloud) ~20-30x Approximate — varies by model and batch size

torch.mps Utilities

PyTorch exposes several MPS-specific utilities for memory management and synchronisation:

import torch

# Synchronise MPS stream (required for accurate timing)
torch.mps.synchronize()

# Memory management
print(torch.mps.current_allocated_memory())   # bytes currently in use
print(torch.mps.driver_allocated_memory())    # bytes allocated by driver
torch.mps.empty_cache()                        # release cached memory

# Set random seed for MPS
torch.mps.manual_seed(42)

Calling torch.mps.empty_cache() between training runs in a notebook session can help reclaim memory when switching between models. Unlike CUDA, MPS memory management is handled largely by the Metal driver — empty_cache() signals that cached allocations can be released, but the driver decides when to actually free them back to the system.

MPS vs CUDA: The Practical Gap in 2026

The performance gap between MPS and CUDA has narrowed with each Apple Silicon generation. The M4 Ultra (released in 2025) delivers GPU performance that is competitive with mid-range NVIDIA datacenter GPUs for memory-bandwidth-limited workloads. For compute-bound workloads — transformer training at scale — CUDA on high-end NVIDIA GPUs (H100, A100) remains 10–20x faster. The gap that matters most in practice is not the peak compute but the memory: a Mac with 192GB unified memory can hold model states that simply do not fit on any consumer GPU, making it uniquely suited for certain large-model workflows like multi-adapter LoRA fine-tuning or multi-modal model inference where the parameter count and KV cache together exceed 80GB. For these workloads, MPS on a maxed-out Mac is not a compromise — it is the only practical local option.

Why MPS Matters for ML Practitioners

Before MPS, Mac users had two options for GPU-accelerated ML: use a cloud GPU for everything (expensive and slow iteration), or use CPU-only training (slow). MPS changes this by making the Mac’s integrated GPU usable for PyTorch workloads. The unified memory architecture is the key differentiator — on a 96GB M2 Ultra or 128GB M3 Ultra, you can train models that would not fit in any consumer discrete GPU’s VRAM. This is practically significant for large language model fine-tuning where VRAM is the binding constraint, not compute throughput.

The MPS speedup relative to CPU is real and consistent: expect 3–8x for CNNs, 2–5x for transformers, and less for workloads dominated by memory-bound operations. Compared to a cloud NVIDIA A100, an M3 Max is roughly 5–8x slower on typical deep learning — but it is always on, always available, and costs nothing per experiment. For the rapid iteration phase of a project — trying architectures, debugging training loops, running small ablations — that convenience is genuinely valuable.

Known Issues and Solutions

The most common issues when first enabling MPS. Operation not implemented for MPS — set PYTORCH_ENABLE_MPS_FALLBACK=1 to silently fall back to CPU for unsupported ops; this reduces speed but keeps training running. Slow first epoch — MPS compiles Metal shaders on first use, similar to CUDA kernel compilation. This is a one-time cost per model architecture per session. num_workers > 0 in DataLoader — macOS multiprocessing has limitations with CUDA/MPS; use num_workers=0 or set persistent_workers=True if you need parallel data loading. Out of memory errors — unified memory is shared with the OS and other apps; close memory-heavy applications, or reduce batch size. NaN gradients with float16 — gradient scaler is required for float16 stability, same as CUDA AMP; the scaler example above handles this correctly.

MPS makes Apple Silicon Macs genuinely useful for ML development. The right workflow: develop and debug on MPS where iteration is fast and free, scale final training runs to cloud GPU when you need maximum throughput. The device-agnostic code pattern — the get_device() function above — ensures your script runs correctly on CUDA, MPS, and CPU without modification. For models up to a few hundred million parameters, MPS on a high-end Mac is fast enough to run complete experiments locally, which changes the economics of ML experimentation for solo practitioners and small teams significantly.

MPS has made Apple Silicon a serious ML development platform — not a substitute for cloud GPUs in production training, but a capable environment for the daily work of research and development where fast iteration matters more than peak throughput.

Leave a Comment