Run one JAX program across multiple hosts, processes, and devices when a single machine is not enough. JAX uses a multicontroller model: launch the same Python program on every host, initialize the cluster, and let each process compute on its local devices while communicating across global devices. section 7.4
A multihost setup requires you to:
- Create and initialize the cluster with jax.distributed.initialize() before running JAX computations. section 7.4
- Use jax.process_index() to identify each process and distinguish local devices from global devices. section 7.4
- Prepare each process's input separately, then use pmap() and collective operations for cross-host computation. section 7.4
Listing 7.30 Program for TPU Pod slice
import jax
import jax.numpy as jnp
from jax import random
jax.distributed.initialize() #1
print('== Running worker: ', jax.process_index()) #2
def dot(v1, v2):
return jnp.vdot(v1, v2)
rng_key = random.PRNGKey(42 + 10*jax.process_index()) #3
vs = random.normal(rng_key, shape=(2_000_000,3)) #3
v1s = vs[:1_000_000,:] #3
v2s = vs[1_000_000:,:] #3
# The total number of TPU cores in the Pod
device_count = jax.device_count() #4
# The number of TPU cores attached to this host
local_device_count = jax.local_device_count() #5
if jax.process_index() == 0: #6
print('-- global device count:', jax.device_count())
#print('global devices:', jax.devices())
print('-- local device count:', jax.local_device_count())
#print('local devices:', jax.local_devices())
print('-- JAX version:', jax.__version__)
v1sp = v1s.reshape(
(local_device_count,
v1s.shape[0]//local_device_count,
v1s.shape[1]))
v2sp = v2s.reshape(
(local_device_count,
v2s.shape[0]//local_device_count,
v2s.shape[1]))
if jax.process_index() == 0:
print('-- v1sp shape: ', v1sp.shape) # (8, 125000, 3)
dots = jax.pmap(jax.vmap(dot))(v1sp,v2sp) #7
if jax.process_index() == 0:
print('-- dots shape: ', dots.shape) # (8, 125000)
global_sum = jax.pmap( #8
lambda x: jax.lax.psum(jnp.sum(x), axis_name='p'), #9
axis_name='p'
)(dots)
if jax.process_index() == 0:
print('-- global_sum shape: ', global_sum.shape) # (8,)
print(f'== Worker {jax.process_index()} global sum: {global_sum}') #10
dots = dots.reshape((dots.shape[0]*dots.shape[1]))
if jax.process_index() == 0:
print('-- result shape: ', dots.shape) # (1000000,)
local_sum = jnp.sum(dots) #11
print(f'== Worker {jax.process_index()} local sum: {local_sum}')
print(f'== Worker {jax.process_index()} done')
The TPU Pod slice program initializes the cluster, gives workers distinct random seeds, reports global and local device counts, and combines per-process results. section 7.4
