Skip to content

← JAX Performance & Distributed step 15 of 25

Hard Primitives

shard_map with psum Collective

Jump to the task

Why this matters

A local sum covers only the values on one device. jax.lax.psum adds values across a named mesh axis and makes the result available to each participant. Combining a local sum with psum produces a global sum of partitioned data.

For example, two blocks [2, -1] and [4, 3] have local sums 1 and 7. The collective produces 8 on both devices. With more mesh axes, reducing over data combines participants along that axis, not necessarily every device in the whole mesh.

Worked mini-example

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, PartitionSpec as P

mesh = Mesh(np.array(jax.devices()[:1]), ('data',))

def sum_block(block):
    local_sum = jnp.sum(block)
    return jax.lax.psum(local_sum, axis_name='data')

sum_all = jax.shard_map(
    sum_block, mesh=mesh,
    in_specs=P('data'), out_specs=P(),
)
x = jnp.array([2.0, -1.0, 4.0, 3.0], dtype=jnp.float32)
print(sum_all(x))  # 8.0

The example selects one device. The collective leaves its local sum unchanged because there are no other participants along data. A larger mesh is needed to observe communication between devices.

Problem

Implement shard_map_psum(x):

  1. Construct a mesh from jax.devices()[:1], with axis name 'data'.
  2. Define a local function that computes jnp.sum(block), then calls jax.lax.psum(local_sum, axis_name='data').
  3. Wrap it with jax.shard_map, using PartitionSpec('data') for input and an empty PartitionSpec() for output.
  4. Call the wrapped function on x and return its result.

x is a one-dimensional float32 JAX array. Return its global sum as a float32 scalar. Numerical tests on this one-device mesh cannot distinguish this implementation from jnp.sum(x) or verify actual communication.

Common pitfalls

  • Use the exact mesh axis name. A collective over an unbound name is invalid.
  • PartitionSpec() declares that the output is replicated. It does not perform the reduction for you. Without psum, different local sums do not satisfy that declaration; replication checking can reject the function.
  • All participants must execute compatible collectives in a consistent order. Device-dependent control flow that skips a collective for some participants is invalid and can fail or hang. This requirement does not mean the Python call always blocks until the device work is complete.
  • psum is a sum, not a mean. Gradient aggregation needs the normalization appropriate to the loss and local batch sizes; summing local mean gradients is not automatically the global mean gradient.
  • A larger mesh needs compatible input lengths and may change floating-point addition order. It does not promise bit-for-bit equality for every input.