--- jupytext: formats: md:myst text_representation: extension: .md format_name: myst format_version: 0.13 jupytext_version: 1.16.4 kernelspec: display_name: Python 3 language: python name: python3 --- (jax-301-sharding-ad)= # Autodiff and sharding In explicit sharding mode ({ref}`jax-201-sharding`), shardings are part of JAX types: `jax.typeof(x)` might print `float32[8@X,4]`, meaning the leading axis is sharded along mesh axis `X`. This page is about what happens when you differentiate such a program. The main idea is that you control how your backward pass is sharded through local reasoning about your forward pass. The rule behind this is that *cotangent types are a function of primal types*: the type of each gradient value in the backward pass (shape, dtype, and sharding) is determined by the type of the corresponding value in the forward pass. Once you know your forward-pass types, you know your backward-pass types, and hence where backward-pass communication happens. There's a second theme running through this page. If every axis of every array were plain *sharded*, autodiff would need nothing new: sharded cotangents for sharded primals, no communication anywhere. *Replication*, meaning a full copy of an array on several devices, is what requires something new. Replication's dual under transposition is *reduction*, i.e. a cross-device sum, and someone has to decide where that sum happens. Demanding efficient autodiff plus local reasoning in the presence of replication is what leads to the two new sharding states introduced below, `unreduced` and `reduced`. Most of this page works in explicit mode; the final section shows the same thing in manual mode, inside {func}`jax.shard_map`. We'll use two CPU devices and a one-axis mesh throughout: ```{code-cell} import jax import jax.numpy as jnp jax.config.update('jax_num_cpu_devices', 2) jax.set_mesh(jax.make_mesh((2,), ('X',))) # explicit mode by default ``` ## Cotangent shardings are a function of primal shardings Here's a data-parallel loss: the batch `x` is sharded along `X`, and the weights `w` are replicated (a full copy on every device): ```{code-cell} x = jax.device_put(jnp.arange(8 * 4.).reshape(8, 4), jax.P('X', None)) w = jax.device_put(jnp.arange(4 * 2.).reshape(4, 2) / 10., jax.P(None, None)) def loss(w, x): return jnp.sum((x @ w) ** 2) dw, dx = jax.grad(loss, argnums=(0, 1))(w, x) print(jax.typeof(w), '->', jax.typeof(dw)) print(jax.typeof(x), '->', jax.typeof(dx)) ``` The gradient with respect to a sharded input is sharded the same way, and the gradient with respect to a replicated input is replicated. This isn't just true of the top-level inputs and outputs of `jax.grad`; it holds for every intermediate in the backward pass. We can see it in the jaxpr, where each cotangent value's sharding matches its primal's: ```{code-cell} print(jax.jit(jax.grad(loss)).trace(w, x).jaxpr) ``` The final two equations compute `dw`: a `dot_general` of two values sharded along `@X` producing a replicated (unsharded) result, then a transpose. We'll come back to what that dot costs. There are two reasons JAX insists that cotangent shardings are determined by primal shardings: 1. **User control.** The goal of explicit mode is that user-written code determines all the shardings in the computation, in an easy-to-predict, local way. The backward pass is part of the computation, but it's written by autodiff rather than by you. Making cotangent shardings a function of primal shardings means your forward-pass sharding decisions *are* your backward-pass sharding decisions. (Compiler-based automatic sharding mode has no analogous guarantee: there, backward-pass shardings can be chosen by the compiler, unrelated to the corresponding primal shardings.) 2. **Ruling out ambiguities.** If a variable is used more than once in the forward pass (fan-out), autodiff generates an addition of cotangents in the backward pass. If cotangent shardings could be unrelated to primal shardings, the two summands might have *different* shardings, and the addition would require communication that nothing in your code specifies. (Similarly, the zero cotangents autodiff generates would have no determined sharding.) When cotangent shardings are a function of primal shardings, both summands automatically agree: they're cotangents for the same primal variable, so they have the same type. This is a special case of a general principle in JAX's autodiff: cotangent types are always a function of the corresponding primal types, meaning shapes, dtypes, and now shardings. That gives us three things: each op's backward rule has a clear type and knows exactly what kind of cotangent it will receive; a well-typed forward program always yields a well-typed backward program; and you can predict backward-pass types by looking only at your forward-pass code, without knowing what the rest of the program looks like. ## Where backward-pass communication comes from In the example above, look at the types in the backward pass of `x @ w`. The cotangent `dw` must be replicated, because `w` is. But the data it's built from is spread across devices: `dw` is (the transpose of) a contraction of two arrays sharded along `@X`, the batch axis. Producing a replicated result from a contraction over a sharded axis requires a cross-device sum, an AllReduce. The partitioner inserts it inside that final `dot_general`. This is the familiar gradient synchronization of data-parallel training, and notice that you can predict it purely locally: `w` is replicated, each device touches only part of the batch, so somewhere in the backward pass the per-device gradient contributions must be summed. The types tell you the communication exists, but in the jaxpr above it's *implicit*, hidden inside an op whose operands have a sharded contracting dimension. The rest of this page is about making that communication explicit: something you can see in the types, move around, and batch up. ## Unreduced: a reduction waiting to happen Consider a matmul whose *contracting* dimension is sharded: ```{code-cell} a = jax.device_put(jnp.arange(4.).reshape(2, 2), jax.P(None, 'X')) b = jax.device_put(jnp.arange(4., 8.).reshape(2, 2), jax.P('X', None)) print(jax.typeof(a)) print(jax.typeof(b)) ``` ```{code-cell} :tags: [raises-exception] a @ b ``` JAX makes us say what we want here, because there's a genuine choice. Look at the data each device holds: ```text a: P(None, 'X') b: P('X', None) (device k has column k) (device k has row k) [ 0 | 1 ] [ 4 5 ] @ --------- [ 2 | 3 ] [ 6 7 ] ``` Device 0 can multiply its column of `a` by its row of `b` without any communication, and likewise device 1. Each local matmul produces a full-shape *partial sum*, and the true answer is the elementwise sum of the two. One option is to finish the job with an AllReduce, by asking for an ordinary output sharding: ```{code-cell} c = jnp.einsum('ij,jk->ik', a, b, out_sharding=jax.P(None, None)) print(jax.typeof(c)) print(c) ``` The other option is to *stop before the reduction*: ```{code-cell} c = jnp.einsum('ij,jk->ik', a, b, out_sharding=jax.P(None, None, unreduced={'X'})) print(jax.typeof(c)) ``` The type `float32[2,2]{U:X}` reads: a 2×2 array, *unreduced* along mesh axis `X`. Each device along `X` holds a full-shape partial sum, and the array's true value is the sum of those pieces: ```{code-cell} for shard in c.addressable_shards: print(f'device {shard.device.id}:\n{shard.data}') ``` An unreduced array is a reduction waiting to happen. To perform it, reshard to an ordinary sharding, which runs the deferred AllReduce: ```{code-cell} print(jax.reshard(c, jax.P(None, None))) ``` Deferring lets you do more work first and pay for fewer, bigger collectives. Only *linear* operations make sense on unreduced arrays. Adding two arrays unreduced along the same axes is fine, since sums of partial sums are partial sums of sums: ```{code-cell} c2 = jnp.einsum('ij,jk->ik', 2. * a, b, out_sharding=jax.P(None, None, unreduced={'X'})) print(jax.typeof(c + c2)) ``` That lets you compute a sum of sharded matmuls (for example, a LoRA-style `x @ W + x @ A @ B`) with a single AllReduce at the end rather than one per matmul. Nonlinear operations, on the other hand, *can't* have a rule for unreduced inputs: the cosine of a sum is not the sum of the cosines, so applying `cos` to each device's partial sum would compute the wrong answer. Such ops raise an error instead: ```{code-cell} :tags: [raises-exception] jnp.cos(c) ``` One caveat: for straight-line code like `c + c2` above, XLA can often merge adjacent AllReduces on its own, so the compiled program may be equally good either way. The type-level guarantee matters when the compiler can't find the merge, above all across the iterations of a loop. The microbatch example below shows that case. ## Reduced: choosing unreduced gradients Now back to autodiff. We said cotangent shardings are a function of primal shardings, and we've seen two entries of that function: sharded primals get sharded cotangents, replicated primals get replicated cotangents. The replicated entry is exactly where backward-pass communication comes from: replication's transpose is a cross-device sum. With `CT(Replicated) = Replicated`, autodiff performs that sum *eagerly*: each backward-pass op that produces a cotangent for a replicated primal does its AllReduce on the spot, implicitly, like the `dw` dot we saw earlier. The natural, before-any-communication state of such a cotangent is a set of per-device partial sums, an *unreduced* array. Suppose we want autodiff to leave it that way and let us decide when to reduce. We need a way to ask for that in the forward pass, and it has to be a new *type*, since cotangent types are a function of primal types. That's what `reduced` is for: ```{code-cell} w_ = jax.reshard(w, jax.P(None, None, reduced={'X'})) print(jax.typeof(w_)) ``` An array that's *reduced* along `X`, written `{R:X}`, is physically identical to a replicated array: a full copy on every device along `X`. (The name is the past tense of "unreduced": it's the state an array is in after its reduction has happened.) In the forward pass it behaves exactly like a replicated array, and the reshard above moves no data. The only difference is how autodiff treats it. The complete cotangent map is: | primal type | cotangent type | |-------------------|-------------------| | sharded `@X` | sharded `@X` | | replicated | replicated | | reduced `{R:X}` | unreduced `{U:X}` | | unreduced `{U:X}` | reduced `{R:X}` | Use `Replicated` and you get replicated gradients; use `Reduced` and you get unreduced gradients. The reshard-to-reduced above is a communication-free cast in the forward pass, and under transposition it becomes a reshard-from-unreduced in the backward pass, which is the AllReduce. A free cast in your forward code pins down where the collective runs in your backward code, so backward-pass communication becomes something you can see and place while writing the forward pass. Compare the backward pass of our data-parallel loss with and without the cast: ```{code-cell} def loss2(w, x): w = jax.reshard(w, jax.P(None, None, reduced={'X'})) return jnp.sum((x @ w) ** 2) print(jax.jit(jax.grad(loss2)).trace(w, x).jaxpr) ``` Where the original jaxpr had a `dot_general` with an implicit AllReduce buried inside, this one shows the dot producing an explicit `f32[2,4]{U:X}` value, and a final `reshard` (the transposed cast) turning it into the replicated `dw`. Same math, same total communication, but now the reduction is a visible, movable object in the program. Once it's visible, you can move it. Suppose the weights are used twice, as with two heads, two microbatches, or a LoRA branch: ```{code-cell} def loss_fanout(w, x1, x2): w = jax.reshard(w, jax.P(None, None, reduced={'X'})) return jnp.sum(x1 @ w) + jnp.sum(x2 @ w) x1 = jax.device_put(jnp.ones((8, 4)), jax.P('X', None)) x2 = jax.device_put(jnp.ones((8, 4)), jax.P('X', None)) print(jax.jit(jax.grad(loss_fanout)).trace(w, x1, x2).jaxpr) ``` Both backward dots produce `{U:X}` contributions, the fan-out addition happens *unreduced* (addition is linear), and a single `reshard` performs one AllReduce for the whole gradient. Without the cast, each contribution would be reduced separately inside its own dot. Here the fusion is guaranteed by the types rather than left to compiler pattern-matching, which matters in the next example. ## Example: microbatch gradient accumulation The classic case is gradient accumulation. A realistic training step has *two* loops that the compiler can't see through: a scan over layers inside the model, and a scan over microbatches accumulating gradients, with one update at the end. The gradient AllReduce should happen once per step, but if the weights are replicated, every microbatch's backward pass synchronizes its own gradient contribution, inside both loops, and XLA cannot hoist collectives out of a loop for you. With `reduced` weights, the gradients come out unreduced, the accumulator stays unreduced across the whole scan, and you reduce once: ```{code-cell} def predict(stacked_ws, xs): # stacked_ws: [layer, features, features] def apply_layer(xs, w): return jnp.tanh(xs @ w), None final_xs, _ = jax.lax.scan(apply_layer, xs, stacked_ws) return final_xs def loss3(stacked_ws, batch): return jnp.sum(predict(stacked_ws, batch) ** 2) @jax.jit def step(stacked_ws, xs): # xs: [microbatch, batch@X, features] def microbatch_step(grad_acc, xs_mb): grads = jax.grad(loss3)(stacked_ws, xs_mb) # ws are reduced, so grads are unreduced -- and we can check it! assert jax.typeof(grads).sharding.spec.unreduced == {'X'} return grad_acc + grads, None grad_acc = jax.reshard(jnp.zeros_like(stacked_ws), jax.P(unreduced={'X'})) grad_acc, _ = jax.lax.scan(microbatch_step, grad_acc, xs) grads = jax.reshard(grad_acc, jax.P()) # the one AllReduce ws = jax.reshard(stacked_ws, jax.P()) # free: full copies already return ws - 0.01 * grads stacked_ws = jax.device_put(jnp.stack([jnp.eye(4) / 2] * 3), jax.P(reduced={'X'})) xs = jax.device_put(jnp.ones((5, 2, 4)), jax.P(None, 'X', None)) new_ws = step(stacked_ws, xs) print(jax.typeof(new_ws)) ``` Everything type-checks locally: the weights are `{R:X}`, so each microbatch's gradients come out `{U:X}` (even though they're computed by a scan over layers); unreduced arrays support addition, so the microbatch scan carry accumulates them; and one reshard to replicated is the step's single AllReduce. Notice the `assert` in the scan body: because shardings are part of JAX types, "the gradients are unreduced" is a property of a value that you can check with `jax.typeof`, inside traced code, at trace time. The updated weights come back replicated; casting them back to `{R:X}` for the next step is free. And the compiled program has exactly one AllReduce, outside both loops: ```{code-cell} import re def print_all_reduces(jitted, *args): hlo = jitted.lower(*args).compile().as_text() for line in hlo.splitlines(): if 'all-reduce(' in line or 'all-reduce-start(' in line: print(re.search(r'op_name="([^"]*)"', line).group(1)) print_all_reduces(step, stacked_ws, xs) ``` Compare the same step written with plain replicated weights: ```{code-cell} @jax.jit def step_replicated(stacked_ws, xs): def microbatch_step(grad_acc, xs_mb): grads = jax.grad(loss3)(stacked_ws, xs_mb) # replicated: AllReduce inside! return grad_acc + grads, None grad_acc = jnp.zeros_like(stacked_ws) grad_acc, _ = jax.lax.scan(microbatch_step, grad_acc, xs) return stacked_ws - 0.01 * grad_acc ws_replicated = jax.reshard(stacked_ws, jax.P()) print_all_reduces(step_replicated, ws_replicated, xs) ``` The op name `while/body/.../while/body` shows that this AllReduce sits inside the transposed layer scan, inside the microbatch scan: it runs once per layer per microbatch. Hoisting the gradient reduction out of both loops, from once per layer per microbatch to once per step, has produced large wins in production LLM training (in one case cutting the per-step time spent in gradient reduction several-fold). The compiler cannot make this transformation on its own, because it can't pattern-match collectives through a loop. ## Why this design? A few design questions come up often. **Why not make the cotangent of Replicated be Unreduced?** It would be coherent, but it takes options away. Lots of code returns replicated values (loss values, for one), and making their cotangents unreduced would introduce surprising communication requirements into existing programs. By keeping `Replicated ↔ Replicated` and adding `Reduced ↔ Unreduced` as a separate pair, you get to choose per-array whether gradients arrive replicated (reduction done for you, eagerly) or unreduced (reduction is yours to place), and the choice is made by a communication-free cast in the forward pass. **Is there a theoretical justification for Unreduced, or is it just a practical trick?** Here's an argument that something like it is forced. Suppose we want the following, none of which is individually exotic: 1. autodiff should preserve communication cost, op by op: a forward op that requires no communication should transpose to a backward op that requires no communication (or at least, we need *some* way to write such a backward pass); 2. replication should be expressible (we don't want *only* sharded axes); 3. multiplying a replicated scalar by a sharded vector, a communication-free forward op, should be expressible; 4. cotangents should have the same shape as their primals. Transposing (3) with respect to the scalar means mapping a sharded cotangent vector to a cotangent for the replicated scalar. Doing that with no communication (1) forces a state that is scalar-shaped (4) but holds only a per-device *piece* of the answer, pending a sum: that's Unreduced. (Without premise 4 you could simulate it by tacking on an extra device-indexed axis; Unreduced is in effect that axis, tracked in the sharding instead of the shape.) ## Manual mode: autodiff and `shard_map` In manual mode ({ref}`jax-201-shard-map`), you write per-device code and communication is explicit: an AllReduce isn't implied by an `out_sharding` or performed by a `reshard`; it's a `jax.lax.psum` you place yourself. All the machinery above has a manual-mode counterpart, governed by the same rule: cotangent types are a function of primal types. Inside a `shard_map`, types track how each value relates to its counterparts on the other devices, along each mesh axis the map is manual over. There are four states per axis, mirroring the four sharding states outside (see {ref}`jax-201-shard-map` for a fuller introduction): | outside (explicit mode) | inside (manual mode) | along `X`, devices hold | |--------------------------|-------------------------|--------------------------------| | sharded, `f32[8@X]` | varying, `f32[4]{V:X}` | different values | | replicated, `f32[4]` | invarying, `f32[4]` | the same value | | unreduced, `f32[4]{U:X}` | unreduced, `f32[4]{U:X}`| partial sums of the true value | | reduced, `f32[4]{R:X}` | reduced, `f32[4]{R:X}` | the same value, with unreduced cotangents | The `in_specs` and `out_specs` mediate the correspondence: a `P('X')` input binds to a varying value inside, `P()` to an invarying one, and `P(unreduced={'X'})` / `P(reduced={'X'})` pass through as themselves. (One difference from the outside types: sharded names the array axis that is split, as in `8@X`, while inside the split has already happened, so varying is a fact about the mesh axis alone.) The cotangent map mirrors the one above: varying and invarying are each their own cotangent type, and unreduced and reduced swap. Here's the matmul from the unreduced section, with the same `a` and `b`, written per-device: ```{code-cell} @jax.shard_map(in_specs=(jax.P(None, 'X'), jax.P('X', None)), out_specs=jax.P(None, None, unreduced={'X'})) def matmul_partial(a, b): out = a @ b # a local matmul of this device's pieces: a partial sum return jax.lax.pcast(out, 'X', to='unreduced') c = matmul_partial(a, b) print(jax.typeof(c)) ``` Each device multiplies its column block of `a` with its row block of `b`, a computation whose result is *varying* over `X`: each device holds a different partial sum. The `pcast` to unreduced is a free varying-to-unreduced cast, which `shard_map` requires here because the body's output type must match the unreduced `out_specs`. Outside, we get exactly the `float32[2,2]{U:X}` value we built with `out_sharding` before. The backward pass: ```{code-cell} def matmul_loss(a, b): return jnp.sum(jax.reshard(matmul_partial(a, b), jax.P(None, None))) print(jax.jit(jax.grad(matmul_loss, argnums=(0, 1))).trace(a, b).jaxpr) ``` Read the two `shard_map` equations. In the forward one, the body is a local `dot_general` on varying operands followed by the free `vary_unreduced_cast`. In the backward one, the cotangent of the unreduced output arrives with `in_specs=P(None, None, reduced={'X'})` (the boundary specs obey the cotangent map) and the body is `reduced_vary_cast` (the free transpose of the forward's free cast) followed by two local dots. Neither body contains any communication: the single AllReduce is the transposed outer `reshard`, exactly where explicit mode put it. More generally, in manual mode every type cast (`jax.lax.pcast`) is free and `psum` is the communication, and transposition pairs them up: | forward | type | transpose | type | |----------------------------------|----------------------|----------------------------------|-----------------------| | `jax.lax.psum` | varying → invarying | `jax.lax.pcast(..., to='varying')` | invarying → varying | | `jax.lax.pcast(..., to='reduced')` | invarying → reduced | `jax.lax.psum` on an unreduced value | unreduced → invarying | | `jax.lax.pcast(..., to='unreduced')` | varying → unreduced | `jax.lax.pcast(..., to='varying')` | reduced → varying | The first row is the classic story (see the [shard_map transposition JEP](https://docs.jax.dev/en/latest/jep/17111-shmap-transpose.html) and the collectives table in {ref}`jax-201-shard-map`): summing transposes to marking-as-varying. The second row is the manual-mode version of this page's main trick: a free cast to reduced in the forward pass transposes to the `psum` that pays for it in the backward pass. And the third row is free in both directions, since declaring values to be partial sums moves no data either way. As in explicit mode, where you place free casts in your forward code determines where the backward-pass `psum`s run. That means the reduced-weights pattern carries over verbatim. Give the weights a reduced type on the way in, and gradients come out unreduced, with no `psum` anywhere in the backward pass: ```{code-cell} w_r = jax.reshard(w, jax.P(None, None, reduced={'X'})) @jax.shard_map(in_specs=(jax.P('X', None), jax.P(None, None, reduced={'X'})), out_specs=jax.P('X', None)) def apply_layer(x, w): return jnp.tanh(x @ w) dw = jax.grad(lambda w, x: jnp.sum(apply_layer(x, w) ** 2))(w_r, x) print(jax.typeof(dw)) ``` So a model whose layers are written with `shard_map` accumulates unreduced gradients across microbatches with a single AllReduce per step, exactly like the example above. For the full table of collectives and their transposes, see {ref}`jax-201-shard-map`.