Sam D. Buchanan

SPMD in JAX #3: Infrastructure Buildout

2025-12-10T00:00:00+00:00

In this third post in a series on programming with JAX for training transformers, we pick up where we left off at the end of the previous post. We’ll extend our previous small transformer model to the NanoGPT scale, incorporating actual text data, distributed training, and a real optimizer, and perform some basic evaluation! We’ll end up with a reasonable substrate for building future experiments on top of.

As in previous posts in the series, we make good use of the Google scaling playbook (Austin et al., 2025) for profiling.

Package Conversion and Improvements
Actual Data
- Checking the Results
Conclusions
Acknowledgments

Package Conversion and Improvements

We’re running in Python 3.13, and we’ve upgraded to JAX 0.8.0 from previous posts. Now that we’ve written a good chunk of model, training, and config code, we’ll also be working on top of a packaged version of our code.

The code is hosted in this GitHub repository: baremetal-gpt. There are instructions to install the repository as a package using uv there.

We’ll start by talking through some refactoring that leads us from the state of our code in the previous post to the state of the code in the GitHub repository.

Config Upgrades

To run experiments with our code in a distributed (i.e., multi-host) setting, we’ll benefit from having a robust way to dispatch different experimental configurations from the command line. We’ll use Hydra for this purpose. Hydra is not the most minimal solution, but it’s the most feature-complete library for managing configs that I know of (e.g., specifying lists as arguments), and for our simple purposes we’ll be able to adopt it with only a small amount of added code and complexity.

Small Modifications to our Config Dataclass

Hydra supports registering config dataclasses as so-called “Structured Configs”. This tutorial gives a good description of the core functionalities we need. We only need to make a few small modifications to our previous dataclass:

Use only supported types: instead of random class objects (e.g., for dtypes or PartitionSpecs), we need to create Enums for passing values, and disambiguate them at call time. There are other minor caveats (e.g., can’t arbitrarily nest lists; restrictions on use of Optional).
Make the config mutable (remove frozen=True).
Re-register the config as static with JAX. Hydra will wrap our dataclass so that it ends up duck-typed as a Config object, but we will need to re-register the wrapped config as static. Since the dataclass is no longer frozen, we can do this with a post-init function, which we’ll now call when we launch our training loop.
Create a function to register the Config dataclass in a Hydra ConfigStore instance. In general, we can create arbitrarily stratified configs for easier management (e.g., separate configs for model, experiment, data, etc.). We’ll call this function in our main method.
Refactor the training loop into a train.py file with a main method that we’ll call when we want to launch an experiment, and add the necessary Hydra setup code.

We will also make a larger refactor, which will help us later with setting up config YAML files for experiments: we will change from a flat config to a hierarchical config, with different parameters split into different sub-configs (e.g., for the optimizer, for data, …). This is not an objectively superior design – for example, it’s arguably a bit easier to write new code and integrate static type checking with a flat config – but we’ll follow it in order to add a bit more encapsulation to distinct config parameters.

Here’s a representative snipped from the new config.py with the modifications discussed above. There are many forward-looking modifications too, which we just show out-of-context selections of.

# config.py

# We omit the sub-configs, and other nitty-gritty details...
# Config is now top-level (orchestration parameters, etc.)


@dataclass(kw_only=True, unsafe_hash=True)
class Config:
    """Overall config class, containing top-level orchestration parameters"""

    seed: int = 1337
    logger_type: LoggerType = LoggerType.WANDB
    project_name: str = "bmgpt-debug"
    run_name: str = ""
    val_log_interval: int = 1000  # log validation metrics every <this many> batches

    train_dataset: DatasetConfig = MISSING
    val_list: list[EvaluationConfig] = field(default_factory=list)  # validation metrics
    eval_list: list[EvaluationConfig] = field(default_factory=list)

    optimizer: OptimizerConfig = field(default_factory=OptimizerConfig)
    model: ModelConfig = field(default_factory=ModelConfig)
    inference: InferenceConfig = field(default_factory=InferenceConfig)
    sharding: ShardingConfig = field(default_factory=ShardingConfig)


def register_configs():
    cs = ConfigStore.instance()
    cs.store(group="optimizer", name="base_optimizer", node=OptimizerConfig)
    cs.store(group="model", name="base_model", node=ModelConfig)
    cs.store(group="inference", name="base_inference", node=InferenceConfig)
    cs.store(group="sharding", name="base_sharding", node=ShardingConfig)
    cs.store(name="base_config", node=Config)


def config_post_init(config: Config):
    """Input validation."""
    # ...

And here’s the relevant parts of train.py that we modify. There is some extra code for distributed initialization, which we discuss in more detail later in the post.

from pathlib import Path

import hydra
from bmgpt.config import Config, config_post_init, register_configs

register_configs()

# ...


@hydra.main(
    version_base=None,
    config_path=str(Path("configs").absolute().resolve()),
    config_name="base_config",
)
def main(config: Config):
    try:
        # Launch distributed and register configs
        jax.distributed.initialize()
        jax.tree_util.register_static(type(config))
    except RuntimeError:
        # This implies the distributed backend has already been initialized
        pass
    config_post_init(config)


if __name__ == "__main__":
    main()

The modifications are quite simple: we register our Config dataclass in the ConfigStore instance as 'config', then decorate our main method in train.py to use that appropriate config. The 'config_path' argument in that decorator specifies a directory where we store different .yaml config files specifying different experimental setups (e.g., different model architectures, different datasets, different distributed training environments, etc.). Then we just call our post-init config method right at the start of main: right now, it just validates some config inputs.

When we want to launch an experiment now, we can run a command like the following (train loop refactoring follows later in the blog):

uv run --extra tpu train model.num_layers=1

Here, we are using uv optional dependencies so that we can run on either CPU or TPU depending on the platform, as well as a build script in our package (see the pyproject.toml). Super easy! Anything that isn’t specified follows our defaults from the Config dataclass definition.

As the experiments we want to run become more diverse and complex, it becomes harder and harder to productively specify all default overrides via the command line. To address this, we can write YAML configs for different experiments as we mentioned above. For example, we can make a file at configs/experiments/text_toy.yaml, with:

# @package _global_

defaults:
  - /dataset/staircase@train_dataset

train_dataset:
  num_steps: 100

eval_list:
  - dataset:
      name: NUMBER_STAIRCASE
      path: ""
      split: TEST
      seq_len: 1  # At inference time, this functions like prompt size
      global_batch_size: 16
    evaluator: AUTOREGRESSIVE_ROLLOUTS

model:
  transformer_type: DISCRETE
  is_causal: True
  num_vocab: 10

See the Hydra docs for more details about how to parse this file – it has some boilerplate to ensure it overrides the config defaults we specify in our config.py. We’ll discuss some of the enums that appear here (you can identify them by the all-caps text) later on. Then to run an experiment, we can simply override as:

uv run --extra tpu train +experiment=text_toy

In addition, any other command line arguments we want to override with can be added!

Model Improvements and Refactors

We’ll add support for the Pallas Splash Attention kernel, which will help us with model scaling later!

Splash Attention Integration

Splash Attention stands for “sparse Flash Attention” (Pagliardini et al., 2023). Flash Attention is a well-known memory-efficient implementation of the attention operation (Dao et al., 2022); it also is often used to refer to the associated CUDA kernel and surrounding PyTorch package for the algorithm. Compared to flash attention, splash attention adds better support for sparse attention masks, which arise frequently in practice: for example, when we want to pack many relatively short sequences into one long sequence, or precisely mask between document boundaries, we need to make sure not to attend across these boundaries.

The types of optimizations that flash/splash attention entail require a level of control over hardware-level execution that is not possible within the standard JAX-XLA programming pipeline. So JAX exposes a domain-specific language for writing these kernels, called Pallas! The authoring of Pallas kernels is out-of-scope for this post, but we will be interested in learning how to call pre-written Pallas kernels so that we can integrate splash attention.

The interfaces for achieving this are naturally a bit rough around the edges, with scant documentation and relatively few examples for how to proceed. With the help of multiple LLMs, I eventually pieced together a solution based on example code available in the JAX Pallas tests. Here’s a high-level overview, with the code to follow:

We need to configure the Pallas kernel itself based on how we’ll use it at runtime: specifically, the sequence length, number of heads, base mask type, and parallelism configuration. The latter has to be enforced manually with shard_map. Once we figure out where everything needs to go, this works more-or-less okay with our current model implementation structure, which uses Explicit mesh axes and vmap as much as possible to simplify ‘inner’ model code. However, it currently doesn’t work in JAX 0.8.1. It should work in the next release, thanks to a bug report I submitted and a fix from Yash Katariya (thanks!).
Once we have the kernel (represented as a tuple of two Callables, one corresponding to the actual kernel and one corresponding to a wrapper with the shard_map configuration), we need to thread it through our top-level _transformer call down to the _attn call. A more ‘elegant’ approach to doing this (at least at the call site) might be to use some metaprogramming, but we will go with the simplest solution.

The code for the kernel setup is in splash_helpers.py:

# splash_helpers.py
from functools import partial

import jax
from bmgpt.config import Config, DatasetConfig
from jax.experimental.pallas.ops.tpu.splash_attention import (
    BlockSizes,
    CausalMask,
    MultiHeadMask,
    make_splash_mha,
)

# ...


def make_splash_kernel(
    config: Config,
    config_data: DatasetConfig,
    cache_capacity: int,
    mesh,
    head_shards: int = 1,
    q_seq_shards: int = 1,
):
    if not config_data.use_splash:
        # None ends up calling jax-xla attention
        # see _attn
        return None
    # s is Q len (seq_len @ train; variable/1 at prefill/decode)
    # t is K len (s + cache_capacity)
    s = config_data.seq_len
    t = s + cache_capacity

    if config.model.is_causal:
        mask = MultiHeadMask(
            [
                CausalMask(shape=(s, t), offset=cache_capacity)
                for _ in range(config.model.num_heads)
            ]
        )
    else:
        mask = MultiHeadMask(
            [FullMask(shape=(s, t)) for _ in range(config.model.num_heads)]
        )
    BLOCK_SIZE = min(config_data.seq_len, 128)
    if config_data.seq_len % 128 != 0 or BLOCK_SIZE % 128 != 0:
        # splash attention kernel requires block size to be a multiple of 128
        raise NotImplementedError("Splash block size needs to be a multiple of 128")
    block_sizes = BlockSizes(
        block_q=BLOCK_SIZE,
        block_kv=BLOCK_SIZE,
        block_kv_compute=BLOCK_SIZE,
        block_q_dkv=BLOCK_SIZE,
        block_kv_dkv=BLOCK_SIZE,
        block_kv_dkv_compute=BLOCK_SIZE,
        block_q_dq=BLOCK_SIZE,
        block_kv_dq=BLOCK_SIZE,
    )
    splash_spec = jax.P(None, None)
    splash_sharding = jax.sharding.NamedSharding(mesh, splash_spec)
    kernel = make_splash_mha(
        mask,
        head_shards=head_shards,
        q_seq_shards=q_seq_shards,
        block_sizes=block_sizes,
    )
    kernel_spec = kernel.manual_sharding_spec(splash_sharding)

    @partial(
        jax.shard_map,
        mesh=mesh,
        in_specs=(kernel_spec, splash_spec, splash_spec, splash_spec, jax.P()),
        out_specs=splash_spec,
        check_vma=False,
    )
    def splash_sharded(kernel, q, k, v, segment_ids):
        return kernel(q, k, v, segment_ids=segment_ids)

    return (splash_sharded, kernel)

The function argument cache_capacity specifies the maximum KV cache size we’ll use (we’ve split this off from max_seq_len which filled a dual function encompassing this previously). This allows us to specify a cache size of 0 for training, for example. Within the function body, we have a new configuration parameter use_splash to denote whether splash attention should be used or not. For example, we always fall back to XLA attention for autoregressive inference, since the splash attention Pallas kernel requires the block size to be a multiple of 128 (and for decode, we almost necessarily have to process a query block size of 1). In this case, the kernel creation factory just returns None, and we check for this in _attn.¹

The “sparse” part of splash attention comes from the proper use of the segment_ids argument: distinct documents/sub-sequences can be assigned different segment_ids. We’ll see an example below, in how to set this up for use during training with respect to the KV cache.

First, at the level of train.py, we now create our kernels when we set up. For training, this leads to a simple kernel = make_splash_kernel(config, config.train_dataset, 0, mesh) statement. We then pass the kernel to any model call we have: for example, logits, _ = jax.vmap( partial(_transformer, config, kernel, params, cache_params=cache_params))(inputs, train_state.kv_cache) in the training loop step. Here, we’ve refactored the way we use the KV cache – a CacheParams class is defined simply as follows:

class CacheParams(NamedTuple):
    enabled: bool
    size: int

This way, as we’ll see momentarily, we can easily trace with respect to the size argument (i.e., for inference) while still omitting all KV cache maintenance code at compile time if we want to train (i.e., by setting enabled=False).

The kernel argument and this new cache_params argument get threaded down to the level of _attn, which now reads as follows (abbreviating unchanged components):

# model.py
def _make_causal_mask(seq_len_q: int, seq_len_k: int, cache_size: int) -> Array:
    """(seq_len_q, seq_len_k) bool array with True for past token positions"""
    q_positions = cache_size + jnp.arange(seq_len_q)
    k_positions = jnp.arange(seq_len_k)
    return q_positions[:, None] >= k_positions[None, :]


def _make_cache_mask(seq_len_q: int, seq_len_k: int, cache_size: int) -> Array:
    """(1, seq_len_k) bool array with True for actual cache+context positions"""
    k_positions = jnp.arange(seq_len_k)
    return k_positions[None, :] < cache_size


def _attn(
    config: Config,
    kernel,
    params: Attn,
    x_seq: jax.Array,
    kv_cache: jax.Array,  # 2 x cache_capacity x n x h
    cache_params: CacheParams,
):
    # s: sequence length
    # d: embedding dim (config.d_model)
    # n: attention heads (config.num_heads)
    # h: head dim (config.d_head)
    # x_seq: s x d

    q, k, v = jnp.einsum(
        "sd,d3nh->3nsh",
        x_seq,
        params.w_qkv,
        out_sharding=jax.P(*config.sharding.att_qkv),
    )
    s = q.shape[1]

    # skipping rope...

    k_cache, v_cache = kv_cache[0], kv_cache[1]
    if cache_params.enabled:
        k_cache_out = jax.lax.dynamic_update_slice(
            k_cache, k, (0, cache_params.size, 0)
        )
        v_cache_out = jax.lax.dynamic_update_slice(
            v_cache, v, (0, cache_params.size, 0)
        )
        kv_cache_out = jnp.concatenate((k_cache_out[None], v_cache_out[None]), axis=0)
    else:
        kv_cache_out = kv_cache
    # Cache read scheme: to enable same mask for the same s value (Q seq len),
    #  we concatenate the full cache to K, and mask empty entries
    # For efficient training, set cache size zero + cache_params.enabled=False
    cache_capacity = k_cache.shape[-2]
    k = jnp.concatenate((k_cache, k), axis=1)
    v = jnp.concatenate((v_cache, v), axis=1)

    # Attention
    t = k.shape[1]  # t = s + cache_capacity
    if kernel:
        q_segment_ids = jnp.zeros((s,))
        kv_mask = _make_cache_mask(s, t, cache_params.size) | (
            ~_make_cache_mask(s, t, cache_capacity)
        )
        kv_mask = ~kv_mask[0]
        kv_segment_ids = kv_mask.astype(jnp.int32)
        segment_ids = SegmentIds(q=q_segment_ids, kv=kv_segment_ids)

        splash_sharded, kernel = kernel
        attn_out = splash_sharded(
            kernel,
            q / config.model.d_head**0.25,
            k / config.model.d_head**0.25,
            v,
            segment_ids,
        )
    else:
        # Make mask
        if config.model.is_causal:
            mask = _make_causal_mask(s, t, cache_capacity)
            cache_mask = _make_cache_mask(s, t, cache_params.size) | (
                ~_make_cache_mask(s, t, cache_capacity)
            )
            mask = mask & cache_mask
        else:
            mask = ~_make_cache_mask(s, t, 0)  # full attention
        mask = mask[None, ...]  # broadcast over heads
        # Scale and causal mask
        logits = jnp.einsum("nsh,nth->nst", q, k).astype(
            config.model.compute_dtype.value
        )
        logits *= 1.0 / config.model.d_head**0.5
        logits = jnp.where(mask, logits, -jnp.inf)
        probs = jax.nn.softmax(logits, axis=2)  # type: ignore[reportArgumentType]
        probs = probs.astype(config.model.param_dtype.value)
        attn_out = jnp.einsum("nst,nth->nsh", probs, v)
    out = jnp.einsum(
        "nsh,hnd->sd",
        attn_out,
        params.w_o,
        out_sharding=jax.P(*config.sharding.res_stream),
    )

    return out, kv_cache_out


# ... skipping to the cache init function


def init_kv_cache(config: Config, global_batch_size: int, cache_capacity: int):
    if not config.sharding.data:
        sharding_batch_layer = [None, None]
    else:
        sharding_batch_layer = config.sharding.data + [None]
    sharding = jax.P(*(sharding_batch_layer + config.sharding.att_qkv))

    return jnp.zeros(
        (
            global_batch_size,
            config.model.num_layers,
            2,
            config.model.num_heads,
            cache_capacity,
            config.model.d_head,
        ),
        dtype=config.model.param_dtype.value,
        out_sharding=sharding,
    )

The splash Pallas kernel expects the shape to be N x S x H, so we refactor some code compared to what we had before. The key change here is that in order to repeatedly reuse the same splash Pallas kernel as we grow the cache, we write the code so that the KV sequence length is always fixed as the cache capacity plus S (the Q sequence length). The cache capacity is specified dynamically at initialization, allowing us to use the same Config object for different training/inference scenarios (e.g., cache size 0 for training, and some fixed positive integer for autoregressive inference). This means we need to construct a SegmentIds object that masks out all the unused portions of the KV cache, which is performed correctly by the above code. Previously, we were just updating the cache with jax.lax.dynamic_update_slice, rather than concatenating, but with the splash kernel structure, this would require us to change the mask after every additional cache update (requiring a new kernel!).

That’s it! We’ve integrated splash attention, allowing us to scale much more effectively to larger model sizes down the road.

Optimizer Improvements

We were using 100% vanilla SGD in the previous post. Let’s upgrade that to a much stronger optimizer, AdamW! (Loshchilov & Hutter, 2017)

In our previous post, we used a very simple tree-mapped update, facilitated by the simple structure of constant-learning-rate SGD:

# ...
loss, grad = jax.value_and_grad(loss_fn)(train_state.params)
new_state = TrainState(
    params=jax.tree.map(lambda x, y: x - config.lr * y, train_state.params, grad),
    opt=None,
)
# ...

For implementing Adam and AdamW (the “W” denotes “decoupled weight decay”), it’s better to use a more flexible API: both because the updates become harder to express in a single concise lambda, and because we generally want to be able to apply different updates to different parameters.² Here’s how we’ll do this, at a basic level:

We’ll define a handfull of optimizer updates that share the same API.
We’ll allow one update to be selected via the Config class (i.e., with the Enum we mentioned above), and then we’ll jax.tree.map it over the model pytree.
We’ll be able to provide a pytree of arguments to the update, as necessary, allowing to do different things for different parameters.

For example, here’s how this looks for the previous constant-learning-rate SGD update. Since AdamW requires us to save state information (i.e., first-order and second-order momentum buffers), we’ll incorporate that into the API.

# Factory for getting the update, used externally
def opt_update_factory(opt_type: OptType):
    match opt_type:
        case OptType.ADAMW:
            return adamw_update
        case OptType.SGD:
            return sgd_update


class OptState(NamedTuple):
    mu: jax.Array  # 1st moment EMA
    nu: jax.Array  # 2nd moment EMA
    step: jax.Array  # step number


def init_adam_state(config: Config, param: jax.Array) -> OptState:
    return OptState(
        mu=jnp.zeros_like(param, dtype=config.optimizer_dtype.value),
        nu=jnp.zeros_like(param, dtype=config.optimizer_dtype.value),
        step=jnp.array(0, dtype=jnp.int32),
    )


def sgd_update(
    config: Config, param: jax.Array, grad: jax.Array, state: OptState, wd_mask: bool
):
    update = -config.lr * grad
    if wd_mask:
        # Apply weight decay
        update = update - config.lr * config.weight_decay * param
    return update, state

The update API takes parameters, gradients, state, and additional flags, and returns the update direction and the updated state. We add the update direction to the parameters externally since this is a useful debugging abstraction (and matches the optax API).

Here’s how it looks to apply the optimizer now. It gets a bit more cumbersome to do this with only pytrees, since the update returns a tuple – meaning the output pytree structure is different from the model/grad pytree structure. The latter is a prefix of the former, though, so we can still extract what we need fairly easily.

# In train.py

# ...
class TrainState(NamedTuple):
    params: Transformer
    opt_state: Any
    kv_cache: jax.Array


@jax.jit
def init_train_state(key, config: Config) -> TrainState:
    model_params = init_transformer(key, config)
    adam_state = jax.tree.map(partial(init_adam_state, config), model_params)
    cache = init_kv_cache(config, config.train_dataset.global_batch_size, 0)
    return TrainState(params=model_params, opt_state=adam_state, kv_cache=cache)


# ...


def main(config: Config):
    # ...
    opt_update = opt_update_factory(config.optimizer.type)

    # ...
    @partial(jax.jit, donate_argnums=2)
    def train_step(config: Config, batch, train_state: TrainState):
        # ...
        loss, grad = jax.value_and_grad(loss_fn)(train_state.params)
        # ...
        update__opt_state = jax.tree.map(
            partial(opt_update, config),
            train_state.params,
            grad,
            train_state.opt_state,
            weight_decay_mask,
        )
        # Transpose the output tree to get update tree and state tree
        update, opt_state = map(
            lambda i: jax.tree.map(lambda x, y: y[i], grad, update__opt_state),
            range(2),
        )
        params = jax.tree.map(lambda x, y: x + y, train_state.params, update)
        new_state = TrainState(
            params=params, opt_state=opt_state, kv_cache=train_state.kv_cache
        )
        # ...

This is pretty straightforward, once we’ve written it!

Now, to add a new optimizer – AdamW in our case – we just need to define the appropriate opt_update and possibly define new state initialization functions. We’ve already done the latter above for Adam. Here’s a simple implementation of AdamW that agrees with the optax implementation (not bitwise):

def adamw_update(
    config: Config, param: jax.Array, grad: jax.Array, state: OptState, wd_mask: bool
):
    beta1 = config.beta1
    beta2 = config.beta2
    lr = config.lr
    eps = config.eps_adam
    weight_decay = config.weight_decay

    mu = beta1 * state.mu + (1 - beta1) * grad.astype(config.optimizer_dtype.value)
    nu = beta2 * state.nu + (1 - beta2) * grad.astype(config.optimizer_dtype.value) ** 2
    new_state = OptState(mu=mu, nu=nu, step=state.step + 1)

    mu_debias = mu / (1 - beta1**new_state.step)
    nu_debias = nu / (1 - beta2**new_state.step)
    update = -lr * mu_debias / (eps + jnp.sqrt(nu_debias))
    if wd_mask:
        # Apply weight decay
        update = update - lr * weight_decay * param
    return update.astype(config.param_dtype.value), new_state

The slightly tricky part remaining is how to pass the right weight decay mask pytree. It’s best to do this with model metadata, since the format we store the parameters in doesn’t necessarily match how they’re used. For example, we store our attention QKV parameters as a D x 3 x N x H array, where D is the model dimension, N the number of attention heads, and H the projected attention head dimension. We use these parameters as D x H matrices, but we can’t safely infer this from only the shape information! Similarly, since we vmap our Blocks initialization over the layer dimension L when constructing the model, our MLP output bias parameters are shaped like L x D. This is a matrix, so when iterating over the model pytree, we can’t just check whether the current leaf has ndim > 1!

This is a situation where we clearly see the advantage of the bigger, standard model libraries (e.g. Equinox): there, metadata and parameters can easily be stored together. For us, we’d need to depart from our appealingly simple NamedTuple model architecture to do this (e.g. register custom pytree leaves). So we’ll instead use a slightly unpleasant approach: we’ll just add a function to our model.py file that returns a model “spec” associated to the architecture, containing all relevant metadata.³

# in model.py
def model_spec(model: Transformer) -> Any:
    # Make the spec (we need some way to pass metadata around)
    def _make_spec_from_str(path: str) -> tuple[int, int] | None:
        param_str = path[-1].__str__()
        matrix_axes_dict = {
            ".w_qkv": (-4, -1),
            ".w_o": (-3, -1),
            ".w_up": (-2, -1),
            ".w_down": (-2, -1),
            ".w": (-2, -1),
            # ...
        }
        return matrix_axes_dict.get(param_str, None)

    spec = jax.tree.map_with_path(lambda p, _: _make_spec_from_str(p), model)
    return spec

The keys in the dictionary used to generate the output pytree correspond to the parameter names in our Transformer model. We can get these from the pytree using jax.tree.map_with_path, which gives the pytree structure (easily converted into a string, like .blocks.norm_attn.gamma) in addition to the leaves.

The returned pytree looks like the following:

Transformer(emb=Embedding(w=(-2, -1)), blocks=Block(norm_attn=LayerNorm(gamma=None, beta=None), attn=Attn(w_qkv=(-4, -1), w_o=(-3, -1)), norm_mlp=LayerNorm(gamma=None, beta=None), mlp=Mlp(w_up=(-2, -1), bias_up=None, w_down=(-2, -1), bias_down=None)), unemb=Unembedding(w=(-2, -1)))

Now, back to the matter at hand: we need to generate a mask pytree that tells us which parameters in the model to apply weight decay to. To apply it only to the matrix parameters in the model, we can do the following:

spec = model_spec(train_state.params)
weight_decay_mask = jax.tree.map(lambda x, s: bool(s), train_state.params, spec)

Since None is not truthy, this will give us the mask we desire! Then we use it as above, where we already provided weight_decay_mask as an argument to the update calculation tree map (but didn’t specify how to generate it).

Here’s the entire updated train step. We add gradient clipping to it, since this doesn’t require too much effort:

# in optimizers.py
def grad_norm_and_clip(
    config: Config, model: Transformer
) -> tuple[Transformer, Transformer, float]:
    # Gradient norms in param precision (NOTE: might want fp32?)
    grad_norms_squared = jax.tree.map(lambda grad: jnp.sum(grad**2), model)
    global_grad_norm = jax.tree.reduce(operator.add, grad_norms_squared) ** 0.5
    truncated_norm = jax.lax.select(
        global_grad_norm >= config.clip_grad,
        global_grad_norm,
        jnp.ones_like(global_grad_norm),
    )
    return (
        jax.tree.map(lambda grad: grad / truncated_norm, model),
        grad_norms_squared,
        global_grad_norm,
    )

It’s quite straightforward overall!

  # Initialize state, configure forward pass and optimization
  with jax.set_mesh(mesh):
      train_state = init_train_state(key_model, config)
  cache_params = CacheParams(enabled=False, size=0)
  kernel = make_splash_kernel(config, config.train_dataset, 0, mesh)
  spec = model_spec(train_state.params)
  opt_update = opt_update_factory(config.optimizer.type)
  weight_decay_mask = jax.tree.map(lambda _, s: bool(s), train_state.params, spec)


  @partial(jax.jit, donate_argnums=2)
  def train_step(config: Config, batch, train_state: TrainState):
      def loss_fn(params: Transformer):
          inputs, targets = batch
          logits, _ = jax.vmap(
              partial(_transformer, config, kernel, params, cache_params=cache_params)
          )(inputs, train_state.kv_cache)
          logits = logits.astype(config.model.compute_dtype.value)
          logprobs = jax.nn.log_softmax(logits, axis=-1)
          return -jnp.take_along_axis(logprobs, targets[..., None], axis=-1).mean()

      loss, grad = jax.value_and_grad(loss_fn)(train_state.params)
      grad_clipped, _, global_grad_norm = grad_norm_and_clip(config, grad)
      update__opt_state = jax.tree.map(
          partial(opt_update, config),
          train_state.params,
          grad_clipped,
          train_state.opt_state,
          weight_decay_mask,
      )
      # Transpose the output tree to get update tree and state tree
      update, opt_state = map(
          lambda i: jax.tree.map(lambda x, y: y[i], grad, update__opt_state), range(2)
      )
      params = jax.tree.map(lambda x, y: x + y, train_state.params, update)
      new_state = TrainState(
          params=params, opt_state=opt_state, kv_cache=train_state.kv_cache
      )

      metrics = {"batch_loss": loss, "grad_norm": global_grad_norm}
      return metrics, new_state

Distributed Training

Next, let’s set up distributed training. On the whole, this is actually very easy to do on Google’s TPU VMs!⁴ This is because:

The VMs are already configured with the IP addresses / ICI information of the other hosts in the slice.
We’re using jax.sharding.AxisType.Explicit sharding, and our model is relatively simple. So communication gets handled by the XLA compiler! We don’t actually need to use pmap or pjit anywhere.

For our purposes, the only part that requires care is data management, particularly data loading. We’ll describe this below for our simple synthetic data distribution, as well as the other changes we need to make for distributed training.

High-Level Paradigm

In JAX distributed training, each host (or “process”) is physically connected to some number of chips via PCIe. On v4 TPUs, which we’ve been using, each host is connected to 4 chips (each chip with 2 cores). When working with TPU VMs, we can think of one host as corresponding to one VM that we can ssh into; in our tests previously, we’ve been working on a single host.

Different hosts are arranged into “slices”. Within a slice, chips are interconnected via high-bandwidth optical links, letting them communicate relatively efficiently! When it comes to data movement from the host to the device, though, each host can only directly communicate with the chips it is physically connected to. These are called the “addressable” devices of that host.

With jax.distributed, we can programmatically handle these differences while only writing a single program! Our program gets run identically on each host, but we have access to certain functions that can allow us to branch and shard differently depending on which host we are. Key functions that we’ll see below are:

jax.distributed.initialize(), which synchronizes information about all interconnected devices among the different hosts;
jax.process_index(), which returns a different integer for each host (allowing branching);
jax.Array.addressable_shards, jax.Array.is_fully_addressable (etc.), which can be used to programmatically determine which devices can directly modify a given array. For example, if x is a jax.Array that is fully replicated across devices (with jax.P() as its partition spec), then x.is_fully_addressable is False in a multi-host setting! We want to avoid having our program modify an array in a way that is not compatible with either its sharding or the devices with respect to which it is addressable.

For more detailed information, see Chapter 2 of the TPU scaling book (Austin et al., 2025) as well as relevant sections of the JAX documentation (JAX Team, 2025).

Setting up the Mesh and Initializing the Distributed Backend

One multi-host footgun I learned about the hard way is that setting a global mesh with jax.set_mesh, as we have been doing so far, makes it hard to create arrays that exist only on the local process (more on that below). So we switch to creating a mesh from our config options at the start of the main function, then instantiate it as a context manager when we need it. This turns out to be pretty easy—we don’t need to pass the mesh to all parts of our model’s forward pass and add it as an additional parameter to out_sharding (etc.)!

Here’s the new setup (which we saw part of above):

# in config.py
def mesh_from_config(config: Config):
    mesh = jax.make_mesh(
        config.sharding.mesh_shape,
        config.sharding.mesh_axis_names,
        len(config.sharding.mesh_shape) * (jax.sharding.AxisType.Explicit,),
    )
    return mesh


# in train.py
@hydra.main(
    version_base=None,
    config_path=str(Path("configs").absolute().resolve()),
    config_name="base_config",
)
def main(config: Config):
    try:
        # Launch distributed and register configs
        jax.distributed.initialize()
        jax.tree_util.register_static(type(config))
    except RuntimeError:
        # This implies the distributed backend has already been initialized
        pass
    config_post_init(config)
    mesh = mesh_from_config(config)
    # ...

That’s it! We just have to be sure to set the config.mesh_shape parameter appropriately when we launch a job (see below), as well as all our desired sharding settings. Our config makes simple data parallel training over the entire list of devices the default.

Modifications for Data Loading

We refactor the data generation code into data.py. At a high level, here’s what we end up doing in train.py to set up the data loading:⁵

# in train.py
# ...
# Randomness
key = jax.random.key(config.seed)
key_model, key_train, key_val, key_eval = jax.random.split(key, 4)

# Data
batch_iter = get_distributed_batch_iter(config, config.train_dataset, key_train, mesh)
# ...

We abstract the previous dataset-specific code into this factory function, which calls the appropriate sub-functions (letting us reuse this for val/eval):

# in data.py


def dataset_dataloader_factory(config: DatasetConfig):
    match config.name:
        case DatasetName.MNIST:
            return (load_mnist(config), dataloader_without_replacement)
        case DatasetName.NUMBER_STAIRCASE:
            return (
                make_number_staircase_data(config),
                dataloader_with_replacement,
            )
        case DatasetName.SHAKESPEARE:
            return (
                load_shakespeare(config),
                dataloader_with_replacement,
            )


# Helper to just get the global batch iterator, if we don't need the local data/loader
def get_distributed_batch_iter(
    config: Config, dataset_config: DatasetConfig, key, mesh
):
    data, dataloader = dataset_dataloader_factory(dataset_config)
    return get_dataset_on_device(config, dataloader(key, dataset_config, data), mesh)

Distributed data loading has a large number of potential footguns; the JAX Advanced Guides listing has a useful entry about this (JAX Team, 2025), and see also the article about explicit sharding we used in previous blogs (JAX Team, 2025). In the present setting, we take advantage of the fact that we can load the entire dataset into memory on each host. The following code does this:

def make_number_staircase_data(config: DatasetConfig):
    # Simple data sequence!
    # Small, so mega replicate for ease
    text = "012345678987654321" * 1024
    ids = [int(c) for c in text]
    data = jnp.array(ids, dtype=jnp.int32)
    Xtr, Xdev, Xte = split_data(data, 0.8, 0.1)
    match config.split:
        case SplitType.TRAIN:
            return Xtr, jnp.array(0)
        case SplitType.TEST:
            return Xte, jnp.array(0)
        case SplitType.VAL:
            return Xdev, jnp.array(0)


def split_data(data: Array, train_fraction: float, dev_fraction: float):
    num_data = len(data)
    Xtr = data[: int(train_fraction * num_data)]
    Xdev = data[
        int(train_fraction * num_data) : int((train_fraction + dev_fraction) * num_data)
    ]
    Xte = data[int((train_fraction + dev_fraction) * num_data) :]
    return Xtr, Xdev, Xte

The tricky part here is to note that as configured, the above code generates the data array so that it is locally addressable by every process: each host creates its own (identical) data array. More precisely, if we evaluate data.is_locally_addressable, it will be True. This must be contrasted with the case of a fully replicated data array, which we don’t want here—our approach will be to piece together a global batch out of local batches on each device, and having a locally addressable data array is (currently) the canonical way to do this!

Next, we need to make sure each host has access to an independent and identically distributed stream of batches. This gets handled by the get_dataset_on_device and dataloader functions.

# in data.py

# Simple next-token prediction dataloader for text training:
# - Each host has the whole dataset (data input)
# - Random sampling without replacement to draw batches (consumes key)
# - Targets are the next token (expect data to be (num_data,) shape)
def dataloader_with_replacement(
    key, config: DatasetConfig, data: tuple[Array, Array]
) -> DataloaderOutputType:
    inputs, _ = data
    num_data = len(inputs)
    key = jax.random.fold_in(key, jax.process_index())
    for step in it.count():
        key = jax.random.fold_in(key, step)
        offsets = jax.random.randint(
            key,
            (config.global_batch_size // jax.process_count(),),
            0,
            num_data - config.seq_len - 1,
        )
        seqs = inputs.at[offsets[:, None] + jnp.arange(config.seq_len)].get()
        targets = inputs.at[1 + offsets[:, None] + jnp.arange(config.seq_len)].get()
        yield seqs, targets


# Helper to map a dataloader with make_array_from_process_local_data
def get_dataset_on_device(
    config: Config, dataloader: DataloaderOutputType, mesh: Mesh
) -> DataloaderOutputType:
    return map(
        lambda batch: jax.make_array_from_process_local_data(
            NamedSharding(mesh, jax.P(*config.sharding.data)), batch
        ),
        dataloader,
    )

We must pass the mesh we created earlier here in order to correctly create the global batch out of the local shards that the dataloader generates. Explicit sharding seems to propagate to outputs, so the data.at operations return locally-addressable batch arrays on each process, which make_array_from_process_local_data pieces together into the global array. To generate independent local batches on each process, we are sure to fold_in the process index to the RNG in order to create a distinct stream of random offsets.

This quick-and-dirty data loader, which does not shuffle nor guarantee sampling without replacement for batches across processes, is fine for our purposes on this synthetic dataset. It should also be okay for use with much larger datasets, when training for a relatively small number of tokens. But it’s worth keeping in mind that something more robust would be needed in general.⁶

Remaining Cleanup

All that’s left is to judiciously insert our mesh context manager at key points of the training loop. We saw above (when discussing optimizers) that we need to add it to the model initialization call. Other than that, we just need to add it to the model call in our training loop:

# in train.py

# ...
for step in range(config.num_steps):
    batch = next(batch_iter)
    with jax.set_mesh(mesh):
        cur_metrics, train_state = train_step(config, batch, train_state)
    # ...

For simplicity’s sake, we also wrap our sampling code with it.

Launching Everything

Given our DIY ethos, we handle launching distributed jobs on our TPU VM with an ad-hoc launch script:

# run.sh

TPU_NAME="$1"
shift
SSH_FLAGS='-A -o ForwardAgent=yes'
COMMANDS="if [ ! -d \"baremetal-gpt\" ]; then git clone git@github.com:sdbuch/baremetal-gpt; fi \
    && export HYDRA_FULL_ERROR=1 \
    && export WANDB_ENTITY='$WANDB_ENTITY' \
    && export WANDB_API_KEY='$WANDB_API_KEY' \
    && export HF_TOKEN='$HF_TOKEN' \
    && cd baremetal-gpt \
    && git fetch \
    && git checkout -f main \
    && git pull \
    && uv sync --extra tpu \
    && uv run train $@"

gcloud compute tpus tpu-vm ssh "$TPU_NAME" \
  --ssh-flag="$SSH_FLAGS" \
  --command="$COMMANDS" \
  --worker=all

For logging purposes, we need the Weights & Biases default entity and your API key to be exported as environment variables locally. We use Github for code synchronization, and provide the tpu extra to uv when building since we’re running on a TPU VM.

Here’s an example call on the local machine, for data-parallel training:

./deploy/run.sh tpu-v4-32 num_vocab=10 num_layers=4 num_steps=300 lr=3e-4 'mesh_shape=[16]'

For this to work well, when I first create the TPU VM,⁷ I make sure the following startup script is run on each host in the configuration:

# startup.sh
curl -LsSf https://astral.sh/uv/install.sh | sh  # Download uv
sudo cp /root/.local/bin/uv /usr/local/bin/uv
sudo cp /root/.local/bin/uvx /usr/local/bin/uvx
ssh-keyscan -H github.com >> /etc/ssh/ssh_known_hosts  # Code copying

This installs uv and adds github.com to the known hosts, so that I can connect to Github via ssh (agent forwarding is enabled for authentication). If we don’t add it to the known hosts list, it will ask to confirm connection and the gcloud command will fail.

That’s it! This gives a pretty easy minimal setup for distributed training with JAX on TPU VMs.

Adding Logging

Custom logging is easy to add, as well. We add an Enum to the config to select the logger, then add two basic loggers: one for printing to the command line, and one for using Weights and Biases.

There are two things to keep in mind here:

In a distributed setting, by default we only want one process (say, the one with index 0) to log metrics.⁸
We want to set up appropriate host-device pipelining, so as we alluded to in the previous post (taking a trick from the JAX training cookbook) we should log the metrics from the previous batch on every step. We can incorporate automatic buffering into our Logger class to do this.

# config.py

# ...
class LoggerType(Enum):
    PRINT = "print"
    WANDB = "wandb"


# ...


# loggers.py
def logger_factory(logger_type: LoggerType):
    match logger_type:
        case LoggerType.PRINT:
            return PrintLogger
        case LoggerType.WANDB:
            return WandbLogger


def get_run_name(base: str) -> str:
    """Prepend current UTC time to config-specified run name"""
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
    run_name = f"{timestamp}-{base}"
    return run_name


class Logger:
    """Logger base class. Implements a depth-1 buffer for device-host pipelining"""

    def __init__(self, config: Config):
        self.project_name = config.project_name
        self.run_name = get_run_name(config.run_name)
        if isinstance(config, DictConfig):
            config_dict = OmegaConf.to_container(config, resolve=True)
        else:
            config_dict = asdict(config)
        self.config = config_dict
        self.prev_log_data = None

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return False

    def log(self, log_dict: dict):
        pass

    def warn(self, message: str):
        pass

    def buffer(self, log_dict: dict) -> dict | None:
        self.prev_log_data, buffered_metrics = log_dict, self.prev_log_data
        return buffered_metrics

    def flush_buffer(self):
        if self.prev_log_data is not None:
            self.log({})
        self.prev_log_data = None


class PrintLogger(Logger):
    def __init__(self, config: Config):
        super().__init__(config)

    def __enter__(self):
        print(f"Project: {self.project_name}")
        print(f"Run: {self.run_name}")
        return super().__enter__()

    def log(self, log_dict: dict):
        if (buffered_dict := self.buffer(log_dict)) is None:
            return
        print(*[f"{metric}: {val}" for metric, val in buffered_dict.items()], sep="\t")

    def warn(self, message: str):
        print(message)


class WandbLogger(Logger):
    def __init__(self, config: Config):
        super().__init__(config)
        self.is_master = jax.process_index() == 0

    def __enter__(self):
        if self.is_master:
            wandb.init(
                project=self.project_name, name=self.run_name, config=self.config
            )
        return super().__enter__()

    def __exit__(self, exc_type, exc_value, traceback):
        if self.is_master:
            wandb.finish()
        return super().__exit__(exc_type, exc_value, traceback)

    def log(self, log_dict: dict):
        if self.is_master:
            if (buffered_dict := self.buffer(log_dict)) is None:
                return
            wandb.log(buffered_dict)

    def warn(self, message: str):
        if self.is_master:
            print(message)

Then we simply wrap our training loop with this context manager:

# train.py
# ...
Logger = logger_factory(config.logger_type)
# ...
with Logger(config) as logger:
    # ...
    for step, batch in enumerate(batch_iter):
        with jax.set_mesh(mesh):
            metrics, train_state = train_step(config, batch, train_state)
        logger.log(metrics | {"step": step})
        # ...
        if step == config.train_dataset.num_steps - 1:
            break
# ...

Actual Data

Now that we have a reasonable backend, let’s get it working on some text data, and try it at GPT-2 scale! Within the context of the above refactors, we just need to write the data preprocessing code and the dataloader. We also need to add evaluation code – we’ve done this in the Github, but not detailed it above.

We’ll do a minimal demo with Andrej Karpathy’s Tiny Shakespeare dataset, available in the nanoGPT repo. We adapt Karpathy’s script for downloading and tokenizing the dataset, making use of the tiktoken library for this purpose:

# data/download_and_tokenize_tiny_shakespeare.py
import os

import jax.numpy as jnp
import requests
import tiktoken

"""From https://github.com/karpathy/nanoGPT/blob/master/data/shakespeare/prepare.py"""

# download the tiny shakespeare dataset
input_file_path = os.path.join(os.path.dirname(__file__), "input.txt")
if not os.path.exists(input_file_path):
    data_url = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
    with open(input_file_path, "w", encoding="utf-8") as f:
        f.write(requests.get(data_url).text)

with open(input_file_path, "r", encoding="utf-8") as f:
    data = f.read()
n = len(data)
train_data = data[: int(n * 0.9)]
val_data = data[int(n * 0.9) :]

# encode with tiktoken gpt2 bpe
enc = tiktoken.get_encoding("gpt2")
train_ids = enc.encode_ordinary(train_data)
val_ids = enc.encode_ordinary(val_data)
print(f"train has {len(train_ids):,} tokens")
print(f"val has {len(val_ids):,} tokens")

# export to bin files
train_ids = jnp.array(train_ids, dtype=jnp.int32)
val_ids = jnp.array(val_ids, dtype=jnp.int32)
jnp.save("data/tiny-shakespeare/train.npy", train_ids)
jnp.save("data/tiny-shakespeare/val.npy", val_ids)

# train.bin has 301,966 tokens
# val.bin has 36,059 tokens

In a distributed setting, we run this on each host, with a command like:

# deploy/download_shakespeare.sh

TPU_NAME="$1"
shift
SSH_FLAGS='-A -o ForwardAgent=yes'
COMMANDS="if [ ! -d \"baremetal-gpt\" ]; then git clone git@github.com:sdbuch/baremetal-gpt; fi \
    && export HF_TOKEN='$HF_TOKEN' \
    && cd baremetal-gpt \
    && git fetch \
    && git checkout -f main \
    && git pull \
    && uv sync --extra tpu \
    && uv run data/tiny-shakespeare/download_and_tokenize_tiny_shakespeare.py $@"

gcloud compute tpus tpu-vm ssh "$TPU_NAME" \
  --ssh-flag="$SSH_FLAGS" \
  --command="$COMMANDS" \
  --worker=all

Then the data is ready for loading. We write a python loader for the data to go with our previously-written factor, using the same dataloader_with_replacement as for the synthetic data:

# data.py
def load_shakespeare(config: DatasetConfig):
    path = Path(config.path)
    data = jnp.load(path / (config.split.value + ".npy"))
    return data, jnp.array(0)

Then running the experiment is simply a matter of setting up the right experimental config. We use the following, which runs almost all the evaluation code we wrote (to test it!):

# configs/experiment/tiny-shakespeare.yaml
# @package _global_

defaults:
  - /dataset/shakespeare@train_dataset

val_log_interval: 250

train_dataset:
  num_steps: 3000

val_list:
  - dataset:
      name: SHAKESPEARE
      path: ${train_dataset.path}
      split: VAL
      seq_len: ${train_dataset.seq_len}  # At inference time, this functions like prompt size
      global_batch_size: ${train_dataset.global_batch_size}
      num_steps: 200
    evaluator: NLL
  - dataset:
      name: SHAKESPEARE
      path: ${train_dataset.path}
      split: VAL
      seq_len: 128  # At inference time, this functions like prompt size
      global_batch_size: 16
      use_splash: False
    evaluator: AUTOREGRESSIVE_ROLLOUTS

eval_list:
  - dataset:
      name: SHAKESPEARE
      path: ${train_dataset.path}
      split: VAL
      seq_len: 128  # At inference time, this functions like prompt size
      global_batch_size: 16
      use_splash: False
    evaluator: AUTOREGRESSIVE_ROLLOUTS

model:
  transformer_type: DISCRETE
  max_seq_len: ${train_dataset.seq_len}  # no longer than train
  is_causal: True
  num_vocab: 50257  # gpt2 tokenizer

optimizer:
  weight_decay: 1e-1


inference:
  tokenizer: GPT2
  max_tokens_to_generate: 128

The tokenizer parameters, as suggested, are just used for rollouts at inference time. We use some Hydra interpolation in the config above to avoid retyping the same values. Here’s the overall training code we run (which is pretty short!):

# train.py

# ...
@hydra.main(
    version_base=None,
    config_path=str(Path("configs").absolute().resolve()),
    config_name="base_config",
)
def main(config: Config):
    try:
        # Launch distributed and register configs
        jax.distributed.initialize()
        jax.tree_util.register_static(type(config))
    except RuntimeError:
        # This implies the distributed backend has already been initialized
        pass
    config_post_init(config)
    mesh = mesh_from_config(config)
    Logger = logger_factory(config.logger_type)

    # Randomness
    key = jax.random.key(config.seed)
    key_model, key_train, key_val, key_eval = jax.random.split(key, 4)

    # Data
    batch_iter = get_distributed_batch_iter(
        config, config.train_dataset, key_train, mesh
    )

    # Initialize state, configure forward pass and optimization
    with jax.set_mesh(mesh):
        train_state = init_train_state(key_model, config)
    cache_params = CacheParams(enabled=False, size=0)
    kernel = make_splash_kernel(config, config.train_dataset, 0, mesh)
    spec = model_spec(train_state.params)
    opt_update = opt_update_factory(config.optimizer.type)
    weight_decay_mask = jax.tree.map(lambda _, s: bool(s), train_state.params, spec)

    @partial(jax.jit, donate_argnums=2)
    def train_step(config: Config, batch, train_state: TrainState):
        def loss_fn(params: Transformer):
            inputs, targets = batch
            logits, _ = jax.vmap(
                partial(_transformer, config, kernel, params, cache_params=cache_params)
            )(inputs, train_state.kv_cache)
            logits = logits.astype(config.model.compute_dtype.value)
            logprobs = jax.nn.log_softmax(logits, axis=-1)
            return -jnp.take_along_axis(logprobs, targets[..., None], axis=-1).mean()

        loss, grad = jax.value_and_grad(loss_fn)(train_state.params)
        grad_clipped, _, global_grad_norm = grad_norm_and_clip(config, grad)
        update__opt_state = jax.tree.map(
            partial(opt_update, config),
            train_state.params,
            grad_clipped,
            train_state.opt_state,
            weight_decay_mask,
        )
        # Transpose the output tree to get update tree and state tree
        update, opt_state = map(
            lambda i: jax.tree.map(lambda x, y: y[i], grad, update__opt_state), range(2)
        )
        params = jax.tree.map(lambda x, y: x + y, train_state.params, update)
        new_state = TrainState(
            params=params, opt_state=opt_state, kv_cache=train_state.kv_cache
        )

        metrics = {"batch_loss": loss, "grad_norm": global_grad_norm}
        return metrics, new_state

    # Training loop
    with Logger(config) as logger:
        do_evals = partial(eval_loop, config, mesh=mesh, logger=logger)
        for step, batch in enumerate(batch_iter):
            with jax.set_mesh(mesh):
                metrics, train_state = train_step(config, batch, train_state)
            logger.log(metrics | {"step": step})
            if (step + 1) % config.val_log_interval == 0:
                # Calculate val metrics
                key_val = do_evals(key_val, config.val_list, train_state.params, step)
            if step == config.train_dataset.num_steps - 1:
                break

        # Run evals (testing)
        key_eval = do_evals(key_eval, config.eval_list, train_state.params, step)


def eval_loop(
    config: Config,
    key,
    eval_list: list[EvaluationConfig],
    params: Transformer,
    step: int,
    logger: Logger,
    mesh,
):
    logger.flush_buffer()
    for evaluation in eval_list:
        kernel = make_splash_kernel(config, evaluation.dataset, 0, mesh)
        key, key_d, key_e = jax.random.split(key, 3)
        batch_iter = get_distributed_batch_iter(config, evaluation.dataset, key_d, mesh)
        evaluation_fn = evaluator_factory(evaluation)
        metrics = evaluation_fn(config, key_e, kernel, mesh, params, batch_iter)
        logger.log(metrics | {"step": step})
    logger.flush_buffer()
    return key

Check out the evaluators.py file in the Github for the full structure of the evaluation code.

Checking the Results

The default model we train has 12 layers, embedding dimension 768, 12 heads per attention layer, and a MLP expansion factor of 4. With GPT2 vocabulary size, this ends up at about 160M parameters. Our Tiny Shakespeare config follows nanoGPT, setting the defaults as:

# configs/dataset/shakespeare.yaml
name: SHAKESPEARE
path: "data/tiny-shakespeare"
split: TRAIN
seq_len: 256
global_batch_size: 128

So we process 32K tokens per batch. The total number of tokens in the training set is only 301,966! So without additional regularization, we expect to overfit this dataset very rapidly with our model.

We test this out, via the launch command

./deploy/run.sh tpu-v4-32 +deploy=v4-16 +experiment=tiny-shakespeare

on a 16-chip TPU v4 VM (4 hosts). Following the experiment config we saw above, we train for 3000 steps, and log val metrics every 250 steps.

Plot of batch loss (negative log-likelihood) versus training step for the first 500 steps of Tiny Shakespeare training. The model overfits the training data.

We see that the model rapidly overfits the training data. It takes about 10 steps per epoch, so the plot shows about 50 epochs worth of training – significantly more than we encounter in standard language model training datasets.

Plot of gradient norm versus training step for the first 500 steps of Tiny Shakespeare training. Training is stable, even without learning rate scheduling.

The model trains stably (we are performing gradient clipping to norm 1.0 by default), even without learning rate scheduling. Scheduling would improve the initial stability; since we are heavily overfitting, there does not seem to be any need for LR cooldown at the end of training.

Plot of validation loss (negative log likelihood) versus training step for the 3000 step Tiny Shakespeare training trajectory. The model is heavily overfit.

Our evaluation code lets us log various metrics, including the validation set negative log likelihood every 250 steps of training. We verify that the model is significantly overfit to the training set – the validation loss actually increases monotonically! Taking a look at some of the trained model’s rollouts verifies that they are indeed not generalizing very well (although they’re quite fun to read):

################################
Prompt:
################################
 father? Gentles, methinks you frown:
And wherefore gaze this goodly company,
As if they saw some wondrous monument,
Some comet or unusual prodigy?

BAPTISTA:
Why, sir, you know this is your wedding-day:
First were we sad, fearing you would not come;
Now sadder, that you come so unprovided.
Fie, doff this habit, shame to your estate,
An eye-sore to our solemn festival!

TRANIO:
And tells us, what occasion of import
Hath all
################################
Generated text:
################################
 the castle wall:
And wheither is come to me, and sit by fortune's side,
The right is not! God's my son,
Immoderate but a little distressed widow;
For joyful wife's son is no weary way.

ROMEO:
What say'st thou, that didst me yet?

MERCUTIO:
Good dream of mine own.

ROMEO:
Give me that mattock and the wrenching iron.
Hold, take this letter; early, and my lord

Till then be urged to my advancement?

Nurse:

The model seems to be regurgitating tokens from Romeo and Juliet, with some random digressions sprinkled in! For these rollouts, we generate up to max_seq_len tokens (based on the training config; for us, it’s 256) off of a 128 token prompt. This avoids out-of-distribution RoPE bugs and KV cache errors.

Conclusions

That’s all for our base infra buildout! Some of the key takeaways include:

Correct configuration of splash attention Pallas kernel;
Flexible infrastructure for training and inference across different datasets and evaluation types;
Basic training improvements (with fairly robust infrastructure, except perhaps for our model spec) that will let us experiment with new optimizers in the future!

As we’ve gone from the notebook-based code in the previous blog to the full repository, we’ve seen along the way that building reasonably general experiment infrastructure requires a bit of a different approach than writing one-off research code. It also requires a decent amount of time! There’s a good chance a lot of this can be automated in a fast-paced environment with LLMs – I have not used LLMs for writing any code in these blogs – but writing the code by hand is hard to beat for learning and for precision of implementation (as Karpathy suggested in his recent conversation with Dwarkesh!). As a small piece of evidence of this, we’ve ended up with a very concise and readable train.py script: not much more than 150 lines!

Some remaining infrastructure-related tweaks we should add in the future:

Model checkpointing: we’ll add this at some point soon, likely just using Orbax.
Further parallelism configurations, in particular basic FSDP.

Look forward to less infrastructure and closer-to-research experimentation in future iterations of this blog series.

Acknowledgments

Thanks to the TRC program for compute support.

The reason we actually check for None is that we generally need to call our model with different configurations (e.g., when doing autoregressive inference, or when evaluating with a different parallelism configuration), and it is suboptimal to create an entirely different Config object for each of these, or to manually edit the global Config object for these as we run. So instead, we check for None, and just pass this as the kernel parameter to the model as we need to. ↩
For example, we typically only apply weight decay to matrix-like parameters in the network (e.g., not to biases). ↩
Note that even this approach is not perfect, as for example the w_out matrix associated to the attention output projection actually does reduce over one of the dimensions marked as non-matrix. We should revisit this for future improvements (i.e., Muon-type optimizers!). ↩
At least in our setting, where we’re not using multislice operation. I think the distributed backend is just as easy to set up for multislice, though! ↩
Creating the training/val/test split on each distributed process is probably not a good idea, even though we seed each host’s RNG with the same seed, and generate the splits with the same key. It won’t matter for our purposes here because the synthetic data distribution is so simple (even if every process had a different training split, it would be okay, because there isn’t a material distinction between the training and test set for this distribution). We’ll fix this with proper processing of data when we move to text data. ↩
We could use Grain for this purpose! ↩
Here, tpu-v4-32, for a four-host v4 setup: the 32 denotes the number of cores, chips = cores / 2, and the number of chips dictates the mesh shape. ↩
For more complex debugging, say in interpretability settings, we need to add a bit more custom code in order to have each process log its respective metrics (say, per-device gradients before all reduce, etc.). ↩

A Faster Manifold Muon with ADMM

2025-10-20T00:00:00+00:00

In a recent blog post, Jeremy Bernstein (Thinking Machines) gave an algorithm for optimizing matrix-valued parameters—say, weights in a transformer—that ensures that both the update directions and the parameters themselves are well-conditioned (Bernstein, 2025). This algorithm, called manifold Muon, has an inner loop that is relatively slow to converge, and requires some hyperparameter tuning to get right.

Comparison of the default manifold Muon inner loop solver (dual ascent) to a solver based on ADMM, for one weight matrix in a neural network being trained. ADMM converges much more rapidly, with minimal hyperparameter tuning!

In this post, we’ll describe a modification of the manifold Muon algorithm that greatly speeds up convergence—by more than a factor of two in wall-clock time in our experiments on 1x H100! To realize this, we apply a venerable algorithm from convex optimization called ADMM (the “alternating direction method of multipliers”) to the optimization problem at the heart of the manifold Muon update.

The rest of the post will re-introduce the manifold Muon subproblem, give some intuition for why the existing implementation leaves some room for improvement on the table, and walk through the ADMM derivation for manifold Muon. We’ll verify the improved algorithm on the same toy model used in the linked Github repository. We hope this improvement will facilitate larger-scale experimentation with manifold Muon!

If you haven’t already, you’re encouraged to read Jeremy’s blog post on manifold Muon before proceeding. It gives an excellent technical overview and derivation of the method, and accessible, well-written intuition on the deep links between optimization and geometry. See also Jianlin Su’s work (苏剑林, 2025), which gives a more thorough technical derivation.

Background: Manifold Muon

Manifold Muon is a first-order method for optimizing a matrix-valued parameter $\vW \in \bbR^{m \times n}$ (we’ll assume $m \geq n$ throughout), such that it satisfies the quadratic constraint $\vW^\top \vW = \vI$. Given a putative update direction $\vG \in \bbR^{m \times n}$—say, back-propagated gradients of some loss with respect to $\vW$—and a step size $\eta > 0$, we choose an update $\vA$ by solving the convex optimization problem¹

\[\begin{equation}\label{eq:manifold-muon} \min_{\vA \in \bbR^{m \times n}}\, \ip{\vA}{\vG} \quad \text{subject to} \quad \norm{\vA} \leq \eta, \quad \vA^\top \vW + \vW^\top \vA = \Zero, \end{equation}\]

then perform the update

\[\vW \mapsto \mathrm{msign} \left( \vW + \vA \right) ,\]

where $\mathrm{msign}(\vX) = \vX (\vX^\top \vX)^{-1/2}$ is the matrix sign function.²

The manifold Muon subproblem \eqref{eq:manifold-muon} does not have a closed-form solution in general, so Jeremy proposed an iterative solver for it, based on the equivalent unconstrained formulation³

\[\begin{equation}\label{eq:manifold-muon-dual} \min_{\vLambda \in \bbR^{n \times n}}\, \norm*{\vG + \vW\left( \vLambda + \vLambda^\top \right)}_*. \end{equation}\]

One can obtain an optimal solution $\vA_\star$ to \eqref{eq:manifold-muon} from an optimal solution $\vLambda_\star$ to \eqref{eq:manifold-muon-dual} via $\vA_\star = -\eta \mathop{\mathrm{msign}}( \vG + \vW(\vLambda_\star + \vLambda_\star^\top))$.

In practice, we solve \eqref{eq:manifold-muon-dual} with subgradient descent: we pick an iteration count $K \in \bbN$ and step sizes $\gamma_k > 0$ for $k = 1, \dots, K$, and repeatedly perform the update

\[\begin{equation}\label{eq:subgradient} \begin{split} \vGamma_{k} &= \mathop{\mathrm{msign}}( \vG + \vW(\vLambda_k + \vLambda_k^\top))^\top \vW + \vW^\top\mathop{\mathrm{msign}}( \vG + \vW(\vLambda_k + \vLambda_k^\top)); \\ \vLambda_{k+1} &= \vLambda_k - \gamma_k \vGamma_k. \end{split} \end{equation}\]

The step sizes $\gamma_k$ are chosen to decay as $k \to K$ to guarantee convergence. The gradient $\vGamma_k$ is computed efficiently on modern accelerators using Newton-Schulz iterations, or optimal algorithms such as the “Polar Express” algorithm (Amsel et al., 2025), as in the Github for Jeremy’s blog.

Drawbacks of Subgradient Descent

The subgradient descent iteration \eqref{eq:subgradient} is guaranteed to converge, but it does so very slowly! In fact, its worst-case rate of convergence is $O(1/\sqrt{k})$ (Wright & Ma, 2022), in contrast to the standard $O(1/k)$ rate for gradient descent on smooth objectives.

Even worse, this is typical behavior for this algorithm, not just worst-case. It stems from the fact that subgradient descent is an algorithm for nonsmooth problems, and nonsmoothness leads to a ‘chattering’ behavior of iterates. Here is a quick simulation on a canonical example, the absolute value function $\abs{\spcdot}$ in dimension one, to illustrate:⁴

from math import sqrt

init = 1.0
lr = 2.0
iters = []
grads = []
it = init
for i in range(100):
    grad = 1 if it >= 0 else -1
    it = it - lr / sqrt(i + 1) * grad
    iters.append(it)
    grads.append(grad)

Subgradient descent on the absolute value function, with $1/\sqrt{k}$ LR decay.

The nuclear norm objective in \eqref{eq:manifold-muon-dual} is nondifferentiable at every point $\vLambda$ for which $\vG + \vW(\vLambda + \vLambda^\top)$ is not full rank. Even if the iterates $\vLambda_k$ never witness such a point of nondifferentiability, the lack of continuity of the gradient near these points leads to the same ‘chattering’ behavior of the iterates as in the toy example above.

Fortunately, we can do better, with a bit of algorithmic cleverness (and without too much more compute)!

A Faster Algorithm from ADMM with Splitting

ADMM is an optimization algorithm for solving problems of the form

\[\begin{equation}\label{eq:admm} \begin{split} \min_{\vx \in \bbR^{m}, \vz \in \bbR^{n}}\, &f(\vx) + g(\vz) \\ \text{subject to}\,\, &\vA \vx + \vB \vz = \vc, \end{split} \end{equation}\]

where $f$ and $g$ are convex. It was originally developed in the 1970s, and was re-popularized among a generation or two of researchers by the excellent monograph of Boyd, Parikh et al. (Boyd et al., 2011), which we recommend as a general reference beyond the rapid overview we give here. The key empirical property that makes ADMM very well suited for use in solving the manifold Muon subproblem \eqref{eq:manifold-muon-dual} is its rapid initial convergence to a solution of reasonable quality—the algorithm’s worst-case convergence rate is only $O(1/k)$, but this initial rapid convergence often allows it to be stopped much earlier.

The ADMM algorithm solves problems of the form \eqref{eq:admm} as follows. For a penalty parameter $\rho > 0$, we define the “augmented Lagrangian”

\[\sL_{\rho}(\vx, \vz, \vlambda) = f(\vx) + g(\vz) + \ip{\vlambda}{\vA \vx + \vB \vz - \vc} + \frac{\rho}{2} \norm*{ \vA \vx + \vB \vz - \vc }_2^2,\]

then iteratively perform the updates

\[\begin{equation}\label{eq:admm-update} \begin{split} \vx_{k+1} &= \argmin_{\vx}\, \sL_{\rho}(\vx, \vz_k, \vlambda_k) \\ \vz_{k+1} &= \argmin_{\vz}\, \sL_{\rho}(\vx_{k+1}, \vz, \vlambda_k) \\ \vlambda_{k+1} &= \vlambda_k + \rho \left( \vA \vx_{k+1} + \vB \vz_{k+1} - \vc \right). \end{split} \end{equation}\]

The algorithm is reminiscent of a block coordinate descent method, with gradient ascent on the dual variable $\vlambda$ (with step size $\rho$). It converges for any setting of the penalty parameter $\rho$. Empirically, this parameter can be tuned to accelerate convergence.

ADMM is ideally suited for problems where the minimization subproblems in the $\vx$ and $\vz$ updates in \eqref{eq:admm-update} have closed-form solutions. For the manifold Muon subproblem \eqref{eq:manifold-muon-dual}, we need to pass two related obstacles:

How do we apply ADMM to \eqref{eq:manifold-muon-dual}? It doesn’t seem to have the ADMM-amenable structure of \eqref{eq:admm}.
Having resolved this, do the minimization subproblems in the ADMM update \eqref{eq:admm-update} applied to \eqref{eq:manifold-muon-dual} have closed-form solutions that can be computed efficiently on GPUs/TPUs?

The Immortal Splitting Trick

To resolve the first issue, we apply a trick known as variable splitting. We add an auxiliary variable $\vX \in \bbR^{m \times n}$ to the manifold Muon subproblem \eqref{eq:manifold-muon-dual} along with an extra constraint to give the equivalent problem

\[\begin{equation}\label{eq:manifold-muon-split} \begin{split} \min_{\vLambda \in \bbR^{n \times n}, \vX \in \bbR^{m \times n}}\, &\norm*{\vX}_* \\ \text{subject to}\,\, &\vX = \vG + \vW(\vLambda + \vLambda^\top). \end{split} \end{equation}\]

This may initially seem counterintuitive—didn’t we work hard to reduce the original manifold Muon problem \eqref{eq:manifold-muon} to an unconstrained form? But the advantage is that this problem is now amenable to ADMM, since it has the form of \eqref{eq:admm} after we re-interpret matrices as vectors!

For this problem, given that we’ve already used $\vLambda$ for a dual variable, we’ll write the ADMM dual variable as $\vOmega \in \bbR^{m \times n}$, giving the augmented Lagrangian

\[\begin{equation}\label{eq:augmented-lagrangian} \begin{split} \sL_{\rho}(\vLambda, \vX, \vOmega) = \norm*{\vX}_* &+ \ip{\vOmega}{\vX - \vW(\vLambda + \vLambda^\top) - \vG} \\ &+ \frac{\rho}{2} \norm*{ {\vX - \vW(\vLambda + \vLambda^\top) - \vG} }_{\frob}^2. \end{split} \end{equation}\]

Now we can work on the ADMM subproblems for this augmented Lagrangian.

ADMM Subproblems for Manifold Muon

In deriving ADMM algorithms for specific problems, it’s often useful to “complete the square” to more easily see the solutions to the minimization subproblems. We rewrite \eqref{eq:augmented-lagrangian} as

\[\begin{equation} \begin{split} \sL_{\rho}(\vLambda, \vX, \vOmega) &= \norm*{\vX}_* + \frac{\rho}{2}\norm*{\frac{1}{\rho}\vOmega}_{\frob}^2 - \frac{\rho}{2}\norm*{\frac{1}{\rho}\vOmega}_{\frob}^2 + \ip{\vOmega}{\vX - \vW(\vLambda + \vLambda^\top) - \vG} \\ &\qquad\qquad+ \,\frac{\rho}{2} \norm*{ {\vX - \vW(\vLambda + \vLambda^\top) - \vG} }_{\frob}^2 \\ &= \norm*{\vX}_* - \frac{\rho}{2}\norm*{\frac{1}{\rho}\vOmega}_{\frob}^2 + \frac{\rho}{2} \norm*{ \frac{1}{\rho} \vOmega + {\vX - \vW(\vLambda + \vLambda^\top) - \vG} }_{\frob}^2. \end{split} \end{equation}\]

Now the dependence on $\vLambda$ is an easy-to-parse quadratic function, and the dependence on $\vX$ is also simple.

$\vLambda$ Subproblem

To solve

\[\argmin_{\vLambda}\, \sL_{\rho}(\vLambda, \vX_k, \vOmega_k) = \argmin_{\vLambda}\, \norm*{ \frac{1}{\rho} \vOmega_k + {\vX_k - \vW(\vLambda + \vLambda^\top) - \vG} }_{\frob}^2,\]

it’s helpful to define $\Sym(\vLambda) = \half ( \vLambda + \vLambda^\top)$. This is a linear operator, and in fact is an orthogonal projection. Then because $\vW^\top \vW = \vI$, we have

\[\argmin_{\vLambda}\, \sL_{\rho}(\vLambda, \vX_k, \vOmega_k) = \argmin_{\vLambda}\, \norm*{ \Sym(\vLambda) - \frac{1}{2}\vW^\top\left( \frac{1}{\rho} \vOmega_k + \vX_k - \vG \right) }_{\frob}^2.\]

It then follows that

\[\frac{1}{2}\Sym\left(\vW^\top\left( \frac{1}{\rho} \vOmega_k + \vX_k - \vG \right)\right) \in \argmin_{\vLambda}\, \sL_{\rho}(\vLambda, \vX_k, \vOmega_k),\]

and that this is the minimum Euclidean norm solution to this problem.

In particular, we can solve the $\vLambda$ subproblem in closed form, using only matrix multiplies and transposes!

$\vX$ Subproblem

We have to solve

\[\begin{split} \argmin_{\vX}\, &\sL_{\rho}(\vLambda_{k+1}, \vX, \vOmega_k) = \\ &\qquad\frac{1}{\rho}\norm*{\vX}_* + \frac{1}{2} \norm*{ \vX + \frac{1}{\rho} \vOmega_k - \vW(\vLambda_{k+1} + \vLambda_{k+1}^\top) - \vG }_{\frob}^2. \end{split}\]

This optimization problem actually has a closed-form solution called singular value thresholding, which can be derived by considering the SVD of everything that doesn’t depend on $\vX$ in the quadratic term. Namely, if we define $\vM = \vW(\vLambda_{k+1} + \vLambda_{k+1}^\top) + \vG - \frac{1}{\rho} \vOmega_k$ for concision and let $\vU$, $\diag(\vs)$, $\vV$ denote an economy SVD of $\vM$, then it can be shown that (Wright & Ma, 2022)

\[\argmin_{\vX}\, \sL_{\rho}(\vLambda_{k+1}, \vX, \vOmega_k) = \vU \diag \left( \max \set*{ \vs - \frac{1}{\rho}, 0 } \right)\vV^\top,\]

with the natural broadcasting interpretation. Thus, the $\vX$ subproblem can be solved in closed-form with a singular value decomposition! Unfortunately, this is not enough for our purposes: SVD algorithms are not as well-suited for scaling up on GPUs/TPUs as algorithms that rely purely on matrix multiplications.

Fortunately, it turns out that we can compute singular value thresholding using only the $\mathrm{msign}$ algorithm, letting us exploit efficient algorithms for this operation that have already been developed for Muon!⁵ There are likely to be better custom solutions for this problem, but this quick-and-dirty approach will get us on our way. Here’s how:

First, we make a simple observation: if $s_i - 1/\rho \geq 0$ for every singular value $s_i$, then
\[\vU \diag \left( \max \set*{ \vs - \frac{1}{\rho}, 0 } \right)\vV^\top = \vM - \frac{1}{\rho} \vU \vV^\top.\]
Since $\vU\vV^\top = \mathrm{msign}(\vM)$, we’re done!
We cannot rely on every singular value being larger than $1/\rho$ in general. But, observe that those singular values that are smaller simply need to be set to $0$, which we can achieve by other means. More precisely, since $n \leq m$, consider the matrix $\mathrm{msign}(\vM^\top \vM - \frac{1}{\rho^2} \vI)$. Let $\vV_{\geq}$ denote the submatrix of (right) singular vectors of $\vM$ associated to singular values that are greater than or equal to $1/\rho$, and $\vV_{<}$ the complementary submatrix. Then one has
\[\mathrm{msign}(\vM^\top \vM - \frac{1}{\rho^2} \vI) = \vV_{\geq} \vV_{\geq}^\top - \vV_{<} \vV_{<}^\top.\]
Each of the (unsigned) factors in this sum is an orthogonal projection matrix. Hence we have
\[\frac{1}{2} \left( \vI + \mathrm{msign}(\vM^\top \vM - \frac{1}{\rho^2} \vI) \right) = \frac{1}{2} \left(\vV_{\geq} \vV_{\geq}^\top + \left(\vI - \vV_{<} \vV_{<}^\top \right) \right) = \vV_{\geq} \vV_{\geq}^\top.\]
The matrix $\vV_{\geq} \vV_{\geq}^\top$ projects onto the right singular vectors of $\vM$ associated to singular values at least as large as $1/\rho$. This is exactly what we need to implement the zero clipping operation in singular value thresholding!

In summary, we’ve shown how to compute singular value thresholding using only matrix sign operations. Zooming out, we’ve shown that the solution to the $\vX$ subproblem in our manifold Muon ADMM algorithm is

\[\argmin_{\vX}\, \sL_{\rho}(\vLambda_{k+1}, \vX, \vOmega_k) = \frac{1}{2}\left( \vM - \frac{1}{\rho} \mathrm{msign}(\vM) \right) \left( \vI + \mathrm{msign}\left(\vM^\top \vM - \frac{1}{\rho^2} \vI\right) \right).\]

This can be computed using only matrix multiplications and matrix sign operations, which themselves are efficiently approximated using matrix multiplications. Almost as hardware-efficient as we could hope for!

ADMM for Manifold Muon: Final Algorithm

Plugging our derivations above into the basic ADMM skeleton \eqref{eq:admm-update} leads to the following final algorithm, here in math form:⁶

\[\begin{equation}\label{eq:admm-update-mfld-muon} \begin{split} \vLambda_{k+1} &= \frac{1}{2}\Sym\left(\vW^\top\left( \frac{1}{\rho} \vOmega_k + \vX_k - \vG \right)\right) \\ \vM_{k+1} &= 2\vW\vLambda_{k+1} + \vG - \frac{1}{\rho} \vOmega_k \\ \vX_{k+1} &= \frac{1}{2}\left( \vM_{k+1} - \frac{1}{\rho} \mathrm{msign}(\vM_{k+1}) \right) \left( \vI + \mathrm{msign}\left(\vM_{k+1}^\top \vM_{k+1} - \frac{1}{\rho^2} \vI\right) \right) \\ \vOmega_{k+1} &= \vOmega_k + \rho \left( \vX_{k+1} - 2\vW\vLambda_{k+1} - \vG \right). \end{split} \end{equation}\]

When we’re done iterating, just as we did for dual ascent, we return $\vA_\star = -\eta \mathop{\mathrm{msign}}( \vG + 2\vW\vLambda_\star)$.

In code form, implementing the ADMM update we have derived above in PyTorch and combining it with the remainder of the manifold Muon algorithm leads to the following. This code is available in my fork of Jeremy’s Github repository, and it’s a modification of Jeremy’s original code. Disclaimer: it has not been heavily optimized, beyond choosing sensible defaults for the hyperparameters via experiments we’ll share below!

@torch.no_grad()
def manifold_muon_admm(W, G, eta=0.1, steps=10, rho=4.0):
    # Ensure that W and G are both tall matrices
    should_tranpose = W.shape[0] < W.shape[1]
    if should_tranpose:
        W = W.T
        G = G.T
    # Initialize the lagrangian, slack, and dual variable
    Lambda = -0.25 * (W.T @ G + G.T @ W)
    X = G + 2 * W @ Lambda
    Omega = torch.zeros_like(X)
    # Solve the dual problem with ADMM to find the update direction A
    for step in range(steps):
        # Update for Lambda (orthonormal least-squares solve)
        P = W.mT @ (1 / rho * Omega + X - G)
        Lambda_upd = 0.25 * (P + P.mT)
        # Update for X (singular value thresholding)
        B = G + 2 * W @ Lambda_upd - 1 / rho * Omega
        eye = torch.eye(B.shape[1], device=B.device, dtype=B.dtype)
        P_pos = 0.5 * (eye + msign(B.mT @ B - 1 / rho**2 * eye))
        X_upd = (B - 1 / rho * msign(B)) @ P_pos
        # Update for Omega (dual ascent)
        Omega_upd = Omega + rho * (X_upd - 2 * W @ Lambda_upd - G)
        Lambda, X, Omega = Lambda_upd, X_upd, Omega_upd
    # Calculate A from final ADMM solution
    # (at convergence, G + 2 * W @ Lambda \approx X)
    A = msign(G + 2 * W @ Lambda)
    # Descend on the primal problem
    new_W = W - eta * A
    # Retract to the manifold
    new_W = msign(new_W)
    # Restore the shape of the solution and return
    return new_W.T if should_tranpose else new_W

In this code, msign calls Jeremy’s implementation of the “Polar Express” algorithm (Amsel et al., 2025), which is a fairly aggressive algorithm for efficiently approximating the matrix sign function. It’s worth keeping this in mind, as a source of numerical inexactness in our subsequent experiments: for ADMM to be useful to us, it has to be able to cope well with these approximations! There are also some ‘reparameterizations’ of learning rates relative to the notation we’ve used above, but the algorithm is equivalent.

Intuition for the Speedup over Dual Ascent

As a final sanity check, we can use our derivations above to visualize a toy case of the manifold Muon ADMM algorithm, and contrast it with the pathological behavior of subgradient descent we visualized before. If we pick $m = n = 1$, $W = 1$ and $G = 0$ and simplify with algebra, the $\Lambda_k$ and $\Omega_k$ variables actually become redundant, and we end up with the very simple equivalent ADMM update

\[X_{k+1} = \sign(X_k) \max \set{\abs{X_k} - 1/\rho, 0}.\]

It should be intuitively clear that this iteration converges stably to zero: we reduce the magnitude of $X_k$ by $1/\rho$ until it drops below zero, then it stays there! On an intuitive level, this is because the $\vX$ update in the ADMM algorithm we derived for manifold Muon is continuous, whereas the subgradient descent update is discontinuous.⁷

ADMM on the absolute value function with splitting (after algebraic simplifications), with $\rho = 16.0$.

Experimental Results

We base our experiments on the experimental setup in the Thinking Machines Github repository for the blog post: a simple three-layer MLP is trained on CIFAR-10 for five epochs, and per-epoch timing information is reported.

Here’s a summary of what we’ll look at below:

We perform a sweep over the ADMM penalty parameter $\rho$ and the number of ADMM iterations for each manifold Muon solve. This gives us a sensible choice of hyperparameters for ADMM for this setup, and lets us verify performance.
Since the model being trained is so small-scale, we also plot full dual ascent and ADMM inner loop loss curves for each iteration of the outer loop and eyeball these to assess how fast each algorithm is converging.⁸ We believe this ranking will be more indicative of properties of the algorithms that persist at scale (since raw test error may not be a suitable proxy, based on what we’ll see below).
Having convinced ourselves of sensible ADMM performance at different choices of inner loop iterations, we look at timing!
We finally look at some other relevant metrics for ADMM, such as feasibility (how well are we satisfying the splitting constraint).

The code for these experiments is available in the admm branch of my fork of Jeremy’s Github repository.

Parameter Sweep

The main takeaway of this experiment is that the model training might be a bit too small scale! In general, for this model, manifold Muon leads to a fairly similar test error for all methods examined, including running zero inner loop iterations. This suggests a need to examine behaviors further at larger scale. Nevertheless, we sweep over values in $\set{0, 5, 10, 25, 50, 100}$ for the number of inner loop iterations to run, and for ADMM, we sweep over $\rho \in \set{2, 4, 8, 16}$. We also run dual ascent as a baseline, with the same initialization as in Jeremy’s repo.⁹

Plot of test accuracy vs. total inner loop step counts for different ADMM penalty parameters (or dual ascent).

From this plot, we see that test performance is not directly related with the number of steps we run the inner optimizer for. To try to have some generalizable signal towards larger-scale experiments, we will focus instead on assessing convergence of the inner loop with different representative hyperparameter values. Based on this plot, we select $(K, \rho) = (10, 4.0)$, $(100, \mathrm{DA})$ (for dual ascent), and as a baseline $(100, 8.0)$. This will give us a sample of ADMM’s speed of convergence, as well as its performance run for a large number of iterations.

Inner Loop Convergence

We compare the inner loop losses for the hyperparameter settings above.

We can choose to visualize the inner loop loss for one of three possible weight matrices (first layer, second layer, third layer), for any one of the outer loop iterations (we do a total of five epochs over the data).

Dual ascent losses for each weight matrix in the network, for the nonzero-steps hyperparameter choices shown in the table above. Traces can be identified based on where they terminate.

Notice that both nonsmoothness and initial scale of the learning rate for dual ascent being too large both seem to play a role in the slow convergence of dual ascent in the plots above, which are evaluated for each weight matrix at the first epoch and first iteration of optimization. Dual ascent needs all 100 iterations to reach a good loss value, whereas the ADMM algorithms converge much more quickly. The ADMM algorithms only require setting the penalty parameter, and are not very sensitive to it in this experiment, as long as it is sufficiently large: we found smaller values of the penalty parameter to lead to instabilities in inner loop loss convergence, although outer loop performance was not significantly affected.¹⁰

For the ADMM iteration counts, ten iterations is aggressive, but reasonable: only for fc3 does it end up a bit far from convergence, and this seems to be due to the setting of the penalty parameter we’ve chosen more than anything else.

These plots are more or less representative for all iterations of training. In later iterations, the loss decrease is often less dramatic than at the first iteration, but trends persist: dual ascent is slower, ten ADMM iterations is aggressive, but not unreasonable (see below).

Dual ascent losses for weights in the first layer of the network, for the nonzero-steps hyperparameter choices shown in the table above. Traces can be identified based on where they terminate.

Timing and Other Diagnostics

We profile the code with the pre-implemented timing commands in the original repository. It would be good to do this in a more rigorous setting, also after optimizing the code a bit more (e.g., adding compilation to unroll and fuse the inner loop where possible). But this test gives us a rough sense of how much we stand to improve.

Here are outputs for the dual ascent algorithm (100 iterations), run on a 1x H100 VM.

ubuntu@flaky-muskox-1:~/thinky-manifolds$ uv run src/main.py
Training with: manifold_muon
Epochs: 5 --- LR: 0.1
Epoch 1, Loss: 1.6895257502186054, Time: 11.8445 seconds
Epoch 2, Loss: 1.4129536565469236, Time: 11.7680 seconds
Epoch 3, Loss: 1.2799630311070656, Time: 11.7151 seconds
Epoch 4, Loss: 1.183229492635143, Time: 11.6794 seconds
Epoch 5, Loss: 1.1131725724862547, Time: 12.1164 seconds
Accuracy of the network on the 10000 test images: 52.58 %
Accuracy of the network on the 50000 train images: 65.936 %

And here are outputs for the ADMM algorithm (10 iterations).

ubuntu@flaky-muskox-1:~/thinky-manifolds$ uv run src/main.py --update manifold_muon_admm
Training with: manifold_muon_admm
Epochs: 5 --- LR: 0.1
Epoch 1, Loss: 1.6916365696459401, Time: 5.1556 seconds
Epoch 2, Loss: 1.4166164446850211, Time: 5.1545 seconds
Epoch 3, Loss: 1.2806039537702287, Time: 5.0617 seconds
Epoch 4, Loss: 1.1837433649569142, Time: 5.0579 seconds
Epoch 5, Loss: 1.1136368123852476, Time: 5.0620 seconds
Accuracy of the network on the 10000 test images: 52.62 %
Accuracy of the network on the 50000 train images: 66.086 %

We see a 2.3x speedup—not bad! We’ve cut the iteration count by a factor of 10, but ADMM uses more FLOPs per iteration because of the extra variables from splitting. With further optimizations (e.g., better fusing), it’s likely the ADMM runtime would better reflect the total iteration count and improve further.

Another natural question is the extent to which we’re obtaining ADMM iterates $(\vLambda_\star, \vX_\star)$ that satisfy the desired ADMM constraint $\vX = \vG + \vW(\vLambda + \vLambda^\top)$. We show a plot below from epoch 0, iteration 0 for the first layer (other layers/steps are similar, or better). It demonstrates that the residual norm is always rather small, but still non-negligible before convergence. Ten iterations of ADMM is not quite converged for this layer, but strikes a reasonable tradeoff between convergence and efficiency.

Feasibility (RMSE) for the ADMM algorithm, for the nonzero-steps hyperparameter choices shown in the table above. Traces can be identified based on where they terminate.

Conclusions

We’ve managed to speed up manifold Muon by a significant amount over the dual ascent baseline by deriving and implementing ADMM for this problem! ADMM converges faster and requires less hyperparameter tuning, allowing us to safely use a much smaller number of inner loop iterations.

You can download the code and try it yourself from my fork. I’m excited to run larger-scale experiments with this algorithm to see how it works!

As a closing thought: here’s a quote from Jeremy’s original blog post:

Manifold Muon increased the wall clock time per step compared to AdamW, although this could be improved by running fewer steps of dual ascent or adding momentum to the algorithm and running dual ascent online. Depending on other systems bottlenecks, the overhead may not be an issue.

I think the switch from dual ascent to ADMM dovetails with all of these considerations! In particular, the results on this experiment (highly clustered performance) suggest that incremental / online manifold Muon might be a natural choice for efficiency’s sake, and our ADMM algorithm should work out-of-the-box in this regime.

Cite

If you found this post or the code useful, please consider citing it:

@misc{buchanan2025mmuonadmm,
  author = {Buchanan, Sam},
  title = {A Faster Manifold Muon with {ADMM}},
  year = 2025,
  howpublished = {\url{https://sdbuchanan.com/blog/manifold-muon/}}
}

Acknowledgments

Thanks to Jeremy Bernstein for helpful discussions. Thanks to the TRC program, Hyperbolic, and Mithril for compute.

Postscript: Other Algorithms than ADMM

After writing this post, it came to my attention that Franz Louis Cesista had also written about an algorithmic improvement to the baseline dual ascent method for manifold Muon using the primal-dual hybrid gradient method (PDHG) (Cesista, 2025). PDHG is sometimes referred to as “linearized ADMM”.

PDHG for manifold Muon has some similar benefits to the ADMM approach we derived above: it involves proximal steps to “smooth out” the nonsmooth aspects of the subproblem \eqref{eq:manifold-muon}.¹¹ At the same time, PDHG is classically most commonly used in situations where the ADMM subproblems are intractable (e.g., in image processing, for total-variation-regularized image denoising), since it involves three hyperparameters that require careful tuning in practice (Goldstein et al., 2015), as opposed to than ADMM’s one.

It would be interesting to run a thorough comparison between PDHG and ADMM for manifold Muon. The JAX code in Cesista’s blog post is missing some implementations, but could likely be integrated into our (PyTorch) Github repo with a bit of effort.

Postscript: Other Quadratic Constraints

Another recent blog post by Ben Keigwin, Dhruv Pai, and Nathan Chen at Tilde Research describes how the essential structure of the manifold Muon algorithm is preserved under replacements of the constraint $\vW^\top \vW = \vI$ by similar constraints (e.g., constraining only the diagonal) (Keigwin et al., 2025). As a corollary, the ADMM algorithm we derived above can be applied to constraints from this family with little extra effort! We leave this as an exercise to the reader.

Here and below, we write $\ip{\vA}{\vB} = \tr(\vA^\top \vB) = \sum_{i, j} A_{ij} B_{ij}$ to denote the Frobenius inner product on matrices, which treats them as if they were vectors, and $\norm{\spcdot}$ to denote the operator norm (or spectral norm, or Schatten $\infty$-norm), which is the largest singular value of its (matrix) argument. ↩
As written, the matrix sign function is only defined for inputs that have no zero singular values. It turns out that as long as $\vA$ satisfies the tangent constraint $ \vA^\top \vW + \vW^\top \vA = \Zero$, the update $\vW -\eta \vA$ never has zero singular values! This can be derived straightforwardly via contradiction, and provides a strong motivation for the need to exactly solve the manifold Muon subproblem. ↩
Here and below, we write $\norm{\spcdot}_*$ to denote the nuclear norm (or Schatten 1-norm), which is the sum of the singular values of its (matrix) argument. ↩
Less naive algorithms for nonsmooth optimization can ameliorate this pathological behavior, sometimes to a significant extent. See, for example, work of Damek Davis and collaborators (Davis et al., 2024). These algorithms definitely merit a closer look in the context of manifold Muon! But we will stick our neck out and suggest that it might be hard to beat ADMM’s solution quality vs. iteration count tradeoff, based on the experiments to come—although we would be happy to be proven wrong. ↩
This is a relatively straightforward argument, but at the same time, I haven’t seen it anywhere in the compressed sensing literature previously (and this literature applied singular value thresholding to almost everything one could). If you know of a reference that has applied this previously, please let me know! ↩
After writing down these updates all together, it becomes clear that $\vLambda_{k+1}$ is always a symmetric matrix. We use this invariant to simplify some steps in the algorithm here and below! ↩
More precisely, it’s straightforward to prove (using the SVD) that singular value thresholding is a $1$-Lipschitz continuous function of the input matrix with respect to the operator norm. ↩
It may be obvious, but it’s worth repeating: the Manifold Muon runs an iterative optimization solver after every forward/backward pass of the model being trained, for every weight matrix in the model. This leads to some interesting debugging issues, such as small differences in the final solution obtained by different manifold Muon solvers at iteration zero leading to significant differences in the behavior of subsequent gradients/inner loop solves (even though everything is convex!). These differences seem to mostly be restricted to inner loop loss, as final val error does not seem to vary much, but it will be interesting to study this behavior at larger scales. ↩
For the sake of comparison, we also use this initialization for our $\vLambda$ in our ADMM algorithm. ↩
This is likely because of the fact that we compute msign with an approximate algorithm, leading to unstable behavior when we aggressively threshold singular values in the $\vX$ update of ADMM. ↩
In contrast to ADMM, it proceeds directly on the Lagrangian for \eqref{eq:manifold-muon}, and hence involves proximal steps on the (characteristic function of the) operator norm ball, rather than on the nuclear norm, as in ADMM. ↩

SPMD in JAX #2: Transformers in Bare-Metal JAX

2025-09-12T00:00:00+00:00

In this second post in a series on programming with JAX for training models like transformers, we write and train a transformer with a decent set of bells and whistles, then benchmark and scale it as much as we can on our TPU v4 half-cube!

As in previous posts in the series, we make good use of the Google scaling playbook (Austin et al., 2025), as well as some useful JAX tutorials, especially (JAX Team, 2025).

Setup
Design Choices
- Model Architecture Overview
Building the Model: First Cut
Sampling from the Model: First Cut
Building the Model: Attention and Positional Encodings
Building the Model: Sampling with a Cache
Conclusions
Accumulated Configs
Acknowledgments

Setup

As usual, we’ll test things out in Python 3.13 and JAX 0.7.1 on a single TPU v4 host.

import jax
import jax.numpy as jnp

jax.devices()

Output:
[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),
 TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0),
 TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0),
 TpuDevice(id=3, process_index=0, coords=(1,1,0), core_on_chip=0)]

Design Choices

A project like this is for learning and pedagogy—it’s inevitably behind the research/infra frontiers. So the code we write is going to lean into this somewhat. Notably,

We’ll implement everything in bare-metal JAX, rather than using neural network libraries like Flax/nnx/Equinox.

This will let us expose and understand all the implementation details, both for infra and architecture, that we might normally overlook. On the flip side, we’ll take some performance hits for this (which we’ll attempt to profile and characterize), and it will lead to some slightly unfortunate code repetition.¹ We’ll also be reinventing the wheel throughout – but for pedagogy’s sake!

Model Architecture Overview

Our approach to implementing the network architecture is a modification of the approach used in the JAX training cookbook (JAX Team, 2025), which uses a data structure like those in Google’s ML Collections library to store the model. At a high level, this approach consists of these standard JAX neural net patterns:

Storing the model parameters as a pytree.
Defining separate functions for initialization and the model’s forward pass which take such a pytree as input (as well as other relevant inputs). This pattern promotes functionally-pure code for the model’s initialization and forward pass, making it easy to jit for performance.

Our implementation tries to move this slightly in the direction of Equinox, without all the powerful features that library entails. In particular, we will:

Use NamedTuples for parameters and layers (i.e., where you’d normally expect to use a nn.Module in Pytorch), enabling good static type checking with e.g. Pyright. NamedTuples are automatically registered as pytrees in JAX!
Write a top-level function for model initialization, and layer-level functions for each separate layer’s operation (which, as above, take parameters and inputs as arguments). This is a ‘bottom-up’ version of the previous design, where we’ll end up with many different functions, one for each ‘layer’, which are combined bottom-up to implement the overall model’s forward pass (more like in a typical neural network library).

Building the Model: First Cut

When writing this post, I wrote and debugged the initial model/training code iteratively in a notebook. Rather than just jumping to the ‘final answer’, we’ll walk through the development process below.

In this part of the post, we’ll build up to the transformer through a few simple sub-modules of the overall model. Our targets will be:

We’ll work with a very simple data distribution, described below, which requires minimal dataloading code (and no tokenization).
We’ll focus on the training loss, leaving sampling for later. Although simple sampling gives us an important correctness test (“did we accidentally give the model access to future tokens?”, for example…), doing it ‘fast’ requires a bit of code, so we’ll save it for later.

On the way to these targets, we’ll build up the following:

Just the MLP, giving us a position-wise model (bigrams with a nonlinearity!).
A training loop with SGD and cross-entropy (next token prediction) loss.
Just causal attention, letting us build a transformer that can solve our toy task (discussed in the next section) with two layers.
Rotary positional encodings, so that we can solve our task with a one-layer model.

We’ll then step back and refactor things with an eye towards training larger models on actual language data in a distributed setting.

If you’re following along and you notice anything I’ve implemented that seems suboptimal, please let me know! Drop me an email or a DM on Twitter (links at the bottom of the page).

A Quick Note on Configuration

For overall convenience, we eventually want to combine all our different configuration options in a large dataclass, then pass our runtime instantiation of the config (with defaults, overrides, etc.) to all our functions, which can extract just the parameters they need.

When writing this code, I tend to define the config options I need locally, then go back later and turn them into attributes of the overall config dataclass. To avoid frontloading with all the complexity of the config, and to avoid too much unnecessary description of the refactoring process, we’ll reference different config options below as attributes of an (as of yet undefined) Config class. Code outputs will also have this class defined. You can Ctrl+F for type definitions and defaults, which will appear later, if it feels necessary.

Brief Setup

We are going to make reference to different sources of randomness below, which we seed from our Config class.

# Test it out
config = Config()
key = jax.random.key(config.seed)
key_params, key_data, key_sampling = jax.random.split(key, 3)

Data

The data will consist of an (infinite) sequence of consecutive digits (from 0 to 9), which ‘reflect’ when they reach 0 or 9. I.e.,

01234567898765432101...

We’ll generate a dataset of fixed-length subsequences of this infinite sequence for training.

text = "012345678987654321" * 1024
seqs = [
    [int(c) for c in text[i : i + config.seq_len + 1]]
    for i in range(len(text) - config.seq_len - 1)
]
data = jnp.array(seqs, dtype=jnp.int32)


key_data, sk = jax.random.split(key_data)
data = jax.random.permutation(sk, data, axis=0)
num_data = len(data)
print(data[128])

Xtr = data[: int(0.8 * num_data)]
Xdev = data[int(0.8 * num_data) : int(0.9 * num_data)]
Xte = data[int(0.9 * num_data) :]
Xtr.shape

Output:
[8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8
 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9
 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8
 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7
 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6
 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5
 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6 5 4 3 2 1 0 1 2 3 4 5 6 7 8 9 8 7 6]
(14540, 257)

We shuffle the sequences above and perform a simple train-dev-test split (although we’ll just use train for now). We take one more character than the sequence length, so that we can easily turn these raw sequences into inputs and targets. We do this with a simple dataloader:

import itertools as it


def dataloader(key, config: Config):
    for step in it.count():
        key = jax.random.fold_in(key, step)
        offsets = jax.random.randint(key, (config.global_batch_size,), 0, num_data)
        yield (Xtr.at[offsets, :-1].get(), Xtr.at[offsets, 1:].get())


key_data, sk = jax.random.split(key_data)
batch = map(
    lambda batch: jax.device_put(batch, config.sharding_data),
    dataloader(sk, config),
)

Here batch is a map iterator which wraps the dataloader with a device_put statement to enforce our desired sharding (e.g., data parallel). When we create our config class, we create a global mesh and set it as the default mesh, so we can pass a PartitionSpec object to device_put in our single-host setting (see the previous post).

The batches being loaded have shape (config.global_batch_size, config.seq_len), so expect config.sharding_data to be for example jax.P('dp') for a “data-parallel” mesh axis 'dp', which leads the batch to be sharded across accelerators.

Defining the Model: MLP

We’ll implement the overarching structure of our model here, described above in the “Design Choices” section, but keep the low-level implementations restricted to just the MLP.

We’ll follow the following conventions:

“Layers” correspond to NamedTuples (which are pytrees). The elements of such a “layer” can be either jax.Arrays (parameters/pytree leaves), or other NamedTuples (for composing layers).
We’ll name our layers with capital letters, and the functions that implement them (i.e., their “forward” method in Pytorch lingo) in lower-case, with a leading underscore.

Here’s our MLP layer following these conventions:

from typing import NamedTuple


class Mlp(NamedTuple):
    w_up: jax.Array
    bias_up: jax.Array
    w_down: jax.Array
    bias_down: jax.Array


def _mlp(config: Config, params: Mlp, x: jax.Array):
    preact = jnp.dot(x, params.w_up, out_sharding=config.sharding_mlp_hidden)
    if config.use_bias_mlp:
        preact += params.bias_up
    act = jax.nn.gelu(preact, approximate=True)
    out = jnp.dot(act, params.w_down, out_sharding=config.sharding_res_stream)
    if config.use_bias_mlp:
        out += params.bias_down
    return out

We would create parameters for such a layer and apply it as follows:

my_mlp = Mlp(
    w_up=jnp.zeros((2, 4)),
    bias_up=jnp.zeros((4,)),
    w_down=jnp.zeros((4, 2)),
    bias_down=jnp.zeros((2,)),
)
_mlp(config, my_mlp, jnp.ones((2,)))

Output:
Array([0., 0.], dtype=float32)

It’s important to note here that we are not able to enforce shape information at this level of implementation, since abstractly speaking an MLP can be implemented for any compatible weight matrix dimensions. These shapes get enforced later, when we initialize an Mlp with settings specified in our config instance.

More importantly, there is some ambiguity in the types of the arguments x and params that _mlp will apply to – e.g. the above function could apply to x being a vector of shape (d_model,), or a matrix of shape (seq_len, d_model), as long as the parameters of Mlp are initialized correctly. In other words, it’s up to us to make sure we apply this implementation correctly.

In our subsequent implementation, we’re going to follow a helpful principle that JAX affords – we’ll operate at the finest level of granularity possible at every stage of the model, and ‘move up’ to the next level of granularity as necessary, using JAX’s powerful primitives like vmap and scan. In particular, we’ll see soon that we’ll only use the above _mlp function in cases where x is a vector of shape (d_model,). Since the MLPs in transformers operate the same on every position in the sequence, and every sequence in the batch, we’ll just vmap our _mlp function when we need to use it at the next level of granularity! This is a pretty empowering design strategy that JAX enables: it lets us write our code’s core numerical functionality without having to plan ahead for exactly how it will be used later, making implementation simpler and separating complexity nicely.

A few other notes (some drawbacks of this bare metal approach, which seem unvaoidable):

There’s no ‘state’ in the _mlp function – it just performs the operation of a MLP given parameters/inputs. This makes it easy to jit.
Initialization is not handled here – we’ll have to do it later, since we want to initialize the entire transformer (top-down) in one shot.
We enforce shardings for the hidden activations/output of the MLP in the _mlp function. We enforce shardings for the parameters of the MLP when we initialize them (see later).

Defining the Model: Embeddings and Placeholders

To have a basic working model, we need token embeddings (mapping the raw integer tokens we generated above to vectors the transformer can operate on), unembeddings (mapping the vector outputs of the transformer to logits, which we can map to a probability distribution over tokens to predict with), and normalization layers. We’ll have attention layers too, eventually; for now, we’ll make a placeholder function for the attention layer.

class Attn(NamedTuple):
    w_qkv: jax.Array
    w_o: jax.Array


def _attn(
    config: Config,
    params: Attn,
    x_seq: jax.Array,
    kv_cache: jax.Array,
    cache_size: int,
):
    return x_seq, kv_cache


class Embedding(NamedTuple):
    w: jax.Array


def _embedding(config: Config, params: Embedding, token: jax.Array):
    emb = params.w.at[token].get(out_sharding=config.sharding_res_stream)
    return emb


class Unembedding(NamedTuple):
    w: jax.Array


def _unembedding(config: Config, params: Unembedding, x: jax.Array):
    logits = jnp.dot(x, params.w, out_sharding=config.sharding_res_stream)
    return logits


class LayerNorm(NamedTuple):
    gamma: jax.Array
    beta: jax.Array


def _layernorm(config: Config, params: LayerNorm, x: jax.Array):
    x_std = jax.nn.standardize(x.astype(config.compute_dtype), epsilon=config.eps_ln)
    out = params.gamma * x_std.astype(config.param_dtype)
    if config.use_bias_ln:
        out += params.beta
    return out

These are standard unoptimized implementations for these different building blocks. Looking ahead slightly, we make sure to provide the ability to compute the layer normalization operation in a higher precision (config.compute_dtype), since the variance operation in standardize can lead to underflow and other losses of precision. The API for _attn involves arguments for utilizing a cache for faster inference (we’d normally leave this out to begin with, and add it only later).

Defining the Model: Building Blocks to Transformer

A transformer consists of an embedding, then a sequence of transformer blocks (which involve normalization + attention and a residual connection, then normalization + MLP and a residual connection), then the unembedding. Accordingly, we make a Block class which combines our low-level building blocks.

from functools import partial


class Block(NamedTuple):
    norm_attn: LayerNorm
    attn: Attn
    norm_mlp: LayerNorm
    mlp: Mlp


def _block(
    config: Config,
    params: Block,
    x_seq: jax.Array,
    cache_in: jax.Array,
    cache_size: int,
):
    att_skip = x_seq
    out = jax.vmap(partial(_layernorm, config, params.norm_attn))(x_seq)
    out, cache_out = _attn(
        config, params.attn, out, kv_cache=cache_in, cache_size=cache_size
    )
    out += att_skip

    mlp_skip = out
    out = jax.vmap(partial(_layernorm, config, params.norm_mlp))(out)
    out = jax.vmap(partial(_mlp, config, params.mlp))(out)
    out += mlp_skip

    return out, cache_out

Notice that we’re following the design principle we mentioned above with the MLP: here we expect x_seq to have two dimensions, the first for the sequence length and the second for the model dimension, and we vmap the _mlp and _layernorm functions to apply to it correctly.

We can now define the transformer. We’ll do this in a slightly refined way versus just using a for loop: we’ll expect the blocks parameter below, of type Block, to have a mapped leading axis corresponding to the layer dimension, and we’ll consume this leading axis using a scan. That means we expect every parameter in the Block pytree we pass to not just have its usual shape, say param.shape (expected by the _mlp, etc. APIs above), but additionally be of shape (num_layers,) + param.shape. This turns out to be very easy to guarantee with a vmap at initialization (see below).

If you aren’t familiar with scan, you can read about the API here. It takes a function that accepts two arguments – a ‘carry’ and an ‘input’ – and produces one output, as well as an ‘initial’ carry, and an initial input, which must have an additional mapped axis somewhere relative to the input shape the function argument expects. It will then iteratively apply the function to the current carry and the current input; the output carry is passed to the next call of the function. In this way, it’s like a vectorized for loop – you can also draw a direct analogy to a recurrent neural network, or other dynamical systems.

In our case, the ‘carry’ for scan is our network’s input and activations (we pass these from layer to layer sequentially), and the ‘input’ for scan is the mapped Block pytree (as well as the mapped KV cache for inference – the updated cache becomes the output – but we can ignore that for now).

class Transformer(NamedTuple):
    emb: Embedding
    blocks: Block  # vmapped at init
    unemb: Unembedding


def _transformer(
    config: Config,
    params: Transformer,
    tokens: Array,
    cache_in: jax.Array,
    cache_size: int = 0,
):
    x_seq = jax.vmap(partial(_embedding, config, params.emb))(tokens)

    def _block_fun(x_seq: Array, params__cache_in: tuple[Block, jax.Array]):
        params, cache_in = params__cache_in
        return _block(config, params, x_seq, cache_in, cache_size)

    out, cache_out = jax.lax.scan(_block_fun, x_seq, (params.blocks, cache_in))

    out = jax.vmap(partial(_unembedding, config, params.unemb))(out)

    return out, cache_out

Using scan leads to a very concise forward pass definition for the transformer. Note that our transformer is still operating on sequences of embeddings – we’ll apply it to a batch of sequences with vmap later!

Defining the Model: Initialization

The last step before testing is to initialize the model. This leads to a long definition, but the intrinsic complexity is not high.

def init_model_params(key, config: Config) -> Transformer:
    def init_embedding(key) -> Embedding:
        emb = config.param_std * jax.random.normal(
            key,
            (config.num_vocab, config.d_model),
            config.param_dtype,
            out_sharding=config.sharding_res_stream,
        )
        return Embedding(w=emb)

    def init_unembedding(key) -> Unembedding:
        unemb = config.param_std * jax.random.normal(
            key,
            (config.d_model, config.num_vocab),
            config.param_dtype,
            out_sharding=config.sharding_res_stream,
        )
        return Unembedding(w=unemb)

    def init_mlp(key) -> Mlp:
        k_w_up, k_w_down = jax.random.split(key, 2)
        w_up = config.param_std * jax.random.normal(
            k_w_up,
            (config.d_model, config.mlp_factor * config.d_model),
            config.param_dtype,
            out_sharding=config.sharding_wup,
        )
        w_down = config.param_std * (
            jax.random.normal(
                k_w_down,
                (config.mlp_factor * config.d_model, config.d_model),
                config.param_dtype,
                out_sharding=config.sharding_wdown,
            )
        )
        bias_up = jnp.zeros(
            (config.mlp_factor * config.d_model,),
            config.param_dtype,
            out_sharding=config.sharding_mlp_hidden,
        )
        bias_down = jnp.zeros(
            (config.d_model,),
            config.param_dtype,
            out_sharding=config.sharding_res_stream,
        )
        return Mlp(w_up=w_up, bias_up=bias_up, w_down=w_down, bias_down=bias_down)

    def init_attn(key) -> Attn:
        k_qkv, k_o = jax.random.split(key, 2)
        w_qkv = config.param_std * jax.random.normal(
            k_qkv,
            (config.d_model, 3, config.num_heads, config.d_head),
            config.param_dtype,
            out_sharding=jax.P(*config.sharding_wqkv),
        )
        w_out = config.param_std * (
            jax.random.normal(
                k_o,
                (config.d_head, config.num_heads, config.d_model),
                config.param_dtype,
                out_sharding=jax.P(*config.sharding_wo),
            )
        )
        return Attn(w_qkv=w_qkv, w_o=w_out)

    def init_layernorm() -> LayerNorm:
        gamma = jnp.ones(
            (config.d_model,),
            config.param_dtype,
            out_sharding=config.sharding_res_stream,
        )
        beta = jnp.zeros(
            (config.d_model,),
            config.param_dtype,
            out_sharding=config.sharding_res_stream,
        )
        return LayerNorm(gamma=gamma, beta=beta)

    def init_block(key) -> Block:
        key_attn, key_mlp = jax.random.split(key)
        return Block(
            norm_attn=init_layernorm(),
            attn=init_attn(key_attn),
            norm_mlp=init_layernorm(),
            mlp=init_mlp(key_mlp),
        )

    # Make the full network
    key_emb, key_blocks, key_unemb = jax.random.split(key, 3)
    keys_blocks = jax.random.split(key_blocks, config.num_layers)
    return Transformer(
        emb=init_embedding(key_emb),
        blocks=jax.vmap(init_block)(keys_blocks),
        unemb=init_unembedding(key_unemb),
    )

Notice above how we produce the mapped axis needed for the scan in the _transformer function above: we just split our PRNGKey into an array of keys, one for each layer, and vmap the init_block function over the array of keys. Easy!

Testing the Basic Model

We can perform a quick test to make sure everything above is at least syntactically correct. There are two things we want to see:

Can we run the model in eager mode?
Can we jit the model?

We really only care about 2 (it implies 1), so let’s test it.

key_params, sk = jax.random.split(key_params)
tf = init_model_params(sk, config)

print("Number of layers in model: ", config.num_layers)
print("Model dimension: ", config.d_model)

cache = None
tokens = jnp.arange(10)


@jax.jit
def model(params, tokens):
    return _transformer(config, params, tokens, cache, 0)


out, out_cache = model(tf, tokens)
out

Output:
Number of layers in model:  2
Model dimension:  768
Array([[0.894531, -0.482422, 1.41406, 0.194336, -2.76562, -1.27344,
        1.21094, 0.417969, -0.460938, -0.304688],
       [0.0610352, 2.8125, 1.08594, -0.351562, -0.326172, -1.46094,
        0.0354004, 0.283203, 0.796875, -1.42969],
       [0.186523, 1.21875, 0.429688, -1.10156, 0.667969, -2.07812,
        -0.078125, 0.410156, 0.0273438, -2.60938],
       [0.201172, 1.46094, 0.124512, -1.27344, -1.05469, -0.667969,
        0.453125, 1.78906, -1.00781, 1.67188],
       [-0.925781, 0.988281, -1.14844, 0.0942383, 0.0898438, 1.70312,
        1.125, 1.03125, -1.17969, -1.69531],
       [-0.710938, -0.679688, -1.53125, 0.800781, 0.443359, -0.71875,
        -2.29688, 0.488281, -0.169922, -1.38281],
       [1.44531, 0.949219, -0.589844, 0.357422, -1.30469, 2.85938,
        -0.388672, 0.46875, -0.695312, -0.425781],
       [-1.85156, 0.679688, 1.25, 0.109863, -0.667969, -0.361328,
        -0.0405273, -0.964844, 3.01562, 0.333984],
       [-0.304688, 1.83594, 0.734375, 1.16406, -0.114746, -0.800781,
        0.486328, 1, -1.00781, -1.01562],
       [-1.66406, 0.925781, -0.0908203, 1, 0.298828, -2.40625, 0.574219,
        0.945312, -2.54688, -0.123047]], dtype=bfloat16)

Looks good! There might be a small issue with the unembedding scaling, as we can see:

jax.nn.softmax(out, axis=-1)

Output:
Array([[0.163086, 0.0410156, 0.275391, 0.0810547, 0.00418091, 0.0186768,
        0.224609, 0.101562, 0.0419922, 0.0493164],
       [0.0390625, 0.613281, 0.108887, 0.026001, 0.0264893, 0.00848389,
        0.0380859, 0.0488281, 0.081543, 0.00872803],
       [0.100098, 0.279297, 0.126953, 0.02771, 0.161133, 0.010376,
        0.0766602, 0.124512, 0.0849609, 0.006073],
       [0.0583496, 0.204102, 0.0539551, 0.0133057, 0.0164795, 0.0244141,
        0.0751953, 0.285156, 0.017334, 0.253906],
       [0.0227051, 0.15332, 0.0181885, 0.0629883, 0.0629883, 0.314453,
        0.176758, 0.160156, 0.0177002, 0.010437],
       [0.0588379, 0.060791, 0.026123, 0.267578, 0.1875, 0.0588379,
        0.012146, 0.195312, 0.101562, 0.0300293],
       [0.141602, 0.0864258, 0.0184326, 0.0476074, 0.00909424, 0.582031,
        0.022583, 0.0534668, 0.0164795, 0.0218506],
       [0.00500488, 0.0629883, 0.112305, 0.0358887, 0.0164795, 0.0224609,
        0.0307617, 0.012207, 0.65625, 0.0444336],
       [0.0395508, 0.335938, 0.111328, 0.171875, 0.0473633, 0.0239258,
        0.0869141, 0.145508, 0.0194092, 0.0194092],
       [0.0145874, 0.193359, 0.0693359, 0.208008, 0.102539, 0.00689697,
        0.135742, 0.196289, 0.00598145, 0.0673828]], dtype=bfloat16)

In particular, the initial outputs are not very uniform, which we would like to have. However, we should still be able to learn – our network is not very deep.

A Minimal Training Loop

Now we have enough to train a model. A simple training loop follows, using SGD. First, we set up the training step, jitting it for performance.

class TrainState(NamedTuple):
    params: Transformer
    opt: None


@jax.jit
def init_train_state(key, config: Config) -> TrainState:
    return TrainState(params=init_model_params(key, config), opt=None)


train_state = init_train_state(key_params, config)
cache = None


@partial(jax.jit, donate_argnums=2)
def train_step(config: Config, batch, train_state: TrainState):
    def loss_fn(params: Transformer):
        inputs, targets = batch
        logits, cache_out = jax.vmap(partial(_transformer, config, params))(
            inputs, cache
        )
        logits = logits.astype(config.compute_dtype)
        logprobs = jax.nn.log_softmax(logits, axis=-1)
        return -jnp.take_along_axis(logprobs, targets[..., None], axis=-1).mean()

    loss, grad = jax.value_and_grad(loss_fn)(train_state.params)
    new_state = TrainState(
        params=jax.tree.map(lambda x, y: x - config.lr * y, train_state.params, grad),
        opt=None,
    )

    metrics = {"loss": loss}
    return metrics, new_state

In the training step, after loading a batch of sequences from our previously-written batch loader, we vmap _transformer over it, then calculate the next-token prediction loss with the standard bare-metal approach (with possible up-casting for numerical stability of the softmax). The donate_argnums argument to jit lets us mark that the train_state buffer can be discarded after the step finishes (since we create and return a new TrainState with the results of the SGD step), letting the XLA compiler reuse that memory for the output state.

The loop is then simple to write:

prev_metrics = None
for step in range(config.num_steps):
    cur_metrics, train_state = train_step(config, next(batch), train_state)
    log_metrics, prev_metrics = prev_metrics, cur_metrics
    if log_metrics and step % 50 == 0:
        log_metrics |= {"step": step}
        print(*[f"{metric}: {val}" for metric, val in log_metrics.items()], sep="\t")

Output:
loss: 1.1560431718826294	step: 50
loss: 0.9151328206062317	step: 100
loss: 1.004440188407898	    step: 150
loss: 0.8609314560890198	step: 200
loss: 0.9039498567581177	step: 250
loss: 0.8058883547782898	step: 300
loss: 0.7688642144203186	step: 350
loss: 0.6995882987976074	step: 400
loss: 0.6380590796470642	step: 450
loss: 0.6304700374603271	step: 500
loss: 0.6276624798774719	step: 550
loss: 0.6255137920379639	step: 600
loss: 0.6255109310150146	step: 650
loss: 0.6271587610244751	step: 700
loss: 0.6273760795593262	step: 750
loss: 0.6262620687484741	step: 800
loss: 0.6236293911933899	step: 850
loss: 0.6241434216499329	step: 900
loss: 0.6261416673660278	step: 950

It’s not great, but it’s indeed training! We don’t expect to be able to solve this task perfectly with our simple model, since a position-wise MLP can only use one token of context to predict the next token, and our task requires at least two. To see why, just note that, for example, the character following 5 could be either 6 or 4, depending on whether we’re at a rising or falling part of the sequence. Two characters of context is enough to estimate this, meaning adding attention should help here!

As an aside, above, we make sure to print metrics for the previous iteration of the training loop at each step. This helps with correctly pipelining host to data transfers: see (JAX Team, 2025) for a detailed description.

Sampling from the Model: First Cut

To get a sense of the behavior of our learned model, we can write our sampler for it now. Normally, efficient sampling requires a cache of past outputs to be maintained (we already have the API for this), but we’ll discuss this later after we implement attention.

@partial(jax.jit, donate_argnums=(4,))
def sample_one_token(
    config: Config,
    key,
    params: Transformer,
    x: jax.Array,
    cache_in: jax.Array,
    cache_size: int,
    temperature: float,
):
    y, cache_out = _transformer(config, params, x, cache_in, cache_size)
    logits = y.astype(config.compute_dtype)
    cache_size = cache_size + x.shape[-1]
    next_token = jnp.array((jax.random.categorical(key, logits[-1] / temperature),))
    return next_token, cache_out, cache_size

config_sampling_args = {"sharding_data": jax.P()}
config_sampling = Config(**config_sampling_args)

output = jnp.array((1,))
cache = None
tokens_to_generate = 63
temperature = 0.1
for step in range(tokens_to_generate):
    key_sampling, sk = jax.random.split(key_sampling)
    next_token, cache, cache_size = sample_one_token(
        config_sampling, sk, train_state.params, output, cache, 0, temperature
    )
    output = jnp.concatenate((output, next_token))
output

Output:
Array([1, 2, 3, 2, 3, 4, 3, 2, 3, 2, 3, 5, 6, 5, 4, 3, 4, 5, 2, 1, 2, 3,
       2, 1, 0, 1, 2, 3, 4, 3, 4, 7, 6, 7, 6, 7, 6, 7, 6, 7, 6, 5, 6, 5,
       6, 5, 1, 2, 0, 1, 0, 2, 3, 2, 3, 2, 3, 2, 3, 2, 3, 4, 3, 2],      dtype=int32)

We set a relatively low temperature in order to better highlight what the model has learned. The generated sequence is in line with our hypothesis – the model has lowered the loss by learning that it needs to predict the number that’s one above or one below the current input, but it can’t do much better than guessing randomly between these two possibilities (hence why our final loss is close to $\log(2)$).

Let’s improve this with attention!

Building the Model: Attention and Positional Encodings

An attention implementation that follows jax.nn.dot_product_attention shape conventions follows, with both manual implementations and one using this function. For now, we omit the cache implementation and positional embedding implementation.

def _attn(
    config: Config,
    params: Attn,
    x_seq: jax.Array,
    kv_cache: jax.Array,
    cache_size: int,
):
    # s: sequence length
    # d: embedding dim (config.d_model)
    # n: attention heads (config.num_heads)
    # h: head dim (config.d_head)
    # x_seq: s x d

    qkv = jnp.einsum(
        "sd,d3nh->3snh",
        x_seq,
        params.w_qkv,
        out_sharding=jax.P(*config.sharding_att_qkv),
    )
    q, k, v = [qkv[i] for i in range(3)]
    s = q.shape[0]

    # Attention computation
    t = k.shape[0]
    mask = (cache_size + jnp.arange(s))[:, None] >= jnp.arange(t)[None, :]
    mask = mask[None, ...]  # broadcast over heads
    if config.use_fa:
        attn_out = jax.nn.dot_product_attention(q, k, v, scale=None, mask=mask)
    else:
        logits = jnp.einsum("snh,tnh->nst", q, k).astype(config.compute_dtype)
        # Scale and causal mask
        logits *= 1.0 / config.d_head**0.5
        logits = jnp.where(mask, logits, -jnp.inf)
        probs = jax.nn.softmax(logits, axis=2)  # type: ignore[reportArgumentType]
        probs = probs.astype(config.param_dtype)
        attn_out = jnp.einsum("nst,tnh->snh", probs, v)
    out = jnp.einsum(
        "snh,hnd->sd",
        attn_out,
        params.w_o,
        out_sharding=jax.P(*config.sharding_res_stream),
    )

    return out, kv_cache

Attention takes the embeddings of the input sequences $\vX \in \bbR^{S \times D}$ as input. It computes $N$ “attention heads”, then combines them to form the output. Each head, indexed by $n \in [N]$, is generated from different projections of the input embedding sequences: we calculate

\[\begin{equation*} \vQ_n = \vX \vW_{Q, n}, \quad \vK_n = \vX \vW_{K, n}, \quad \vV_n = \vX \vW_{V, n}, \end{equation*}\]

where the output dimension of the projections is $H$, then generate the head output as

\[\begin{equation*} \vA_n = \mathrm{softmax}(\vQ_n \vK_n^\top + \vM) \vV_n, \end{equation*}\]

with the softmax being taken along rows, and $\vM$ denoting a causal mask:

\[M_{ij} = \begin{cases} 0 & i \leq j \\ -\infty & \ow. \end{cases}\]

The heads are then concatenated together along the embedding dimension, producing an $S \times NH$ dimensional result (usually $NH = D$), and then an output projection $\vW_O$ is applied to produce the attention output, which has the same shape as the input $\vX$. The code above implements these operations concisely with jnp.einsum.

We can replace our previous placeholder attention implementation with a complete one and then rerun our training loop, since we already added the proper initialization code above.

Output:
loss: 0.6685549020767212	step: 50
loss: 0.58647620677948	step: 100
loss: 0.5948932766914368	step: 150
loss: 0.5258455872535706	step: 200
loss: 0.5094199776649475	step: 250
loss: 0.40264642238616943	step: 300
loss: 0.42422279715538025	step: 350
loss: 0.48532211780548096	step: 400
loss: 0.45087528228759766	step: 450
loss: 0.3302219808101654	step: 500
loss: 0.39742058515548706	step: 550
loss: 0.3917175531387329	step: 600
loss: 0.3619387745857239	step: 650
loss: 0.34305423498153687	step: 700
loss: 0.2451237142086029	step: 750
loss: 0.27099400758743286	step: 800
loss: 0.4500095248222351	step: 850
loss: 0.26931577920913696	step: 900
loss: 0.2725009322166443	step: 950

Output:
loss: 0.15929250419139862	step: 50
loss: 0.14883831143379211	step: 100
loss: 0.165802463889122	step: 150
loss: 0.14946813881397247	step: 200
loss: 0.15420110523700714	step: 250
loss: 0.14391759037971497	step: 300
loss: 0.2680056691169739	step: 350
loss: 0.14069046080112457	step: 400
loss: 0.27090030908584595	step: 450
loss: 0.14207065105438232	step: 500
loss: 0.13676664233207703	step: 550
loss: 0.13658639788627625	step: 600
loss: 0.13079527020454407	step: 650
loss: 0.39731210470199585	step: 700
loss: 0.1367562860250473	step: 750
loss: 0.13199977576732635	step: 800
loss: 0.1323813498020172	step: 850
loss: 0.12948714196681976	step: 900
loss: 0.13398176431655884	step: 950

Two back-to-back output traces are posted above. We can get a sense that this model should be able to learn to solve the task, but the poor optimizer (constant learning rate SGD) is holding it back. If we rerun the sampling code above with the new model at low temperature (temperature=0.1), we do get correct outputs:

Output:
Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 1, 2, 3, 4,
       5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 1, 2, 3, 4, 5, 6, 7, 8,
       9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 8],      dtype=int32)

Yet the final loss is suggestive that the model we’ve learned is not very good. We can improve the model further by incorporating positional encodings. We’ll use the modern choice, rotary positional encodings (RoPE). There are already many great blogs explaining the intuition behind these encodings, so I’ll just focus here on the algebraic ideas.

The basic idea is to consider the dot product of a query-key pair forming one head’s attention matrix $\vA$, as we defined it above (dropping the subscript $n$ for concision). Ignoring the mask, for the $ij$-th entry of this attention matrix, we consider the dot product (“attention logits”) $\vq_i \vk_j^\top$. A notable property of attention is that the only dependence of its output on the position $i$ is through the causal mask $\vM$: if we drop the mask, then permuting the input positions produces a corresponding permutation of the output positions, because every position receives the same projections $\vW_{QKV}$.²

This is arguably undesirable: the model benefits in many cases from the ability to learn a sequence-to-sequence mapping that is more position sensitive. For example, in our toy data setting, the model would be able to learn a good mapping faster if it were able to selectively focus on only a small number of past tokens (since this is enough to determine whether we’re in a “rising” or “falling” part of the sequence).

A nice way to incorporate this positional information is through a relative positional embedding. In such a method, we modify the attention operation so that the attention logits incorporate the magnitude of the difference in positions $\abs{i - j}$. In RoPE, this is done by inducing a quadratic form

\[\begin{equation*} \vq_i \vR_{\abs{i-j}} \vk_j^\top, \end{equation*}\]

where $\vR_{\abs{i-j}}$ is a specific rotation matrix.³ RoPE’s choice of this matrix has the useful property that for fixed-norm keys and queries, the magnitude of the encoded attention logits decays as a function of $\abs{i-j}$, giving a priori less weight to logits of positions that are further apart.

It also has the essential property that the encoded logits can be computed efficiently. It turns out that there are (real-valued) sequences $\vc_i$, $\vs_i$ corresponding to each position, such that the following holds. Let $\vPi \in \bbR^{H \times H}$ denote the permutation matrix that reverses the top $H/2$ coordinates of a vector with the bottom $H/2$ coordinates (assume $H$ is even).⁴ Then one has

\[\begin{equation*} \vq_i \vR_{\abs{i-j}} \vk_j^\top = (\vq_i \circ \vc_i + \vPi \vq_i \circ \vs_i) (\vk_j \circ \vc_j + \vPi \vk_j \circ \vs_j)^\top, \end{equation*}\]

where matrix multiplication binds before elementwise multiplication $\circ$. In particular, we can compute the encoded logits via a small number of elementwise multiplications, adds, and concatenations involving the raw keys and queries, which leads to a minimal-complexity implementation. The sequences $\vc_i, \vs_i$ depend only on position, and can be precomputed and reused.

We give a simple implementation of RoPE below, and its integration into our attention layer we implemented previously.

def _precompute_rope_sincos(config: Config):
    freqs = jnp.exp(
        -jnp.log(config.rope_theta).astype(config.compute_dtype)
        * 2
        / config.d_head
        * jnp.arange(config.d_head // 2, out_sharding=jax.P())
    )
    positions = jnp.arange(config.max_seq_len, out_sharding=jax.P())
    cycles = positions[:, None] * freqs[None, :]
    return jnp.cos(cycles), jnp.sin(cycles)


def _apply_rope(
    config: Config,
    cos: jax.Array,
    sin: jax.Array,
    positions: jax.Array,
    x: jax.Array,  # S x H
):
    x_1, x_2 = x[:, : config.d_head // 2], x[:, config.d_head // 2 :]
    c, s = cos.at[positions].get(), sin.at[positions].get()
    return jnp.concatenate(
        (c * x_1 - s * x_2, c * x_2 + s * x_1), axis=-1, dtype=config.param_dtype
    )


def _attn(
    config: Config,
    params: Attn,
    x_seq: jax.Array,
    kv_cache: jax.Array,
    cache_size: int,
):
    # s: sequence length
    # d: embedding dim (config.d_model)
    # n: attention heads (config.num_heads)
    # h: head dim (config.d_head)
    # x_seq: s x d

    qkv = jnp.einsum(
        "sd,d3nh->3snh",
        x_seq,
        params.w_qkv,
        out_sharding=jax.P(*config.sharding_att_qkv),
    )
    q, k, v = [qkv[i] for i in range(3)]
    s = q.shape[0]

    # Apply RoPE
    if config.use_rope:
        if not config.update_cache:
            cache_size = 0  # ignore passed value
        with jax.ensure_compile_time_eval():
            rope_cos, rope_sin = _precompute_rope_sincos(config)
        positions = jnp.arange(s, out_sharding=jax.P())
        _apply_rope_one_head = partial(
            _apply_rope, config, rope_cos, rope_sin, positions + cache_size
        )
        _apply_rope_all_heads = jax.vmap(_apply_rope_one_head, in_axes=1, out_axes=1)
        q, k = _apply_rope_all_heads(q), _apply_rope_all_heads(k)

    # Attention computation
    t = k.shape[0]
    mask = (cache_size + jnp.arange(s))[:, None] >= jnp.arange(t)[None, :]
    mask = mask[None, ...]  # broadcast over heads
    if config.use_fa:
        attn_out = jax.nn.dot_product_attention(q, k, v, scale=None, mask=mask)
    else:
        logits = jnp.einsum("snh,tnh->nst", q, k).astype(config.compute_dtype)
        # Scale and causal mask
        logits *= 1.0 / config.d_head**0.5
        logits = jnp.where(mask, logits, -jnp.inf)
        probs = jax.nn.softmax(logits, axis=2)  # type: ignore[reportArgumentType]
        probs = probs.astype(config.param_dtype)
        attn_out = jnp.einsum("nst,tnh->snh", probs, v)
    out = jnp.einsum(
        "snh,hnd->sd",
        attn_out,
        params.w_o,
        out_sharding=jax.P(*config.sharding_res_stream),
    )

    return out, kv_cache

The changes to the _attn function are all within the code block marked with # Apply RoPE (everything else is unchanged); the _apply_rope function takes care of the necessary permutation/concatenation operation. We use JAX’s ensure_compile_time_eval context manager to ensure the RoPE multiplying sequences are precomputed. The multiplying sequences themselves are chosen to have geometrically-varying time values, with a rate parameter rope_theta set to 1e4 by default (in general, set based on the target context length, here 1024).

With RoPE, we find it much easier to learn a strong model for our toy data. After rerunning training with these modifications, we observe the following loss progression:

Output:
loss: 0.6326758861541748	step: 50
loss: 0.4087381660938263	step: 100
loss: 0.14004486799240112	step: 150
loss: 0.08325601369142532	step: 200
loss: 0.06802305579185486	step: 250
loss: 0.05945558100938797	step: 300
loss: 0.05321500450372696	step: 350
loss: 0.052829645574092865	step: 400
loss: 0.048602353781461716	step: 450
loss: 0.04696694016456604	step: 500
loss: 0.04576549679040909	step: 550
loss: 0.04564153775572777	step: 600
loss: 0.04341501370072365	step: 650
loss: 0.043622132390737534	step: 700
loss: 0.04205687716603279	step: 750
loss: 0.04401983320713043	step: 800
loss: 0.04256521165370941	step: 850
loss: 0.041673485189676285	step: 900
loss: 0.0421551875770092	step: 950

Much better! For a completely overconfident model, we might expect to have loss around $\log(2) / 256 \approx 0.002$, corresponding to randomly guessing for the first token in each sequence of the batch (we train with sequence length 256). This is within about an order of magnitude.

Building the Model: Sampling with a Cache

The implementation of sampling we gave above is extremely slow: although we jit the sample_one_token function, we are re-forwarding the entire updated sequence through the model after every generation, which is very wasteful. A quick test on our previously-implemented sampling function (on 1x TPU v4 host) takes about 50 seconds to sample 63 tokens, including the time to jit compile the sampling function.

This makes sense, given our naive implementation. Indeed, if we recall the attention definition, at the $i$-th token each head computes

\[\begin{equation*} \mathrm{softmax}( \vq_i \vK^T ) \vV, \end{equation*}\]

and the output projection just accumulates these independent heads and ‘lifts’ the head dimension $H$ to the output dimension $D$. In other words, if we were to save the previous key and value tokens $(\vk_j, \vv_j)_{j=1}^{i-1}$ in a cache, we could reuse these to compute each head’s attention output with no wasted operations: we’d just need to compute $\vk_i$ and $\vv_i$ for the current position, then load the previous positions’ keys and values from the cache to compute the full attention output.⁵

The data structure that stores these past keys and values is called a KV cache, and it’s essential for fast inference, since autoregressive operation implies we can only compute one token at a time (in contrast to training, where we forward an entire batch of sequences at once). A back-of-the-envelope analysis suggests that the naive sampling strategy (re-forward the entire sequence on every new token generation) takes $O(S^2D^2 + DS^3)$ FLOPs to generate an $S$-length sequence starting from an empty prompt, whereas the KV caching strategy takes only $O(SD^2 + DS^2)$ FLOPs. This is a significant savings!

The catch is that the KV caching strategy requires significantly more memory accesses and usage compared to the naive approach. In JAX, we also need to be careful to implement the KV cache ‘correctly’ – since the cache grows after every generation, naive implementations can entail inputs of different sizes to _attn (triggering a re-compilation every time we call it, when it’s jitted!), or dynamically-sized arrays (not allowed within a jit context!).⁶

Our simple implementation of the KV cache is below. Since we left the API for the cache in the code we’ve written already, we just need to implement its functionality in the _attn function, and add a function to initialize an empty cache. To play nice with jit, we will implement a static-sized cache with a fixed size config.max_seq_len (the same parameter we used for RoPE), ensuring we can jit the _attn function only once, and maintain an auxilary cache_size variable to keep track of how much we’ve written to the cache. To avoid having dynamically-sized arrays in the _attn function, we are a bit wasteful and zero-initialize the cache, then just multiply the current query $\vq_i$ against the entire cache on every attention operation.⁷ To avoid having these extra keys influence the attention computation (since softmaxing them produces a nonzero output!), we augment the attention mask to mask them out.

The implementation is below.

def _attn(
    config: Config,
    params: Attn,
    x_seq: jax.Array,
    kv_cache: jax.Array,  # 2 x config.max_seq_len x n x h (see below)
    cache_size: int,
):
    # s: sequence length
    # d: embedding dim (config.d_model)
    # n: attention heads (config.num_heads)
    # h: head dim (config.d_head)
    # x_seq: s x d

    qkv = jnp.einsum(
        "sd,d3nh->3snh",
        x_seq,
        params.w_qkv,
        out_sharding=jax.P(*config.sharding_att_qkv),
    )
    q, k, v = [qkv[i] for i in range(3)]
    s = q.shape[0]  # we save k shape later, after possibly prepending cache

    # Apply RoPE
    if config.use_rope:
        if not config.update_cache:
            cache_size = 0  # ignore passed value
        with jax.ensure_compile_time_eval():
            rope_cos, rope_sin = _precompute_rope_sincos(config)
        positions = jnp.arange(s, out_sharding=jax.P())
        _apply_rope_one_head = partial(
            _apply_rope, config, rope_cos, rope_sin, positions + cache_size
        )
        _apply_rope_all_heads = jax.vmap(_apply_rope_one_head, in_axes=1, out_axes=1)
        q, k = _apply_rope_all_heads(q), _apply_rope_all_heads(k)

    # Read/update/concatenate the cache
    if config.update_cache:
        k_cache, v_cache = kv_cache[0], kv_cache[1]
        k = jax.lax.dynamic_update_slice(k_cache, k, (cache_size, 0, 0))
        v = jax.lax.dynamic_update_slice(v_cache, v, (cache_size, 0, 0))
        kv_cache_out = jnp.concatenate((k[None], v[None]), axis=0)
    else:
        kv_cache_out = kv_cache
        cache_size = 0  # ignore passed value

    # Attention computation
    t = k.shape[0]
    mask = (cache_size + jnp.arange(s))[:, None] >= jnp.arange(t)[None, :]
    # with static kv cache, must mask unused memory!
    mask = mask & (jnp.arange(t)[None, :] < cache_size + s)
    mask = mask[None, ...]  # broadcast over heads
    if config.use_fa:
        attn_out = jax.nn.dot_product_attention(q, k, v, scale=None, mask=mask)
    else:
        logits = jnp.einsum("snh,tnh->nst", q, k).astype(config.compute_dtype)
        # Scale and causal mask
        logits *= 1.0 / config.d_head**0.5
        logits = jnp.where(mask, logits, -jnp.inf)
        probs = jax.nn.softmax(logits, axis=2)  # type: ignore[reportArgumentType]
        probs = probs.astype(config.param_dtype)
        attn_out = jnp.einsum("nst,tnh->snh", probs, v)
    out = jnp.einsum(
        "snh,hnd->sd",
        attn_out,
        params.w_o,
        out_sharding=jax.P(*config.sharding_res_stream),
    )

    return out, kv_cache_out


def init_kv_cache(config: Config):
    return jnp.zeros(
        (
            config.global_batch_size,
            config.num_layers,
            2,
            config.max_seq_len,
            config.num_heads,
            config.d_head,
        ),
        dtype=config.param_dtype,
        out_sharding=config.sharding_data + jax.P(None) + config.sharding_att_qkv,
    )

Running this block and then rerunning our training script shows that our changes haven’t broken anything. We can test how much it speeds up sampling, as well. We modify the sampler above slightly – we enable the cache in the config, and forward one token per sampling step to the model, rather than the entire output, as the cache will store all the context.

config_sampling_args = {"sharding_data": jax.P(), "update_cache": True}
config_sampling = Config(**config_sampling_args)

next_token = jnp.array((1,))
output = next_token
cache = init_kv_cache(config_sampling)[0]
cache_size = 0
tokens_to_generate = 63
temperature = 0.1
for step in range(tokens_to_generate):
    key_sampling, sk = jax.random.split(key_sampling)
    next_token, cache, cache_size = sample_one_token(
        config_sampling,
        sk,
        train_state.params,
        next_token,
        cache,
        cache_size,
        temperature,
    )
    output = jnp.concatenate((output, next_token))
output

Output:
Array([1, 2, 3, 4, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 1, 2, 3, 4,
       5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 1, 2, 3, 4, 5, 6, 7, 8,
       9, 8, 7, 6, 5, 4, 3, 2, 1, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 8],      dtype=int32)

Correct outputs! The runtime is about 2 seconds including the jit compilation time on the first run, and 0.5 seconds on subsequent runs (since the static cache and one-token-input-per-sampling-step means the jitted sampler sees the same input shapes on every call). Quite a speedup!

Conclusions

In this blog post, we’ve built up:

A simple jittable transformer model with RoPE;
A basic training loop with SGD and toy data (digit alphabet);
A minimal jittable sampler with a static KV cache.

All in bare-metal JAX.

Since the length is already nontrivial, we’ll save further extensions – refactoring into a python package, improving the config scaffolding for running experiments and adding distributed training support, improving the optimizer, and training on actual language data – to the next post. I hope this exposition is helpful to learners of JAX and transformers – writing it has been very helpful to me!

Accumulated Configs

For reference, here’s the implementation of the Config class we’ve used throughout the discussion above.

# Architecture: config
@jax.tree_util.register_static
@dataclass(kw_only=True, frozen=True)
class Config:
    # Experiment orchestration params
    mesh_axis_names: tuple[str, ...] = ("dp",)
    mesh_shape: tuple[int, ...] = (4,)
    seed: int = 1337

    # Data and training params
    seq_len: int = 256
    global_batch_size: int = 128
    num_steps: int = 10**3
    lr: float = 1e-2

    # Model architecture params
    num_vocab: int = 10
    d_model: int = 768
    num_heads: int = 12
    d_head: int = 64
    mlp_factor: int = 4
    num_layers: int = 2
    param_std: float = 0.02
    rope_theta: float = 10000.0
    max_seq_len: int = 1024

    # Model dtypes
    param_dtype = jnp.bfloat16  # weights, activations
    compute_dtype = jnp.float32  # layernorm, attn logits, rope
    optimizer_dtype = jnp.float32  # optimizer state

    # Model call-time params
    eps_ln: float = 1e-6
    use_bias_ln: bool = False
    use_fa: bool = True
    use_bias_mlp: bool = False
    use_rope: bool = True
    update_cache: bool = False  # default training

    # Model sharding params
    sharding_data: jax.sharding.PartitionSpec = jax.P("dp")
    sharding_wqkv: jax.sharding.PartitionSpec = jax.P()
    sharding_wo: jax.sharding.PartitionSpec = jax.P()
    sharding_wup: jax.sharding.PartitionSpec = jax.P()
    sharding_wdown: jax.sharding.PartitionSpec = jax.P()
    sharding_mlp_hidden: jax.sharding.PartitionSpec = jax.P()
    sharding_res_stream: jax.sharding.PartitionSpec = jax.P()
    sharding_att_qkv: jax.sharding.PartitionSpec = jax.P()

    def __post_init__(self):
        # Set up and register mesh
        mesh = jax.make_mesh(
            self.mesh_shape,
            self.mesh_axis_names,
            len(self.mesh_shape) * (jax.sharding.AxisType.Explicit,),
        )
        jax.sharding.set_mesh(mesh)

        # Checks
        assert self.d_head % 2 == 0, (
            "Head dimension needs to be divisible by 2 for RoPE"
        )

Acknowledgments

Thanks to the TRC program for compute support.

Namely, because we won’t have a clean way to combine our model parameter definitions with the actual application of the model. ↩
The causal mask actually is enough to enable sequence-to-sequence mappings that depend on position to be learned in multi-layer transformers. The recent trend of “NoPE” (no positional encoding) in open-source large models is evidence of this. ↩
It’s actually a block-diagonal rotation matrix, with $2 \times 2$ blocks. This is what leads to the efficient approach to computing its induced quadratic form below! ↩
In code, for a list v, this is just v[H//2:] + v[:H//2]. Moreover, the choice of the permutation $\vPi$ is not essential (because the coordinates of the embedded keys and queries follow a learnable linear transformation), and can be changed for either mathematical clarity or code clarity (we choose based on the latter; the original RoPE paper emphasized the former). ↩
Using some sort of a cache turns out to be essential to avoid using as many FLOPs as the naive strategy (calculated below). If we try to adaptively recompute the keys and values necessary for generating the current token’s output, it can be seen that we need the previous queries in order to do so. Caching instead the past keys and values avoids any recomputation. ↩
Actually, our naive sampling implementation above, which re-forwards the entire sequence after concatenating, runs so slow because the input to the transformer keeps growing in size – which triggers a recompilation after every generation! Avoiding these pitfalls is a general challenge in JAX. ↩
Although this feels very wasteful, it actually only loses (asymptotically) a factor of $2$ in FLOPs versus the optimal implementation. ↩

SPMD in JAX #1: Sharding

2025-08-20T00:00:00+00:00

This post will be the first in a series on programming with JAX, for training models like transformers. I’m experimenting with these as an alternative to my usual scratch paper or LaTeX notes while learning, in the hope that it will help me with recall and perhaps be useful to others learning this material.

The focus in this post is on building a mental model (mathematical) for sharding and communication based on linear algebra, in particular block matrices, and on studying some low-level code for different communication primitives. The notes are based on two tutorials on sharding: (Austin et al., 2025) and (JAX Team, 2025), as well as some ‘original research’.

Setup
Sharding
- Valid Shardings
Example: Matrix Multiplication
Types of Inter-Device Communication with Sharded Computation
- Communication Primitives through Matrix Multiplication
- AllToAll Communication
Example: Communication Primitives in JAX
Conclusion
Acknowledgments

Setup

Here’s the environment we will be working in—using Python 3.13 and JAX 0.7.1, on a single TPU v4 host.

from functools import partial

import jax
import jax.numpy as jnp

jax.devices()

Output:
[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0),
 TpuDevice(id=1, process_index=0, coords=(1,0,0), core_on_chip=0),
 TpuDevice(id=2, process_index=0, coords=(0,1,0), core_on_chip=0),
 TpuDevice(id=3, process_index=0, coords=(1,1,0), core_on_chip=0)]

Sharding

When we want to perform a high-throughput computation involving some data and some mathematical operations (e.g., passing data through layers of a neural network) using multiple hardware accelerators (e.g., TPUs), we need to, at a minimum:

Have a scheme for indexing hardware accelerators relative to their physical (spatial) layout.
Specify what data and what parameters (inputs) go on which accelerators.
Specify where the outputs of the computation should go.

In JAX, these three steps are abstracted into creation of a Mesh specifying device layout and a PartitionSpec specifying how data is split across devices relative to this mesh. Such a splitting is called a sharding.

Here’s how the first two steps look in code:

mesh = jax.make_mesh((2, 2), ("x", "y"))  # 1x v4 tpu host
sharding = jax.NamedSharding(mesh, jax.P("x", "y"))

A = jnp.zeros((1024, 1024), device=sharding)
jax.debug.visualize_array_sharding(A)

Output:
┌──────────┬──────────┐
│          │          │
│  TPU 0   │  TPU 1   │
│          │          │
│          │          │
├──────────┼──────────┤
│          │          │
│  TPU 2   │  TPU 3   │
│          │          │
│          │          │
└──────────┴──────────┘

Above,

The Mesh indexes devices based on spatial axes. JAX will optimize this indexing based on the actual device topology.
The NamedSharding creates a sharding specification relative to the mesh, which is defined by the PartitionSpec we provide. Here, we specify that a 2D array has its first dimension sharded along the x dimension of the mesh and its second dimension along y.
For the PartitionSpec, we give None for dimensions that should not be sharded. That means this dimension will be present on all devices.
If a mesh axis does not appear in the sharding, then the sharded array is fully replicated over that axis. In the extreme case, a NamedSharding of all Nones puts the full array on every device.
The compiler will complain about invalid shardings.

We can also move data around after creation on, say, one host.

key = jax.random.key(42)
B = jax.random.normal(key, (1024, 1024))
B_sharded = jax.device_put(B, sharding)
jax.debug.visualize_array_sharding(B)
jax.debug.visualize_array_sharding(B_sharded)

Output:
┌───────────────────────┐
│                       │
│                       │
│                       │
│                       │
│         TPU 0         │
│                       │
│                       │
│                       │
│                       │
└───────────────────────┘
┌──────────┬──────────┐
│          │          │
│  TPU 0   │  TPU 1   │
│          │          │
│          │          │
├──────────┼──────────┤
│          │          │
│  TPU 2   │  TPU 3   │
│          │          │
│          │          │
└──────────┴──────────┘

In addition, we can ‘combine’ physical device axes into one or more aggregated axes. This is a common operation, for example in parallelizing a batch axis across all devices. Since the Mesh is our link between partition indices and physical devices, we do this through the mesh:

mesh_flat = jax.make_mesh((4,), ("xy",))
sharding_flat = jax.NamedSharding(mesh_flat, jax.P("xy", None))
C = jnp.zeros((1024, 1024), device=sharding_flat)
jax.debug.visualize_array_sharding(C)

Output:
┌───────────────────────┐
│         TPU 0         │
├───────────────────────┤
│         TPU 2         │
├───────────────────────┤
│         TPU 1         │
├───────────────────────┤
│         TPU 3         │
└───────────────────────┘

When we perform computations on data, we can either let the JAX compiler infer the output sharding, or we can specify it ourselves.

@jax.jit
def f_contract(x):
    return x.sum(axis=0)


@partial(
    jax.jit,
    out_shardings=jax.NamedSharding(
        mesh,
        jax.P(
            None,
        ),
    ),
)
def f_contract_replicated(x):
    return x.sum(axis=0)


result = f_contract(B_sharded)
result_replicated = f_contract_replicated(B_sharded)
print("Inferred sharding:")
jax.debug.visualize_array_sharding(result)
print("Manual sharding (replicated):")
jax.debug.visualize_array_sharding(result_replicated)

Output:
Inferred sharding:
┌───────┬───────┐
│TPU 0,2│TPU 1,3│
└───────┴───────┘
Manual sharding (replicated):
┌───────────┐
│TPU 0,1,2,3│
└───────────┘

Passing a NamedSharding to jax.jit’s out_shardings keyword argument lets us specify here that the result of the computation should be propagated to all devices. This entails some communication!

Given a valid input-output sharding for a more complex computation, the JAX compiler can heuristically optimize the inner bits of the computation (intermediate computations’ shardings, and necessary communication) to try to make it run as efficiently as possible. In cases where this doesn’t work, JAX has ways to hint at the compiler, as well as more advanced programming paradigms, to explicitly specify intermediate sharding and communications.

Valid Shardings

Given a mesh with $M$ axes, and an array with $N$ axes, we can enumerate all valid shardings as follows. First note:

A sharding corresponds to, for each of the $N$ array axes, the selection of one of the $M$ mesh axes, or None.
Each mesh axis can appear at most once in a valid sharding.

Then for each $k = 0, \dots, \min \set{M, N}$, pick $\binom{M}{k}$ of the $M$ mesh axes and $\binom{N}{k}$ of the $N$ array axes, then consider all possible permutations of those $k$ selected axes. This is a total of

\[\sum_{k=0}^{\min \set{M, N}} \binom{M}{k} \binom{N}{k} k!\]

possible shardings. This is actually not too large—roughly on the order of $N^M$ when $N$ is large.¹

Example: Matrix Multiplication

The most common operations in neural networks are related to matrix multiplications, and so thinking in terms of sharding data with two axes goes a long way. Say we have matrices $\vA \in \bbR^{m \times n}$ and $\vB \in \bbR^{n \times p}$. A sharding of a 2D array (matrix) corresponds to two things: a partitioning of the matrix into sub-matrices, and a splitting of those sub-matrices across devices. In JAX, these two operations are coupled through the Mesh/NamedSharding interface, as we described above:

If an array dimension is sharded across a mesh axis, that dimension is partitioned, and the partitioned components will be split across that mesh axis.
If an array dimension is not sharded, that dimension is not partitioned, and it appears across all devices.
If a mesh axis does not appear in a sharding specification, the sharded data is copied across that mesh axis.

Let’s consider the special case of a 2D mesh, in particular $2 \times 2$. For simplicity, let’s place the first mesh axis in correspondence with the first array dimension, and the second mesh axis in correspondence with the second array dimension.² Then every partitioning of $\vA$ or $\vB$, for example

\[\begin{equation}\label{eq:partitioning-0} \begin{bmatrix} \vA_{00} & \vA_{01} \\ \vA_{10} & \vA_{11} \end{bmatrix}, \end{equation}\]

corresponds to an arrangement of the matrix across devices. In the above example, we have $\vA_{ij} \in \bbR^{(m/2) \times (n/2)}$, and the coordinates $(i, j)$ of the submatrix correspond to the coordinates of the device it is stored on. For a different partitioning, for example

\[\begin{bmatrix} \vA_{0} \\ \vA_{1} \end{bmatrix},\]

we need to take copying into account: here we have $\vA_{i} \in \bbR^{(m/2) \times n}$, and the physical device layout is

\[\begin{equation}\label{eq:partitioning-1-dev} \begin{bmatrix} \vA_{0} & \vA_{0} \\ \vA_{1} & \vA_{1} \end{bmatrix}. \end{equation}\]

Notice that $\eqref{eq:partitioning-0}$ and $\eqref{eq:partitioning-1-dev}$ are incomparable: the second one uses twice as much memory! But there is a simple ‘algorithm’ to go from a partitioning to a device layout: if an axis is not partitioned, simply copy the array along that axis to reach the full mesh size (and remember that this leads to additional memory usage).

Now, what happens if we want to compute the product $\vA \vB$ when $\vA$ is sharded in some way, for example as in $\eqref{eq:partitioning-0}$? First, thinking abstractly: by properties of block matrix multiplication, we have

\[\vA \vB = \begin{bmatrix} \vA_{11} & \vA_{12} \\ \vA_{21} & \vA_{22} \end{bmatrix} \begin{bmatrix} \vB_{11} & \vB_{12} \\ \vB_{21} & \vB_{22} \end{bmatrix} = \begin{bmatrix} \vA_{11}\vB_{11} + \vA_{12} \vB_{21} & \vA_{11}\vB_{12} + \vA_{12} \vB_{22} \\ \vA_{21}\vB_{11} + \vA_{22} \vB_{21} & \vA_{21}\vB_{12} + \vA_{22} \vB_{22} \\ \end{bmatrix}.\]

So no matter how $\vB$ is sharded, we need to perform communication such that we compute each of the matrix products appearing above. We can think about two cases:

No communication involving $\vA$: Here, we can use the RHS of the above display to tell us what needs to be computed. Processor $(0, 0)$ needs the first (block) row of $\vB$, processor $(0, 1)$ needs the second (block) row of $\vB$, etc.—so the sharding of $\vB$ had better support this, or communication will be necessary. Regardless, some communication between devices is necessary to accumulate the block matrix products: e.g., we need to add results from processor $(0, 0)$ and $(0, 1)$ to get the top-left block of the product.
We can re-structure how $\vA$ is sharded: It might be possible, given a specific sharding for $\vB$, to have a more efficient computation by sharding $\vA$ differently (and vice versa). Thinking in terms of block partitions, this is the same as observing that there are many ways of partitioning a matrix into blocks, each of which gives a different decomposition for the matrix product. For example, this is also a compatible partitioning (with suitable definitions of the blocks):
\[\begin{equation}\label{eq:nice-sharding} \vA \vB = \begin{bmatrix} \vA_{1} \\ \vA_{2} \end{bmatrix} \begin{bmatrix} \vB_{1} & \vB_{2} \end{bmatrix} = \begin{bmatrix} \vA_{1}\vB_{1} & \vA_{1}\vB_{2} \\ \vA_{2}\vB_{1} & \vA_{2}\vB_{2} \end{bmatrix}, \end{equation}\]
and so is
\[\vA \vB = \begin{bmatrix} \vA_{1} & \vA_{2} \end{bmatrix} \begin{bmatrix} \vB_{1} \\ \vB_{2} \end{bmatrix} = \vA_1 \vB_1 + \vA_2 \vB_2.\]
Re-sharding requires communication, though, and so this may or may not be more efficient.

As an exercise, verify that $\eqref{eq:nice-sharding}$ corresponds to a zero-communication sharding for a $2 \times 2$ output sharding.³

Now, if $\vA$ and $\vB$ both have their own sharding, in general these correspond to (possibly incompatible) block matrix partitions of each matrix. For each individual partition, we can, as above, write out the blockwise matrix products and accumulation operations needed to compute the product, and try to intuit what computation/communication should happen—in particular involving re-sharding one or both arrays. In JAX, the XLA compiler will be recruited to perform this optimization; as algorithm designers, we can try to provide sensible shardings up front in order to make sure we end up with an efficient system.

Types of Inter-Device Communication with Sharded Computation

The matrix multiplication example suggests that different forms of communication naturally arise in distributed computation. We will describe the ones that are implemented in XLA here. The JAX scaling book (Austin et al., 2025) gives a nice way to think about matrix multiplication of sharded arrays using a variant of named-axis notation. Since this is closer to code, we’ll describe this briefly below.

Communication Primitives through Matrix Multiplication

We saw in the previous section how to think about sharding of 2D arrays in terms of block matrices:

Sharded array axes are partitioned; non-sharded array axes aren’t.
Any mesh axis that isn’t used means the partitioned array is copied across that axis. This gives us a scheme to map a block matrix to the corresponding arrangement of that matrix across physical devices.

With named index notation, things are more compact. Consider a 2D mesh with axes $x$ and $y$, and let us write $A[I, J]$ for the 2D array $\vA$ (likewise for $\vB$) we looked at in the previous section. We write $A[I_x, J]$ (etc.) to denote the sharding of $A$ along its first dimension with respect to the mesh axis $x$. The use of $I$, $J$, etc.\ for ‘dummy’ indices is in analogy to einsum-type notation, which is very useful for tensor operations.

Now consider a matrix product of sharded matrices $A[I_x, J] \cdot B[J, K_y]$. Suppose we want the output, say $\vC$, to be sharded as $C[I_x, K_y]$. Then this matrix product can be computed with no communication: to see why, think in terms of the correspondence to block matrix partitions and their associated map to device arrangements in the previous section. However, for an unsharded output $C[I, K]$, it would be necessary to perform some communication—either calculating $C[I_x, K_y]$ and then sharing the local results among all devices, or something else.

The following is a convenient way to think about this, following the previous example:

Given a sharding of matrix multiplicands $A$, $B$, there is a ‘natural’ output sharding whenever the contracting dimension is unsharded. (As above, this comes by ‘removing’ the contracting axis, just like in einsum notation.) If this sharding is invalid, communication is necessary.
If the contracting dimension is sharded, different types of communication are necessary to turn the multiplication into one with an unsharded contracting dimension. We can identify typical ‘best practices’ for what communication type to use based on how the contracting dimension is sharded.
If an output sharding is specified and it is not the ‘natural’ one, some communication is required to achieve it. We can use the techniques from either of the two previous bullets to do this.

We’ll describe these different cases below.

Case 1: Unsharded Contracting Dimension, Valid Output

As above, the ‘natural’ output sharding is just an einsum-style replacement:

\[A[I_x, J] \cdot B[J, K] \to C[I_x, K].\]

If this natural output sharding is invalid, we need to do something else (see below).

Case 2: One Multiplicand Sharded Along Contracting Dimension (AllGather)

For a product like

\[A[I, J_x] \cdot B[J, K] \to C[I, K],\]

it isn’t possible to directly perform the multiplication, because in order to contract, we either need to perform a series of block-matrix multiplies, then accumulate the results (see below), or have the entire contracting dimension of both arrays available on all devices (see above).

We use AllGather for this, which removes a specified sharded axis:

\[\mathrm{AllGather}_{x}(A[I, J_x]) \to A[I, J].\]

For a compound mesh, AllGather can also just remove one sub-axis:

\[\mathrm{AllGather}_{x}(A[I, J_{xy}]) \to A[I, J_y].\]

AllGather can also be used to satisfy certain output shardings that aren’t naturally compatible with the input shardings. For example, in the product

\[A[I_x, J] \cdot B[J, K] \to C[I, K],\]

one removes the sharded axis either before or after the computation in order to achieve the desired output sharding.

AllGather Cost

On TPUs with enough chips to form a cube, AllGather can be performed using algorithms that exploit the toroidal interconnectivity of TPU devices. For example, to AllGather a single axis, the following procedure can be used:

Initialize a local buffer on each device with its local shard.
Send the buffer contents to the next device; overwrite it with the contents of the previous device.
Store the received shard in memory.
Repeat.

This takes at most $N/2$ rounds of communication if the mesh axis has size $N$. For an array of size $V$ bytes, each round sends $V/N$ bytes over each link. Given an ICI bandwidth of $W$ (unidirectional),⁴ the total time this takes is no more than

\[\frac{N}{2} \cdot \frac{V}{NW} = \frac{V}{2W} \text{ s},\]

which is independent of the size of the mesh. However, one must also take into account the intrinsic overhead of ICI communication. Each operation takes around $T_{\min} = 10^{-6} \text{ s}$ (Austin et al., 2025), so the total time is actually only bounded by

\[\max \set*{\frac{V}{2W}, \frac{NT_{\min}}{2}}.\]

In particular, there are latency-bound (small array size) and throughput-bound (large array size) regimes of communication.

AllGathering a compound axis (e.g., $xy$) is similar: one can utilize some per-device state to efficiently communicate data across the entire multidimensional mesh. Since each device can then communicate along each mesh axis simultaneously, the bandwidth increases by a factor proportional to the number of mesh axes, but the maximum possible latency also increases, since there are more devices. The total time is no larger than

\[\max \set*{\frac{V}{2n_{\mathrm{axes}}W}, \frac{N_{\mathrm{total}}T_{\min}}{2}},\]

where $N_{\mathrm{total}}$ is the product of the axes lengths.

Case 3: Both Multiplicands Sharded Along Contracting Dimension (AllReduce)

In this case, we want to compute a matrix product like

\[A[I, J_x] \cdot B[J_x, K] \to C[I, K].\]

Thinking in terms of block matrices, we know that computing this matrix product entails a sum of the block matrix multiplications corresponding to the shards of both $A$ and $B$, and since the contracting axis is sharded, we can:

Compute per-device multiplications of the sharded matrices.
Add these up.

The communication primitive that performs this reduction is called AllReduce. It is similar to AllGather, but it sums up the components being communicated, and distributes the resulting sum to all devices.

We will use the notation from (Austin et al., 2025): we write

\[C[I, K]\{ U_x \}\]

to denote an array $C$ which is “unreduced” over the mesh axis $x$, i.e. the partial products are sharded over the $x$ mesh axis, which gives

\[\mathrm{AllReduce}_y\left( A[I_x, J]\{ U_y \} \right) \to A[I_x, J]\]

as the signature for AllReduce.

Decomposition in Terms of ReduceScatter

Generally, an AllReduce is about two times as expensive as an AllGather. One way to see why is to decompose it into another useful communication primitive: the ReduceScatter operation.

ReduceScatter has signature

\[\mathrm{ReduceScatter}_{y, J}\left( A[I_x, J]\{ U_y \} \right) \to A[I_x, J_y].\]

It takes an array unreduced along a mesh axis, then shards the array along a specified array axis with respect to that mesh axis. To compute an AllReduce, we can compose a ReduceScatter with an AllGather:

\[\mathrm{AllReduce}_y\left( A[I_x, J]\{ U_y \} \right) = \mathrm{AllGather}_{y}\left( \mathrm{ReduceScatter}_{y, J}\left( A[I_x, J]\{ U_y \} \right) \right).\]

For multidimensional arrays, there is a possibility to choose the array axis that is ReduceScattered to optimize performance.

ReduceScatter can be performed by an algorithm similar to the approach we discussed for AllGather. The process is somewhat difficult to describe algorithmically, but there is a very good pictorial representation in (Austin et al., 2025).

The upshot is that this similarity persists to the latency and throughput analysis of the algorithm, making it identical to that of AllGather. This demonstrates the cost of AllReduce as twice that of AllGather.

Aside: The Backward Pass

There is a more fundamental relationship between the AllGather and ReduceScatter operations, suggested by the representation of AllReduce above: they can be represented as adjoints of one another, which means that an AllGather in the forward pass will trigger a ReduceScatter in the backward pass!

To see this, let’s go back to our more visual block matrix example. As above, consider a simplified setting where a matrix is sharded along a 1D mesh corresponding to the first array dimension. With mesh size $3$ for simplicity, we represent this as the block matrix

\[\vA = \begin{bmatrix} \vA_0 \\ \vA_1 \\ \vA_2 \end{bmatrix}.\]

Remember that the partitioning of the block matrix corresponds to the device layout in this model. Now, if we perform an AllGather, we have the ($3\times$ larger) block matrix output

\[\mathrm{AllGather}(\vA) = \begin{bmatrix} \vA \\ \vA \\ \vA \end{bmatrix} = \begin{bmatrix} \begin{bmatrix} \vA_0 \\ \vA_1 \\ \vA_2 \end{bmatrix} \\ \begin{bmatrix} \vA_0 \\ \vA_1 \\ \vA_2 \end{bmatrix} \\ \begin{bmatrix} \vA_0 \\ \vA_1 \\ \vA_2 \end{bmatrix} \end{bmatrix}.\]

Let’s imagine we perform some subsequent computations on the AllGathered output, and we want to compute a backward pass. In general, each individual device will have different gradients, which we can write in partitioned form as

\[\begin{bmatrix} \Delta \vA^0 \\ \Delta \vA^1 \\ \Delta \vA^2 \end{bmatrix} = \begin{bmatrix} \begin{bmatrix} \Delta \vA_0^0 \\ \Delta \vA_1^0 \\ \Delta \vA_2^0 \end{bmatrix} \\ \begin{bmatrix} \Delta \vA_0^1 \\ \Delta \vA_1^1 \\ \Delta \vA_2^1 \end{bmatrix} \\ \begin{bmatrix} \Delta \vA_0^2 \\ \Delta \vA_1^2 \\ \Delta \vA_2^2 \end{bmatrix} \end{bmatrix}.\]

Then because AllGather behaves in the sense of linear algebra as a copying operation, its backward pass involves its adjoint operation, which is a sum. We end up with⁵

\[\mathrm{dAllGather}^* \left( \begin{bmatrix} \Delta \vA^0 \\ \Delta \vA^1 \\ \Delta \vA^2 \end{bmatrix} \right) = \begin{bmatrix} \sum_{i=1}^3\Delta \vA^i_0 \\ \sum_{i=1}^3 \Delta \vA^i_1 \\ \sum_{i=1}^3 \Delta \vA^i_2 \end{bmatrix}.\]

This is exactly what we’d get if we:

Interpret the per-device gradients as unreduced partial sums with respect to the mesh axis of the AllGather;
Perform a ReduceScatter with respect to this mesh axis and the array axis that was sharded before the AllGather.

So, an AllGather in the forward pass produces a corresponding ReduceScatter in the backward pass!

Similarly, if there is a ReduceScatter in the forward pass that produces an array sharded along one of its axes, then in the backward pass, there will be an AllGather operation along that axis in the backward pass. And finally, because an AllReduce can be expressed as the composition of an AllGather with a ReduceScatter, the backward pass of an AllReduce is an AllReduce!

Case 4: Sharded Non-Contracting Dimensions, Same Axis

A direct reduction to the natural sharding for this case would lead to an invalid sharding:

\[A[I_x, J] \cdot B[J, K_x] \to C[I_x, K_x]\]

A mesh axis cannot appear multiple times in a sharding specification.

To resolve this, we just AllGather one of the axes in the multiplicands:

\[\begin{split} \mathrm{AllGather}_x(A[I_x, J]) \to A[I, J] \\ A[I, J] \cdot B[J, K_x] \to C[I, K_x]. \end{split}\]

We can select which depending on the downstream shardings needed.

AllToAll Communication

This primitive can be thought of as a resharding operation. It has signature

\[\mathrm{AllToAll}_{x, J}( A[I_x, J] ) \to A[I, J_x].\]

AllToAll is actually significantly cheaper than an AllGather, by a factor of 4. This is another one where the excellent picture from (Austin et al., 2025) describes where this efficiency is coming from well.

Notice that the message that is being sent over each link is composed of pieces of different lengths from each shard. The communication is balanced, but later messages are shorter, which may cause issues with latency-boundedness of the communication.

Example: Communication Primitives in JAX

Let’s test out what we learned above in the JAX sandbox we started this post with. For slightly more granular control of sharding, we’ll switch to Explicit mode sharding (see e.g. (JAX Team, 2025)).

mesh = jax.make_mesh(
    (2, 2), ("x", "y"), axis_types=(jax.sharding.AxisType.Explicit,) * 2
)  # 1x v4 tpu host
jax.set_mesh(mesh)

A = jnp.zeros((2048, 8192), dtype=jnp.bfloat16, device=jax.NamedSharding(mesh, jax.P("x", "y")))

jax.typeof(A)

Output:
ShapedArray(bfloat16[2048@x,8192@y])

AllGather

First, a simple AllGather implementation along the y mesh axis—we compile the resharding operation and inspect the “high-level operations” (HLO) to see what actually gets implemented on the device.

@jax.jit
def all_gather(x):
    return jax.sharding.reshard(x, jax.NamedSharding(mesh, jax.P("x")))


lowered = all_gather.lower(A)
compiled = lowered.compile()
B = compiled(A)

print("Output type: ", jax.typeof(B))
print()
print("StableHLO:\n", lowered.as_text())
print()
print("XLA HLO:\n", compiled.as_text())

Output:
Output type:  bfloat16[2048@x,8192]

StableHLO:
 module @jit_all_gather attributes {mhlo.num_partitions = 4 : i32, mhlo.num_replicas = 1 : i32} {
  sdy.mesh @mesh = <["x"=2, "y"=2]>
  func.func public @main(%arg0: tensor<2048x8192xbf16> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {"y"}]>}) -> (tensor<2048x8192xbf16> {jax.result_info = "result", sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) {
    %0 = sdy.sharding_constraint %arg0 <@mesh, [{"x"}, {}]> : tensor<2048x8192xbf16>
    return %0 : tensor<2048x8192xbf16>
  }
}

XLA HLO:
 HloModule jit_all_gather, is_scheduled=true, entry_computation_layout={(bf16[1024,4096]{1,0:T(8,128)(2,1)})->bf16[1024,8192]{1,0:T(8,128)(2,1)}}, num_partitions=4
ENTRY %main.5_spmd (param: bf16[1024,4096]) -> bf16[1024,8192] {
  %param = bf16[1024,4096]{1,0:T(8,128)(2,1)} parameter(0), sharding={devices=[2,2]<=[4]}, metadata={op_name="x"}
  %all-gather = bf16[1024,8192]{1,0:T(8,128)(2,1)S(3)} all-gather(%param), channel_id=1, replica_groups=[2,2]<=[4], dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(all_gather)/reshard" source_file="/tmp/ipykernel_794398/1925679859.py" source_line=3}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"collective_algorithm_config":{"emitter":"1DAllGatherNonMajorDim","debug":"\ngroup_size = 2 \nhas_reordering_map: false \nper_stride_size = 65536 bytes \nshard_size = 8388608 bytes "},"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"0"}],"retry_config":{"retry_count":"0"}}
  ROOT %copy.3 = bf16[1024,8192]{1,0:T(8,128)(2,1)} copy(%all-gather), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16","64"],"input_window_bounds":[],"estimated_cycles":"58232","iteration_bounds":["8","1"]},"megacore_config":{"megacore_split_dim":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"8388608"}],"retry_config":{"retry_count":"0"}}
}

The first representation above (output of lower) is in XLA’s StableHLO language. It is slightly imposing, but there is a clearly written spec that helps with parsing it after a bit of studying. It’s still too high level to clearly show what communication primitives are being generated, though.

The second representation (output of compile) is the result of the XLA compiler’s first device-independent compilation stage; it’s in its own dialect of HLO, which is different from StableHLO and doesn’t have as easily-accessible of a public spec.⁶ Here’s how to parse it above:

Types are relatively self-explanatory, aside from the text in braces—the first sequence of integers here is most important, and it specifies the memory layout of the array (see layout.h in the XLA source). Axis indices are listed in minor-to-major order (i.e., 1, 0 denotes row-major layout).
The all-gather semantics are clear at a high level: we gather along dimension $1$ of our array. The replica_groups argument specifies which devices communicate in the gather (see tile_assignment.h in the XLA source)—compare to our input %param’s sharding specified above to see that these agree.
Independent of these lower-level details, we can see that the input and output shapes are as we expect.

AllToAll

We can test out AllToAll with a similar high-level API. Let’s start by transposing the sharding axis of the output of the previous AllGather.

@jax.jit
def all_to_all(x):
    return jax.sharding.reshard(x, jax.NamedSharding(mesh, jax.P(None, "x")))


lowered = all_to_all.lower(B)
compiled = lowered.compile()
C = compiled(B)

print("Output type: ", jax.typeof(C))
print()
print("StableHLO:\n", lowered.as_text())
print()
print("XLA HLO:\n", compiled.as_text())

The HLO that we get is more complex:

Output:
Output type:  bfloat16[2048,8192@x]
StableHLO:
 module @jit_all_to_all attributes {mhlo.num_partitions = 4 : i32, mhlo.num_replicas = 1 : i32} {
  sdy.mesh @mesh = <["x"=2, "y"=2]>
  func.func public @main(%arg0: tensor<2048x8192xbf16> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) -> (tensor<2048x8192xbf16> {jax.result_info = "result", sdy.sharding = #sdy.sharding<@mesh, [{}, {"x"}]>}) {
    %0 = sdy.sharding_constraint %arg0 <@mesh, [{}, {"x"}]> : tensor<2048x8192xbf16>
    return %0 : tensor<2048x8192xbf16>
  }
}
XLA HLO:
 HloModule jit_all_to_all, is_scheduled=true, entry_computation_layout={(bf16[1024,8192]{1,0:T(8,128)(2,1)})->bf16[2048,4096]{1,0:T(8,128)(2,1)}}, num_partitions=4
%fused_computation (param_0.1: bf16[128,2,8,4096]) -> bf16[1024,2,4096] {
  %param_0.1 = bf16[128,2,8,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0)
  %copy.4 = bf16[128,2,8,4096]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.1), metadata={op_name="jit(all_to_all)/reshard" source_file="/tmp/ipykernel_874047/2971478236.py" source_line=3}
  ROOT %bitcast.5 = bf16[1024,2,4096]{2,0,1:T(8,128)(2,1)S(3)} bitcast(%copy.4), metadata={op_name="jit(all_to_all)/reshard" source_file="/tmp/ipykernel_874047/2971478236.py" source_line=3}
}
ENTRY %main.5_spmd (param: bf16[1024,8192]) -> bf16[2048,4096] {
  %param = bf16[1024,8192]{1,0:T(8,128)(2,1)} parameter(0), sharding={devices=[2,1,2]<=[4] last_tile_dim_replicate}, metadata={op_name="x"}
  %bitcast.7 = bf16[128,2,8,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%param)
  %copy_bitcast_fusion = bf16[1024,2,4096]{2,0,1:T(8,128)(2,1)S(3)} fusion(%bitcast.7), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(all_to_all)/reshard" source_file="/tmp/ipykernel_874047/2971478236.py" source_line=3}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["1","32","32"],"input_window_bounds":[],"estimated_cycles":"78212","iteration_bounds":["2","4","1"]},"megacore_config":{"megacore_split_dim":"1"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"8388608"}],"retry_config":{"retry_count":"0"}}
  %all-to-all = bf16[1024,2,4096]{2,0,1:T(8,128)(2,1)S(3)} all-to-all(%copy_bitcast_fusion), channel_id=1, replica_groups=[2,2]<=[2,2]T(1,0), dimensions={1}, metadata={op_name="jit(all_to_all)/reshard" source_file="/tmp/ipykernel_874047/2971478236.py" source_line=3}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"GLOBAL","id":"-1"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"262144"}],"retry_config":{"retry_count":"0"}}
  %bitcast.6 = bf16[2048,4096]{1,0:T(8,128)(2,1)S(3)} bitcast(%all-to-all)
  ROOT %copy.6 = bf16[2048,4096]{1,0:T(8,128)(2,1)} copy(%bitcast.6), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32","32"],"input_window_bounds":[],"estimated_cycles":"58232","iteration_bounds":["8","1"]},"megacore_config":{"megacore_split_dim":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"8388608"}],"retry_config":{"retry_count":"0"}}
}

Some amount of the above amounts to memory optimizations that are necessary to actually implement the communication in physical memory.⁷ We can more easily interpret what’s going on by thinking in terms of block matrices again: recall that we started with a $2048 \times 8192$ array $\vA$, sharded two ways:

\[\vA = \begin{bmatrix} \vA_{00} & \vA_{01} \\ \vA_{10} & \vA_{11} \end{bmatrix},\]

and we performed an AllGather along the second axis, leading to (we’ll use $\oplus_i$ for concatenation along dimension $i$)

\[\vB = \begin{bmatrix} \vA_{00} \oplus_1 \vA_{01} & \vA_{00} \oplus_1 \vA_{01} \\ \vA_{10} \oplus_1 \vA_{11} & \vA_{10} \oplus_1 \vA_{11} \end{bmatrix}.\]

The output of the AllToAll reshards along the trailing axis, but with respect to the 'x' mesh axis, which gives

\[\vC = \begin{bmatrix} \vA_{00} \oplus_0 \vA_{10} & \vA_{00} \oplus_0 \vA_{10} \\ \vA_{01} \oplus_0 \vA_{11} & \vA_{01} \oplus_0 \vA_{11} \end{bmatrix}.\]

To go from $\vB$ to $\vC$, we can do the following, which is essentially what the HLO above is doing (modulo memory operations): unbind the concatenated arrays, to give (in a rough approximation of tensor notation)

\[\begin{bmatrix} (\vA_{00}, \vA_{01}) & (\vA_{00}, \vA_{01}) \\ (\vA_{10}, \vA_{11}) & (\vA_{10}, \vA_{11}) \end{bmatrix},\]

then communicate ‘vertically’, swapping $01$ and $10$ on the left and right, to get

\[\begin{bmatrix} (\vA_{00}, \vA_{10}) & (\vA_{00}, \vA_{10}) \\ (\vA_{01}, \vA_{11}) & (\vA_{01}, \vA_{11}) \end{bmatrix}.\]

If we concatenate these in the right order, we end up with $\vC$! We can see that this is what the HLO is doing:

After some memory manipulations, an extra axis of size $2$ is split off from the second dimension, making the local array shaped like $1024 \times 2 \times 4096$.
The all-to-all function has to be interpreted relative to the input sharding. Here {devices=[2,1,2]<=[4] last_tile_dim_replicate} denotes a [2, 1] tile size, replicated twice; <=[4] is equivalent to a device list {0, 1, 2, 3}. This is as we expect—the input is sharded over 'x'.
Now, in all-to-all, replica_groups=[2,2]<=[2,2]T(1,0) specifies how the communication is done across devices (relative to the sharding previously specified). This corresponds to a specification of {{0, 2}, {1, 3}} (see hlo_sharding.h in the XLA source), which means that each replica’s array gets split into two parts along dimension $1$ (the axis of length $2$), then communication occurs across the specified groups. This gets us what we wanted above.

It’s slightly interesting to note that different input sharding and an AllToAll leads to different memory operations in the HLO:

@jax.jit
def all_to_all(x):
    return jax.sharding.reshard(x, jax.NamedSharding(mesh, jax.P("x")))


lowered = all_to_all.lower(C)
compiled = lowered.compile()
D = compiled(C)

print("Output type: ", jax.typeof(D))
print()
print("StableHLO:\n", lowered.as_text())
print()
print("XLA HLO:\n", compiled.as_text())

Output:
Output type:  bfloat16[2048@x,8192]
StableHLO:
 module @jit_all_to_all attributes {mhlo.num_partitions = 4 : i32, mhlo.num_replicas = 1 : i32} {
  sdy.mesh @mesh = <["x"=2, "y"=2]>
  func.func public @main(%arg0: tensor<2048x8192xbf16> {sdy.sharding = #sdy.sharding<@mesh, [{}, {"x"}]>}) -> (tensor<2048x8192xbf16> {jax.result_info = "result", sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {}]>}) {
    %0 = sdy.sharding_constraint %arg0 <@mesh, [{"x"}, {}]> : tensor<2048x8192xbf16>
    return %0 : tensor<2048x8192xbf16>
  }
}
XLA HLO:
 HloModule jit_all_to_all, is_scheduled=true, entry_computation_layout={(bf16[2048,4096]{1,0:T(8,128)(2,1)})->bf16[1024,8192]{1,0:T(8,128)(2,1)}}, num_partitions=4
ENTRY %main.5_spmd (param: bf16[2048,4096]) -> bf16[1024,8192] {
  %param = bf16[2048,4096]{1,0:T(8,128)(2,1)} parameter(0), sharding={devices=[1,2,2]<=[4] last_tile_dim_replicate}, metadata={op_name="x"}
  %bitcast.5 = bf16[2,1024,4096]{2,1,0:T(8,128)(2,1)} bitcast(%param)
  %all-to-all = bf16[2,1024,4096]{2,1,0:T(8,128)(2,1)S(3)} all-to-all(%bitcast.5), channel_id=1, replica_groups=[2,2]<=[2,2]T(1,0), dimensions={0}, metadata={op_name="jit(all_to_all)/reshard" source_file="/tmp/ipykernel_874047/814736672.py" source_line=3}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"GLOBAL","id":"-1"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"0"}],"retry_config":{"retry_count":"0"}}
  %copy.3 = bf16[2,1024,4096]{1,2,0:T(8,128)(2,1)S(3)} copy(%all-to-all), metadata={op_name="jit(all_to_all)/reshard" source_file="/tmp/ipykernel_874047/814736672.py" source_line=3}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["1","48","8"],"input_window_bounds":["1","128","3"],"estimated_cycles":"47696","iteration_bounds":["1","11","1"]},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"3145728"}],"retry_config":{"retry_count":"0"}}
  %bitcast.4 = bf16[1024,8192]{0,1:T(8,128)(2,1)S(3)} bitcast(%copy.3)
  ROOT %copy.4 = bf16[1024,8192]{1,0:T(8,128)(2,1)} copy(%bitcast.4), metadata={op_name="jit(all_to_all)/reshard" source_file="/tmp/ipykernel_874047/814736672.py" source_line=3}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["64","6"],"input_window_bounds":["96","4"],"estimated_cycles":"67782","iteration_bounds":["1","11"]},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"3182592"}],"retry_config":{"retry_count":"0"}}
}

The HLO for this operation is more concise, but neither is able to avoid a copy operation.

AllReduce and ReduceScatter

To trigger an AllReduce/ReduceScatter, let’s follow our matrix multiplication example above—we’ll multiply two arrays that are both sharded along the contracting dimension. In the AllToAll section, we’ve generated C, which is sharded along 'x' in its 1st dimension. We can try computing a Gram matrix associated to C, and specifying the out_sharding (singular!) in the dot operation to avoid the compiler complaining:

@jax.jit
def gram(x):
    return jnp.dot(x, x.mT, out_sharding=jax.NamedSharding(mesh, jax.P()))


lowered = gram.lower(C)
compiled = lowered.compile()
E = compiled(C)


print("Output type: ", jax.typeof(E))
print()
print("StableHLO:\n", lowered.as_text())
print()
print("XLA HLO:\n", compiled.as_text())

Output:
Output type:  bfloat16[2048,2048]
StableHLO:
 module @jit_gram attributes {mhlo.num_partitions = 4 : i32, mhlo.num_replicas = 1 : i32} {
  sdy.mesh @mesh = <["x"=2, "y"=2]>
  func.func public @main(%arg0: tensor<2048x8192xbf16> {sdy.sharding = #sdy.sharding<@mesh, [{}, {"x"}]>}) -> (tensor<2048x2048xbf16> {jax.result_info = "result", sdy.sharding = #sdy.sharding<@mesh, [{}, {}]>}) {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<2048x8192xbf16>) -> tensor<8192x2048xbf16>
    %1 = sdy.sharding_constraint %0 <@mesh, [{"x"}, {}]> : tensor<8192x2048xbf16>
    %2 = stablehlo.dot_general %arg0, %1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2048x8192xbf16>, tensor<8192x2048xbf16>) -> tensor<2048x2048xbf16>
    %3 = sdy.sharding_constraint %2 <@mesh, [{}, {}]> : tensor<2048x2048xbf16>
    return %3 : tensor<2048x2048xbf16>
  }
}
XLA HLO:
 HloModule jit_gram, is_scheduled=true, entry_computation_layout={(bf16[2048,4096]{1,0:T(8,128)(2,1)})->bf16[2048,2048]{1,0:T(8,128)(2,1)}}, num_partitions=4
%add.clone (x.1: bf16[], y.1: bf16[]) -> bf16[] {
  %y.1 = bf16[]{:T(256)} parameter(1)
  %x.1 = bf16[]{:T(256)} parameter(0)
  ROOT %add.1 = bf16[]{:T(256)} add(%x.1, %y.1)
}
%bitcast_fusion (bitcast_input: bf16[2048,4096]) -> bf16[2048,4096] {
  %bitcast_input = bf16[2048,4096]{1,0:T(8,128)(2,1)S(3)} parameter(0)
  ROOT %bitcast = bf16[2048,4096]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input)
}
%bitcast_fusion.1 (bitcast_input.1: bf16[2048,4096]) -> bf16[2048,4096] {
  %bitcast_input.1 = bf16[2048,4096]{1,0:T(8,128)(2,1)S(3)} parameter(0)
  ROOT %bitcast.1 = bf16[2048,4096]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.1)
}
%fused_computation (param_0: bf16[2048,4096]) -> bf16[2048,2048] {
  %param_0 = bf16[2048,4096]{1,0:T(8,128)(2,1)S(3)} parameter(0)
  %fusion.1 = bf16[2048,4096]{1,0:T(8,128)(2,1)} fusion(%param_0), kind=kLoop, calls=%bitcast_fusion
  %fusion.2 = bf16[2048,4096]{1,0:T(8,128)(2,1)} fusion(%param_0), kind=kLoop, calls=%bitcast_fusion.1
  ROOT %convolution.1 = bf16[2048,2048]{1,0:T(8,128)(2,1)S(3)} convolution(%fusion.1, %fusion.2), dim_labels=bf_oi->bf, metadata={op_name="jit(gram)/dot_general" source_file="/tmp/ipykernel_874047/670048905.py" source_line=3}
}
ENTRY %main.8_spmd (param: bf16[2048,4096]) -> bf16[2048,2048] {
  %param = bf16[2048,4096]{1,0:T(8,128)(2,1)} parameter(0), sharding={devices=[1,2,2]<=[4] last_tile_dim_replicate}, metadata={op_name="x"}
  %copy-start = (bf16[2048,4096]{1,0:T(8,128)(2,1)S(3)}, bf16[2048,4096]{1,0:T(8,128)(2,1)}, u32[]{:S(2)}) copy-start(%param), cross_program_prefetch_index=0
  %copy-done = bf16[2048,4096]{1,0:T(8,128)(2,1)S(3)} copy-done(%copy-start)
  %fusion = bf16[2048,2048]{1,0:T(8,128)(2,1)S(3)} fusion(%copy-done), kind=kOutput, calls=%fused_computation, metadata={op_name="jit(gram)/dot_general" source_file="/tmp/ipykernel_874047/670048905.py" source_line=3}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":["64","32"],"output_window_bounds":["64","4"],"input_window_bounds":["64","32"],"estimated_cycles":"149352","iteration_bounds":["4","4","1"]},"megacore_config":{"megacore_split_dim":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"5734400"}],"retry_config":{"retry_count":"0"},"convolution_algorithm_config":{"emitter":"EmitAllBatchInSublanes"}}
  ROOT %all-reduce = bf16[2048,2048]{1,0:T(8,128)(2,1)} all-reduce(%fusion), channel_id=1, replica_groups=[2,2]<=[2,2]T(1,0), use_global_device_ids=true, to_apply=%add.clone, metadata={op_name="jit(gram)/dot_general" source_file="/tmp/ipykernel_874047/670048905.py" source_line=3}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[{"memory_space":"0","offset":"0","size":"67108864"}],"collective_algorithm_config":{"emitter":"RotatedPincerEmitter","strategy":"UniDirection1DRingStrategy","debug":"\nUniDirection1DRingStrategy{colors:2 phases:1 cores:{2},{2} nophase0:0 reserved_sflags:0 cross_module_on_2d_plane:0 has_reordering_map:0 use_routing_table_indices:0}"},"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"12582912"}],"retry_config":{"retry_count":"0"}}
}

The above is relatively clear! A copy operation accounts for the fact that we’re multiplying the matrix with itself to calculate the Gram matrix, then the actual multiplication is performed with the XLA convolution operation. It’s performed on each local shard, and then the local shards are accumulated and shared with an AllReduce.

We expect to be able to trigger a ReduceScatter instead of an AllReduce by asking the computation output to be sharded in a particular way (rather than fully replicated, as above). Let’s try to verify this:

@jax.jit
def dot_custom(x, y):
    return jnp.dot(x, y, out_sharding=jax.NamedSharding(mesh, jax.P('x', 'y')))

F = jnp.zeros((8192, 4096), dtype=jnp.bfloat16, out_sharding=jax.NamedSharding(mesh, jax.P('y')))

lowered = dot_custom.lower(A, F)
compiled = lowered.compile()
G = compiled(A, F)


print("Output type: ", jax.typeof(G))
print()
print("StableHLO:\n", lowered.as_text())
print()
print("XLA HLO:\n", compiled.as_text())

Output:
Output type:  bfloat16[2048@x,4096@y]
StableHLO:
 module @jit_dot_custom attributes {mhlo.num_partitions = 4 : i32, mhlo.num_replicas = 1 : i32} {
  sdy.mesh @mesh = <["x"=2, "y"=2]>
  func.func public @main(%arg0: tensor<2048x8192xbf16> {sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {"y"}]>}, %arg1: tensor<8192x4096xbf16> {sdy.sharding = #sdy.sharding<@mesh, [{"y"}, {}]>}) -> (tensor<2048x4096xbf16> {jax.result_info = "result", sdy.sharding = #sdy.sharding<@mesh, [{"x"}, {"y"}]>}) {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<2048x8192xbf16>, tensor<8192x4096xbf16>) -> tensor<2048x4096xbf16>
    %1 = sdy.sharding_constraint %0 <@mesh, [{"x"}, {"y"}]> : tensor<2048x4096xbf16>
    return %1 : tensor<2048x4096xbf16>
  }
}
XLA HLO:
 HloModule jit_dot_custom, is_scheduled=true, entry_computation_layout={(bf16[1024,4096]{1,0:T(8,128)(2,1)}, bf16[4096,4096]{1,0:T(8,128)(2,1)})->bf16[1024,2048]{1,0:T(8,128)(2,1)}}, allow_spmd_sharding_propagation_to_parameters={false,false}, num_partitions=4
%add.1.clone (x.3: bf16[], y.3: bf16[]) -> bf16[] {
  %y.3 = bf16[]{:T(256)} parameter(1)
  %x.3 = bf16[]{:T(256)} parameter(0)
  ROOT %add.3 = bf16[]{:T(256)} add(%x.3, %y.3)
}
%all-reduce-scatter (input: bf16[1024,4096]) -> bf16[1024,2048] {
  %input = bf16[1024,4096]{1,0:T(8,128)(2,1)S(3)} parameter(0)
  %all-reduce.2 = bf16[1024,4096]{1,0:T(8,128)(2,1)} all-reduce(%input), channel_id=3, replica_groups={{0,1},{2,3}}, use_global_device_ids=true, to_apply=%add.1.clone, frontend_attributes={from-cross-replica-sharding="true"}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[]}
  %constant.13 = u32[] constant(0)
  %constant.shard_id_table = u32[4]{0:T(128)} constant({0, 1, 0, 1})
  %partition-id.1 = u32[] partition-id()
  %dynamic-slice.3 = u32[1]{0:T(128)} dynamic-slice(%constant.shard_id_table, %partition-id.1), dynamic_slice_sizes={1}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"0","ones":"0","bitwidth":"32"}],"is_index_aligned":[false]},"used_scoped_memory_configs":[]}
  %bitcast.1 = u32[]{:T(128)} bitcast(%dynamic-slice.3)
  %constant.14 = u32[] constant(2048)
  %multiply.3 = u32[]{:T(128)} multiply(%bitcast.1, %constant.14)
  ROOT %dynamic-slice.4 = bf16[1024,2048]{1,0:T(8,128)(2,1)} dynamic-slice(%all-reduce.2, %constant.13, %multiply.3), dynamic_slice_sizes={1024,2048}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294961151","ones":"0","bitwidth":"32"}],"is_index_aligned":[true,false]},"used_scoped_memory_configs":[]}
}
%bitcast_fusion (bitcast_input: bf16[1024,4096]) -> bf16[1024,4096] {
  %bitcast_input = bf16[1024,4096]{1,0:T(8,128)(2,1)} parameter(0)
  ROOT %bitcast.2 = bf16[1024,4096]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input)
}
%bitcast_fusion.1 (bitcast_input.1: bf16[4096,4096]) -> bf16[4096,4096] {
  %bitcast_input.1 = bf16[4096,4096]{1,0:T(8,128)(2,1)S(3)} parameter(0)
  ROOT %bitcast.3 = bf16[4096,4096]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.1)
}
%fused_computation (param_0: bf16[1024,4096], param_1: bf16[4096,4096]) -> bf16[1024,4096] {
  %param_0 = bf16[1024,4096]{1,0:T(8,128)(2,1)} parameter(0)
  %fusion.2 = bf16[1024,4096]{1,0:T(8,128)(2,1)} fusion(%param_0), kind=kLoop, calls=%bitcast_fusion
  %param_1 = bf16[4096,4096]{1,0:T(8,128)(2,1)S(3)} parameter(1)
  %fusion.3 = bf16[4096,4096]{1,0:T(8,128)(2,1)} fusion(%param_1), kind=kLoop, calls=%bitcast_fusion.1
  ROOT %convolution.1 = bf16[1024,4096]{1,0:T(8,128)(2,1)S(3)} convolution(%fusion.2, %fusion.3), dim_labels=bf_io->bf, metadata={op_name="jit(dot_custom)/dot_general" source_file="/tmp/ipykernel_874047/871715024.py" source_line=3}
}
ENTRY %main.7_spmd (param: bf16[1024,4096], param.1: bf16[4096,4096]) -> bf16[1024,2048] {
  %param.1 = bf16[4096,4096]{1,0:T(8,128)(2,1)} parameter(1), sharding={devices=[2,1,2]<=[2,2]T(1,0) last_tile_dim_replicate}, metadata={op_name="y"}
  %copy-start = (bf16[4096,4096]{1,0:T(8,128)(2,1)S(3)}, bf16[4096,4096]{1,0:T(8,128)(2,1)}, u32[]{:S(2)}) copy-start(%param.1), cross_program_prefetch_index=0
  %param = bf16[1024,4096]{1,0:T(8,128)(2,1)} parameter(0), sharding={devices=[2,2]<=[4]}, metadata={op_name="x"}
  %copy-done = bf16[4096,4096]{1,0:T(8,128)(2,1)S(3)} copy-done(%copy-start)
  %fusion.1 = bf16[1024,4096]{1,0:T(8,128)(2,1)S(3)} fusion(%param, %copy-done), kind=kOutput, calls=%fused_computation, metadata={op_name="jit(dot_custom)/dot_general" source_file="/tmp/ipykernel_874047/871715024.py" source_line=3}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":["512","4"],"output_window_bounds":["32","4"],"input_window_bounds":["32","32"],"estimated_cycles":"144544","iteration_bounds":["8","4","1"]},"megacore_config":{"megacore_split_dim":"1"},"scoped_memory_configs":[],"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"5591040"}],"retry_config":{"retry_count":"0"},"convolution_algorithm_config":{"emitter":"EmitAllBatchInSublanes"}}
  ROOT %fusion = bf16[1024,2048]{1,0:T(8,128)(2,1)} fusion(%fusion.1), kind=kCustom, calls=%all-reduce-scatter, metadata={op_name="jit(dot_custom)/dot_general" source_file="/tmp/ipykernel_874047/871715024.py" source_line=3}, backend_config={"flag_configs":[],"scoped_memory_configs":[{"memory_space":"0","offset":"0","size":"67108864"}],"collective_algorithm_config":{"emitter":"SingleInputAllReduceScatterFusion","strategy":"StrategyRing","debug":"\nStrategyRing{colors:1 phases:1 cores:{2} nophase0:0 reserved_sflags:0 cross_module_on_2d_plane:0 has_reordering_map:0 use_routing_table_indices:0}\nStrategyRing{colors:1 phases:1 cores:{2} nophase0:0 reserved_sflags:0 cross_module_on_2d_plane:0 has_reordering_map:0 use_routing_table_indices:0}\nType: 1D phase_count: 1; color_count: 1; sharded_partitions: 4; original_shape: bf16[1024,4096]{1,0:T(8,128)(2,1)S(3)}; per_color_shard_counts: 2; color_dim: -1; sharding_dim: 1; sharding_type: minor; convert_all_gather_output_to_bf16: 0; formatting steps: ()\nall-reduce-scatter fusion ND: span_size:8192*k512Byte, shard count:2, span_count:2, total_size:16384*k512Byte, valid_granules:16384*k512Byte"},"used_scoped_memory_configs":[{"memory_space":"1","offset":"0","size":"12582912"}],"retry_config":{"retry_count":"0"}}
}

I actually haven’t found a way to generate a ReduceScatter operation in the HLO—if you know how to do this, please let me know! Above, I’ve picked different shardings and arrays, and we end up triggering an operation similar in effect to a ReduceScatter; but the HLO implementation instead uses all-reduce and dynamic-slice to split up the reduced array.

Conclusion

This turned out to be on the longer side! But we have two valuable takeaways:

A linear-algebraic mental model for sharding in JAX—sharded dimensions partition data, and unsharded dimensions lead to replication in the underlying mesh! All of this can be understood in terms of block matrices and their algebra.
The ability to read HLO and parse the communications primitives that arise therein when we manipulate sharded arrays.

We also recapped information from the TPU Scaling book about how to interpret performance tradeoffs between these different forms of communication.

Acknowledgments

Thanks to the TRC program for compute support.

One way to generalize this to simultaneously considering different meshes is to consider that the mesh simply indexes different physical devices axes, or ‘merges’ of them. Hence if there are $M$ physical device axes (e.g., devices with 3D interconnects), one can count shardings across meshes by allowing the $k$ selected axes to be merged together. This is a noncommutative merging, as the traversal order for the devices will be different based on the ordering of the merged axes. ↩
To go beyond this restriction, it suffices to think of the mapping between block indices and devices as being performed by a pre-configured mesh. Algebraically, the mesh axis corresponds to a sort of ‘block permutation matrix’ (i.e., a tensor product of a permutation matrix with the identity), which left-multiplies for sharding the first dimension of the matrix, and right-multiplies for the second. ↩
Note that, algebraically, if we perform copying to represent a partitioned matrix in terms of devices, then elementwise matrix multiplication corresponds to the within-device matrix products!) ↩
E.g., about 42 GB/s for v4 TPUs. ↩
Since AllGather is linear, its derivative doesn’t involve any cached activations from the forward pass, so we omit these from the notation. ↩
But the translator is open source, and the admissible operations can be gleaned from the source code (use Google’s LLM for help with this!). A high-level specification of operations can be found here. ↩
Some high-level takeaways about these operations: think of bitcast like a view of the array, at least in this context, whereas copy doesn’t change the underlying tensor, but does change the memory layout. (Notice that the copy operation involves a transposition of the layout axes.) ↩