JAX sharding - a distributed-array taster

JAX sharding - a distributed-array taster#

Track 07 - GPU · Notebook 08 · Runtime: ≈2 min on CPU

Prerequisites: 03_training/02 (DDP vs FSDP2) helps but isn’t required.

References:

JAX docs: Distributed arrays and automatic parallelization.

Xu et al. 2021, GSPMD: General and Scalable Parallelization for ML Computation Graphs (2105.04663).

What#

PyTorch’s FSDP2 shards parameters and all-gathers them on demand. JAX takes the idea further: every array is always a jax.Array with a sharding attached, and the compiler (XLA) plans all the all-gathers, reduce-scatters, and all-reduces automatically based on the shardings of inputs and outputs.

We run JAX on CPU and pretend we have 8 “devices” by setting XLA_FLAGS=--xla_force_host_platform_device_count=8 before JAX is imported - the XLA backend consumes this at startup and presents the single CPU as eight logical devices. We then build a 1-D mesh, shard a parameter, and verify:

You can place a parameter matrix so it’s sharded along rows; a matmul against an un-sharded activation produces a sharded output and XLA inserts the right all-reduce.
Re-sharding (jax.device_put with a new sharding) has the expected effect on the global array layout.
A two-dim mesh (data x model) is the natural structure for combined data- and tensor-parallel inference.

from llm_systems_cookbook.nb import bootstrap

import os

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

s = bootstrap("07_gpu_08_jax_sharding_pipeline")

Import JAX, set up the mesh#

With 8 CPU devices simulated, we build a 1-D mesh of shape (8,) with a single axis named "model". Any tensor we shard along "model" is split evenly across the 8 devices.

(On GPU you’d get real devices; the mesh API is identical.)

jnp = None
jax = None
have_jax = False
try:
    import jax
    import jax.numpy as jnp
    from jax.sharding import Mesh, NamedSharding, PartitionSpec

    devices = jax.devices()
    print(f"jax version = {jax.__version__}  devices = {len(devices)}")
    have_jax = len(devices) >= 2
except Exception as e:  # noqa: BLE001
    print(f"jax unavailable: {type(e).__name__}: {e}")
    have_jax = False

mesh = None
if have_jax:
    import numpy as np
    mesh = Mesh(np.array(jax.devices()).reshape(-1), axis_names=("model",))
    print(f"mesh = {mesh}")
    s.check(
        "mesh_has_multiple_devices",
        lambda: len(mesh.devices) >= 2,
        msg=f"mesh devices = {len(mesh.devices)}",
    )
else:
    s.skip("mesh_has_multiple_devices", "JAX not available or < 2 devices")

Shard a matrix across the mesh#

Create a (1024, 1024) matrix and place it with PartitionSpec("model", None) - first axis partitioned across the mesh, second axis replicated. Each device then owns (1024/8, 1024) = (128, 1024).

W_global = None
if have_jax and mesh is not None:
    W_local = (jnp.arange(1024 * 1024, dtype=jnp.float32).reshape(1024, 1024) / (1024 * 1024))
    sharding = NamedSharding(mesh, PartitionSpec("model", None))
    W_global = jax.device_put(W_local, sharding)

    n_shards = len(W_global.addressable_shards)
    print(f"global shape = {W_global.shape}   n_shards = {n_shards}")
    print(f"shard[0].data.shape = {W_global.addressable_shards[0].data.shape}")

    s.assert_close(
        "shard_row_count_matches_mesh",
        actual=float(W_global.addressable_shards[0].data.shape[0]),
        expected=1024.0 / len(mesh.devices),
        rtol=1e-9,
    )
    s.check(
        "shard_col_count_is_full",
        lambda: W_global.addressable_shards[0].data.shape[1] == 1024,
        msg=f"shard col dim = {W_global.addressable_shards[0].data.shape[1]}",
    )
else:
    s.skip("shard_row_count_matches_mesh", "JAX not available")
    s.skip("shard_col_count_is_full",      "JAX not available")

Matmul with automatic collective insertion#

jnp.dot(x, W) where x is replicated and W is sharded along its first axis: XLA sees that x @ W requires a partial-reduce along the sharded axis and inserts an all-reduce. The output is replicated.

We verify the numerical output matches the single-device computation up to FP rounding.

if have_jax and mesh is not None:
    x_local = jnp.ones((16, 1024), dtype=jnp.float32)
    x_replicated = jax.device_put(x_local, NamedSharding(mesh, PartitionSpec(None, None)))
    y_sharded = jnp.dot(x_replicated, W_global)
    y_single = jnp.dot(x_local, W_local)
    err = float(jnp.max(jnp.abs(y_sharded - y_single)))
    print(f"max abs err, sharded vs single-device: {err:.3e}")
    # Both computations use the same underlying XLA kernels; with
    # well-scaled inputs the absolute error should sit in FP32 noise.
    y_scale = float(jnp.max(jnp.abs(y_single)))
    rel_err = err / max(y_scale, 1e-12)
    print(f"relative err = {rel_err:.3e}  (output scale = {y_scale:.3g})")
    s.check("sharded_matmul_matches_single_device",
             lambda: rel_err < 1e-4,
             msg=f"rel err = {rel_err:.3e}")
else:
    s.skip("sharded_matmul_matches_single_device", "JAX not available")

2-D mesh: data × model parallel#

Reshape the 8 devices into a (2, 4) mesh. A tensor with PartitionSpec("dp", "mp") is sharded along both axes: batch across dp, features across mp. This is the layout most production JAX LLM-training codebases use.

if have_jax:
    import numpy as np
    mesh_2d = Mesh(np.array(jax.devices()).reshape(2, 4), axis_names=("dp", "mp"))
    B, D = 8, 1024
    x2d = jnp.arange(B * D, dtype=jnp.float32).reshape(B, D)
    x2d_sharded = jax.device_put(x2d, NamedSharding(mesh_2d, PartitionSpec("dp", "mp")))

    n_shards_2d = len(x2d_sharded.addressable_shards)
    shard0 = x2d_sharded.addressable_shards[0].data
    print(f"2-D mesh: {n_shards_2d} shards; each shape = {shard0.shape}")
    s.assert_close(
        "shard_batch_equals_global_over_dp",
        actual=float(shard0.shape[0]),
        expected=B / 2,
        rtol=1e-9,
    )
    s.assert_close(
        "shard_features_equal_global_over_mp",
        actual=float(shard0.shape[1]),
        expected=D / 4,
        rtol=1e-9,
    )
else:
    s.skip("shard_batch_equals_global_over_dp",       "JAX not available")
    s.skip("shard_features_equal_global_over_mp",     "JAX not available")

Per-device shard sizes across strategies#

Three layouts on the same 8-device fabric: fully replicated (every device holds the full tensor), 1-D model-parallel (rows split across 8), and 2-D data x model (batch split across 2, features split across 4). Plot the per-device element count; the picture makes the memory savings of sharding concrete.

import matplotlib.pyplot as plt

if have_jax and mesh is not None:
    n_dev = len(jax.devices())
    B, D = 8, 1024
    total = B * D

    layouts = {
        "replicated":         [total] * n_dev,
        "1-D model-parallel": [sh.data.size for sh in W_global.addressable_shards],
        "2-D (dp x mp)":      [sh.data.size for sh in x2d_sharded.addressable_shards],
    }

    fig, ax = plt.subplots(figsize=(7.5, 3.4))
    width = 0.26
    xs = list(range(n_dev))
    for i, (name, sizes) in enumerate(layouts.items()):
        offs = (i - 1) * width
        # Pad / truncate so bars align to n_dev positions.
        sizes = (sizes + [0] * n_dev)[:n_dev]
        ax.bar([x + offs for x in xs], sizes, width, label=name)
    ax.set_xticks(xs)
    ax.set_xticklabels([f"dev {i}" for i in xs])
    ax.set_ylabel("elements per device")
    ax.set_title(f"shard size per device on a {n_dev}-device mesh")
    ax.legend()
    ax.grid(True, axis="y", alpha=0.3)
    fig.tight_layout()
    plt.show()
else:
    print("skipped - JAX not available.")

Exercises#

Replace the matmul with a full Transformer block and declare PartitionSpecs for QKV projections and the MLP. This is the core pattern of a tensor-parallel inference server in JAX.
jax.jit the whole forward. The XLA compiler reports how many collectives it inserted; check it matches the minimum for the sharding you specified.
Compare with pjit. The newer shard_map API gives explicit collective control; try expressing the matmul via shard_map and see where the user has to insert jax.lax.psum manually.

References#

JAX distributed-arrays tutorial.
GSPMD paper §3 for the axis-partition programming model.
JAX’s LLaMA training example as a production reference.

s.summary()
s.save()