Data placement

Data placement#

Every array lives somewhere: in the memory of one device, or spread across the memories of many. The unit of placement in JAX is the mesh, a set of devices arranged in a grid with named axes. A large program might place different arrays on different meshes and move data between them; the everyday case, an array sitting on one device, is just placement on a single-device mesh. This page covers how placement is decided by default, how to control it with jax.device_put(), what it means for an array to be committed to a mesh, and how to move data between meshes.

How data is laid out across a mesh’s devices (shardings, partition specs, and parallelism) is the subject of the next page, Distributed arrays and automatic parallelization. Here we only need the idea of a mesh itself.

We’ll simulate eight devices on CPU:

import jax

jax.config.update('jax_num_cpu_devices', 8)

import jax.numpy as jnp
import numpy as np
from jax.sharding import NamedSharding

A mesh over all eight devices, arranged along one named axis:

full = jax.make_mesh((8,), ('x',))
full

Mesh(axis_sizes=(8,), axis_names=('x',), axis_types=(Explicit,))

And, using jax.make_mesh()’s devices argument, some smaller meshes over subsets of the same hardware — two halves, and a couple of single-device meshes:

mesh_a = jax.make_mesh((4,), ('x',), devices=full.devices[:4])
mesh_b = jax.make_mesh((4,), ('x',), devices=full.devices[4:])

m0 = jax.make_mesh((1,), ('x',), devices=full.devices[:1])
m1 = jax.make_mesh((1,), ('x',), devices=full.devices[1:2])

Default placement, and commitment#

When you create an array without saying where it should go, JAX places it somewhere sensible and leaves it uncommitted:

x = jnp.arange(4.0)
print(x.committed)
print(x.sharding)

False
SingleDeviceSharding(device=CpuDevice(id=0), memory_kind=device)

Uncommitted means the array isn’t attached to its location: it’s free to follow. If it’s used in a computation together with data that is committed somewhere, JAX will move it there implicitly.

To place an array deliberately, use jax.device_put() with a sharding, which is a mesh plus a description of how the array is laid out over it:

data = np.arange(8.0)
xa = jax.device_put(data, NamedSharding(mesh_a, jax.P('x')))
print(xa.committed)
print(xa.sharding.mesh == mesh_a)

True
True

Now the array is committed to mesh_a and it stays put. Computations follow their data: an operation on xa runs on mesh_a’s devices.

Commitment is a promise JAX enforces. Combining arrays committed to different meshes in one operation is an error:

a = jax.device_put(1.0, NamedSharding(m0, jax.P()))
b = jax.device_put(2.0, NamedSharding(m1, jax.P()))

a + b  # committed to different meshes!

ValueError: Received incompatible devices for jitted computation. Got argument x of add with shape float32[] and with device ids [0] on platform CPU and argument y of add with shape float32[] and with device ids [1] on platform CPU

This is deliberate: transfers are expensive, so JAX makes you say where data should move rather than shuffling it behind your back. Uncommitted values, by contrast, follow along freely:

c = jnp.float32(3.0)   # uncommitted
print((a + c).sharding)  # runs on m0, where `a` is committed

NamedSharding(mesh=Mesh('x': 1, axis_types=(Explicit,)), spec=P(), memory_kind=device)

In older code you’ll see jax.device_put(x, some_device), committing an array to a single device. That’s the same model: committing to a single-device mesh.

Moving data between meshes#

The same function that places data also moves it: calling jax.device_put() on an already-placed array with a sharding on a different mesh transfers it there.

xb = jax.device_put(xa, NamedSharding(mesh_b, jax.P('x')))
print(xb.sharding.mesh == mesh_b)

True

There are two related operations:

  • Changing layout within a mesh (same devices, different partitioning) is resharding, covered in Distributed arrays and automatic parallelization (see jax.reshard, and device_put with a new spec on the same mesh).

  • Changing meshes, meaning moving data to a different set of devices, is this page’s focus, and jax.device_put is the tool.

One more case: because a mesh is an arrangement of devices, keeping an array’s data on the same set of devices but changing the device order in the sharding is also a change of mesh, and therefore a device_put rather than a reshard. The same partition spec over differently-ordered meshes assigns different data to each device, so real data movement is required.

mesh_rev = jax.make_mesh((8,), ('x',), devices=full.devices[::-1])

x = jax.device_put(np.arange(8.0), NamedSharding(full, jax.P('x')))
y = jax.device_put(x, NamedSharding(mesh_rev, jax.P('x')))  # same devices, new order

def data_by_device(arr):
  return [float(s.data[0]) for s in sorted(arr.addressable_shards,
                                           key=lambda s: s.device.id)]

print(data_by_device(x))  # what each of devices 0, 1, ..., 7 holds
print(data_by_device(y))

[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]
[7.0, 6.0, 5.0, 4.0, 3.0, 2.0, 1.0, 0.0]

As mathematical values, x and y are identical, but every element has moved: device 0 held 0.0 before and holds 7.0 after. (And since x and y are committed to different meshes, combining them in one operation would be an error, as above.)

Mesh changes are runtime-level operations: everything in this section happens at the top level of your program. A compiled computation belongs to one fixed mesh: under jax.jit (or inside any staged-out control flow like jax.lax.scan), you can change an array’s layout on the mesh, but you cannot move data to a different mesh. Attempting a cross-mesh device_put inside jit is an error:

@jax.jit
def f(x):
  return jax.device_put(x, NamedSharding(mesh_b, jax.P('x')))

f(xa)  # error: can't move to a different mesh inside jit!

ValueError: Received incompatible devices for jitted computation. Got argument x of f with shape float32[8@x] and with device ids [0, 1, 2, 3] on platform CPU and device_put inside jit with device ids [4, 5, 6, 7] on platform CPU at /tmp/ipykernel_1658/3373542785.py:3:9 (f)

Moving between meshes is the runtime’s job, done from Python between compiled computations, which is why the pipeline example below drives each stage-to-stage transfer from the top level.

Like everything else in JAX’s runtime, device_put is asynchronous (Asynchronous dispatch): it returns immediately, and the transfer proceeds concurrently with whatever Python does next. As a consequence, computations and transfers on different meshes run in parallel whenever their inputs are ready, because each device works through its own queue independently. That’s enough to express pipeline parallelism at the runtime level, with no compiler involvement: stages committed to disjoint meshes, device_put moving each microbatch from stage to stage:

stage_meshes = [jax.make_mesh((2,), ('x',), devices=full.devices[2*i:2*i+2])
                for i in range(4)]

@jax.jit
def stage(x):
  return jnp.sin(x) + 1.0   # stand-in for real per-stage computation

microbatches = [np.full((256, 256), float(i)) for i in range(8)]

results = []
for mb in microbatches:
  for m in stage_meshes:
    mb = jax.device_put(mb, NamedSharding(m, jax.P('x')))
    mb = stage(mb)
  results.append(mb)

print(results[-1].sharding.mesh == stage_meshes[-1])

True

Each microbatch is enqueued through the stages sequentially, but the queues are independent: while stage 1 works on microbatch 0, stage 0 has already started microbatch 1, with the transfers overlapping the computation. (For a richer treatment of pipelining, see the microbatched gradient-accumulation example in Autodiff with refs.)

Everything on this page generalizes to multiple processes: meshes can span hosts, can cover just a subset of the devices in a cluster, and jax.device_put can transfer between meshes across processes over high-speed interconnects. That story is in the multi-process systems docs; see Introduction to multi-controller JAX (aka multi-process/multi-host JAX).

Notes#

  • jax.device_put(x) with no target is (at most) a transfer to the default placement, and the result is uncommitted.

  • jax.device_put accepts donate=True to reuse the input’s buffers when possible, in the same spirit as jit’s buffer donation (Buffer donation).

  • Placement round-trips through the host too: np.asarray(x) copies device data back to host memory as a NumPy array.

Next steps#

The next page, Distributed arrays and automatic parallelization, covers what happens within a mesh: how an array’s data is partitioned over the mesh’s devices, and how JAX turns one program into coordinated per-device computation.