Chapter 3: Distributed Computing with JAX

Course: "Advanced JAX: Performance, Optimization, and Scale", Chapter 3 · Original: https://apxml.com/courses/advanced-jax/chapter-3-distributed-computing-jax Environment: Python 3.14 · JAX / jaxlib 0.12.0.1 · CPU only (8 devices simulated with XLA_FLAGS; the multi-host demo uses 2 subprocesses × 4 devices)

Learning goals

1. Tell apart data parallelism, model (tensor) parallelism, pipeline parallelism, and ZeRO / FSDP, along with the memory and communication cost of each;
2. Understand device objects, array placement, and the rules for committed vs. uncommitted arrays;
3. Read example code written with pmap, and rewrite each example in the style 0.12 recommends: jit + NamedSharding or jax.shard_map;
4. Know exactly what psum / pmean / pmax / all_gather / psum_scatter / all_to_all / ppermute do and how much data each one moves;
5. Use PartitionSpec, AxisType, with_sharding_constraint and reshard on a 2D mesh (data × model), and inspect sharding types with jax.typeof;
6. Understand the multi-host programming model, and actually run a 2-process demo on this machine (a cross-process psum and data-parallel training);
7. Write distributed data-parallel training that matches the single-device result bit for bit, and read the gradient-sync communication volume from the compiled output.

Sections

How to read this chapter

• The example code follows the structure of each course page. Variable names, shapes, hyperparameters, in_axes / axis_name and so on all match the course.
• All code uses the style 0.12 recommends (jax.random.key, jax.tree.map, x.device, etc.). Outdated or broken patterns go into # ❌ Old style cells, together with what they output, followed right away by the recommended version.
• The course uses jax.pmap. In 0.12 that is a legacy API: the first sentence of its docstring is "Old way of doing parallel map. Use jax.shard_map instead." (Section 3.3 checks this with an assert). So every major pmap example is followed by a # ✅ Modern style cell that asserts both give the same result.

Links to the Basics part (cross-references only; not repeated here)

• Basics 05: Multi-device parallelism with pmap already covers: basic pmap usage, in_axes/out_axes, the fact that pmap is implemented with shard_map in 0.12, make_mesh using Explicit axes by default, vma (varying/invariant) types, a hand-written ring all-reduce with ppermute, and the "gradients come out N times too large inside shard_map" pitfall. This chapter assumes you know all of that and builds on it with 2D partitioning, all_to_all/psum_scatter, FSDP, explicit sharding types, multi-host, and communication-volume analysis.
• Basics 02: JIT compilation: tracing and compilation with jit. Every "modern style" in this chapter is built on jit.
• Basics 03: Automatic differentiation with grad: Section 3.7 uses one property: the transpose (backward pass) of all_gather is psum_scatter.
• Basics 04: Automatic vectorization with vmap: vmap also accepts axis_name, so you can test collectives on a single device (extra in Section 3.6).

⚠️ Requirement: the first code cell must set XLA_FLAGS before import jax, so please restart the kernel and run from the top. If jax has already been imported in this kernel, the device count will be 1 and the assert in the first cell will fail.

3.1 Introduction to parallelism¶

Modern machine learning often outgrows a single device: either the dataset is huge, or the model has billions of parameters, so the work has to be spread over several devices. Parallel processing means cutting a task into smaller, independent or semi-independent pieces and running them on several processing units at the same time. Which strategy you pick for training and inference depends on whether the bottleneck is data volume or model size.

Data parallelism¶

This is the most common way to speed up training. The course breaks it into five steps:

1. Replicate the model: copy the same model structure and current parameters to every device (GPU or TPU core);
2. Shard the data: split the global batch into smaller batches, one per device, with no overlap;
3. Parallel forward and backward pass: each device uses its own model copy and mini-batch to compute local gradients;
4. Aggregate gradients: combine the local gradients by averaging (pmean) or summing (psum). The result equals the gradient computed on the global batch;
5. Update parameters: every device applies the same aggregated gradient, so all replicas stay in sync.

Why step 4 works can be written as a formula. Split the global batch $B$ evenly into $N$ parts $B_1,\dots,B_N$, each of size $|B|/N$, and take the loss as the batch mean $L_B(\theta)=\frac{1}{|B|}\sum_{x\in B}\ell(x;\theta)$. Then

$$ \nabla_\theta L_B(\theta) \;=\; \frac{1}{N}\sum_{i=1}^{N} \nabla_\theta L_{B_i}(\theta), $$

In other words, each device computes the gradient of its local mean loss, and then the devices average them (pmean), which gives the global gradient. This effectively increases the total batch per step, which usually uses the accelerators better. Data parallelism fits best when "the model fits on one device and the dataset is large". The course uses pmap to implement this kind of SPMD data parallelism.

Model parallelism¶

If the parameters, activations or intermediate state don't fit on a single accelerator, you have to split the model itself across devices. Data flows through the pieces in order: the forward pass sends activations between devices and the backward pass sends gradients, which means a lot of communication. This lets you train very large models, but partitioning and managing the communication is fairly complex. Two common approaches are tensor parallelism (splitting individual weight matrices) and pipeline parallelism.

Pipeline parallelism¶

This is a more refined form of model parallelism whose goal is to keep devices busy. In naive model parallelism, a batch flows through the devices one after another, so at any moment most devices are idle. Pipeline parallelism works like this:

1. Stages: split the model into several sequential stages and give each stage to one device (or a group of devices);
2. Micro-batches: split the input batch further into smaller micro-batches;
3. Pipelined execution: once device 1 finishes stage 1 of micro-batch 1 and moves on to micro-batch 2, device 2 can start stage 2 of micro-batch 1. That way more devices work at the same time.

It reduces the idle "bubbles" of naive model parallelism, but dependencies, scheduling and state (such as the activations the backward pass needs) all have to be managed carefully.

Focus of this chapter¶

This chapter is mostly about data parallelism: pmap maps a function over several devices, distributes the data automatically, and provides collectives such as gradient aggregation. Model parallelism and pipeline parallelism are only background here; Chapter 6 of the course (very large models) comes back to them.

🔍 Extra: memory and communication budgets for the three kinds of parallelism¶

🟢 Must-read

The above is qualitative. When you actually choose a strategy, you need two numbers: how much memory each device must hold, and how much data has to be communicated per step. Let the model have $\Phi$ parameters and let $N$ be the number of devices.

Memory: with mixed-precision Adam, each parameter takes about 16 bytes: 2 bytes for the fp16 parameter, 2 bytes for the fp16 gradient, plus 4 bytes each for the fp32 master parameter, first moment and second moment.

Communication: with ring algorithms, each device sends/receives the following ($S$ is the tensor size in bytes):

$$ \text{all-reduce}: 2\tfrac{N-1}{N}S,\qquad \text{reduce-scatter}: \tfrac{N-1}{N}S,\qquad \text{all-gather}: \tfrac{N-1}{N}S . $$

One ring all-reduce is exactly one reduce-scatter plus one all-gather (Section 3.5 checks this with an assert). So when ZeRO-1/2 replace DP's "all-reduce the gradients" with "reduce-scatter the gradients + all-gather the updated parameters", the communication volume stays the same, while the state each device keeps around shrinks a lot.

ZeRO-3 / FSDP shard the parameters too, so the long-lived parameters, gradients and optimizer state all drop to $1/N$. But to avoid holding the full parameters between the forward and backward pass, the backward pass usually all-gathers them again, so communication is about 1.5× that of DP ($3\tfrac{N-1}{N}S$ vs. $2\tfrac{N-1}{N}S$). If the full parameters gathered in the forward pass are kept until the backward pass (as the simplified implementation in Section 3.7 does), communication is the same as DP, at the cost of temporarily holding the full parameters during compute. That is the core idea of ZeRO / FSDP. Section 3.7 actually implements it and checks it with memory_analysis.

Pipeline bubble: for a GPipe-style schedule with $p$ stages and $m$ micro-batches, the idle fraction is

$$ \text{bubble} = \frac{p-1}{m+p-1}. $$

The more micro-batches, the smaller the bubble. So the key to pipeline parallelism is making the number of micro-batches $m \gg p$.

3.2 Device management in JAX¶

Before using pmap, you need to know how JAX finds and manages accelerators. JAX automatically detects the CPUs, GPUs and TPUs in the system (or the ones a distributed environment assigns to it), and provides tools to inspect and work with them.

Listing available devices¶

• jax.devices(): all devices JAX can use. In a multi-host setup this is the global device list across all processes; in a single process it is just the local devices;
• jax.device_count(): the global device count;
• jax.local_devices() / jax.local_device_count(): the devices directly attached to the current Python process, and how many there are;
• Device object attributes: platform ('cpu', 'gpu', 'tpu'), id (unique within a platform), process_index (which process it belongs to), device_kind.

The course specifically points out that in a multi-host setup (such as a TPU Pod) you need to distinguish jax.devices() (global view) from jax.local_devices() (only this process's devices). On a single-host GPU or CPU setup, the two are usually the same. Section 3.8 shows the difference in a real 2-process setup.

Default device placement¶

By default, arrays are created and computed on the first device in jax.local_devices() (cpu:0, gpu:0 or tpu:0). That is convenient for a single device; in distributed settings you usually want explicit control.

⚠️ Note: device is an attribute, so write x.device. The x.device() you often see in older material raises TypeError: 'Device' object is not callable. There are two other ways to check: x.devices() (a method that returns a set) and x.sharding.

Explicit placement with jax.device_put¶

The signature is jax.device_put(x, device=None). x is the data to place (a NumPy array, Python scalar, list or pytree); the second argument is the target device (taken from jax.devices() or jax.local_devices()), and None means the default device. In 0.12 the second argument can also be a Sharding, which is the main way to place data on multiple devices; it is used a lot from Section 3.4 on.

The course lists three uses of device_put:

1. Data transfer: it is the main way to move data from host memory to accelerator memory;
2. Preparing for parallel compute: prepare shards on the host, and explicitly place initial parameters when needed;
3. Controlling initial placement: to avoid unnecessary transfers later.

⚠️ Note: it is easy to think "where computation runs is decided by jit; device_put only affects where data starts out". In fact, the device a jit function runs on is decided by where its input arrays live: "computation follows data". If you device_put the input to cpu:1, the jit computation and its result are both on cpu:1. The cell in the committed extra below, and the ✅ cell after the jax.jit(device=...) note in Section 3.7, both check this with asserts. So device_put does more than decide where data starts out; it also indirectly decides where later computation runs.

Device data residency¶

Once an array is on a device, JAX tries to keep the computations that use it on that device too, to reduce transfers. So if the arrays in an operation live on different devices, does JAX move data automatically, or raise an error? The rule is actually deterministic; see the extra below.

🔍 Extra: committed vs. uncommitted arrays¶

🟢 Must-read

Basics 01 §1.6 "Cross-device operations" already covered this rule; here we review it in the multi-device setting. Whether JAX "moves the data" or "raises an error" depends on whether the arrays are committed (bound to a device):

• Arrays created without a device, like jnp.ones(...), are uncommitted; they automatically follow the other inputs;
• jax.device_put(x, device_or_sharding) gives you a committed array;
• If the inputs of an operation have only one committed device, the computation runs on that device and uncommitted inputs are moved there;
• If two inputs are committed to different devices, you get ValueError: Received incompatible devices right away. JAX will not move the data for you.

Environment variables for device visibility¶

Before importing JAX, you can use environment variables to control which devices it sees:

• CUDA_VISIBLE_DEVICES=1 python my_jax_script.py: inside the script, jax.local_devices() lists just 1 GPU device, with id 0. Visible devices are renumbered from 0, so this id=0 is physically GPU 1;
• CUDA_VISIBLE_DEVICES="" usually hides all GPUs and forces the CPU;
• JAX_PLATFORMS=cpu (or gpu, tpu): only initialize the backend for that platform. Use it to force one when several backends are available.

To sum up: device management is the foundation of distributed computing. You need to know which devices are available, where data lives (x.device), and how to control placement (jax.device_put). These are prerequisites for using pmap.

🔍 Extra: two ways to simulate multiple devices, and "must be set before initialization"¶

⚪ Optional: read this if you want to simulate multiple devices on your machine, or if you see "device count is 1"

There is no GPU on this machine, so this chapter uses XLA_FLAGS=--xla_force_host_platform_device_count=8 to simulate 8 CPU devices. 0.12 also has a config option jax_num_cpu_devices (used in Basics 01 §1.6) that does the same thing. But like the environment variable, it must be set before the backend is initialized (before any JAX operation runs for the first time), or you get a RuntimeError. The backend in this kernel is already initialized, so the demo below uses a subprocess.

3.3 Single Program, Multiple Data (SPMD) with pmap¶

SPMD (Single Program, Multiple Data) is a very common parallel computing model, and it suits GPUs and TPUs especially well: the same program runs on many processors at once, and each instance handles a different part of the data. The alternative is MPMD, where different processors run different programs.

In the course, the main tool for SPMD is jax.pmap. It turns a function written for one device into one that runs in parallel on many devices, handling replication and data sharding automatically. Think of Python's built-in map: pmap "maps the function over the devices in parallel". Each device runs the same compiled function, but gets its own slice of the input.

How pmap implements SPMD (with 4 TPU cores as an example)¶

1. Write single-device code: write the function as if for one device, e.g. one network layer or one loss computation;
2. Wrap it in jax.pmap;
3. Arrange the data shards: the leading axis of the input corresponds to devices. For example, data of shape (128, 50) (batch size 128, 50 features) on 4 devices must be reshaped to (4, 32, 50); the 4 in dimension 0 is the device axis;
4. Run: JAX compiles the original function with XLA when needed, copies the compiled result to the 4 devices, and sends each device its slice (device 0 gets slice 0, and so on). All devices compute their slices at the same time, and the results are collected and stacked along a new leading axis.

So the key point of pmap is: you only write single-device logic; pmap maps the input's leading axis onto devices and handles replication and parallelism. Underneath, XLA optimizes for the target hardware.

🔍 Extra: three "views" of the same SPMD program¶

🟢 Must-read

All three styles compile to an SPMD program. The difference is the shape the function body sees:

Below we print jax.typeof during tracing. On an Explicit mesh, the jit view's type shows @data directly; the shard_map view shows {V:data}, meaning the value differs across devices (varying); the pmap view only sees a single slice.

3.4 Data parallelism with pmap¶

The basic recipe for data parallelism: replicate the model on several devices, and have each device process a different slice of the input batch with the same parameters. The course uses a simple prediction function to show how pmap does both:

• Model replication: the pmapped function (usually the forward pass or the whole training step) runs on every device. Parameters are broadcast with in_axes=None, so each device gets an identical copy;
• Data sharding: an input whose leading axis length equals the number of devices is split along that axis; the first slice goes to the first device, and so on.

What in_axes means: None means broadcast, the same value goes to every device, which is typical for model parameters; 0 means map over axis 0, so the data is split along axis 0, which is how input data is usually handled in data parallelism.

⚠️ Note: pmap's in_axes defaults to 0, and arguments without a mapped leading axis are not broadcast automatically. You must write in_axes=(None, 0) explicitly, otherwise the parameters will also be split along axis 0 (the ❌ cell below fails outright). The example here writes in_axes=(None, 0), so it runs; the ❌ cell in the gradient example of Section 3.6 leaves it out, so it doesn't.

⚠️ Note: jax.random.PRNGKey is the legacy API; 0.12 recommends jax.random.key. Also, JAX's convention is use each key only once: don't use the same key(0) to generate w, b and the data; split it first. The code below does exactly that. Basic key usage is in Basics 01, and a deeper discussion is in Advanced 10.

To sum up the data flow: the global batch is split across devices, the parameters are copied to every device, each device computes independently, and the outputs are stacked along the device axis. sharded_predictions[i] is what device i computed from shard sharded_data[i] and the replicated parameters.

In real training, predict is just one part of a bigger train_step that also computes the loss and gradients. Before updating the parameters, the gradients from all devices have to be combined across devices (usually averaged). That needs the collective primitives in the next section.

🔍 Extra: three practical points for data parallelism¶

⚪ Optional: for reference when you actually build a data-parallel training loop, tune the learning rate, and deal with the last partial batch

1. Keep parameters resident on device. In a training loop, device_put the parameters as replicated (P()) once. Each step outputs parameters with the same sharding, and you just pass them back in for the next step. With pmap, the common approach is to give the parameters a device axis and pass them with in_axes=0 (see Section 3.9), so they stay on the devices.
2. Global batch and learning rate. When the device count goes from 1 to $N$ with the per-device batch size fixed, the global batch grows $N$ times and gradient noise drops. A common rule of thumb is to scale the learning rate linearly ($\eta \to N\eta$) together with warmup. For very large batches linear scaling often stops working, and people usually switch to layer-wise adaptive optimizers like LARS / LAMB. This doesn't change the conclusion that "after pmean you get the global gradient"; it only affects convergence-related hyperparameters.
3. The sharded dimension must divide evenly. P('data') requires axis 0 to be divisible by the size of the mesh axis. The last incomplete batch is either dropped, or padded and used with a mask (Basics 05, extra exercise 4, covers how to average correctly over batches of unequal length).

3.5 Collective communication primitives (psum, pmean, etc.)¶

Inside pmap, each device runs the same program on its own slice of the data. But distributed algorithms can't do without communication and synchronization between devices. In data-parallel training, for example, once each device has computed gradients on its local batch, those gradients must be combined across all devices (usually averaged) before the parameters can be updated. That step relies on collective operations.

The collectives in jax.lax let arrays spread across devices (along the pmap's mapped axis) take part in one computation together, and the result is usually replicated back to every participating device. Collectives must be called inside a transformation that binds the corresponding axis name. Besides pmap, in 0.12 that can also be shard_map or a vmap with axis_name (the extra in Section 3.6 demos vmap).

Use axis_name to pick the communication axis: a collective needs axis_name to say which axis to communicate along, and the name must match the one given by the enclosing pmap(..., axis_name=...). Explicit names avoid ambiguity, which matters especially for multi-dimensional parallelism.

Core reductions: psum, pmean, pmax, pmin¶

Suppose there are $N$ devices along axis 'i', and device $j$ holds $x_j$. The formulas below are all element-wise, and every device gets the same result:

$$ \texttt{psum}: y = \sum_{j=0}^{N-1} x_j,\qquad \texttt{pmean}: y = \frac{1}{N}\sum_{j} x_j,\qquad \texttt{pmax}: y = \max_j x_j,\qquad \texttt{pmin}: y = \min_j x_j . $$

• psum: commonly used to sum data-parallel gradients before the optimizer update. For example, if 4 devices hold 2, 3, 1 and 4, after psum every device gets $2+3+1+4=10$;
• pmean: the standard way to average gradients in data parallelism;
• pmax / pmin: useful for syncing flags (e.g. "did any device hit an error?") or finding the largest loss in the whole batch.

⚠️ Note: DeviceArray was removed long ago. Example outputs in older material often look like DeviceArray([6., 6., 6., 6.], dtype=float32). Today every array is a unified jax.Array and prints as Array([...], dtype=float32).

⚠️ Note: it's easy to assume that pmean is more efficient than psum followed by a manual division by the device count, because the backend can optimize it specially. In fact pmean is psum plus one division. The jaxpr below only contains psum_invariant and div; there is no dedicated pmean primitive, and both versions compile to exactly the same communication. pmean is recommended only because it reads better and you won't forget to divide by $N$.

Other collectives¶

Three more are commonly used:

• jax.lax.all_gather: collects the inputs from all devices and concatenates them along the mapped axis, so every device gets the full array. Keep in mind that the output size grows linearly with the device count, so it can take a lot of memory;
• jax.lax.ppermute: exchanges data according to the (source, destination) pairs in the perm argument. It is more general and is the building block for ring all-reduce and model-parallel communication, but the indices need careful handling (Basics 05, extra exercise 2, already used it to write a ring all-reduce by hand);
• jax.lax.axis_index(axis_name): strictly speaking, not a communication primitive. It returns the integer index of the current device along the given axis, so devices at different positions can run different logic.

Overall, collective primitives are the basic building blocks for coordinating devices inside pmap: gradient aggregation, sync points and data exchange all rely on them. JAX and XLA use NVLink or inter-chip interconnect (ICI) on GPUs/TPUs to optimize execution, but communication always has a cost.

🔍 Extra: exact semantics of the four "rearranging" collectives¶

🟢 Must-read

psum_scatter and all_to_all are at the heart of ZeRO/FSDP, MoE and sequence parallelism. Suppose there are $N$ devices on the axis, device $j$ holds $x_j$, and $x_j$ is split evenly along some dimension into $N$ blocks $x_j^{(0)},\dots,x_j^{(N-1)}$:

There is also an identity: all-reduce = reduce-scatter + all-gather. Below we verify each one on 8 devices (an 8×8 matrix split by rows, device i holds row i).

🔍 Extra: halo exchange with ppermute (stencil computations)¶

⚪ Optional: for domain decomposition in stencil computations such as convolutions or finite differences

Basics 05 already used ppermute to write a ring all-reduce. Here is another equally common case: boundary exchange in stencil computations, such as 1D convolution or finite differences. Split an array of length 64 into 8 segments and compute $y_k = x_{k-1} + x_k + x_{k+1}$ (zero padding at both ends). Besides its own segment, each device needs the last element of its left neighbor and the first element of its right neighbor.

The key is that ppermute's perm doesn't have to be a full permutation: a device that receives nothing gets all zeros, which is exactly the zero-padding boundary condition.

🔍 Extra: values after all_gather are still varying¶

⚪ Optional: read this when out_specs=P() in shard_map raises a replication error

shard_map defaults to check_vma=True and tracks whether each value differs across devices (Basics 05 covered this). The results of psum / pmean / pmax are invariant (identical on every device), so you can output a single copy with out_specs=P(). The trap: in 0.12's type system, the result of all_gather is still marked varying ({V:i}), so writing out_specs=P() right after an all_gather raises an error. There are two ways out: keep out_specs=P('i') (from outside you see N copies concatenated); or, if you really need a replicated output, switch to jit + sharding and let the compiler handle it.

3.6 Handling axis names in pmap¶

When you write a function to parallelize with pmap, especially one that uses collectives like lax.psum or lax.pmean, you have to say clearly which group of devices the communication happens across. You do that by naming the axis pmap maps over with axis_name, and then referring to that name in the collective.

⚠️ Note: axis_name is a required argument of psum. It's easy to think that in simple cases you can leave it out and psum will implicitly sum over the mapped axis; in fact leaving it out raises TypeError: psum() missing 1 required positional argument: 'axis_name' right away. JAX does not "infer which axis a collective runs along".

The course lists three benefits of named axes:

• Clarity: lax.psum(..., axis_name='batch') tells you at a glance that the sum runs over the batch dimension;
• Fits multi-dimensional parallelism: when there are several levels of parallelism, each level has its own name, and the collective uses the name to say which level it acts on. For example, with two levels 'data_shards' and 'model_replicas': a collective along 'data_shards' happens between devices that "hold different data shards but the same model replica"; one along 'model_replicas' happens between devices that "hold different model replicas but the same data shard". In 0.12 this corresponds to the two axes of a 2D mesh (see Section 3.7);
• Composability: a function that reduces along a 'spatial' axis can be reused in any parallel context that maps that name.

The second example is gradient averaging in data parallelism: each device computes gradients on its local shard, and before updating the parameters uses pmean to average across all devices.

⚠️ Note: in this example both leaves of params are scalars (shape ()). If you call pmap(parallel_update_step, axis_name='data_parallel_axis') without in_axes, it uses the default in_axes=0, and a scalar has no axis 0 to map over, so you get IndexError: tuple index out of range. Write in_axes=(None, 0) instead. Also, when taking device 0's parameters at the end, you can't use jax.tree_map: it has been removed in 0.12 (AttributeError). Use jax.tree.map.

So writing axis_name explicitly is more than a recommendation: as shown above, collectives can't leave out the axis name at all. Clear axis names make distributed code easier to read, less prone to subtle bugs, and easier to maintain and compose.

🔍 Extra: where axis names "live" in 0.12¶

⚪ Optional: when you need reductions over several axes, or want to test communication logic on a single device with vmap

1. Axis names live on the mesh. jax.make_mesh((4, 2), ('data', 'model')) defines the axis names; PartitionSpec uses those names to describe sharding; collectives inside shard_map refer to them too. pmap's axis_name is just a "one-dimensional, temporary" mesh axis.
2. You can reduce over several axes at once: lax.psum(x, ('data', 'model')) sums over the whole 2D mesh; lax.axis_index('data') returns this device's coordinate along that axis.
3. Test collective logic on a single device: vmap also accepts axis_name, and psum / all_gather / ppermute are all defined under vmap, with the batch dimension playing the role of "devices". When writing a complex communication pattern, you can unit-test it on the CPU with vmap first, and then switch to shard_map.

3.7 Nested pmap and advanced partitioning¶

A single pmap is good at spreading work over a one-dimensional row of devices. More complex parallelism (such as data parallelism + model parallelism) needs a multi-dimensional grid of accelerators. The course's approach is nested pmap: take 4 GPUs arranged in a 2×2 grid, with rows doing data parallelism and columns doing model parallelism; a function under the outer pmap(axis_name='data') calls a function under the inner pmap(axis_name='model').

Three things to watch for when nesting:

1. Axis scope: an axis name defined by the inner pmap is only valid within the inner call and the devices it involves; the outer axis_name refers to collectives between the outer devices;
2. Data mapping: the input shape is (outer size, inner size, ...); axis 0 corresponds to the outer 'data' axis and axis 1 to the inner 'model' axis;
3. Collectives must specify axis_name, which decides which group of devices communicates.

⚠️ Note: this nested-pmap code fails in several places in JAX 0.12:

• from jax.experimental import maps: the module has been removed (ModuleNotFoundError);
• jnp.ones(16.): shapes must be integers, so this raises TypeError; write jnp.ones(16);
• jax.device_put(...).devices: devices is a method, and without parentheses you get a bound method object. The variable name device_mesh is also misleading: the code never builds a mesh. jnp.arange(num_devices).reshape(2, 2) only works with exactly 4 devices and fails with 8. The fallback branch jax.local_devices()[:1] * 4 repeats the same device 4 times, which can't be used for parallelism either;
• The most fundamental problem: nested pmap raises ValueError outright in 0.12 (pmap is now implemented on top of shard_map, and the inner mesh conflicts with the outer one). Use a 2D mesh + shard_map instead.

⚠️ Note: for the simple example (data_simple = arange(4).reshape(2,2), params = 0.5), it's easy to compute output_simple[0] as 12.0, reasoning that "the outer psum(1, 1, 5, 5) = 12". Running the equivalent 2D mesh + shard_map implementation above actually gives [6., 6.]. The reason is that the outer psum(..., 'data') only sums along the 'data' axis, which involves just the 2 devices in the same column, not all 4:

$$ \text{inner}_{i} = 2\sum_{j} x_{ij}p_{ij} = \begin{cases}2(0+0.5)=1 & i=0\\ 2(1+1.5)=5 & i=1\end{cases}, \qquad \text{outer} = \sum_i \text{inner}_i = 1+5 = 6 . $$

The "psum(1, 1, 5, 5)" reasoning sums the inner result once more along both axes. But the inner result is already replicated along 'model' (both devices in a row hold 1, or both hold 5), so summing along 'model' again counts the same value twice. 0.12's vma type check rejects this outright (verified above). The 4D example works the same way: the correct result is output[i, j] = data.sum(axis=(0, 1)), cross-checked against NumPy above.

Advanced partitioning strategies¶

pmap still follows SPMD, and the partitioning is decided implicitly by how inputs are sharded along the mapped axes. Two common patterns:

• Data parallelism + model parallelism: psum the gradients on the outer axis, all_gather the activations on the inner axis;
• 2D/3D tensor parallelism: split a large matrix multiplication over a 2D or 3D device grid.

Picture a 2×2 grid: the outer pmap maps over the "data" axis (rows), the inner one over the "model" axis (columns).

To assign devices by hand, you can pre-place the inputs with jax.device_put, or restrict the devices with pmap(devices=...). For more complex cases, people used to use jax.jit(device=...), jax.experimental.shard_map, or frameworks built on jax.Array.

⚠️ Note: jax.jit(device=...) is deprecated in 0.12 and emits a DeprecationWarning; the recommendation is to device_put the inputs instead. jax.experimental.shard_map has been deprecated since v0.8.0; use jax.shard_map.

🔍 Extra: combining a 2D mesh with PartitionSpec¶

🟢 Must-read

Here is the full toolkit for multi-dimensional partitioning in 0.12. First arrange the 8 devices into a (data=4, model=2) mesh. Entry k of a PartitionSpec says which mesh axis (or axes) dimension k of the array is split along:

jax.typeof(x) shows the sharding type, e.g. float32[8@data,4@model]; jax.debug.visualize_array_sharding draws the layout.

🔍 Extra: AxisType — Explicit / Auto / Manual¶

🟢 Must-read

In 0.12, jax.make_mesh creates AxisType.Explicit axes by default (Basics 05 mentioned this). How the three axis types differ:

The upside of Explicit mode is that sharding is fixed at trace time, so what you see is what you get. The cost is that you must be explicit when there's ambiguity. The classic example is a matmul where the contracting dimension is sharded on both sides: should the output be all-reduced to replicated, or reduce-scattered to sharded? JAX raises ShardingTypeError and asks you to say so with out_sharding=.

🔍 Extra: with_sharding_constraint and reshard¶

⚪ Optional: when you need to pin or change the sharding of an intermediate value inside jit

jax.lax.with_sharding_constraint(x, sharding) constrains the sharding of an intermediate value inside jit. It behaves differently under the two axis types:

• Auto axes: it is a hint to the compiler. The compiler redistributes the data into the given sharding at that point, inserting communication if needed. It's commonly used to pin the sharding of activations so the compiler doesn't pick an inefficient plan;
• Explicit axes: it becomes an assertion. If the current sharding doesn't match the given one, you get an AssertionError that suggests jax.sharding.reshard. In Explicit mode, the right tool for changing sharding is reshard.

🔍 Extra: data × model 2D parallelism (a Megatron-style MLP block)¶

⚪ Optional: for tensor parallelism; full training is in Advanced 06 §6.8

The classic way to split for tensor parallelism comes from Megatron-LM. In a two-layer MLP $y = \mathrm{relu}(xW_1)W_2$, split $W_1$ by columns (P(None, 'model')) and $W_2$ by rows (P('model', None)). Then: