Skip to content

Chapter 15: Device Management (CPU, GPU, TPU)

Overview

Managing data and computation across different devices with JAX.

Topics

1. Checking Available Devices

import jax

# List all devices
print(jax.devices())
# Output: [gpu(id=0), gpu(id=1), ...]

# Check default device
print(jax.default_backend())

# Device count
print(jax.device_count())

2. Device Placement

import jax
import jax.numpy as jnp

# Place data on specific device
x = jnp.array([1, 2, 3])
x_gpu = jax.device_put(x, jax.devices()[0])

print(x_gpu.device())  # Which device

# Bring back to CPU
x_cpu = jax.device_get(x_gpu)

3. Data Parallelism with pmap

import jax
from jax import pmap, random

def per_device_forward(x, params):
    return model(params, x)

# Map to all devices
n_devices = jax.device_count()
forward_pmap = pmap(per_device_forward)

# Split batch across devices
x_split = x.reshape((n_devices, -1) + x.shape[1:])
params_split = jax.tree_map(lambda p: p[None], params)

# Run on all devices in parallel
output = forward_pmap(x_split, params_split)

4. Collective Operations

from jax.experimental.pjit import pjit
from jax import lax

# All-reduce across devices
def allreduce_sum(x):
    return lax.psum(x, 'i')

# All-gather
def allgather(x):
    return lax.all_gather(x, 'i')

Summary

  • Check devices with jax.devices()
  • Place data with device_put/device_get
  • Use pmap for data parallelism
  • Collective ops (psum, all_gather) for sync
  • TPU support automatic with same code