How One Python Dictionary Replaces 2,000 Lines of Distributed Communication Code

If you’ve ever looked inside a Megatron-LM or Picotron-style PyTorch codebase to see how Tensor Parallelism (TP) and Sequence Parallelism (SP) are implemented, you know the drill:

1. You replace nn.Linear with ColumnParallelLinear and RowParallelLinear.
2. You write custom torch.autograd.Function classes (_CopyToModelParallelRegion, _ReduceFromModelParallelRegion, _GatherFromSequenceParallelRegion, _ReduceScatterToSequenceParallelRegion) that manually fire dist.all_reduce, dist.all_gather, and dist.reduce_scatter_tensor in the forward and backward passes.
3. If you want FSDP too, you wrap the whole thing in another layer of pre-forward and post-backward hooks.

It works, and it’s a great exercise to write once. But it tangles your model definition with your communication schedule.

In JAX, GSPMD (General and Scalable Parallelization for ML Computation Graphs) flips this around. You write standard single-device math using plain jnp.einsum, give each tensor dimension a logical name, and map those logical names to physical mesh axes in a single Python dictionary. The XLA compiler figures out which collectives to insert for both the forward and backward passes.

Let’s look at how jaxotron/parallel/rules.py implements Data Parallelism, FSDP (ZeRO-3), Tensor Parallelism, Sequence Parallelism, Pipeline Parallelism, and Context Parallelism in about 60 lines of code.

The One File That Decides Everything

First, install Jaxotron from GitHub if you want to run the commands alongside this post:

pip install "jaxotron @ git+https://github.com/rishiraj/jaxotron"

Now open jaxotron/parallel/rules.py. Every parameter and activation in the transformer has named dimensions:

from jax.sharding import PartitionSpec as P

# How parameters are laid out across the mesh.
WEIGHT_RULES: dict[str, object] = {
    "layers": "pp",      # Pipeline parallel: each stage owns a contiguous block of layers
    "embed": "fsdp",     # FSDP / ZeRO-3: hidden dim of every weight is sharded across fsdp
    "vocab": "tp",       # Tensor parallel: vocab split (Megatron VocabParallelEmbedding)
    "heads": "tp",       # Tensor parallel: attention heads split (column-parallel q/k/v)
    "kv_heads": "tp",
    "mlp": "tp",         # Tensor parallel: MLP hidden split (column-parallel up/gate, row-parallel down)
}

# How activations are laid out across the mesh.
ACTIVATION_RULES: dict[str, object] = {
    "batch": ("dp", "fsdp"),  # Data parallel: batch split across dp and fsdp
    "seq": "cp",              # Context parallel: sequence split
    "res_seq": "cp",          # Residual stream sequence axis: "cp", or ("cp", "tp") with --sp
    "kv_seq": None,           # Keys/values are all-gathered across cp (all-gather context parallelism)
    "stage": "pp",            # Pipeline parallel: the pipeline state buffer has one slot per stage
    "heads": "tp",
    "kv_heads": "tp",
    "mlp": "tp",
    "vocab": "tp",
    "act_embed": None,        # Residual hidden dim is replicated over tp (Megatron style)
}

ACTIVATION_RULES_SP: dict[str, object] = {
    **ACTIVATION_RULES,
    "res_seq": ("cp", "tp"),  # Megatron-style Sequence Parallelism: shard residual stream over tp too
}

And right below that is the shape map for the model's weights:

PARAM_LOGICAL_AXES: dict[str, tuple] = {
    "embed": ("vocab", "embed"),
    "lm_head": ("embed", "vocab"),
    "final_norm": (None,),
    "attn_norm": ("layers", None),
    "mlp_norm": ("layers", None),
    "wq": ("layers", "embed", "heads", None),
    "wk": ("layers", "embed", "kv_heads", None),
    "wv": ("layers", "embed", "kv_heads", None),
    "wo": ("layers", "heads", None, "embed"),
    "w_gate": ("layers", "embed", "mlp"),
    "w_up": ("layers", "embed", "mlp"),
    "w_down": ("layers", "mlp", "embed"),
    ...
}

Take a look at wq: its shape is [layers, embed, heads, head_dim] and its logical axes are ("layers", "embed", "heads", None). When we translate those logical names through WEIGHT_RULES, wq automatically gets PartitionSpec("pp", "fsdp", "tp", None).

• If pp=1, fsdp=8, tp=1, wq is sharded 8 ways along its hidden dimension (embed).
• If pp=1, fsdp=1, tp=4, wq is sharded 4 ways across attention heads.
• If all three are > 1, wq is sharded across all three axes simultaneously!

You can see this directly in your terminal with --show_sharding:

jaxotron-train --cpu_devices 8 --fsdp 8 --seq_len 128 --global_batch_size 16 \
    --micro_batch_size 2 --dry_run --show_sharding

parameter           global shape         partition spec                MB/device
--------------------------------------------------------------------------------
.embed              (1024, 256)          PartitionSpec('tp', 'fsdp')   0.12
.blocks.attn_norm   (4, 256)             PartitionSpec('pp', None)     0.00
.blocks.wq          (4, 256, 8, 32)      PartitionSpec('pp', 'fsdp', 'tp', None)  0.12
...
TOTAL (params only)                                                    0.82
TOTAL (params + AdamW state)                                           2.46

How the Forward Pass Stays Clean

Now let's look at an entire transformer decoder layer in jaxotron/models/llama.py. There isn't a single all_reduce or all_gather call in the file—just ctx.shard(...), which wraps jax.lax.with_sharding_constraint:

def block_forward(p: Block, x: jax.Array, cos, sin, ctx: ParallelContext, eps: float) -> jax.Array:
    """One decoder layer. Sharding annotations name each activation's logical axes."""
    # --- Attention -----------------------------------------------------------
    h = ctx.shard(rms_norm(x, p.attn_norm, eps), "batch", "res_seq", "act_embed")
    q, k, v = _qkv(p, h, ctx, eps)
    q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
    # Context parallelism: each cp rank keeps its slice of queries, but needs
    # every key/value. Replicating K/V over "cp" makes XLA all-gather them.
    k = ctx.shard(k, "batch", "kv_seq", "kv_heads", None)
    v = ctx.shard(v, "batch", "kv_seq", "kv_heads", None)
    attn = ctx.shard(causal_attention(q, k, v, ctx.attention_impl), "batch", "seq", "heads", None)
    # Row-parallel output projection: XLA inserts the tensor-parallel all-reduce (or reduce-scatter with --sp).
    x = x + ctx.shard(jnp.einsum("bthk,hkd->btd", attn, p.wo), "batch", "res_seq", "act_embed")

    # --- SwiGLU MLP ----------------------------------------------------------
    h = ctx.shard(rms_norm(x, p.mlp_norm, eps), "batch", "res_seq", "act_embed")
    gate = ctx.shard(jnp.einsum("btd,df->btf", h, p.w_gate), "batch", "seq", "mlp")
    up = ctx.shard(jnp.einsum("btd,df->btf", h, p.w_up), "batch", "seq", "mlp")
    out = jnp.einsum("btf,fd->btd", jax.nn.silu(gate) * up, p.w_down)
    return x + ctx.shard(out, "batch", "res_seq", "act_embed")

Watching --sp Turn all-reduce Into all-gather + reduce-scatter

Here is the coolest part to show live. Look at the SwiGLU MLP above:

• gate and up have logical axes ("batch", "seq", "mlp"), where "mlp" maps to "tp". Each TP rank computes 1/tp of the intermediate columns.
• p.w_down has logical axes ("layers", "mlp", "embed"). When we contract over f ("mlp") in jnp.einsum("btf,fd->btd", ..., p.w_down), each TP rank only has a partial sum over its slice of f.
• What happens next depends entirely on what "res_seq" maps to!

Without Sequence Parallelism (--tp 4), "res_seq" maps to "cp" (replicated across tp). To produce a replicated output from a partial sum across tp, XLA must insert an all-reduce over tp. Let’s verify with --show_collectives:

jaxotron-train --cpu_devices 4 --tp 4 --seq_len 128 --global_batch_size 16 \
    --micro_batch_size 4 --dry_run --show_collectives

Communication inserted by XLA:
collective   mesh axes   ops   bytes (static)
---------------------------------------------
all-reduce   tp           17         1.31 MiB
all-gather   tp            2        16.00 KiB

Now add --sp (Sequence Parallelism). The only thing --sp changes is that "res_seq" now maps to ("cp", "tp") in ACTIVATION_RULES_SP. Now the residual stream x and the rms_norm outputs are sharded along the sequence dimension across tp!

• Before w_gate/w_up (which need the full sequence "seq": "cp"), XLA inserts an all-gather over tp.
• After w_down (which has a partial sum over tp and needs to land in "res_seq": ("cp", "tp")), XLA inserts a reduce-scatter over tp.

Let's run the exact same command with --sp:

jaxotron-train --cpu_devices 4 --tp 4 --sp --seq_len 128 --global_batch_size 16 \
    --micro_batch_size 4 --dry_run --show_collectives

Communication inserted by XLA:
collective       mesh axes   ops   bytes (static)
-------------------------------------------------
all-gather       tp            9       544.00 KiB
reduce-scatter   tp            6       192.00 KiB

Zero lines of model code changed. By changing one dictionary entry ("res_seq": ("cp", "tp")), XLA replaced the tensor-parallel all-reduce ops with all-gather + reduce-scatter and sharded the RMSNorm activations across the TP group.