GPU Support

GPU Support#

PyTorch can run tensor operations on CPUs and hardware accelerators. Accelerators are often faster for sufficiently large, parallel workloads, but transfer overhead, unsupported operations, precision, and workload size determine the actual speedup. Small computations can be faster on a CPU.

import torch

Backends#

PyTorch can run computations on different devices. By default, tensors and models are placed on the CPU, but PyTorch also supports GPUs and other hardware accelerators. Common device types include:

  • CUDA — NVIDIA GPUs supported through the CUDA platform, primarily on Windows and Linux.

  • MPS — Apple GPUs supported through the Metal Performance Shaders framework on macOS.

  • Other accelerators — PyTorch also supports additional device types and hardware-specific backends.

Checking for CUDA#

Let’s check whether you have a Nvidia GPU available on your Windows/Linux computer.

torch.cuda.is_available()

True

If the above command returns False, the reason could be one of the following.

  • You don’t have a Nvidia GPU.

  • You don’t have the correct version of CUDA installed.

  • You don’t have the correct version of PyTorch installed.

Refer to the PyTorch website for the correct version of PyTorch and CUDA to install, in case you have a Nvidia GPU.

if not torch.cuda.is_available():
    if not torch.backends.cuda.is_built():
        print("PyTorch was not built with CUDA support")
    else:
        print("No Nvidia GPU on this machine and/or CUDA is not installed properly")

Checking for MPS#

Let’s check whether you have a MPS GPU available on your MacOS computer.

torch.mps.is_available()

True

If the above command returns False, the reason could be one of the following.

  • You don’t have a MPS GPU.

  • PyTorch was not built with MPS enabled.

  • The current MacOS version is lower than 12.3.

if not torch.mps.is_available():
    if not torch.backends.mps.is_built():
        print("PyTorch was not built with MPS enabled.")
    else:
        print("The current MacOS version is not 12.3+ and/or you do not have an MPS-enabled device on this machine.")

Using a GPU#

By default, all tensors you create are stored on the CPU. You can push a tensor to the GPU by using the function .to(...), where the argument is the device you want to use. It is a good practice to define a device object in your code which points to the GPU if you have one, and otherwise to the CPU. Then, you can write a device-agnostic code that will run on the correct device based on the availability. Let’s try it below.

def get_device():

    if torch.cuda.is_available():
        return torch.device("cuda")
    
    if torch.mps.is_available():
        try:
            # A backend may be reported as available even when a machine
            # cannot allocate MPS tensors (for example, without unified memory).
            torch.empty(1, device="mps")
            return torch.device("mps")
        
        except RuntimeError as error:
            print(f"MPS unavailable at runtime: {error}")

    return torch.device("cpu")

device = get_device()

print(f"Using device:", device)

Using device: mps

Now let’s create a tensor and push it to the device.

x = torch.zeros(2, 3)
x = x.to(device)

print(x)

tensor([[0., 0., 0.],
        [0., 0., 0.]], device='mps:0')

Alternatively, you can also use the device argument while creating a tensor.

y = torch.ones(2, 3, device=device)

print(y)

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='mps:0')

In case you have a GPU, you should see device='cuda:0' or device='mps:0'. Operations generally require all participating tensors and model parameters to be on the same device; a device-mismatch error is usually fixed by moving the whole batch and model to the chosen device. Move a result to CPU before converting it to NumPy, for example tensor.detach().cpu().numpy().

Accelerator support is backend-dependent: an operation or dtype available on CUDA may be unavailable or behave differently on MPS. Reduced-precision dtypes can improve speed and memory use but also change numerical accuracy. On CUDA, torch.cuda.memory_allocated() and torch.cuda.memory_summary() help diagnose memory use; deleting references allows PyTorch to reuse cached memory.

Setting the seed on all devices#

A seed makes random-number streams repeatable in a fixed environment, but it does not guarantee identical results across PyTorch versions, platforms, devices, or nondeterministic kernels. The function torch.manual_seed() seeds PyTorch’s random generator on CPU and available accelerators. Deterministic algorithms can be requested separately, sometimes at a performance cost or with an error when no deterministic implementation exists.

# Seed PyTorch on CPU and available accelerators
torch.manual_seed(42)

# Optional deterministic CUDA configuration
if torch.cuda.is_available():
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

Speed comparison#

Let’s compare a matrix multiplication on the CPU and the selected accelerator. Accelerator operations are asynchronous, so timing must synchronize before reading the clock. We also warm up each device and report the median of several repetitions. This microbenchmark measures only matrix multiplication; it does not include data-transfer time and does not predict every training workload.

import statistics
import time

def synchronize(device):
    if device.type == "cuda":
        torch.cuda.synchronize(device)
    elif device.type == "mps":
        torch.mps.synchronize()

def benchmark_matmul(x, repeats=5):
    _ = x @ x  # warm-up
    synchronize(x.device)
    timings = []
    for _ in range(repeats):
        start = time.perf_counter()
        _ = x @ x
        synchronize(x.device)
        timings.append(time.perf_counter() - start)
    return statistics.median(timings)

x_cpu = torch.randn(2000, 2000)
print(f"CPU median: {benchmark_matmul(x_cpu):.5f}s")

CPU time: 0.80048s

if device.type == "cuda":
    x_cuda = x_cpu.to(device)
    print(f"CUDA median: {benchmark_matmul(x_cuda):.5f}s")

GPU time: 0.02454s

if device.type == "mps":
    x_mps = x_cpu.to(device)
    print(f"MPS median: {benchmark_matmul(x_mps):.5f}s")

MPS time: 0.09917s

Conclusion#

In this tutorial, we selected an available accelerator with a CPU fallback, moved tensors between devices, discussed reproducibility and backend limitations, and measured matrix multiplication with synchronization and repeated timings. Accelerator speedups are workload- and hardware-dependent.