Placing an array on purpose

Place the same values twice on the same mesh under opposite specs, add them, and nothing complains. The result takes the first operand's spec, and swapping the operands swaps the answer, so a + b comes back P("data", "model") while b + a comes back P("model", "data").

Order is not the rule, though, which is why guessing here is a bad habit. Add a replicated array to a sharded one and the sharded spec wins in both orders. Read the result's spec rather than working out who should have won.

What the operands may not disagree about is which devices they sit on. An array spread over all eight plus an array placed on device 3 raises, and the message prints both device id lists, [0, 1, 2, 3, 4, 5, 6, 7] against [3]. A layout difference gets resolved. A device-set difference does not.

The flag underneath that refusal, committed, and the placement rules it carries belong to the arrays chapter, which reads its docstring directly. What a mesh adds is the second thing two operands can now differ by.

Layouts negotiate. Device sets do not.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
u = jnp.ones((8, 16))
a = jax.device_put(u, NamedSharding(mesh, P("data", "model")))
b = jax.device_put(u, NamedSharding(mesh, P("model", "data")))
r = jax.device_put(u, NamedSharding(mesh, P()))
one = jax.device_put(u, jax.devices()[3])

print((a + b).sharding.spec, (b + a).sharding.spec)
print((a + r).sharding.spec, (r + a).sharding.spec)
try:
    a + one
except ValueError as err:
    print(err)

# PartitionSpec('data', 'model') PartitionSpec('model', 'data')
# PartitionSpec('data', 'model') PartitionSpec('data', 'model')
# Received incompatible devices for jitted computation. Got argument x of add with
# shape float32[8,16] and device ids [0, 1, 2, 3, 4, 5, 6, 7] on platform CPU and
# argument y of add with shape float32[8,16] and device ids [3] on platform CPU

device_put takes pytrees, and it takes shardings in two shapes. Pass a single sharding and every leaf gets it. Pass a tree of shardings with the same structure and each leaf gets its own, which is how a parameter tree ends up with weights and biases split differently.

The single-sharding form goes further than it looks because specs are positional and short specs pad. P("data") sends an (8, 16) weight and a (16,) bias to the same mesh axis without either leaf having to know its own rank.

The shard shapes in the second call are where the two forms separate. The weight splits both ways down to (2, 8), while the bias splits only across model, down to (8,). Same call, same mesh, two different answers per leaf.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
tree = {"w": jnp.ones((8, 16)), "b": jnp.zeros(16)}

one = jax.device_put(tree, NamedSharding(mesh, P("data")))
print({k: (v.sharding.spec, v.addressable_shards[0].data.shape) for k, v in one.items()})

each = jax.device_put(tree, {"w": NamedSharding(mesh, P("data", "model")),
                            "b": NamedSharding(mesh, P("model"))})
print({k: (v.sharding.spec, v.addressable_shards[0].data.shape) for k, v in each.items()})

# {'b': (PartitionSpec('data',), (4,)), 'w': (PartitionSpec('data',), (2, 16))}
# {'b': (PartitionSpec('model',), (8,)), 'w': (PartitionSpec('data', 'model'), (2, 8))}

Handing device_put an array that is already sharded is legal, and outside jit it is the way data moves between layouts. Shard shapes go from (2, 8) to (4, 4), every value stays where it belongs in the global array, and an elementwise comparison of the two arrays is True everywhere.

Afterwards the original is untouched, still (8, 16) under its old spec, because chapter 1's rule about arrays being values does not stop applying when there are eight devices. Nothing moved in place. A second array exists, and the first one lives until nothing refers to it.

That is the cost worth holding on to. A reshard is bytes across links plus two copies resident at once, so a reshard inside a training step is paying both, every step, for a layout decision that could have been made before the loop.

There is no move. There is a new array, and the old one until you drop it.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
x = jax.device_put(jnp.arange(128.).reshape(8, 16), NamedSharding(mesh, P("data", "model")))
y = jax.device_put(x, NamedSharding(mesh, P("model", "data")))

print(x.addressable_shards[0].data.shape, y.addressable_shards[0].data.shape)
print(bool(jnp.all(x == y)))
print(x.sharding.spec, y.sharding.spec)

# (2, 8) (4, 4)
# True
# PartitionSpec('data', 'model') PartitionSpec('model', 'data')

P() splits no array axis, so every device holds the whole thing. shard_shape returns the global shape, is_fully_replicated is True, and there are still eight shards; they are eight identical copies rather than eight pieces.

Counting those copies is the memory argument. A replicated array costs eight times what the same array costs under P("data", "model") on this mesh, and the next lesson measures exactly that on a pair of matrices big enough for the number to matter.

On this machine global_shards and addressable_shards both return eight, because one process owns every device, and is_fully_addressable says so directly. They separate on multi-host runs, which chapter 12 owns; here the two are interchangeable.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
r = jax.device_put(jnp.arange(128.).reshape(8, 16), NamedSharding(mesh, P()))

print(r.is_fully_replicated, r.addressable_shards[0].data.shape)
print(len(r.addressable_shards), len(r.global_shards), r.is_fully_addressable)

# True (8, 16)
# 8 8 True