This assignment serves as a companion for the JAX-LM blog. For the ambitious readers who want to take a stab at implementing our language model in JAX end-to-end, this assignment provides step-by-step instructions and a test suite to guide your implementation.

This assignment is built on top of Assignment 1 (basics): Building a Transformer LM of Stanford’s CS336: Language Modeling From Scratch course.

The main differences are:

1. We convert the test cases from PyTorch to JAX
2. We add tests and instructions for distributed training in JAX
3. Some sections are organized a little differently from the original assignment

However, the modules can mostly be completed in the same order.

Getting Started

The blank starter code is provided on the assignment branch. To clone only the starter code:

git clone --single-branch --branch assignment https://github.com/chuyishang/jax-lm.git

Repository Structure The repository is a fork of stanford-cs336/assignment1-basics, and contains:

• jax_tests/: the core test suite for JAX-based implementations.
  • Suggested workflow: go through the sections below, passing all non-distributed tests (everything other than test_sharding). Then go back through to implement sharding.
  • adapters.py: This file starts as function stubs. You will edit it as you implement more components to get the relevant tests to run.
• jax_impl/: the directory where your code will live. Feel free to create files and organize your code as you see fit.
• data/tokenizer.py and data/train_bpe.py: BPE implementation and training script are provided by default because that code does not differ from the original PyTorch assignment.
• uv.lock and pyproject.toml which provide dependencies including JAX

By default, uv run and uv sync will assume a Linux + CUDA setup and be able to run out of the box. If you are using a CPU-only setup, or on a macOS device, you should first run:

uv sync --no-group cuda

Running Tests Tests can be run using uv run pytest jax_tests/test_file_name.py or uv run pytest -k test_file_name. As a base, you should be able to run

uv run pytest jax_tests/test_train_bpe.py jax_tests/test_tokenizer.py

and pass these tests right away since we’ve provided byte-pair tokenization implementations in data/train_bpe.py and data/tokenizer.py.

1. Basic Building Blocks: NNX Modules

If you’re following along on the CS336 spec, this section is the JAX-equivalent to Section 3: Transformer Language Model Architecture.

In this section, we will build our first nnx.Module classes that will become the building blocks for our Transformer LM. Specifically, we will implement

• Linear
• Embedding
• RMSNorm
• SwiGLU
• RoPE: RoPE does not use random weight initialization

In your implementation, these modules can be defined anywhere under jax_impl/ as long as jax_tests/adapters.py is updated to call them.

Random Number Generation

NNX modules that require random weight initialization should be initialized with an additional rngs parameter. This is because in NNX, random number generation uses explicit states passed in as nnx.Rngs objects.

Thus, your class signature should look something like:

from flax import nnx
import jax.numpy as jnp

class Linear(nnx.Module):
	def __init__(self, rngs: nnx.Rngs, in_features: int, out_features: int, dtype: jnp.dtype=jnp.float32):
		pass
	def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
		pass

Another note is that the forward pass of an nnx.Module is implemented directly within the __call__ method, instead of a separate forward() method like in PyTorch.

Device Management

Recall that we also don’t need to pass in devices explicitly. Later on, we will modify this section to support sharding annotations, but we don’t need to worry about that now.

1.1 Linear

Let’s get started with the first module: a linear layer. Again, the class signature for your module should look something like:

class Linear(nnx.Module):
	def __init__(self, rngs: nnx.Rngs, in_features: int, out_features: int, dtype: jnp.dtype=jnp.float32):
		pass
	def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
		pass

The Linear module performs the following transformation: \(y = xW\) We first need need to initialize our weight matrix $W \in \mathbb R^{d_{\text{in}} \times d_{\text{out}}}$. Following the original assignment, we will initialize our weights to a truncated normal distribution $W \sim \mathcal N(\mu=0, \sigma^2 = \frac{2}{d_\text{in} + d_\text{out}})$ truncated at $[-3\sigma, 3\sigma]$. This is made easy thanks to nnx.initializers.truncated_normal. You can also play around with more sophisticated initializations offered by the library, although it may not pass the autograder tests.

Once the weight data is initialized, it should be wrapped in an nnx.Param object and stored as a class attribute.

When the function is called, we just apply the linear transformation.

TODOs

1. Implement the Linear module using only JAX arrays and methods.
   • Note: don’t forget to call the superclass constructor.
2. Update run_linear in adapters.py to call your method.
   • Note: to pass the weights to your linear layer, call nnx.update with the parameters formatted in a State dictionary. We do not include a bias term in the weights, and your implementation shouldn’t either to work with the tests and adapter as is.
3. Run uv run pytest -k test_linear to check against the test case.

If you’re stuck, we provide an in-depth walkthrough for implementing a Linear layer in Implementing Our Model. We only do this for select sections. To just see a completed implementation of all sections, refer to the reference implementation!

1.2 Embedding

Next up is the Embedding module, which takes as input a vector of integer token IDs and converts each of them to their vector representations. Its interface is very similar to Linear:

class Embedding(nnx.Module):
	def __init__(self, rngs: nnx.Rngs, n_embeddings: int, embedding_dim: int, dtype: jnp.dtype=jnp.float32):
		pass
	def __call__(self, token_ids: jnp.ndarray) -> jnp.ndarray:
		pass

To make things a bit clearer:

• n_embeddings = vocabulary size
• We maintain a weight matrix $W \in \mathbb R^{\text{n_embeddings} \times \text{embedding_dim}}$ where the $i$-th row gives the embedding for token ID $i$.
  • This can be initialized to $W \sim \mathcal N(0, 1)$ truncated at $[-3, 3]$
• The output will itself be a matrix of size $(\text{token_ids.size}, \text{embedding_dim})$

TODOs

1. Implement the Embedding module using only JAX arrays and methods.
2. Update run_embedding in adapters.py to call your method.
3. Run uv run pytest -k test_embedding to check against the test case.

1.3 RMSNorm

RMSNorm is a basic layer normalization module. For convenience, the formula given a $d$-dimensional vector of activations as input:

class RMSNorm(nnx.Module):
	def __init__(self, rngs: nnx.Rngs, d_model: int, eps: float = 1e-5, dtype: jnp.dtype=jnp.float32):
		pass
	def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
		pass

and the formula

where

Here, our gain vector $g \in \mathbb R^{1 \times \text{d_model}}$ consists of our learnable parameters and $\epsilon$ is a hyperparameter.

TODOs

1. Implement the RMSNorm module using only JAX arrays and methods.
2. Update run_rmsnorm in adapters.py to call your method.
3. Run uv run pytest -k test_rmsnorm to check against the test case.

1.4 Feed-Forward Network

Now we get into implementing the actual feed-forward network that will be used as part our Transformer LM. The assignment calls for a SwiGLU (Swish-Gated Linear Unit) architecture, which consists of a Gated Linear Unit with a Swish (or SiLU) activation. Concretely:

The dimensions here are:

• $x \in \mathbb R^{d_\text{model}}$
• $W_1, W_3 \in \mathbb{R}^{d_{\text{ff}} \times d_{\text{model}}}$
• $W_2 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{ff}}}$ $d_{\text{ff}}$ is the hidden size of the feed-forward block, while $d_\text{model}$ is the hidden size of the token representation. $d_{\text{ff}}$ should be set to approximately $\frac{8}{3} \cdot d_{\text{model}}$.

The SiLU activation function is \(\text{SiLU}(x) = x \cdot \sigma(x) = \frac{x}{1 + e^{-x}}\)

TODOs

1. Implement the SiLU activation function.
   • Update adapters.run_silu and call uv run pytest -k test_silu to test your implementation.
2. Implement the SwiGLU module using only JAX arrays and methods.
   • Note: use your own Linear class!
3. Update run_swiglu in adapters.py to call your method.
4. Run uv run pytest -k test_swiglu to check against the test case.

1.5 Relative Positional Embeddings

The last module for this section will implement Rotary Positional Embeddings (RoPE) to inject information about each token’s relative position in the sequence. Here is the class interface:

class RotaryPositionalEmbedding(nnx.Module):
	def __init__(self, d_k: int, theta: float, max_seq_len: int):
		pass
	def __call__(self, x: jnp.ndarray, token_positions: jnp.ndarray) -> jnp.ndarray:
		pass

Notably, this is the first module we’ve implemented that doesn’t include rngs in the class signature. This is because this layer has no learnable parameters and thus does not need any randomized initialization.

Because of this, it’s also possible to optimize RoPE by caching the sine and cosine values using nnx.Cache at initialization.

TODO: put the definition of RoPE

TODOs

1. Implement the RotaryPositionalEmbedding module using only JAX arrays and methods.
2. Update run_rope in adapters.py to call your method.
3. Run uv run pytest -k test_rope to check against the test case.

2. Attention & Full Transformer Block

The modules we have so far together form almost the entire Transformer, with the exception of the actual self-attention part. That’s what we’ll implement now.

In this section, we will implement

• Scaled Dot-Product Attention
• Causal Multi-Head Self-Attention
• Transformer Block
• Transformer Language Model

  2.1 Scaled Dot-Product Attention

  The SDPA operation takes as input the query, key, and value matrices to perform the attention operation, defined as}”” \(\text{Attention}(Q, K, V) = \text{softmax}\bigg(\frac{Q K^\top}{\sqrt{ d_{k} }} \bigg) \cdot V\) where $Q \in \mathbb R ^{n \times d_{k}}$, $K \in \mathbb R^{m \times d_{k}}$, and $V \in \mathbb R^{m \times d_{v}}$. Here, $m$ is the sequence length and $n$ is the batch size.

For convenience, the numerically-stable softmax activation function is: \(\text{softmax}(x_{i}) = \frac{e^{x_{i} - \max(x)}}{\sum_{j} e^{x_{j} - \max(x)}}\)

TODOs

1. Implement the softmax activation function.
   • Update adapters.run_softmax and call uv run pytest -k test_softmax_matches_pytorch to test your implementation.
2. Implement the scaled_dot_product_attention using only JAX arrays and methods.
   • Note: include an optional mask parameter which accepts an array of shape Float[Array, " ... queries keys"]
3. Update run_scaled_dot_product_attention in adapters.py to call your method.
4. Run uv run pytest -k test_4d_scaled_dot_product_attention to check against the test case.

Hint: if you’ve done the original PyTorch assignment, these methods will look nearly identical.

2.2 Causal Multi-Head Self-Attention

We provide an in-depth walkthrough for implementing Causal Multi-Head Self-Attention in Implementing Our Model.

class MultiHeadSelfAttention(nnx.Module):
	def __init__(
		self,
		rngs: nnx.Rngs,
		d_model: int,
		num_heads: int,
		rope_theta: float = 1e4,
		max_seq_len: int = 1024,
	):
		pass
	
	def __call__(self, x: Array, use_rope: bool = False):
		pass

TODOs

1. Implement the MultiHeadSelfAttention module using only JAX arrays and methods.
2. Update run_multihead_self_attention in adapters.py to call your method.
3. Run uv run pytest -k test_multihead_self_attention to check against the test cases.

2.3 Transformer Block

The last piece of the puzzle is the Transformer block. Conveniently, we’ve already implemented all the submodules we’ll need. Refer to Figure 2 of the below image for the architecture, taken from the original assignment:

Here’s the TransformerBlock interface:

class TransformerBlock(nnx.Module):
	def __init__(
		self,
		rngs: nnx.Rngs,
		d_model: int,
		num_heads: int,
		d_ff: int,
		max_seq_len: int,
		theta: float,
	):
		pass
		
	def __call__(self, x: jnp.ndarray):
		pass

TODOs

1. Implement the TransformerBlock module using only JAX arrays and methods.
2. Update run_transformer_block in adapters.py to call your method.
3. Run uv run pytest -k test_transformer_block to check against the test case.

2.4 Transformer LM

Now its finally time to put the pieces together to create the final Transformer language model. The structure is pasted in Figure 1 in the previous subsection. Similar to the Transformer block, we’ve already implemented all the layers we’ll be needing.

class TransformerLM(nnx.Module):
	def __init__(
		self,
		rngs: nnx.Rngs,
		d_model: int,
		num_heads: int,
		d_ff: int,
		theta: float,
		vocab_size: int,
		context_length: int,
		num_layers: int,
	):
		pass

	def __call__(self, x: jnp.ndarray):
		pass

TODOs

1. Implement the TransformerLM module using only JAX arrays and methods.
2. Update run_transformer_lm in adapters.py to call your method.
3. Run uv run pytest -k test_transformer_lm to check against the test cases.

And there you have it: your own Transformer language model from scratch, completely in JAX. We’ll upgrade this code in Section 4 to add sharding functionality, but this should’ve gotten you comfortable coding in JAX and NNX.

3. Training Loop

Now that we have all the pieces of the model, it’s time to build out the training infrastructure. This will include:

• Cross entropy loss function
• Optax gradient transformations:
  • AdamW
  • Gradient clipping
  • LR schedule
• Train loop with checkpointing

3.1 Cross Entropy Loss

Cross entropy loss is defined as \(\ell_{i}(o, x) = -\log \text{softmax}(o_{i})_{x_{i+1}}\) where the Transformer outputs logits $o \in \mathbb R^{m \times \text{vocab size}}$ for each sequence $x$ of sequence length $m$. Then, the total cross entropy loss for a batch of size $B$ would be \(\mathcal{L}(o, x) = -\frac{1}{B}\sum_{i=1}^{B} \log \big(\mathrm{softmax}(o_i)\big)_{x_i}\) For numerical stability, you should implement the log-sum-exp trick (a helpful blog post about it here).

Note: the test cases can be used as a helpful reference for understanding the expected dimensions of the inputs.

TODOs

1. Implement the cross entropy loss function.
2. Update run_cross_entropy in adapters.py to call your method.
3. Run uv run pytest -k test_cross_entropy to check against the test case.

Hint: if you’ve done the original PyTorch assignment, these methods will look nearly identical.

3.2 Optimizer

In JAX, the optimizer is treated as a series of Optax transformations. We can first construct our pipeline by adding gradient clipping followed by AdamW. Note that we pass in our learning rate schedule lr_schedule as an argument to AdamW during initialization. Then, we can create an nnx.Optimizer using our model, the pipeline we just created, and the wrt argument that specifies what to optimize.

This pipeline could look something like

def build_optimizer_transform(
    optimizer_config: dict,
    gradient_clip: float,
    lr_schedule: optax.Schedule,
) -> optax.GradientTransformation:
    ...

TODOs

1. Implement the get_lr_cosine_schedule method for cosine annealing LR scheduling using only JAX arrays and methods.
2. Update run_get_lr_cosine_schedule in adapters.py to call your method.
3. Run uv run pytest -k test_get_lr_cosine_schedule to check against the test case.

3.4 Training Loop

An in-depth walkthrough for is provided in Implementing the Training Loop.

4. Sharding

Now we’ll extend the model and data pipeline to support distributed execution with JAX sharding. The test suite here checks physical sharding (not just logical intent) of:

• input batches,
• model parameters, and
• optimizer states.

We recommend implementing this section after Sections 1-3 are passing.

For the specific sharding specs we want to do for each mode, you can refer to the Implementing Sharding section of the original blog post.

4.1 Running Sharding Tests

The sharding suite is in jax_tests/test_sharding.py. It assumes we have at least 8 devices. If we don’t have that, it simulates 8 devices using CPU with --xla_force_host_platform_device_count=8.

4.2 Mesh + Mode Validation

Sharding is organized around a 2D mesh with axis names data and tensor.

The tests cover four sharding modes: dp, fsdp, tp, and fsdp_tp

The tests also check for valid device meshes for each mode. Invalid mode/mesh combinations should raise ValueError, and unknown mode names should also raise ValueError.

TODOs

1. Implement mesh construction (for example via create_mesh(mesh_shape, mesh_axis_names)), which is called by adapters.get_mesh.
2. Implement get_sharding_config_for_mode(mode) and make it raise ValueError for unsupported modes.
3. Implement validate_mesh_for_mode(mesh, mode) and enforce valid mode/mesh pairings.
4. Run:
   • uv run pytest jax_tests/test_sharding.py -k test_invalid_mode_name_raises -v
   • uv run pytest jax_tests/test_sharding.py -k test_invalid_mode_mesh_combinations_raise -v

4.3 Sharded Batch Pipeline

The sharding tests construct a tiny dataset wrapper with a .data field and call your sharded batch path through adapters. Your implementation should return (inputs, targets) with shape (batch_size, context_length), and both arrays should be sharded according to the mode-specific batch PartitionSpec.

In particular, the tests call:

• adapters.get_expected_batch_sharding_spec(mode)
• adapters.get_sharded_batch(...)

adapters.get_expected_batch_sharding_spec(mode) already forwards to your mode config via:

• model.get_batch_sharding_for_mode(mode)

adapters.get_sharded_batch(...) is still a TODO adapter path and should route to your sharded training/data batch utility.

TODOs

1. Implement the batch sharding policy for each mode (e.g. via get_batch_sharding_for_mode(mode)).
2. Implement sharded batch loading/sampling and ensure returned arrays are placed with the requested PartitionSpec.
3. Implement adapters.get_sharded_batch.
4. Run uv run pytest jax_tests/test_sharding.py -k test_batch_sharding -v.

4.4 Sharded Model + Optimizer Initialization

Next, shard model parameters and optimizer state. The tests call your model/optimizer initialization path in both:

• sharded mode (mesh + sharding config), and
• no-sharding mode (mesh=None, sharding_config=None).

No-sharding mode should still initialize correctly and keep arrays local (single addressable shard per array). The sharding tests also check that the optimizer state arrays are physically sharded.

TODOs

1. Implement create_model_and_optimizer(rngs, model_config, optimizer_config, sharding_config, mesh) so it supports both sharded and unsharded paths.
2. Make sure optimizer states are initialized with sharding consistent with their corresponding parameters.
3. Implement adapters.get_sharded_model_and_optimizer.
4. Run:
   • uv run pytest jax_tests/test_sharding.py -k test_model_sharding -v
   • uv run pytest jax_tests/test_sharding.py -k test_optimizer_sharding -v
   • uv run pytest jax_tests/test_sharding.py -k test_no_sharding_still_supported -v

Once this section passes, your implementation supports both local and distributed initialization paths and validates that data/model/optimizer arrays are actually distributed across devices.

Congratulations on finishing this assignment!

Introduction

In this blog, we’ll implement a language model from scratch in JAX, then scale it up with distributed training across multiple GPUs/TPUs.

We draw heavily upon 2 great resources. The first is Assignment 1 from Stanford’s CS336 course, which goes over how to implement language models from scratch in PyTorch. The second resource is the How to Scale Your Model textbook, which explains the theory behind parallelism methods and distributed training.

The goal of this blog is to connect these two perspectives by:

1. Implementing a language model from scratch in JAX instead of PyTorch
2. Extending this implementation with distributed training in both JAX and PyTorch
3. Providing a simple, customizable codebase for empirically validating theoretical compute rooflines.

We also provide an assignment with test cases so you can implement all of this in JAX yourself!

The associated code can be found at https://github.com/chuyishang/jax-lm/. If you want to try coding it yourself, you can use the assignment branch of the repo as a starter. The code follows the same flow as CS336 Assignment 1, but with custom test cases for JAX and distributed training.

This blog is split into 4 parts. First, we will provide a brief introduction to JAX and NNX. Next, we will implement the basic components of a language model in JAX, while providing comparisons to PyTorch. Then, we will introduce various distributed training methods (DP, TP, FSDP, FSDP+TP) and show how to implement them in JAX. Finally, we will use our code to empirically validate theoretical compute rooflines (coming soon!).

With that being said, let’s dive in!

Preliminaries

How to Use this Resource

This post is designed to be read front-to-back, but feel free to jump to whatever is most relevant to you. If you’re already comfortable with JAX and NNX, you can go straight to Implementing Our Model. If you’re familiar with distributed training concepts (DP, FSDP, TP) and just want to see the JAX implementation, skip ahead to Implementing Distributed Training and the code implementation.

In addition to this blog post, we also release:

1. An assignment for you to implement a JAX-based language model from scratch. This is based on Assignment 1 of Stanford’s CS336 course, but we add tests in JAX and tests for distributed training (DP, FSDP, TP, FSDP+TP). The assignment can be found here, and the starter code can be found here.
2. We provide code implementations of a simple, non-distributed language model in both JAX and PyTorch. You can compare these two implementations side-by-side to get a feel of the differences. These can be found here:
   • JAX Implementation
   • PyTorch Implementation
3. We extend both the JAX and PyTorch implementations to support distributed training with DP, FSDP, TP, and combined FSDP+TP. These can be found here:
   • JAX Distributed Implementation
   • PyTorch Distributed Implementation

This blog post is a living document and may contain errors. If you spot anything, please leave a comment below!

Recommended Background

As mentioned, this blog post borrows a lot of code from the CS336 Assignment 1, an excellent resource for learning how to build a language model from scratch with PyTorch. If you feel shaky on your PyTorch or Language Modeling fundamentals, we recommend you start there.

This blog post also references the How to Scale Your Model book, which describes distributed training and sharding with JAX in much more detail. We recommend reading that book in parallel with this blog post to gain a deeper understanding of distributed training.

With that being said, let’s dive in!

Intro to JAX and NNX

So what is JAX, and why would we want to use it at all? In this section, we’ll try to answer these questions by providing a brief overview of JAX and its neural network API, NNX. We will also touch upon some key concepts we’ll need for implementing our language model from scratch.

Intro to JAX

JAX is an open-source Python library for high-performance numerical computing developed by Google. It combines a NumPy-like syntax with the ability to run on accelerators (GPUs/TPUs), and also provides an automatic differentiation system for efficiently computing gradients. While JAX was initially developed to be compatible with Google’s Tensor Processing Units (TPUs), it is also compatible with GPUs and CPUs and can provide significant speedups on these devices.

So what makes JAX unique?

1. Firstly, JAX uses a functional programming model. Unlike PyTorch’s stateful, object-oriented approach, JAX adopts pure functions and immutable arrays. As an example, there is no model object that holds parameters. Instead, in core JAX, the model is represented as a nested collection of arrays (a PyTree) and passed explicitly into functions. This makes the execution model more transparent and composable.

2. Secondly, JAX provides a set of very powerful composable function transformations. The core idea is that we can write a plain Python or NumPy function and transform it into a more powerful function. Some common transforms are as follows:
   1. jax.grad, which differentiates a function
   2. jax.jit, which compiles a function via XLA for fast execution
   3. jax.vmap, which vectorizes a function over some batch dimension automatically
   4. jax.pmap, which is similar to jax.vmap but parallelizes a function across devices These functions can be composed arbitrarily (such as jit(vmap(grad(f)))) and can often provide huge speedups in practice.
3. Lastly, JAX uses the XLA Compiler as its backend, which optimizes across the entire computation graph instead of operator by operator as PyTorch does. This allows XLA to perform aggressive fusion and optimizations that span multiple operations and even devices, which can unlock large performance gains.

Putting this all together, the use of pure, composable functions allows the JAX compiler to trace entire functional units into an intermediate representation, called the jaxpr, instead of executing each line in the Python interpreter step-by-step. That’s where the speedup comes from.

Let’s now take a look at some of the code behind these function transformations!

jax.grad

At the heart of any deep learning library is automatic differentiation (autodiff), and JAX’s approach is built around jax.grad. jax.grad is a transform that takes in a function f and returns a new function that computes its gradient.

This is a slightly different mental model from PyTorch. In PyTorch, the computation graph is built incrementally as we perform each operation in the forward pass. To get the backwards pass computation, PyTorch walks through the graph in reverse.

On the other hand, differentiation in JAX is expressed directly as a function transformation. If f computes a scalar output, then jax.grad(f) is a new function with the same inputs that returns the gradient of f with respect to whichever arguments we specify.

Code

To see this in code form, let’s look at the following example of a simple loss function:

import jax
import jax.numpy as jnp

def loss(params, x, y):
    pred = params @ x
    return jnp.mean((pred - y) ** 2)

grad_fn = jax.grad(loss, argnums=(0,1))
grads = grad_fn(params, x, y)
dparams, dx = grads  # for clarity

We can define our loss function normally, using the NumPy-like operation jnp.mean to take the mean of our vector. Then, we can call jax.grad on our loss function to transform it into a new function grad_fn that computes the gradients of its inputs.

One thing to note is that the new function grad_fn will be a wrapper around the original loss function. When we call grad_fn , JAX will run the forward pass (loss) and backwards pass (getting the gradients) in the same call. As a result, grad_fn must have the same function signature as the original loss function.

To specify which arguments to differentiate w.r.t., we can pass in argnums. Here, argnums=(0,1) corresponds to params and x, so grad_fn returns a tuple of 2 gradients, dparams and dx.

If we look carefully, the last two lines can actually be combined into the following form:We drop the `argnums` here for simplicity; it defaults to 0 which corresponds to `params`.

grads = jax.grad(loss)(params, x, y)

where jax.grad acts as a higher order function. If you are familiar with Python decorators, we see that this is actually the exact pattern of a decorator! So an equivalent way of writing the entire code would just be:

import jax
import jax.numpy as jnp

@jax.grad
def loss(params, x, y):
    pred = params @ x
    return jnp.mean((pred - y) ** 2)
    
grads = loss(params, x, y)

Since all JAX transforms are higher-order functions of the same form, they can all be used as decorators in this way.

But while this is a nice illustration of the JAX programming model, for gradients specifically this is usually not the most convenient form. This is because after decoration with @jax.grad, our loss function would only return the gradients instead of the actual loss!

Since we need both the loss in addition to the gradients during training, we typically use the jax.value_and_grad transform instead. The usage is identical, except the returned function now returns a (loss, grads) tuple instead of just grads.

loss_and_grad_fn = jax.value_and_grad(loss)
loss_value, grads = loss_and_grad_fn(params, x, y)

jax.jit

jax.jit is a transform that compiles a function using the XLA compiler.The JIT in `jax.jit` stands for Just-In-Time compilation, which refers to the fact that functions are compiled at runtime the first time they are called. This is in contrast to Ahead-Of-Time compilation, where everything is compiled before the program runs. JAX compiles the function the first time it is called and reuses the compiled version for future calls.

JIT compilation is powerful because XLA sees the whole function and optimizes it its entirety. This allows it to fuse operations, optimize memory layout, and parallelize in ways that line-by-line execution cannot.

So what actually happens when we compile a function? On the first call, JAX traces the function by running it with abstract placeholders. You can think of these as something similar to types: instead of passing in x = jnp.array([1.0, 2.0, 3.0]), JAX traces with something like x = float32[3]. The idea is that XLA doesn’t need to know exact values of arrays to build the computation graph for optimization, it only needs to know the shapes and types of data.

After tracing the function to construct a computation graph, JAX compiles it to optimized XLA bytecode. On subsequent calls, we can just run the compiled binaries instead of the Python code.

In practice, this can speed up code by a lot. The speedup usually gets larger the more operations we fuse in one jax.jit call, which is why many implementations wrap the entire training step (forward + backward + optimizer update) into a single jax.jit block, which we show below.

Code

Here is what jax.jit might look like in code form:

@jax.jit
def f(x):
    return jnp.sum(x ** 2)

x = jnp.ones((3,))
output = f(x)

A more realistic training step might look something like this:

@jax.jit
def train_step(params, x, y):
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    params = params - lr * grads
    return params, loss

for batch_x, batch_y in data:
    params, loss = train_step(params, batch_x, batch_y)

Note that we do the entire forwards, backwards, and parameter updates in a single function compiled by jax.jit!

Common Pitfalls

Note that JAX recompiles functions if the shapes of the input change. Take a look at the following example:

@jax.jit
def f(x):
    return jnp.sum(x ** 2)

f(jnp.ones((3,)))  # compiles for float32[3]
f(jnp.zeros((3,)))  # uses cached compiled function since same shape
f(jnp.ones((5,)))  # !! recompiles for float32[5] !!
f(jnp.ones((7,)))  # !! recompiles for float32[7] !!
f(jnp.zeros((3,)))  # uses cached compiled function from the first call

JAX compilation takes into account the input shapes to best optimize the program, so it saves a separate binary per shape. In our example, JAX compiles f for the first time with a float32[3] input and can reuse this for the second call. However, we trigger a costly recompilation each time our input shape is different!

Secondly, since JAX traces Python control flow at compile time, ordinary if statements and loops can only depend on values known at trace time. Static loops are fine, but if a branch depends on a runtime value, a Python if inside jax.jit will fail. In those cases, we should use JAX control-flow primitives such as lax.cond, lax.scan, lax.fori_loop, or lax.while_loop.

For example, the following looks natural in Python, but fails under jax.jit because the branch depends on the runtime value of x:

import jax

@jax.jit
def bad_f(x):
    count = 0
    for _ in range(3):
        count += 1  # static loop: OK

    if x > 0:      # not OK: value-dependent Python branch under jit
        return x + count
    else:
        return -x + count

# bad_f(1.0)  # raises TracerBoolConversionError

To express the same logic in a JIT-compatible way, we can replace the Python if with lax.cond, which takes a scalar predicate and two branch functions:

import jax
from jax import lax

@jax.jit
def f(x):
    count = 0
    for _ in range(3):
        count += 1  # static loop: OK

    return lax.cond(
        x > 0,
        lambda x: x + count,
        lambda x: -x + count,
        x,
    )

Finally, any Python-specific side effects (such as print statements) only happen at trace time and not during executing the compiled program. As a result, a print statement inside a JIT-compiled function may appear only when JAX first traces the function and on retraces. Here’s what this might look like:

@jax.jit
def f(x):
    print(x)

>>> f(1)
JitTracer<~int32[]>
>>> f(2) # same shape, no print since we use cached compiled function

>>> f(jnp.array([3, 4])) # different shape, retrace
JitTracer<int32[2]>

Note that these prints happen only at trace time, so JAX does not yet know the concrete value, only the abstract tracer type!

jax.vmap

The final core JAX transforms are for vectorization and parallelization. jax.vmap stands for “vectorized map” and transforms a function written for a single example into one that operates over a batch.

import jax
import jax.numpy as jnp

def f(x):
    return jnp.dot(x, x)

batched_f = jax.vmap(f)
batched_f(jnp.ones((8, 4)))

Equivalently, we can also have this in decorator form:

@jax.vmap
def f(x):
    return jnp.dot(x, x)

f(jnp.ones((8, 4)))

An illustration of the above jax.vmapcode.

The idea here is that we can define a dot product function f that operates on a single vector, but then use jax.vmap to make it batched. By default, jax.vmap maps over the first dimension, but we can also specify which dimensions we want to batch over by passing in the in_axes argument.

jax.pmap

jax.pmap stands for parallel map. It’s similar to vmap but operates over devices (different GPUs/TPUs) by running each slice of the batch on a separate accelerator. We describe pmap for background and historical completeness, but it’s worth noting that newer approaches often prefer mesh-based sharding instead, which we’ll discuss later in the parallelism section.

n_devices = jax.device_count()

def loss_fn(params, x):
    return jnp.sum(params * x)

parallel_loss = jax.pmap(loss_fn, in_axes=(None, 0))

xs = jnp.ones((n_devices, 4))  # one slice per device. note that the batch dimension must equal `n_devices`
parallel_loss(params, xs)       # runs shard i on device i

In this example, we pass in in_axes=(None, 0). This means that we want the first argument params to replicated across all devices (None), and the second argument to be split across the device axis 0. Since we don’t specify a device mesh, our devices form a single (n_devices,) mesh, so this effectively sends row i to device i.

import torch
import torch.distributed as dist

def loss_fn(params, x):
    return torch.sum(params * x)

def main():
    dist.init_process_group(backend='nccl')
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    params = torch.ones(4).cuda(rank)
    xs = torch.ones((world_size, 4)).cuda(rank)  # full batch, same shape as JAX

    x_local = xs[rank]          # manually shard: each process takes its slice
    loss = loss_fn(params, x_local)
    print(f"rank {rank}, loss: {loss}")

if __name__ == "__main__":
    main()

NNX

Now that we’ve talked about some of the core concepts behind JAX, let’s talk a bit about NNX, which is what we will be using for the rest of this blog.

NNX is the neural network API for JAX. If you are familiar with PyTorch, you can think of NNX as playing a role similar to torch.nn. NNX was introduced in 2024 as an evolution of Flax Linen, with the goal of keeping the parts of Linen that worked well while making the programming model more Pythonic and easier to debug.

In particular, NNX relaxes the purely functional approach of core JAX by introducing a module-based, object-oriented design similar to PyTorch. Instead of passing model states explicitly in and out of functions, NNX now allows us to define stateful module objects. NNX then provides object-aware versions of common JAX transforms so we can still leverage JAX’s power within this more familiar programming model. For example, @jax.jit becomes @nnx.jit, and jax.value_and_grad becomes nnx.value_and_grad.

To learn a bit more about why you should use NNX, you can refer to the official explanation here.

Implementing Our Model

In this section, we will implement a basic version of a Transformer language model in JAX and provide a version of the same code written in PyTorch for comparison. We will begin with a simple version without worrying too much about sharding and distributed training, then layer in that functionality afterwards.

Let’s start by implementing the simplest component of any modern day neural network: the Linear layer.

The Linear Layer

The following is how you might implement a linear layer (without biases) in PyTorch:

class Linear(nn.Module):
    def __init__(
        self,
        in_features: int,
        out_features: int,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None
    ):
        super().__init__()
        weights = torch.empty(out_features, in_features, device=device, dtype=dtype)
        std = (2 / (in_features + out_features)) ** 0.5
        nn.init.trunc_normal_(weights, mean=0.0, std=std, a=-3.0 * std, b=3.0 * std)
        self.weights = nn.Parameter(weights)

    def forward(self, x: Tensor) -> Tensor:
        return einsum(x, self.weights, "... d_in, d_out d_in -> ... d_out")

And here is how you would do it in NNX:

class Linear(nnx.Module):
    def __init__( self,
        rngs: nnx.Rngs,
        in_features: int,
        out_features: int,
        dtype: jnp.dtype = jnp.float32,
    ):
        super().__init__()
        std = (2 / (in_features + out_features)) ** 0.5
        init_fn = nnx.initializers.truncated_normal(stddev=std, lower=-3.0 * std, upper=3.0 * std)
        weights_data = init_fn(rngs.params(), (in_features, out_features), dtype)
        self.weights = nnx.Param(weights_data)

    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        return jnp.einsum("...i,io->...o", x, self.weights.get_value())

As you can see, the code looks very similar. Aside from minor differences in syntax and storage convention, there are a few key changes that illustrate some of the philosophical differences between JAX and PyTorch.

1. Initialization: In PyTorch, initialization functions usually operate in-place on an already-allocated tensor. This is why we create an empty weights tensor in the desired shape and on the desired device before initializing it. On the other hand, initialization functions in JAX typically return a tensor (weights_data in this case), which we then wrap with nnx.Param() to make it a parameter object so that JAX knows how to use it for transforms and state management. The actual weight tensor of an nnx.Param actually lives in its .value attribute.

   This example highlights a core philosophical difference between PyTorch and JAX. PyTorch is designed in an object-oriented style centered around stateful modules, where parameters live inside modules and workflows mutate tensors in place. JAX treats parameters as explicit data and often represents common workflows as functions that produce new values.

2. Device Placement: In PyTorch, device placement is explicit. Typically, we create parameters directly on a target device with device=... or move them afterwards with .to(device). Meanwhile, in JAX, device placement is implicit and determined automatically by the compiler and any sharding configs we specify.

3. RNG: In JAX, RNGs need to be passed in explicitly, and are split/consumed deterministically. On the other hand, RNG is largely global and ambient in PyTorch (i.e. we can set a global RNG seed) and we don’t need to pass it explicitly to each function.

   In the example above, we use rngs.params() in the init_fn, which generates an random number from the params stream. There are 2 main PRNG key stream names used by Flax NNX’s built-in layers, which are params and dropout.

Some other minor differences:

• PyTorch uses the forward convention, which is wrapped by the __call__ method implemented by the parent class nn.Module. JAX typically uses __call__ directly.

Multi-head Self Attention

Let’s take a look at another example before moving on. For those following along on the assignment, note that the other building blocks of our model have minor syntax differences compared to PyTorch, but follow the same ideas discussed here.

This is what code for multi-head self attention (MHA) in PyTorch would look like:

class MultiHeadSelfAttention(nn.Module):
    def __init__(
        self,
        d_model: int,
        num_heads: int,
        rope_theta: float = 1e4,
        max_seq_len: int = 1024,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.Q_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.K_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.V_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.O_proj = Linear(d_model, d_model, device=device, dtype=dtype)
        self.register_buffer(
            "causal_mask",
            torch.tril(torch.ones(max_seq_len, max_seq_len, device=device)),
            persistent=False,
        )
        self.rope = RotaryPositionalEmbedding(self.d_k, theta=rope_theta, max_seq_len=max_seq_len, device=device)

    def forward(self, x: Tensor, use_rope: bool = False) -> Tensor:
        batch_size, seq_len, _ = x.shape
        q = self.Q_proj(x).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        k = self.K_proj(x).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)
        v = self.V_proj(x).reshape(batch_size, seq_len, self.num_heads, self.d_k).transpose(1, 2)

        if use_rope:
            positions = torch.arange(seq_len, device=x.device)
            q = self.rope(q, positions)
            k = self.rope(k, positions)

        attn_out = sdpa(q, k, v, self.causal_mask[:seq_len, :seq_len])
        attn_out = attn_out.transpose(1, 2).contiguous().reshape_as(x)
        return self.O_proj(attn_out)

And here is what it looks like in JAX:

class MultiHeadSelfAttention(nnx.Module):
    def __init__(
        self,
        rngs: nnx.Rngs,
        d_model: int,
        num_heads: int,
        rope_theta: float = 1e4,
        max_seq_len: int = 1024,
        dtype: jnp.dtype = jnp.float32,
    ):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
        self.num_heads = num_heads
        self.d_k = d_model // num_heads
        self.Q_proj = Linear(rngs=rngs, in_features=d_model, out_features=d_model, dtype=dtype)
        self.K_proj = Linear(rngs=rngs, in_features=d_model, out_features=d_model, dtype=dtype)
        self.V_proj = Linear(rngs=rngs, in_features=d_model, out_features=d_model, dtype=dtype)
        self.O_proj = Linear(rngs=rngs, in_features=d_model, out_features=d_model, dtype=dtype)

        self.causal_mask = jnp.tril(jnp.ones((max_seq_len, max_seq_len)))
        self.rope = RotaryPositionalEmbedding(d_k=self.d_k, theta=rope_theta, max_seq_len=max_seq_len)

    def __call__(self, x: Array, use_rope: bool = False):
        batch_size, seq_len, _ = x.shape
        q = self.Q_proj(x).reshape(batch_size, seq_len, self.num_heads, self.d_k).swapaxes(1, 2)
        k = self.K_proj(x).reshape(batch_size, seq_len, self.num_heads, self.d_k).swapaxes(1, 2)
        v = self.V_proj(x).reshape(batch_size, seq_len, self.num_heads, self.d_k).swapaxes(1, 2)

        if use_rope:
            positions = jnp.arange(seq_len)
            q = self.rope(q, positions)
            k = self.rope(k, positions)

        attn_out = sdpa(q, k, v, self.causal_mask[:seq_len, :seq_len])
        attn_out = attn_out.swapaxes(1, 2).reshape(x.shape)
        return self.O_proj(attn_out)

Aside from the syntax and RNG differences described earlier, the two implementations look basically the same!

However, one interesting difference to note here is how causal masks are handled. In PyTorch, it’s common to register the mask as a buffer at initialization:

self.register_buffer(
    "causal_mask",
    torch.tril(torch.ones(max_seq_len, max_seq_len, device=device)),
    persistent=False,
)

Unlike parameters, buffers are not learnable and not tracked by the optimizer. However, they are automatically moved with the model during device placement.

In JAX, we can just assign it directly:

self.causal_mask = jnp.tril(jnp.ones((max_seq_len, max_seq_len)))

This is because device placement is automatic in JAX, so there’s no need for a special registration mechanism to ensure the mask follows the model to the correct device. As a bonus, since jax.jit treats any closed-over array as a constant in the compiled graph, the mask is never recomputed across calls.

The other components of our language model (SwiGLU, RMSNorm, Embedding, LM Head) follow much of the same design. We recommend taking a look at the full model implementation in JAX, which you can find here.

Implementing the Training Loop

Now that we have our model implemented in JAX, let’s take a look at how we can train it!

There are a few major differences in the JAX training loop compared to PyTorch. We’ll discuss each one and what it might reveal about the JAX paradigm.

The Training Step

In PyTorch, our train step looks something like this:

logits = model(inputs)
B, S, V = logits.shape
loss = model_module.cross_entropy_loss(logits.reshape(B * S, V), targets.reshape(B * S))
optimizer.zero_grad()
loss.backward()
optimizer.step()

Meanwhile, in JAX, it looks something like:

loss, grad_state = train_step(model, optimizer, inputs, targets)

Where the train_step function looks something like:

@nnx.jit
def train_step(model: nnx.Module, optimizer: nnx.Optimizer, inputs: Array, targets: Array):
    def loss_fn(model):
        logits = model(inputs)
        B, S, V = logits.shape
        return cross_entropy_loss(logits.reshape(B * S, V), targets.reshape(B * S))

    loss, grads = nnx.value_and_grad(loss_fn)(model)

    optimizer.update(model, grads)
    return loss, grads

Train Step Function: In PyTorch, we don’t need to wrap our training step into a function. Most of our operations operate in place. However, we need to create a train_step function in JAX so that we can wrap it with nnx.jit.

Gradient Handling: In PyTorch, gradients are implicit, since loss.backward() accumulates gradients into each parameter’s .grad field, and optimizer.step() reads those .grad buffers and updates parameters in place. As a result, our gradient states are effectively distributed across {p.grad for p in model.parameters()}, instead of being unified in a single gradient object grad like in JAX. Gradient clipping in PyTorch follows the same philosophy by scaling the .grad buffers in place.

On the other hand, gradients are treated as explicit values in JAX. To get our grad object, we use nnx.value_and_grad, which is a transformation that turns the loss_fn into a function that also returns gradients. Since JAX transforms act as higher order functions, we need to define a loss function and pass it in as input to nnx.value_and_grad instead of just calling loss.backward() like we do in PyTorch.

Now that we have the basic training step, let’s add gradient clipping and a learning rate schedule.

Gradient Clipping and LR Schedule

If we wanted to add gradient clipping and a learning rate schedule in PyTorch, they would be called separately from our optimizer. In JAX, they would be processed by the optimizer internally.

PyTorch Here’s what it would look like in PyTorch:

def build_lr_scheduler(
    optimizer: torch.optim.Optimizer,
    min_learning_rate: float,
    warmup_iters: int,
    cosine_annealing_iters: int,
    warmup_start_factor: float = 0.1,
) -> torch.optim.lr_scheduler.LRScheduler:
    cosine_decay_iters = cosine_annealing_iters - warmup_iters
    warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
        optimizer,
        start_factor=warmup_start_factor,
        end_factor=1.0,
        total_iters=warmup_iters,
    )
    cosine_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=cosine_decay_iters,
        eta_min=min_learning_rate,
    )
    scheduler: torch.optim.lr_scheduler.LRScheduler = torch.optim.lr_scheduler.SequentialLR(
        optimizer,
        schedulers=[warmup_scheduler, cosine_scheduler],
        milestones=[warmup_iters],
    )
    return scheduler

and the new train_step would look something like this:

logits = model(inputs)
B, S, V = logits.shape
loss = model_module.cross_entropy_loss(logits.reshape(B * S, V), targets.reshape(B * S))
optimizer.zero_grad()
loss.backward()
if gradient_clip > 0:  # gradient clipping
    torch.nn.utils.clip_grad_norm_(model.parameters(), float(gradient_clip))
optimizer.step()
if scheduler is not None:  # scheduler handling
    scheduler.step()

In the function build_lr_scheduler, we create a scheduler object that we chains together a linear warmup schedule with a cosine annealing schedule. To use this, we call scheduler.step(), which computes new LR values and mutates the optimizer’s param_groups’ lr attributes in place.

JAX In contrast, gradient clipping and the learning rate scheduler are used as inputs during optimizer construction in JAX. Our optimizer construction would look something like:

lr_schedule = optax.warmup_cosine_decay_schedule(
    init_value=0.0,
    peak_value=a_max,
    warmup_steps=warmup_iters,
    decay_steps=decay_iters,
    end_value=a_min
)

def build_optimizer_transform(
    optimizer_config: dict,
    gradient_clip: float,
    lr_schedule: optax.Schedule,
) -> optax.GradientTransformation:
    transforms = []
    if gradient_clip > 0:
        transforms.append(optax.clip_by_global_norm(float(gradient_clip)))
    transforms.append(
        adamw(
            lr=lr_schedule,
            betas=tuple(optimizer_config["betas"]),
            weight_decay=float(optimizer_config["weight_decay"]),
            eps=float(optimizer_config["eps"]),
        )
    )
    return optax.chain(*transforms)

optimizer = nnx.Optimizer(
    model,
    build_optimizer_transform(
        optimizer_config,
        gradient_clip=gradient_clip,
        lr_schedule=lr_schedule,
    ),
    wrt=nnx.Param,
)

In JAX, the optimizer is treated as a series of Optax transformations. We can first construct our pipeline by adding gradient clipping followed by AdamW. Note that we pass in our learning rate schedule lr_schedule as an argument to AdamW during initialization. Then, we can create an nnx.Optimizer using our model, the pipeline we just created, and the wrt argument that specifies what to optimize.

Finally, we can call this without changing our code by using optimizer.update(model, grads), which will handle the gradient clipping and LR updates internally. Note that unlike PyTorch, we don’t need to call scheduler.step() or clip_grad_norm_() explicitly in each step.

You can find the full training script here.

Intro to Distributed Training

So far, we have assumed that both the model and data batch fit on a single device. However, once the model state or activations become too large, or once we want more throughput than a single device can provide, we need to spread (a.k.a shard) the work across multiple GPUs/TPUs. Doing so can give us more total memory and compute, but it also introduces a new bottleneck: communication. As a result, a core problem in distributed training is deciding which tensors stay local to each device, which tensors are sharded, and which tensors should be communicated between devices.

Different parallelism strategies answer this question in different ways. In this section, we will focus on 4 common strategies: Data Parallelism (DP), Fully Sharded Data Parallelism (FSDP), Tensor Parallelism (TP), and mixed FSDP+TP. A useful way to think about them is:

• DP splits the data batch
• FSDP splits the persistent model state
• TP splits the hidden-state computation inside a layer
• FSDP+TP does both at once

Don’t worry if this doesn’t make much sense yet! The goal of this section is to build some intuition for these strategies, and in the next section we will show you how to implement them in JAX.

Notation

We will denote variables in the form $\text{VarName}[d_1, d_2, \ldots]$ giving its name and shape, respectively. To keep the notation simple, we will approximate a Transformer block as a pair of linear layers:

One important note here is that $\text{In}[B, D]$ represents the input activations, not necessarily of the raw input batch itself. In language modeling, we typically have a tokenized input batch of shape $[B, T]$ where $T$ is the number of tokens per sample. This input batch would then be fed through our embedding table to get our activations $[B, T, D]$, which is what we refer to here as $\text{In}$. Following common practice, we abuse our notation a bit to roll the $B$ and $T$ dimensions into one $B$ that represents the total number of tokens in a batch.

For sharding, we will use subscripts like $B_X$ and $D_Y$ to denote sharding over device-mesh axes $X$ and $Y$, respectively. For example, $B_X$ means the batch dimension is split across the $X$-axis of the device mesh, while $D_Y$ means that the hidden dimension is split across the $Y$-axis. We’ll introduce the device mesh in the next section.

In our illustrations, we will mostly show weight gradients and omit input gradients when they aren’t the main point of the conversation.

Before we dive into parallelism schemes, let’s take a look at some important preliminaries: the device mesh, and communication collectives.

The Device Mesh

To talk about sharding, we need to have a way to refer to our devices. The standard abstraction is a device mesh, which you can imagine as arranging all of our devices into a matrix. For example, if we had 8 devices, we might arrange them as a $4 \times 2$ mesh, a $8 \times 1$ mesh, a $2 \times 4$ mesh, and so on. Each dimension of the mesh is called a device axis. In this section, we’ll refer to these axes abstractly as $X$ and $Y$. Later, in JAX, we will assign them names like "data" "tensor".

Communication Collectives

When we talk about parallelism, there are also 3 communication operations that will pop up repeatedly. These are the All-Gather, Reduce-Scatter, and All-Reduce.

All-Gather

During an all-gather operation, each node sends its shard of data to all other nodes, so each device can reconstruct the whole. When we have shards of a tensor on different devices, we can use an all-gather to build a copy of the full tensor on each device.

Reduce-Scatter

Each device computes the result on their own shard of data, but a “reduce” operation is needed to compute the overall result incorporating all shards. Then, the final result must be re-sharded and distributed back to the corresponding device—this is the “scatter” operation. Reduce-scatter combines these operations into a single step.

It does this using a ring structure, illustrated below. Devices are organized in a loop and each device passes a segment to its neighbor and receives one from the other neighbor, repeating for $N-1$ steps. By the end, each device is holding their distinct, fully reduced segment of the data.

All-Reduce

An all-reduce combines partial results across devices (often a sum) and leaves every device with a copy of the full result. Equivalently, you can think of it as a reduce-scatter (combine and distribute one shard per device) followed by an all-gather (collect all shards so every device has the complete result).

The beauty of JAX is that we never need to implement or code these operations ourself. This is all handled implicitly once we define the device mesh. Now, with preliminaries out of the way, let’s talk about training!

Standard Training

As a starting point, let’s take a look at the non-distributed case. We can represent this with the following:

On a single device, the full input, weights, activations, gradients, and optimizer state all live together. The forward pass computes $\text{Out}$, the backward pass computes gradients $\text{d}W_1$, $\text{d}W_2$, $\text{d}\text{In}$, and the optimizer uses the gradients to update the weights $W_1, W_2$. This is the simplest case, and it gives us a reference point for the distributed strategies below.

Data Parallel (DP)

Data parallelism is the simplest distributed strategy. We shard the input across the batch dimension and keep a full copy of the weights and optimizer state on every device.

Each device therefore runs the same forward and backward computation, but on a different slice of the batch. Because every device has the full weights locally, the forward pass requires no communication.

However, after the backward pass, each device only has a local contribution to the weight gradients from its own slice of the batch. These are displayed as Grads1 and Grads2 in the diagram below illustrating DP on a 2x1 mesh. To recover the gradient for the full batch, we need to all-reduce those gradients across devices. Each device then receives the same summed gradient Grads and applies the same update to the same weights, which ensures that our weights are always in sync across devices.

To summarize, the main ideas of DP are:

1. Cheap forward pass
2. Replicated model state
3. Synced gradients in the backward pass

Fully-Sharded Data Parallel (FSDP)

You might notice a disadvantage of DP is that each device must hold a complete copy of the weights, gradients, and optimizer state. FSDP tries to solve this by sharding the model state across devices as well:

Now, each device only permanently stores its own shard of the weights and optimizer state. The tradeoff here is that a device usually can’t perform the forward pass matmul using only its local weight shard, since our weights are sharded across the contracting dimension $D$. As a result, we need to all-gather the parameter shards so that each device can materialize the full weight matrix before the layer runs. After this forward pass computation, the temporary full copy can be discarded.

In the backward pass, each device will only have the gradients of the local slice of the input, similar to DP. However, instead of all-reducing the full gradient like DP, we can reduce-scatter the gradient so that each device keeps only the gradient shard corresponding to the parameters it owns. The key idea is that the FSDP lowers the persistent per-device memory footprint, even though each device still temporarily materializes the full layer weights when the layer is executing.

In this illustration, the crosshatched/dashed weight shards are received from the all-gather and discarded after use.

Tensor Parallel (TP)

One important observation is that FSDP mostly reduces persistent parameter memory but not necessarily the peak per-device memory during the forward pass. As a result, if our model or sequence length is so large that our activations don’t fit in a single forward pass, FSDP won’t help us much. Tensor parallelism (TP) takes a different approach by sharding the computation itself, rather than only sharding stored parameters. Instead of moving weights across devices, we shard the weights and distribute the computation of a single layer across devices. Here is how we can represent TP:

Here, we shard along the feature dimension (denoted by $Y$), so each device holds only a slice of the intermediate activations and weights. In this setup, each device computes only its local contribution to the output. Importantly, since each device can compute the gradient for its local weight shard directly, no communication is required for weight gradients.

However, communication is still required elsewhere. After the forward pass, each device holds only a partial output from its local weight shards, so we must all-reduce our output across devices to obtain the full output needed for the loss. In the backward pass, the same issue appears for the input gradients: each device computes only a partial contribution to $\text{dIn}$, so we again need an all-reduce to combine these contributions (although input gradients are not shown in the diagram for simplicity).

FSDP + TP

Can we take FSDP and TP a step further? Turns out we can actually combine them by assigning each to different device axes! Along the $X$ axis, we use FSDP to shard model state, while along the $Y$ axis, we use TP to shard the hidden-state computation:

This combination is useful because the two strategies help each other. FSDP shards the batch across the $X$-axis, which reduces the amount of activation data that TP must move. TP shards the weights across $Y$-axis, which reduces the amount of parameter data that FSDP must move. The tradeoff is that we now pay both activation communication and weight communication, so the implementation is more complex and communication costs are higher as well. For a discussion on how to balance FSDP and TP, and when to use it, we recommend reading the related section in the How to Scale Your Model book.

This image might be a bit hard to read due to the size, but you can find the full-size image here.

In this example, we expand the number of devices we use to 4, since we need to use at least 2 devices to perform FSDP and another 2 to perform TP. We also omit the arrows showing the initial splitting of weights, but each device will hold a shard of both the input and the weights, corresponding to its position in the mesh.

Before the matmul, TP all-gathers the inputs across $Y$ to get the full data dimension $D$. Then, we can use FSDP to gather our weight matrices across the $X$ axis. Then, we perform a TP-like All-Reduce for our output shards across the $Y$ axis, then finally reduce scatter our gradients across $X$ in an FSDP-like way. One thing to note here is that since we perform an all-reduce of the partial $\text{Out}[B_X, D]$ tensors across the $Y$-axis, each row will actually have the same loss.

Implementing Distributed Training

Now that we have a conceptual picture of distributed training, implementing it in JAX is mostly about expressing the same layouts explicitly. We define a device mesh, specify which tensor dimensions are sharded over which mesh axes, and let JAX/XLA insert the necessary collective communications.

In this section, we will go over:

1. Initializing our device mesh
2. Sharding annotations
3. How to use sharding annotations to initialize the model
4. How to define our sharding annotations for each parallelism mode

Let’s begin!

Initializing our Device Mesh

In the previous section, we referred to our mesh axes as $X$ and $Y$. We can now give these axes names, which are conventionally named "data" and "tensor". The "data" axis is used for DP/FSDP-style sharding, while the "tensor" axis is used for TP-style sharding.

We can define our device mesh as following:

from jax.sharding import Mesh

def create_mesh(mesh_shape: list[int], mesh_axis_names: list[str]) -> Mesh:
    auto_mesh = jax.make_mesh(tuple(mesh_shape), tuple(mesh_axis_names))
    return auto_mesh

mesh = create_mesh((4, 2), ("data", "tensor"))

Here, we use the jax.make_mesh() function to create our Mesh object. We pass in (4, 2) as our mesh shape to create a 4x2 device mesh, with the "data" being the X device axis and "tensor" being the Y device axis.

Attaching Sharding Annotations to our Model

Now that we have our device mesh, our next step is to specify how we want our data and model to be sharded on this mesh. We can use the PartitionSpec class to do this, which is a tuple abstraction that defines our sharding annotations.

For example, after importing PartitionSpec using from jax.sharding import PartitionSpec as P, as is customary, we can define:

• P("data", None) , which means that we shard the 1st array dimension across the “data” axis and leave the second dimension unsharded
• P(None, "tensor") means that we shard the 2nd array dimension across the “tensor” axis and leave the first dimension unsharded
• P(None, None) means that we do not shard either array dimension. Essentially, this means that the tensor is replicated across devices. To make this concrete, if I had a device mesh of shape $4 \times 2$ and a tensor $X[32, 8]$, then P("data", None) will result in a $X[32/4, 8] = X[8, 8]$ tensor that is replicated across the $Y$ axis. If I use P("data", "tensor") instead, each device would have a $X[32/4, 8/2] = X[8, 4]$ slice of the original tensor.

The diagram below illustrates several sharding configurations for our 4×2 mesh. Note that the notation Device A,B,C indicates that slice of the matrix is replicated across devices $A$, $B$, and $C$.

For now, you don’t need to worry about what specific PartitionSpec to use, as it depends on what parallelism strategy we choose.

Initializing our Model with Sharding

The next step is to attach these sharding annotations to the tensors that we want to shard. In NNX, the easiest way to do this is to wrap the initializer with nnx.with_partitioning(init_fn, sharding), which will attach sharding annotations to the variable created with init_fn. Typically, these sharding annotations should be passed in at model initialization so that JAX knows how to distribute tensors across devices.

Here’s what it might look like in code:

class Linear(nnx.Module):
    def __init__(self, rngs: nnx.Rngs, in_features: int, out_features: int,
    sharding: Sharding | None = None, dtype: jnp.dtype = jnp.float32,
    ):
        super().__init__()
        std = (2 / (in_features + out_features)) ** 0.5
        # We use `nnx.with_partitioning` here!
        init_fn = nnx.with_partitioning(
            nnx.initializers.truncated_normal(stddev=std, lower=-3.0 * std, upper=3.0 * std), sharding
        )
        weights_data = init_fn(rngs.params(), (in_features, out_features), dtype)
        self.weights = nnx.Param(weights_data)

    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        return jnp.einsum("...i,io->...o", x, self.weights.get_value())

The only differences from a standard Linear layer are that:

1. We wrap the original nnx.initializers.truncated_normal() call with the nnx.with_partitioning() function
2. We pass in the sharding as input The process is the same for other components of the model, all we need to do is just wrap the initializer with nnx.with_partitioning.

Now that we’ve annotated each component of our model with nnx.with_partitioning, we need to actually initialize our model so that our tensors are created on the right devices. The cleanest initialization pattern is to create the model inside a mesh context manager and JIT-compile the initialization for performance.

Here’s what it might look like in code:

import optax

@jax.jit
def init_model_and_optimizer(rngs: nnx.Rngs, model_config, optimizer_config, sharding_config):
    model = TransformerLM(
        rngs=rngs,
        model_config=model_config,
        sharding_config=sharding_config,
    )
    optimizer = nnx.Optimizer(
        model,
        build_optimizer_transform(
            optimizer_config,
            gradient_clip=optimizer_config["gradient_clip"],
            lr_schedule=optimizer_config["lr_schedule"],
        ),
        wrt=nnx.Param,
    )
    return model, optimizer

with jax.set_mesh(mesh):
    model, optimizer = init_model_and_optimizer(
        nnx.Rngs(0),
        model_config,
        optimizer_config,
        sharding_config,
    )

Sharding our Data

Beyond the model itself, we also need to shard the input data. The approach depends on the parallelism strategy: in TP, each device needs the full batch, while in DP, FSDP, and combined FSDP+TP we split the batch across the data axis so each device processes a different slice.

We can capture this logic with a simple helper:

def get_batch_sharding_for_mode(mode: str) -> Sharding:
    if mode == "tp":
        return (None, None)          # replicate across all devices
    if mode in ("dp", "fsdp", "fsdp_tp"):
        return ("data", None)        # shard batch over data axis, replicate over tensor axis
    raise ValueError("invalid mode!")

batch_partition_spec = P(*model_module.get_batch_sharding_for_mode(sharding_mode))

Then, at data loading time, we use jax.device_put to place the batch according to this spec:

def get_batch_from_memmap(
    rngs: nnx.Rngs,
    dataset: DatasetLike,
    batch_size: int,
    context_length: int,
    mesh: Mesh | None = None,
    batch_partition_spec: P | None = None,
) -> tuple[Array, Array]:
    inputs, targets = model_module.get_batch(
        rngs,
        dataset.data,
        batch_size,
        context_length,
    )

    # Place batch on devices according to the partition spec
    if mesh is not None:
        with jax.set_mesh(mesh):
            inputs = jax.device_put(inputs, batch_partition_spec)
            targets = jax.device_put(targets, batch_partition_spec)
    return inputs, targets

Defining Our Sharding

Before discussing each parallelism strategy separately, it’d be helpful to first identify the major tensors in the model and the sharding pattern each one will use. Our model contains the following components:

Rather than assigning every tensor a completely separate sharding rule, we can group tensors by the role they play in the computation.

In particular, many weight matrices fall into one of two common patterns:

1. Projection out of the model dimension $D$ into some larger or different dimension. This includes $W_{Q}, W_{K}, W_{V}$ and $W_{1}, W_{3}$, which have shape $[D, *]$. We can refer to this pattern as dense_in
2. Projection back into the model dimension $D$ from some intermediate dimension. This includes $W_{O}$ and $W_{2}$. We will refer to this pattern as dense_out. This naming is useful because these matrices often share the same partitioning logic under a given parallelism strategy, even if their exact shapes differ.

With this grouping, we can assign each component a sharding label:

These names only serve as categories for now. In the next step, for each parallelism strategy, we will map each one to an actual sharding specification: the tuple of strings describing the 2D partition over the tensor’s axes. For 1D tensors such as RMSNorm, the corresponding specification will be 1D rather than 2D.

Now, let’s take a look at how we can define what each of these sharding names should be under data parallelism, FSDP, tensor parallelism, and FSDP+TP!

Shardings for Each Strategy

DP

In Data Parallel, our basic formula looks something like the following:

We want to split our batches across the $X$, or “data” axis, and replicate our weights otherwise.

Remember that None replicates the dimension of the tensor across that axis!

Here’s what it might look like in code:

def get_sharding_config_for_mode(mode: str) -> ShardingConfig:
    if mode == "dp":
        return ShardingConfig(
            dense_in=(None, None),
            dense_out=(None, None),
            embedding=(None, None),
            lm_head=(None, None),
            norm=(None,),
        )
    ...

And that’s all we need to do! Simple, right?

FSDP

In FSDP, we have the following:

We want to shard both our input activations and weights along the data axis.

And here’s what it would look like in code:

def get_sharding_config_for_mode(mode: str) -> ShardingConfig:
    ...
    if mode == "fsdp":
        return ShardingConfig(
            dense_in=("data", None),
            dense_out=(None, "data"),
            embedding=(None, None),
            lm_head=(None, None),
            norm=(None,),
        )
    ...

TP

In TP, we have the following:

where we want to shard our weights along the “tensor” axis.

And here’s what the code would look like:

def get_sharding_config_for_mode(mode: str) -> ShardingConfig:
    ...
    if mode == "tp":
        return ShardingConfig(
            dense_in=(None, "tensor"),
            dense_out=("tensor", None),
            embedding=(None, "tensor"),
            lm_head=("tensor", None),
            norm=(None,),
        )
    ...

FSDP+TP

Finally, in FSDP+TP, we have the following:

And here’s what the code looks like:

def get_sharding_config_for_mode(mode: str) -> ShardingConfig:
    ...
    if mode == "fsdp_tp":
        return ShardingConfig(
            dense_in=("data", "tensor"),
            dense_out=("tensor", "data"),
            embedding=(None, "tensor"),
            lm_head=("tensor", None),
            norm=(None,),
        )
    ...

As you can see, implementing distributed training at this point is mostly plumbing: 1) initialize the model inside a mesh scope, 2) place the external token batch with the appropriate batch sharding, and 3) let JAX/XLA infer the collective communication implied by those layouts. That’s the main appeal of JAX’s sharding model—rather than manually writing all-gathers and reduce-scatters at the Python level, we can simply describe the layout we want and let the compiler handle these communications implicitly.

We also provide a PyTorch implementation of these distributed training methods (DP, TP, FSDP, FSDP+TP) for comparison. You can find it here. Note how much more complex it is!

Conclusion

Hopefully this was helpful! In this blog post, we went over some of the basic JAX/NNX concepts, demonstrated how we can build a language model from scratch using NNX, and discussed how we can implement distributed training.

Now, you can try applying your knowledge with our assignment or refer to our reference implementations.

This blog is still a work in progress! Some of the things that we plan to release soon:

• Introduction to Roofline Analysis
• Empirical Results and Comparison
• Mixture of Experts (MoE) implementation
• Additional LLM Tricks

We aim for this to be an easily customizable and useful resource for tinkering around in JAX.

Acknowledgements

Thank you to Henry Ko for providing helpful feedback on this blog post! We would also like to thank Machine Learning @ Berkeley for providing compute resources and feedback on an early version.

References

Austin et al., “How to Scale Your Model”, Google DeepMind, online, 2025.

Hashimoto et al., “CS336: Language Modeling from Scratch”, Stanford NLP, online, 2026.

https://pytorch.org/blog/overview-of-pytorch-autograd-engine/

https://docs.jax.dev/en/latest/

https://flax.readthedocs.io/en/stable/