DEV Community: Lewis Won

How does low-rank adaptation for large language models work

Lewis Won — Sat, 13 Sep 2025 12:50:55 +0000

Why study LoRA? The challenge of fine-tuning massive models
Conceptual overview: Full Fine-Tuning
Link back to the Math Equation on Full Fine-tuning
Conceptual Overview: Low-Rank Adaptation (LoRA)
Mathematical Walk-through of Forward Pass with LoRA
Mathematical Walk-through of the Backward Pass
Link Back to the Math Equation on LoRA
LoRA in the Self-Attention Module
Appendix A: The "No Additional Inference Latency" Trick
Appendix B: Why Does the Low-Rank Hypothesis Make Sense?

Why study LoRA? The challenge of fine-tuning massive models

Large Language Models (LLMs), because of their large size, present a significant challenge for training and deployment. For example, a model like GPT-3 has 175 billion parameters, would consume 350GB of storage (assuming FP16 quantization), and can consume up to 1.2TB of VRAM during training. (Note: I use GPT-3 as an example throughout this article to stay close to the original paper on LoRA.)

Fine-tuning is a way to adapt these pre-trained models to specific downstream tasks. Full fine-tuning involves updating all of the model's weights, which means that for every new task, you create a new, massive version of the model. Imagine a company wanting to offer 100 different specialized models to its customers. With full fine-tuning, this would require storing 100 separate 350GB models, consuming 35 terabytes of storage. This could make it difficult to deploy and manage customized LLMs at scale.

Low-Rank Adaptation (LoRA), introduced in the paper "LoRA: Low-Rank Adaptation of Large Language Models" by Hu et al., allows for the adaptation of massive models with a tiny fraction of the trainable parameters, significantly reducing storage costs and training overhead while maintaining satisfactory performance.

This article provides an intuitive walkthrough of how LoRA achieves this. This article was written with the assistance of Google Gemini 2.5 Pro.

Conceptual overview: Full Fine-Tuning

To understand the innovation of LoRA, we first need to look at the standard fine-tuning process. A neural network, at its core, is composed of layers, many of which perform matrix multiplication using weight matrices.

Let's imagine a single weight matrix in a pre-trained model, which we'll call $W_0$ . This matrix might have thousands of rows and columns.

\text{Pre-trained Weights } (W_0)

When we fine-tune this model on a new task, we update these weights based on the new data. The process learns a "delta" or change matrix, $ΔW\Delta W$ , which is added to the original weights.

The new, fine-tuned weight matrix, $W_{ft}$ , is:

W_{ft} = W_0 + \Delta W

The crucial point here is that the change matrix $ΔW\Delta W$ has the exact same dimensions as the original matrix $W_0$ . If $W_0$ has 100 million parameters, then we are training a $ΔW\Delta W$ that also has 100 million parameters. To save the fine-tuned model, we must save the entire $W_{ft}$ matrix, which is just as large as the original.

A Simple Example

Imagine our pre-trained weight matrix is a simple 2x3 matrix:

W_0 = \begin{bmatrix} 0.8 & 0.1 & 0.3 \newline 0.2 & 0.7 & 0.5 \end{bmatrix}

After fine-tuning, we might learn the following update matrix:

\Delta W = \begin{bmatrix} 0.1 & -0.2 & 0.05 \newline -0.05 & 0.15 & 0.1 \end{bmatrix}

The new, fully fine-tuned weight matrix would be:

W_{ft} = W_0 + \Delta W = \begin{bmatrix} 0.9 & -0.1 & 0.35 \newline 0.15 & 0.85 & 0.6 \end{bmatrix}

To make this change, we had to train and store 6 new values for $ΔW\Delta W$ . For a model like GPT-3, this means training and storing 175 billion new values for every single task.

Link back to the Math Equation on Full Fine-tuning

The paper provided the equation below the describe the standard process of full fine-tuning:

\max_{\Phi} \sum_{(x,y) \in Z} \sum_{t=1}^{|y|} \log(P_{\Phi}(y_t | x, y_{<t}))

What it means: This formula states that we want to find the best possible set of model parameters, denoted by $Φ\Phi$ , that maximizes the probability of generating the correct output sequences (y) given the input sequences (x) and all the preceding correct tokens $y_{<t}$ . We are optimizing over the entire set of parameters in the model.

Let's dissect each component:

$max⁡Φ\max_{\Phi}$ : This means our goal is to maximize the following expression by changing the model's parameters, denoted by the set $Φ\Phi$ . In full fine-tuning, $Φ\Phi$ represents every single weight and bias in the entire model. For GPT-3, this is a set of ~175 billion parameters.
$∑(x,y)∈Z\sum_{(x,y) \in Z}$ : This tells us to sum the results over our entire training dataset, which is a set $Z$ of context-target pairs $(x, y)$ . For example, in a summarization task, x would be a long article and y would be its short summary.
$∑t=1∣y∣\sum_{t=1}^{|y|}$ : This is for the autoregressive nature of language models. For each target sequence y, we sum over every single token from the beginning (t=1) to the end (|y|). The model tries to predict each token correctly, one by one.
$log⁡(⋅)\log(\cdot)$ : We use the logarithm of the probability. This is a standard technique that makes the math more stable and turns a long product of probabilities into a more manageable sum (the "log-likelihood"). Maximizing the log-probability is the same as maximizing the probability itself.
$PΦ(yt∣x,y<t)P_{\Phi}(y_t | x, y_{<t})$ : It represents the probability assigned by the model $P$ (with its current parameters $Φ\Phi$ ) to the correct next token $y_t$ , given the input context $x$ and all the preceding correct tokens $y_{<t}$ .

In simple terms, this equation says: Adjust all 175 billion parameters ( $Φ\Phi$ ) to make the model as good as possible at predicting the next correct word in the sequence for all the examples in our training data. The model starts with pre-trained weights $Φ0\Phi_0$ and learns an update $ΔΦ\Delta\Phi$ , resulting in final weights of $Φ0+ΔΦ\Phi_0 + \Delta\Phi$ . The problem is that $ΔΦ\Delta\Phi$ is just as large as $Φ0\Phi_0$ .

Link to this article: This equation is the mathematical representation of the concept described in the section "Conceptual overview: Full Fine-Tuning".
- The initial pre-trained weights, $Φ0\Phi_0$ , correspond to our simple matrix example $W_0$ .
- The final, optimized parameters, $Φ\Phi$ , are equivalent to our fine-tuned matrix $W_{ft}$ .
- The update, $ΔΦ\Delta\Phi$ , which is what we learn during training, corresponds to $ΔW\Delta W$ .
- The paper's key point that |ΔΦ| equals |Φ₀| is exactly what we illustrated: the update matrix $ΔW\Delta W$ has the same large dimensions as the original matrix $W_0$ , making it expensive to train and store.

Conceptual Overview: Low-Rank Adaptation (LoRA)

LoRA is built on a key insight: the update matrix $ΔW\Delta W$ does not need to have full rank to be effective. The authors hypothesize that the change in weights during adaptation has a low "intrinsic rank". This means that the large $ΔW\Delta W$ matrix can be approximated with high fidelity by multiplying two much smaller matrices.

Instead of learning $ΔW\Delta W$ directly, LoRA learns two smaller matrices, which we'll call A and B.

\Delta W \approx B \cdot A

This is a low-rank decomposition. The "rank" (r) is a small number we choose (like 1, 2, 8, or 64) that determines the inner dimension of these thin matrices.

If $W_0$ is a $\times k$ matrix, then $ΔW\Delta W$ is also $\times k$ .
With LoRA, matrix A will have dimensions $\times k$ , and matrix B will have dimensions $\times r$ .

The number of trainable parameters is now the sum of the parameters in A and B ( $\times r) + (r \times k)$ ), which is dramatically smaller than the $\times k$ parameters in $ΔW\Delta W$ , especially when $r$ is much smaller than $d$ and $k$ .

Crucially, during training with LoRA, the original weights $W_0$ are frozen and do not receive gradient updates. We only train the much smaller A and B matrices.

Figure 1 from the paper "LoRA: Low-Rank Adaptation of Large Language Models" by Hu et al. The pretrained weights W are frozen. Only the low-rank matrices A and B are trained.

The Same Example with LoRA

Let's use our 2x3 weight matrix $W_0$ again.

W_0 = \begin{bmatrix} 0.8 & 0.1 & 0.3 \newline 0.2 & 0.7 & 0.5 \end{bmatrix}

Here, $d = 2$ and $k = 3$ . Let's choose a tiny rank, $r = 1$ .

Matrix B will be $\times r \rightarrow 2 \times 1$ .
Matrix A will be $\times k \rightarrow 1 \times 3$ .

Suppose after training, we learn the following A and B:

\begin{bmatrix} 0.4 \newline 0.2 \end{bmatrix} ,\quad A = \begin{bmatrix} 0.25 & -0.5 & 0.1 \end{bmatrix}

Now, let's compute our low-rank approximation of $ΔW\Delta W$ :

\begin{aligned} \Delta W &\approx B \cdot A \newline &= \begin{bmatrix} 0.4 \newline 0.2 \end{bmatrix} \cdot \begin{bmatrix} 0.25 & -0.5 & 0.1 \end{bmatrix} \newline &= \begin{bmatrix} 0.1 & -0.2 & 0.04 \newline 0.05 & -0.1 & 0.02 \end{bmatrix} \end{aligned}

The key comparison is the number of parameters we had to train:

Full Fine-Tuning: The $ΔW\Delta W$ matrix had $\times 3 = 6$ parameters.
LoRA (r=1): We trained a $\times 1$ matrix B and a $\times 3$ matrix A. The total is $\times 1) + (1 \times 3) = 5$ parameters.

While the savings seem small here, for a large matrix in a real model, the difference is significant. The paper notes that for GPT-3, LoRA can reduce the number of trainable parameters by 10,000 times (from 175B to about 17M). This means the storage for each new task drops from 350GB to just a few dozen megabytes.

Mathematical Walk-through of Forward Pass with LoRA

The standard forward pass for a layer is $\cdot x$ , where h is the output, W is the weight matrix, and x is the input.

With LoRA, the output of the frozen pre-trained weights is computed as usual, and the output of the LoRA matrices is added to it.

The modified forward pass is:

W_0 \cdot x + B \cdot A \cdot x

This can be written as $(W_0 + B \cdot A) \cdot x$ . Let's use our example matrices to walk through the calculation.

1. Define Inputs

Pre-trained weights $W_0$ : $W0=[0.80.10.30.20.70.5]W_0 = \begin{bmatrix} 0.8 & 0.1 & 0.3 \newline 0.2 & 0.7 & 0.5 \end{bmatrix}$
Trained LoRA matrices $A$ and $B$ (with r=1): $\begin{bmatrix} 0.4 \newline 0.2 \end{bmatrix} ,\quad A = \begin{bmatrix} 0.25 & -0.5 & 0.1 \end{bmatrix}$
An input vector $x$ : $\begin{bmatrix} 10 \newline 20 \newline 30 \end{bmatrix}$

2. Calculate the original path
First, we compute the output from the frozen, pre-trained weights.

\begin{aligned} h_0 = W_0 \cdot x &= \begin{bmatrix} 0.8 & 0.1 & 0.3 \newline 0.2 & 0.7 & 0.5 \end{bmatrix} \cdot \begin{bmatrix} 10 \newline 20 \newline 30 \end{bmatrix} \newline &= \begin{bmatrix} (8 + 2 + 9) \newline (2 + 14 + 15) \end{bmatrix} \newline &= \begin{bmatrix} 19 \newline 31 \end{bmatrix} \end{aligned}

3. Calculate the LoRA path
Next, we compute the update from our trained LoRA matrices. It's more efficient to multiply $\cdot x$ first.

\begin{aligned} A \cdot x &= \begin{bmatrix} 0.25 & -0.5 & 0.1 \end{bmatrix} \cdot \begin{bmatrix} 10 \newline 20 \newline 30 \end{bmatrix} \newline &= [ (2.5 - 10 + 3) ] \newline &= [-4.5] \end{aligned}

Now multiply that result by B:

\begin{aligned} \Delta h &= B \cdot (A \cdot x) \newline &= \begin{bmatrix} 0.4 \newline 0.2 \end{bmatrix} \cdot [-4.5] \newline &= \begin{bmatrix} -1.8 \newline -0.9 \end{bmatrix} \end{aligned}

4. Combine the outputs
Finally, we add the two results together to get the final output h.

\begin{aligned} h &= h_0 + \Delta h \newline &= \begin{bmatrix} 19 \newline 31 \end{bmatrix} + \begin{bmatrix} -1.8 \newline -0.9 \end{bmatrix} \newline &= \begin{bmatrix} 17.2 \newline 30.1 \end{bmatrix} \end{aligned}

This process—keeping $W_0$ frozen and only passing gradients through the $\cdot A$ path—is how LoRA achieves its remarkable parameter efficiency during training.

Mathematical Walk-through of the Backward Pass

The "backward pass," or backpropagation, is the core mechanism by which a neural network learns. Its goal is to calculate how much each trainable parameter in the model contributed to the final error (or "loss"). Once we know this, we can adjust the parameters slightly to reduce that error.

In LoRA, the key efficiency gain comes from the fact that we only need to calculate these adjustments for the tiny A and B matrices. The massive pre-trained weight matrix, W_0, is frozen, so we can completely ignore it during the backward pass, saving immense amounts of computation and memory.

Let's walk through how the gradients for A and B are calculated.

1. The Setup

First, let's recall the forward pass equation:

W_0 \cdot x + B \cdot A \cdot x

The backward pass starts with a gradient signal coming from the next layer of the network. This signal, which we'll call $grad_h\text{grad\_h}$ , tells us how the final loss $(L)$ would change with respect to a small change in our output, $h$ . Mathematically, this is the derivative $∂L∂h\frac{\partial L}{\partial h}$ . Our task is to use this incoming gradient to figure out $∂L∂A\frac{\partial L}{\partial A}$ and $∂L∂B\frac{\partial L}{\partial B}$ .

We will use the same matrices and input vector from our forward pass example:

$\begin{bmatrix} 0.25 & -0.5 & 0.1 \end{bmatrix}$
$\begin{bmatrix} 0.4 \newline 0.2 \end{bmatrix}$
$\begin{bmatrix} 10 \newline 20 \newline 30 \end{bmatrix}$

Let's assume the incoming gradient $grad_h\text{grad\_h}$ is:

grad_h=∂L∂h=[0.5−0.2] \begin{aligned} \text{grad\_h} &= \frac{\partial L}{\partial h} \newline &= \begin{bmatrix} 0.5 \newline -0.2 \end{bmatrix} \end{aligned}

2. The Chain Rule in Action

The gradient for the LoRA path $(Δh=B⋅A⋅x)(\Delta h = B \cdot A \cdot x)$ is the same as the gradient for the total output $h$ , because $W0⋅xW_0 \cdot x$ is treated as a constant. The gradient flows back only through the parts of the computation that involve our trainable parameters.

3. Calculating the Gradient for `B` (`grad_B`)

To find how B affects the loss, we use the chain rule. The gradient of the loss with respect to B is found by multiplying the gradient of the output (grad_h) by how B affects the output.

The update term is $Δh=B⋅(A⋅x)\Delta h = B \cdot (A \cdot x)$ . The derivative of this with respect to B involves the term it was multiplied by, which is $\cdot x)$ .

The formula is:

∂L∂B=grad_h⋅(A⋅x)T \frac{\partial L}{\partial B} = \text{grad\_h} \cdot (A \cdot x)^T

Let's calculate this:

First, we need the term $\cdot x)$ . We already computed this in the forward pass: $\cdot x = [-4.5]$
The transpose, $\cdot x)^T$ , is $[- 4.5]$ .
Now we multiply: $grad_B=∂L∂B=[0.5−0.2]⋅[−4.5]=[0.5×−4.5−0.2×−4.5]=[−2.250.9]\begin{aligned} \text{grad\_B} &= \frac{\partial L}{\partial B} \newline &= \begin{bmatrix} 0.5 \newline -0.2 \end{bmatrix} \cdot [-4.5] \newline &= \begin{bmatrix} 0.5 \times -4.5 \newline -0.2 \times -4.5 \end{bmatrix} \newline &= \begin{bmatrix} -2.25 \newline 0.9 \end{bmatrix} \end{aligned}$ This grad_B matrix, which has the same shape as B, tells us how to adjust each element in B to reduce the loss.

4. Calculating the Gradient for `A` (`grad_A`)

Similarly, we find how A affects the loss. This time, the gradient has to pass back through B first.

The formula is:

∂L∂A=(BT⋅grad_h)⋅xT \frac{\partial L}{\partial A} = (B^T \cdot \text{grad\_h}) \cdot x^T

Let's calculate this step-by-step:

First, we need the transpose of B: $BT=[0.40.2]B^T = \begin{bmatrix} 0.4 & 0.2 \end{bmatrix}$
Next, we multiply $B^T$ by grad_h: $BT⋅grad_h=[0.40.2]⋅[0.5−0.2]=[(0.4×0.5)+(0.2×−0.2)]=[0.2−0.04]=[0.16]\begin{aligned} B^T \cdot \text{grad\_h} &= \begin{bmatrix} 0.4 & 0.2 \end{bmatrix} \cdot \begin{bmatrix} 0.5 \newline -0.2 \end{bmatrix} \newline &= [ (0.4 \times 0.5) + (0.2 \times -0.2) ] \newline &= [0.2 - 0.04] \newline &= [0.16] \end{aligned}$
Finally, we multiply this result by the transpose of the input, $x^T$ : $grad_A=∂L∂A=[0.16]⋅[102030]=[1.63.24.8]\begin{aligned} \text{grad\_A} &= \frac{\partial L}{\partial A} \newline &= [0.16] \cdot \begin{bmatrix} 10 & 20 & 30 \end{bmatrix} \newline &= \begin{bmatrix} 1.6 & 3.2 & 4.8 \end{bmatrix} \end{aligned}$ This grad_A matrix, which has the same shape as A, tells us how to adjust A.

5. Updating the Weights

After calculating the gradients, the optimizer performs a weight update. Using a simple learning rate (lr), the update rule is:

Anew=Aold−lr⋅grad_A A_{\text{new}} = A_{\text{old}} - \text{lr} \cdot \text{grad\_A}

Bnew=Bold−lr⋅grad_B B_{\text{new}} = B_{\text{old}} - \text{lr} \cdot \text{grad\_B}

And that's it! The crucial part is that W_0 is never updated. It is not part of the gradient calculation and requires no memory to store its gradients or optimizer states (like momentum). This is the source of LoRA's efficiency during the training process.

Link Back to the Math Equation on LoRA

The equation below from the paper describes the optimization objective of LoRA:

\max_{\Theta} \sum_{(x,y) \in Z} \sum_{t=1}^{|y|} \log(P_{\Phi_0 + \Delta\Phi(\Theta)}(y_t | x, y_{<t}))

The key changes of this equation compared to the previous equation for full-model fine-tuning:

$max⁡Θ\max_{\Theta}$ : This is the most important change. Instead of optimizing over the massive set of parameters $Φ\Phi$ , we are now optimizing over a much smaller set of parameters, denoted by $Θ\Theta$ . As the paper notes, the size of $Θ\Theta$ can be as small as 0.01% of the size of the original parameters ( $∣Θ∣≪∣Φ0∣|\Theta| \ll |\Phi_0|$ ). In our LoRA explainer, $Θ\Theta$ represents the collection of all the trainable values in our small A and B matrices.
$PΦ0+ΔΦ(Θ)P_{\Phi_0 + \Delta\Phi(\Theta)}$ : This shows how the model's weights are constructed.
- $Φ0\Phi_0$ : These are the original, pre-trained weights of the large model. They are treated as a fixed constant and are not trained. This corresponds to our frozen $W_0$ matrix.
- $ΔΦ(Θ)\Delta\Phi(\Theta)$ : This represents the weight update, but it's no longer a huge matrix of trainable parameters. Instead, it's a function that generates the large update matrix from the small set of parameters $Θ\Theta$ . For LoRA, this function is the matrix multiplication of our small matrices: $ΔΦ(Θ)=B⋅A\Delta\Phi(\Theta) = B \cdot A$ . The parameters $Θ\Theta$ are the entries of B and A.
What it means:
1. The optimization is no longer over the massive parameter set $Φ\Phi$ , but over a much smaller set of parameters denoted by $Θ\Theta$ . The paper states $∣Θ∣≪∣Φ0∣|\Theta| \ll |\Phi_0|$ .
2. The update to the weights, $ΔΦ\Delta\Phi$ , is now a function of this small parameter set, written as $ΔΦ(Θ)\Delta\Phi(\Theta)$ . The original weights $Φ0\Phi_0$ remain frozen.
Link to this article: This equation is the mathematical foundation for the section "Conceptual Overview: Low-Rank Adaptation (LoRA)".
- The small, trainable parameter set $Θ\Theta$ represents the collection of all the elements in our low-rank matrices, A and B.
- The function that generates the large update, $ΔΦ(Θ)\Delta\Phi(\Theta)$ , corresponds directly to the matrix multiplication $\cdot A$ . This operation takes the small number of parameters in A and B and uses them to produce the full-sized update matrix $ΔW\Delta W$ .
- The term $Φ0+ΔΦ(Θ)\Phi_0 + \Delta\Phi(\Theta)$ in the equation is precisely what we illustrated in the forward pass: $(W_0 + B \cdot A) \cdot x$ .

LoRA in the Self-Attention Module

The self-attention mechanism is a cornerstone of the Transformer architecture. For each input token, it computes Query (Q), Key (K), and Value (V) vectors. These are generated by multiplying the input embedding (x) with three distinct weight matrices: $W_q, W_k, W_v$ . After the attention scores are calculated and applied to the Value vectors, the result is passed through a final output projection matrix, $W_o$ , to produce the layer's output.

While LoRA can be applied to any weight matrix, the paper's authors focus their study on only the weight matrices in the self-attention module ( $W_q, W_k, W_v, W_o$ ). They find that for maximum parameter-efficiency, it is often sufficient to adapt only a subset of these matrices. For example, adapting only the query ( $W_q$ ) and value ( $W_v$ ) matrices can yield strong performance. For the purpose of a complete illustration, we will describe the process as if LoRA is applied to all four:

For $W_q$ , we add $Bq⋅AqB_q \cdot A_q$
For $W_k$ , we add $Bk⋅AkB_k \cdot A_k$
For $W_v$ , we add $Bv⋅AvB_v \cdot A_v$
For $W_o$ , we add $Bo⋅AoB_o \cdot A_o$

All eight of these LoRA matrices (A_q, B_q, A_k, B_k, etc.) are the only parameters that are trained. The original W matrices remain frozen. The calculations for each of the four paths are independent and can be performed in parallel.

Mathematical Walk-through of the Forward Pass

To keep the example clear, we will focus on the generation of a single Query vector. The exact same logic applies simultaneously to the Key and Value vectors.

1. Define Inputs

Let's assume our model has an embedding dimension (d_model) of 4, and the dimension of the Query/Key vectors (d_q or d_k) is 3.

Input Embedding (x), for a single token:

$\begin{bmatrix} 2 \newline 1 \newline 0.5 \newline 1.5 \end{bmatrix} \quad (\text{shape } 4 \times 1)$
Pre-trained Query Weight Matrix (W_q):

$W_q = \begin{bmatrix} 0.1 & 0.8 & 0.2 & 0.4 \newline 0.5 & 0.3 & 0.7 & 0.1 \newline 0.9 & 0.2 & 0.3 & 0.6 \end{bmatrix} \quad (\text{shape } 3 \times 4)$
Trainable LoRA Matrices for Query (A_q, B_q), with rank r=1:

$B_q = \begin{bmatrix} 0.5 \newline -0.2 \newline 0.1 \end{bmatrix} \quad (\text{shape } 3 \times 1) \quad A_q = \begin{bmatrix} 0.2 & -0.1 & 0.4 & 0.3 \end{bmatrix} \quad (\text{shape } 1 \times 4)$

2. Calculate the Original Path (Frozen)

First, we compute the output from the pre-trained W_q matrix.

q_0 = W_q \cdot x = \begin{bmatrix} 0.1 & 0.8 & 0.2 & 0.4 \newline 0.5 & 0.3 & 0.7 & 0.1 \newline 0.9 & 0.2 & 0.3 & 0.6 \end{bmatrix} \cdot \begin{bmatrix} 2 \newline 1 \newline 0.5 \newline 1.5 \end{bmatrix} = \begin{bmatrix} 0.2+0.8+0.1+0.6 \newline 1.0+0.3+0.35+0.15 \newline 1.8+0.2+0.15+0.9 \end{bmatrix} = \begin{bmatrix} 1.7 \newline 1.8 \newline 3.05 \end{bmatrix}

3. Calculate the LoRA Path (Trainable)

Next, we compute the update from our LoRA matrices.

A_q \cdot x = \begin{bmatrix} 0.2 & -0.1 & 0.4 & 0.3 \end{bmatrix} \cdot \begin{bmatrix} 2 \newline 1 \newline 0.5 \newline 1.5 \end{bmatrix} = [0.4 - 0.1 + 0.2 + 0.45] = [0.95]

\Delta q = B_q \cdot (A_q \cdot x) = \begin{bmatrix} 0.5 \newline -0.2 \newline 0.1 \end{bmatrix} \cdot [0.95] = \begin{bmatrix} 0.475 \newline -0.19 \newline 0.095 \end{bmatrix}

4. Combine for the Final Query Vector

Finally, we add the outputs of the two paths.

q_{final} = q_0 + \Delta q = \begin{bmatrix} 1.7 \newline 1.8 \newline 3.05 \end{bmatrix} + \begin{bmatrix} 0.475 \newline -0.19 \newline 0.095 \end{bmatrix} = \begin{bmatrix} 2.175 \newline 1.61 \newline 3.145 \end{bmatrix}

This q_final is the Query vector that will be used in the attention score calculation.

5. The Bigger Picture

This exact process happens in parallel for W_k (with A_k, B_k) and W_v (with A_v, B_v) to produce k_final and v_final. Those three vectors then proceed to the standard attention calculation. The final output of the attention mechanism is then passed through the W_o layer, which has its own LoRA update ( $hout=(Wo+BoAo)⋅hinh_{\text{out}} = (W_o + B_o A_o) \cdot h_{\text{in}}$ ).

Mathematical Walk-through of the Backward Pass

During backpropagation, the gradients flow backward from the loss function. The attention mechanism will provide an incoming gradient for q_final, k_final, and v_final. We'll demonstrate the process using the incoming gradient for our Query vector, grad_q.

1. The Setup

Let's assume the incoming gradient from the rest of the network for our Query vector is:

grad_q=∂L∂q_final=[0.80.1−0.4] \text{grad\_q} = \frac{\partial L}{\partial q\_{final}} = \begin{bmatrix} 0.8 \newline 0.1 \newline -0.4 \end{bmatrix}

Our goal is to use grad_q to calculate the gradients for A_q and B_q. The gradient does not flow back into W_q.

2. Calculating the Gradient for `B_q`

The gradient for B_q is found by multiplying the incoming gradient grad_q by the term that B_q was multiplied by in the forward pass: (A_q * x).

grad_B_q=∂L∂Bq=grad_q⋅(Aq⋅x)T \text{grad\_B\_q} = \frac{\partial L}{\partial B_q} = \text{grad\_q} \cdot (A_q \cdot x)^T

We already know from the forward pass that A_q * x = [0.95].

grad_B_q=[0.80.1−0.4]⋅[0.95]=[0.760.095−0.38] \text{grad\_B\_q} = \begin{bmatrix} 0.8 \newline 0.1 \newline -0.4 \end{bmatrix} \cdot [0.95] = \begin{bmatrix} 0.76 \newline 0.095 \newline -0.38 \end{bmatrix}

This is the adjustment signal for the B_q matrix.

3. Calculating the Gradient for `A_q`

The gradient for A_q requires the gradient to pass back through B_q first.

grad_A_q=∂L∂Aq=BqT⋅grad_q⋅xT \text{grad\_A\_q} = \frac{\partial L}{\partial A_q} = B_q^T \cdot \text{grad\_q} \cdot x^T

First, let's compute B_q^T * grad_q:

BqT⋅grad_q=[0.5−0.20.1]⋅[0.80.1−0.4]=[0.4−0.02−0.04]=[0.34] B_q^T \cdot \text{grad\_q} = \begin{bmatrix} 0.5 & -0.2 & 0.1 \end{bmatrix} \cdot \begin{bmatrix} 0.8 \newline 0.1 \newline -0.4 \end{bmatrix} = [0.4 - 0.02 - 0.04] = [0.34]

Now, multiply this by x^T:

grad_A_q=[0.34]⋅[210.51.5]=[0.680.340.170.51] \text{grad\_A\_q} = [0.34] \cdot \begin{bmatrix} 2 & 1 & 0.5 & 1.5 \end{bmatrix} = \begin{bmatrix} 0.68 & 0.34 & 0.17 & 0.51 \end{bmatrix}

This is the adjustment signal for the A_q matrix.

4. Updating All LoRA Weights

The optimizer will use these computed gradients (grad_B_q, grad_A_q) to update the weights of B_q and A_q.

Crucially, this same backward pass logic is applied independently and simultaneously for the other LoRA pairs. The incoming gradient grad_k is used to update A_k and B_k, grad_v is used to update A_v and B_v, and the gradient from the next layer is used to update A_o and B_o. The massive W_q, W_k, W_v, W_o matrices are completely bypassed during the gradient computation, leading to massive savings in time and memory.

Appendix A: The "No Additional Inference Latency" Trick

A common drawback of other adaptation methods (like adding "adapter" layers) is that they introduce extra layers or computations that permanently increase the model's inference latency. LoRA cleverly avoids this.

During training, the forward pass involves two paths, as we saw: $W_0 \cdot x + B \cdot A \cdot x$ . This does add a small amount of extra computation.

However, once training is complete, we can prepare the model for deployment (inference). Since matrices A and B are now fixed, we can perform their matrix multiplication once to get our final update matrix $ΔW\Delta W$ .

\Delta W = B \cdot A

Then, we can add this final update matrix directly to the original pre-trained weights to create a new, fully merged weight matrix, $W_{ft}$ .

W_{ft} = W_0 + \Delta W

Using our example numbers:

\Delta W = \begin{bmatrix} 0.1 & -0.2 & 0.04 \newline 0.05 & -0.1 & 0.02 \end{bmatrix}

\begin{aligned} W_{ft} &= W_0 + \Delta W \newline &= \begin{bmatrix} 0.8 & 0.1 & 0.3 \newline 0.2 & 0.7 & 0.5 \end{bmatrix} + \begin{bmatrix} 0.1 & -0.2 & 0.04 \newline 0.05 & -0.1 & 0.02 \end{bmatrix} \newline &= \begin{bmatrix} 0.9 & -0.1 & 0.34 \newline 0.25 & 0.6 & 0.52 \end{bmatrix} \end{aligned}

When the model is deployed, we just use this merged $W_{ft}$ matrix. The forward pass becomes $W_{ft} \cdot x$ , which is a single matrix multiplication. This has the exact same computational cost and latency as the original, non-fine-tuned model.

This also makes switching between tasks incredibly fast. As the paper notes, to switch from Task 1 (with matrices $B_1, A_1$ ) to Task 2 (with $B_2, A_2$ ), you can compute $W0+B2⋅A2W_0 + B_2 \cdot A_2$ , which is a very fast operation compared to loading an entirely new 350GB model from disk.

Appendix B: Why Does the Low-Rank Hypothesis Make Sense?

The idea that a massive, over-parameterized model can be adapted by changing only a small number of parameters might seem counter-intuitive, but it is supported by research into the dynamics of deep learning.

The LoRA paper builds on work by Aghajanyan et al. (2020) and others, which showed that pre-trained language models have a low "intrinsic dimension." This suggests that even though they exist in a very high-dimensional parameter space (e.g., 175 billion dimensions), they can learn new tasks effectively by moving along a much smaller, lower-dimensional manifold within that space.

LoRA takes this idea a step further by hypothesizing that the change in weights during adaptation ( $ΔW\Delta W$ ) also has a low "intrinsic rank." This means that the adjustments needed for a new task are not a complex, high-rank transformation but a simpler, low-rank one. The remarkable empirical success of LoRA, even with ranks as low as 1 or 2, provides strong evidence for this hypothesis. By decomposing the update into two small matrices, LoRA is essentially forcing the model to learn an update within this low-rank subspace, which proves to be a very effective and efficient constraint.

FlashAttention by hand

Lewis Won — Sat, 06 Sep 2025 13:42:58 +0000

Table of content

Introduction to FlashAttention
Vanilla self-attention mechanism without FlashAttention
Conceptual Overview: Fused Tiled Attention
FlashAttention by hand
Walkthrough of the FlashAttention diagram
Appendix A - How to derive activation matrices from weight matrices
Appendix B - The Paper's pseudocode vs. this walkthrough

Introduction to FlashAttention

FlashAttention builds upon the principles of online softmax to create an efficient single-pass algorithm for the entire self-attention mechanism. Its key innovation is to compute the final output O (where $\text{softmax}(QK^T)V$ ) directly, without ever forming the full attention matrix A. This is achieved by fusing the matrix multiplications ( $QK^T$ and $A V$ ) and the online softmax calculation into a single GPU kernel.

By avoiding the need to write and read the large $\times L$ attention matrix (where L represents the number of input tokens) to and from global memory (DRAM), FlashAttention significantly reduces memory access, which is often the primary bottleneck in attention calculations. This makes it substantially faster and more memory-efficient, especially for long sequences.

This explanation will walk through the FlashAttention algorithm by hand, using the same tiled approach from the online softmax example in my previous article. This walkthrough is based on "From Online Softmax to FlashAttention" by Ye Zihao (2023).

This article was written with the assistance of Google Gemini 2.5 Pro.

Vanilla self-attention mechanism without FlashAttention

In order to better appreciate the efficiency gains with FlashAttention, we first begin with a vanilla implementation of self-attention. Let's illustrate this with our familiar 1 x 6 example used in the previous article on online softmax. Imagine this vector represents the dot products of a single query vector q with six key vectors k.

Dot Products (Logits): $qK^T = \begin{bmatrix} 1 & 2 & 3 & 6 & 2 & 1 \end{bmatrix}$

We also need a corresponding V matrix. For simplicity, let's assume V has 6 rows (one for each key) and a dimension of 2.

Value Matrix: $\begin{bmatrix} v_1 \newline v_2 \newline v_3 \newline v_4 \newline v_5 \newline v_6 \end{bmatrix} = \begin{bmatrix} 1 & 1 \newline 2 & 2 \newline 3 & 3 \newline 4 & 4 \newline 5 & 5 \newline 6 & 6 \end{bmatrix}$

The standard process is broken down into three distinct, sequential stages:

Calculate Logits: Compute $X = QK^T$ .
Calculate Attention Scores: Compute $\text{softmax}(X)$ . This creates a dense attention score matrix.
Calculate Final Output: Compute $\cdot V$ .

Since the logits X are already given, our walkthrough will start from Step 2.

(Note: Q, K and V here are activation matrices which are result of multiplying the input embeddings by the weight matrices W_Q, W_K and W_V. These activation matrices are different for every input sequence and are the actual inputs to the attention calculation kernel that FlashAttention replaces. Refer to Appendix A to understand how these activation matrices are derived from the weight matrices.)

Step 1: Calculate the Full Attention Score Matrix (`A = softmax(X)`)

This step computes the softmax function over the entire logit vector X. This is a multi-pass process, requiring a pass to find the maximum, a pass to compute the denominator, and a pass to compute the final probabilities. The key difference from FlashAttention is that we must compute and store this entire attention vector A before we can even begin to use the V matrix.

1a. Find the Global Maximum (`m`)

First, we scan the entire logit vector X to find its maximum value for numerical stability (the "safe softmax" trick).

\max(\begin{bmatrix} 1 & 2 & 3 & 6 & 2 & 1 \end{bmatrix}) = 6

1b. Compute the Exponentials and the Denominator (`d`)

Next, we subtract the maximum from each logit, exponentiate the result, and sum them all up to get the denominator.

Subtract max: $\begin{bmatrix} 1-6 & 2-6 & 3-6 & 6-6 & 2-6 & 1-6 \end{bmatrix} = \begin{bmatrix} -5 & -4 & -3 & 0 & -4 & -5 \end{bmatrix}$
Exponentiate: $e(X−m)=[e−5e−4e−3e0e−4e−5]≈[0.00670.01830.049810.01830.0067]e^{(X-m)} = \begin{bmatrix} e^{-5} & e^{-4} & e^{-3} & e^{0} & e^{-4} & e^{-5} \end{bmatrix} \approx \begin{bmatrix} 0.0067 & 0.0183 & 0.0498 & 1 & 0.0183 & 0.0067 \end{bmatrix}$
Sum to get the denominator d: $\sum e^{(X-m)} \approx 0.0067 + 0.0183 + 0.0498 + 1 + 0.0183 + 0.0067 = 1.0998$

1c. Normalize to get the Final Attention Vector `A`

Finally, we divide the exponentiated values by the denominator d to get the final probabilities. The result is the complete attention score vector A.

\begin{aligned} A &= \frac{e^{(X-m)}}{d} \newline &\approx \frac{\begin{bmatrix} 0.0067 & 0.0183 & 0.0498 & 1 & 0.0183 & 0.0067 \end{bmatrix}}{1.0998} \newline &\approx \begin{bmatrix} 0.0061 & 0.0167 & 0.0453 & 0.9092 & 0.0167 & 0.0061 \end{bmatrix} \end{aligned}

At this point, the 1 x 6 vector A is fully computed and stored in memory.

Step 2: Multiply Attention Scores by the Value Matrix ( $\cdot V$ )

Now that we have the attention scores, we can perform the final step: a matrix multiplication between A and V. The attention scores in A act as weights for the corresponding value vectors in V. The output O is the weighted sum of the value vectors.

The operation is $\times 6) \cdot (6 \times 2) \rightarrow (1 \times 2)$ .

\cdot V \approx \begin{bmatrix} 0.0061 & 0.0167 & 0.0453 & 0.9092 & 0.0167 & 0.0061 \end{bmatrix} \cdot \begin{bmatrix} 1 & 1 \newline 2 & 2 \newline 3 & 3 \newline 4 & 4 \newline 5 & 5 \newline 6 & 6 \end{bmatrix}

Let's calculate the weighted sum:

\begin{aligned} O \approx & (0.0061 \cdot \begin{bmatrix}1 & 1\end{bmatrix}) + \newline & (0.0167 \cdot \begin{bmatrix}2 & 2\end{bmatrix}) + \newline & (0.0453 \cdot \begin{bmatrix}3 & 3\end{bmatrix}) + \newline & (0.9092 \cdot \begin{bmatrix}4 & 4\end{bmatrix}) + \newline & (0.0167 \cdot \begin{bmatrix}5 & 5\end{bmatrix}) + \newline & (0.0061 \cdot \begin{bmatrix}6 & 6\end{bmatrix}) \end{aligned}

Now we compute each term and then sum them up:

\begin{aligned} O \approx & \begin{bmatrix}0.0061 & 0.0061\end{bmatrix} + \newline & \begin{bmatrix}0.0334 & 0.0334\end{bmatrix} + \newline & \begin{bmatrix}0.1359 & 0.1359\end{bmatrix} + \newline & \begin{bmatrix}3.6368 & 3.6368\end{bmatrix} + \newline & \begin{bmatrix}0.0835 & 0.0835\end{bmatrix} + \newline & \begin{bmatrix}0.0366 & 0.0366\end{bmatrix} \end{aligned}

Summing the vectors component-wise:

\approx \begin{bmatrix}3.9323 & 3.9323\end{bmatrix}

This result, [3.932, 3.932].

Comparison and Key Differences

Intermediate Matrix: The standard method materialized the full 1 x 6 attention vector A. In a real-world scenario with a sequence length L, this would be a large L x L matrix. This is the main bottleneck. I will demonstrate how FlashAttention computes the final output without ever creating or storing this matrix.
Memory Access: The standard method requires at least two major memory operations: writing the entire A matrix to memory (DRAM), and then reading it all back in to multiply with V. I will show how FlashAttention fuses these operations, keeping tiles of Q, K, and V in fast SRAM and avoiding the slow roundtrip to DRAM.
Computation Flow: The process is strictly sequential. You cannot start the $\cdot V$ multiplication until the entire softmax calculation for A is complete. We will show how FlashAttention integrates these steps, updating a running output vector as it iterates through tiles of the K and V matrices.

Conceptual Overview: Fused Tiled Attention

FlashAttention processes the input matrices Q, K, and V in a tiled manner. For each row of the output matrix O, it iterates through the corresponding rows of K and V in blocks. At each step, it calculates the attention scores for just that block, updates the running statistics (the maximum and the denominator, just like in online softmax), and immediately applies these scores to the corresponding block of V to update a running output vector.

The core idea is to maintain three running statistics for each row of the output:

m_running: The running maximum of the dot products (q · k).
d_running: The running denominator of the softmax.
o_running: The running output vector, which is a weighted sum of the V vectors, scaled by the current (and incomplete) softmax probabilities.

We will use the exact same input data as before with the 1 x 6 example. This vector represents the dot products of a single query vector q with six key vectors k.

Dot Products (Logits): $qK^T = \begin{bmatrix} 1 & 2 & 3 & 6 & 2 & 1 \end{bmatrix}$

We also need a corresponding V matrix. For simplicity, let's assume V has 6 rows (one for each key) and a dimension of 2.

Value Matrix: $\begin{bmatrix} v_1 \newline v_2 \newline v_3 \newline v_4 \newline v_5 \newline v_6 \end{bmatrix} = \begin{bmatrix} 1 & 1 \newline 2 & 2 \newline 3 & 3 \newline 4 & 4 \newline 5 & 5 \newline 6 & 6 \end{bmatrix}$

We will process this with a tile size of 3, breaking X and V into two tiles:

$T_1$ : Logits $[123]\begin{bmatrix} 1 & 2 & 3 \end{bmatrix}$ and Values $[v1v2v3]\begin{bmatrix} v_1 \newline v_2 \newline v_3 \end{bmatrix}$
$T_2$ : Logits $[621]\begin{bmatrix} 6 & 2 & 1 \end{bmatrix}$ and Values $[v4v5v6]\begin{bmatrix} v_4 \newline v_5 \newline v_6 \end{bmatrix}$

The algorithm follows the logic of "Algorithm FlashAttention (Tiling)" on page 6 of "From Online Softmax to FlashAttention" by Ye Zihao (2023).

(Note: For intuitive clarity, this walkthrough maintains an un-normalized running output o_running and performs a single normalization at the end. The paper's pseudocode maintains a normalized running output o' at each step. The final result is mathematically identical. An explanation of why this is true is in Appendix B.)

FlashAttention by hand

Before the main loop begins, the running statistics are initialized.

m_running = -∞
d_running = 0
o_running = [0, 0] (a zero vector of the same dimension as a v vector)

Link back to the FlashAttention algorithm

These initial values correspond to the state before the for loop, where $m0=−∞m_0 = -\infty$ , $d0′=0d^\prime_0 = 0$ , and $o0′=0⃗o^\prime_0 = \vec{0}$ .

Step 1: Process Tile 1 (`i=1`)

We now begin the first iteration of the for i ← 1, #tiles do loop.

1a. Find the New Maximum

First, we calculate the dot products (logits) for this tile and find the maximum value.

Logits for Tile 1: $x1=[123]x_1 = \begin{bmatrix} 1 & 2 & 3 \end{bmatrix}$
Local max of Tile 1: $mT1=max⁡([123])=3m_{T_1} = \max(\begin{bmatrix} 1 & 2 & 3 \end{bmatrix}) = 3$
New overall maximum: $mnew=max⁡(mrunning,mT1)=max⁡(−∞,3)=3m_{new} = \max(m_{running}, m_{T_1}) = \max(-\infty, 3) = 3$

Link back to the FlashAttention algorithm

The calculation of logits corresponds to the first line inside the loop:

x_i \leftarrow Q[k,:] K^T[:, (i-1)b:ib]

Finding the local max corresponds to the second line:

m_i^{(\text{local})} = \max_{j=1..b}(x_i[j])

Updating the running max corresponds to the third line:

m_i \leftarrow \max(m_{i-1}, m_i^{(\text{local})})

(Here, $m_{i-1}$ is our m_running from before the step).

1b. Calculate Local Denominator and Local Output

Next, we compute the un-normalized attention scores for this tile using m_new, and then use them to calculate a local denominator d_local and a local weighted output o_local.

Un-normalized scores: $P1=[e1−3e2−3e3−3]=[e−2e−1e0]≈[0.13530.36791]P_1 = \begin{bmatrix} e^{1-3} \newline e^{2-3} \newline e^{3-3} \end{bmatrix} = \begin{bmatrix} e^{-2} \newline e^{-1} \newline e^{0} \end{bmatrix} \approx \begin{bmatrix} 0.1353 \newline 0.3679 \newline 1 \end{bmatrix}$
Local denominator: $dT1=∑P1≈0.1353+0.3679+1=1.5032d_{T_1} = \sum P_1 \approx 0.1353 + 0.3679 + 1 = 1.5032$
Local output (un-normalized weighted sum of V vectors): $oT1=(0.1353⋅v1)+(0.3679⋅v2)+(1⋅v3)≈(0.1353⋅[11])+(0.3679⋅[22])+(1⋅[33])≈[0.13530.1353]+[0.73580.7358]+[33]=[3.87113.8711]\begin{aligned} o_{T_1} &= (0.1353 \cdot v_1) + (0.3679 \cdot v_2) + (1 \cdot v_3) \newline &\approx (0.1353 \cdot \begin{bmatrix}1 & 1\end{bmatrix}) + (0.3679 \cdot \begin{bmatrix}2 & 2\end{bmatrix}) + (1 \cdot \begin{bmatrix}3 & 3\end{bmatrix}) \newline &\approx \begin{bmatrix}0.1353 & 0.1353\end{bmatrix} + \begin{bmatrix}0.7358 & 0.7358\end{bmatrix} + \begin{bmatrix}3 & 3\end{bmatrix} \newline &= \begin{bmatrix}3.8711 & 3.8711\end{bmatrix} \end{aligned}$

Link back to the FlashAttention algorithm

The calculation of $d_{T_1}$ corresponds to the summation part of the denominator update rule:

\sum_{j=1}^b e^{x_i[j] - m_i}

The calculation of $o_{T_1}$ corresponds to the numerator of the second term in the output update rule (before normalization):

∑j=1bexi[j]−miV[j+(i−1)b,:]\sum_{j=1}^b e^{x_i[j]-m_i} V[j + (i-1)b, :]

1c. Update Running Statistics

Now, we update our global running statistics.

\begin{aligned} d_{new} &= d_{old} \cdot e^{m_{old} - m_{new}} + d_{local} \newline o_{new} &= o_{old} \cdot e^{m_{old} - m_{new}} + o_{local} \end{aligned}

m_old = -∞, d_old = 0, o_old = [0, 0]
m_new = 3
d_local = 1.5032, o_local = [3.8711, 3.8711]

\begin{aligned} d_{running} &= (0 \cdot e^{-\infty - 3}) + 1.5032 = 1.5032 \newline o_{running} &= (\begin{bmatrix}0 & 0\end{bmatrix} \cdot e^{-\infty - 3}) + \begin{bmatrix}3.8711 & 3.8711\end{bmatrix} = \begin{bmatrix}3.8711 & 3.8711\end{bmatrix} \end{aligned}

After processing the first tile, our statistics are: m_running = 3, d_running ≈ 1.5032, o_running ≈ [3.8711, 3.8711].

Link back to the FlashAttention algorithm

This entire step corresponds to the final two lines of the loop body for i=1.

Denominator update:

d^\prime_i \leftarrow d^\prime_{i-1}e^{m_{i-1}-m_i} + \sum_{j=1}^b e^{x_i[j]-m_i}

Since 
 $d0′=0d^\prime_0 = 0$ 
, the first term is zero, and 
 $d1′d^\prime_1$ 
 becomes just the local sum, matching our result.

Output update:

o^\prime_i \leftarrow o^\prime_{i-1}\frac{d^\prime_{i-1}e^{m_{i-1}-m_i}}{d^\prime_i} + \frac{\sum e^{x_i[j]-m_i}V[\dots]}{d^\prime_i}

Similarly, since 
 $o0′o^\prime_0$ 
 and 
 $d0′d^\prime_0$ 
 are zero, the first term vanishes. Our un-normalized `o_running` is equivalent to 
 $o1′⋅d1′o^\prime_1 \cdot d^\prime_1$ 
, which is simply the numerator of the second term, matching our calculation of 
 $o_{T_1}$ 
.

Step 2: Process Tile 2 (`i=2`)

We now proceed to the second and final iteration of the loop.

2a. Find the New Maximum

Logits for Tile 2: $x2=[621]x_2 = \begin{bmatrix} 6 & 2 & 1 \end{bmatrix}$
Local max of Tile 2: $mT2=max⁡([621])=6m_{T_2} = \max(\begin{bmatrix} 6 & 2 & 1 \end{bmatrix}) = 6$
New overall maximum: $m_{new} = \max(m_{running}, m_{T_2}) = \max(3, 6) = 6$

Link back to the FlashAttention algorithm

This again maps to the first three lines of the loop body for i=2. This time, $m_{i-1} = m_1 = 3$ , so the new maximum is correctly found as $max⁡(3,6)=6\max(3, 6) = 6$ .

2b. Calculate Local Denominator and Local Output

We repeat the process for the second tile's data, using the new max, 6.

Un-normalized scores: $P2=[e6−6e2−6e1−6]=[e0e−4e−5]≈[10.01830.0067]P_2 = \begin{bmatrix} e^{6-6} \newline e^{2-6} \newline e^{1-6} \end{bmatrix} = \begin{bmatrix} e^{0} \newline e^{-4} \newline e^{-5} \end{bmatrix} \approx \begin{bmatrix} 1 \newline 0.0183 \newline 0.0067 \end{bmatrix}$
Local denominator: $dT2=∑P2≈1+0.0183+0.0067=1.025d_{T_2} = \sum P_2 \approx 1 + 0.0183 + 0.0067 = 1.025$
Local output: $oT2=(1⋅v4)+(0.0183⋅v5)+(0.0067⋅v6)≈(1⋅[44])+(0.0183⋅[55])+(0.0067⋅[66])≈[44]+[0.09150.0915]+[0.04020.0402]=[4.13174.1317]\begin{aligned} o_{T_2} &= (1 \cdot v_4) + (0.0183 \cdot v_5) + (0.0067 \cdot v_6) \newline &\approx (1 \cdot \begin{bmatrix}4 & 4\end{bmatrix}) + (0.0183 \cdot \begin{bmatrix}5 & 5\end{bmatrix}) + (0.0067 \cdot \begin{bmatrix}6 & 6\end{bmatrix}) \newline &\approx \begin{bmatrix}4 & 4\end{bmatrix} + \begin{bmatrix}0.0915 & 0.0915\end{bmatrix} + \begin{bmatrix}0.0402 & 0.0402\end{bmatrix} \newline &= \begin{bmatrix}4.1317 & 4.1317\end{bmatrix} \end{aligned}$

Link back to the FlashAttention algorithm

This again corresponds to the summation parts of the update rules for i=2.

2c. Update Running Statistics

Now we perform the final update, including the crucial rescaling step.

m_old = 3, d_old ≈ 1.5032, o_old ≈ [3.8711, 3.8711]
m_new = 6
d_local ≈ 1.025, o_local ≈ [4.1317, 4.1317]

First, update the denominator:

\begin{aligned} d_{running} &= d_{old} \cdot e^{m_{old} - m_{new}} + d_{local} \newline &\approx (1.5032 \cdot e^{3 - 6}) + 1.025 \newline &\approx (1.5032 \cdot 0.04979) + 1.025 \newline &\approx 0.0748 + 1.025 = 1.0998 \end{aligned}

Next, update the output vector:

\begin{aligned} o_{running} &= o_{old} \cdot e^{m_{old} - m_{new}} + o_{local} \newline &\approx (\begin{bmatrix}3.8711 & 3.8711\end{bmatrix} \cdot e^{3 - 6}) + \begin{bmatrix}4.1317 & 4.1317\end{bmatrix} \newline &\approx (\begin{bmatrix}3.8711 & 3.8711\end{bmatrix} \cdot 0.04979) + \begin{bmatrix}4.1317 & 4.1317\end{bmatrix} \newline &\approx \begin{bmatrix}0.1927 & 0.1927\end{bmatrix} + \begin{bmatrix}4.1317 & 4.1317\end{bmatrix} \newline &= \begin{bmatrix}4.3244 & 4.3244\end{bmatrix} \end{aligned}

Link back to the FlashAttention algorithm

This step corresponds to the full update rules for i=2.

Denominator update:

d^\prime_2 \leftarrow d^\prime_1 e^{m_1-m_2} + \sum_{j=1}^b e^{x_2[j]-m_2}

The term 
 $d1′em1−m2d^\prime_1 e^{m_1-m_2}$ 
 is the crucial rescaling factor, which perfectly matches the 
 $1.5032 \cdot e^{3 - 6})$ 
 part of our calculation.

Output update:

o^\prime_2 \leftarrow o^\prime_1\frac{d^\prime_1e^{m_1-m_2}}{d^\prime_2} + \frac{\sum e^{x_2[j]-m_2}V[\dots]}{d^\prime_2}

Again, our un-normalized `o_running` is equivalent to 
 $o2′⋅d2′o^\prime_2 \cdot d^\prime_2$ 
. If you multiply the paper's update rule by 
 $d2′d^\prime_2$ 
, you get 
 $sum(o^\prime_1 d^\prime_1) e^{m_1-m_2} + \text{local sum}$ 
. This exactly matches our formula:

oold⋅emold−mnew+olocalo_{old} \cdot e^{m_{old} - m_{new}} + o_{local}

Final Result

After the loop finishes, we have the final un-normalized output and the final denominator.

d_final ≈ 1.0998
o_final ≈ [4.3244, 4.3244]

The last step is to normalize the output vector by dividing by the final denominator.

\begin{aligned} O_{final} &= \frac{o_{final}}{d_{final}} \newline &\approx \frac{\begin{bmatrix}4.3244 & 4.3244\end{bmatrix}}{1.0998} \newline &\approx \begin{bmatrix}3.932 & 3.932\end{bmatrix} \end{aligned}

Link back to the FlashAttention algorithm

This final normalization step is implicitly the result of the algorithm. The final output of the loop, $oN/b′o^\prime_{N/b}$ , is the correctly normalized output row. Our method simply defers this division to the very end for clarity. The final output vector $O [k, :]$ in the algorithm is this final, normalized value.

\leftarrow o^\prime_{N/b}

Walkthrough of the FlashAttention diagram

Having completed the step-by-step walkthrough of the FlashAttention algorithm, we can now walkthrough the FlashAttention schematics from the original paper ["FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"(https://arxiv.org/abs/2205.14135) by Tri Dao et. al. (2022).

High-Level Overview

The diagram shows a tiled computation where the goal is to compute one block of the final output matrix, softmax(QK^T)V, at a time. The loops control which tiles (blocks) of the input matrices (Q, K, V) are loaded from slow global memory (HBM) into fast local memory (SRAM) to perform a piece of the calculation.

Our walkthrough, where we computed a single output row vector, corresponds to one full pass of the "Inner Loop" in this diagram.

Mapping the Diagram Components to Our Walkthrough

Let's look at each element of the diagram and connect it to our example.

1. The Loops

Inner Loop (Blue Arrows): This loop iterates over the queries (rows of Q) and the corresponding output rows. In our example, we only had a single query vector q that produced a single output row o. Therefore, our entire walkthrough represents a single iteration of the Inner Loop.
- The Q matrix block being copied is our single query q.
- The Output to HBM block at the bottom is our o_running vector $[o1o2]\begin{bmatrix}o_1 & o_2\end{bmatrix}$ , which is being progressively built.
Outer Loop (Red Arrows): This loop iterates over the key-value pairs (columns of K^T and rows of V). This loop is the core of the FlashAttention mechanism and maps directly to the steps of our walkthrough.
- Iteration 1 of the Outer Loop corresponds to our "Step 1: Process Tile 1".
- Iteration 2 of the Outer Loop corresponds to our "Step 2: Process Tile 2".

2. The Memory Hierarchy (SRAM vs. HBM)

HBM (High-Bandwidth Memory): This represents the GPU's main global memory (DRAM). This is where the full, large Q, K, V, and final O matrices reside.
SRAM (Fast On-Chip Memory): This is the small but extremely fast "workbench" memory. The diagram shows that we only ever copy small blocks (orange squares) into SRAM to work on them. In our walkthrough, the SRAM would hold:
- Our query vector q.
- The current block of keys and values we are processing (e.g., $k_1, k_2, k_3$ and $v_1, v_2, v_3$ for Tile 1).
- The running statistics: m_running, d_running, and o_running.

3. Putting it all Together: Tracing Our Walkthrough on the Diagram

Let's trace the flow for our single query q.

Initialization:

Before the loops start, the Output to HBM block (our o_running vector) is initialized to [0, 0]. The running statistics m_running = -∞ and d_running = 0 are initialized in SRAM.

Inner Loop Begins (One and only one iteration for our example):

Copy from Q: Our query vector q is loaded from HBM into SRAM. It will stay in SRAM for the entire duration of the Outer Loop.

Outer Loop - Iteration 1 (Our "Process Tile 1"):

Copy from K^T and V: The first block of K^T (corresponding to logits 1, 2, 3) and the first block of V (v_1, v_2, v_3) are loaded from HBM into SRAM.
Compute Block on SRAM: This is the central computation.
- The dot product $\cdot K^T_{\text{block 1}}$ is calculated.
- The local max, local denominator, and local output are computed.
- The running statistics m_running, d_running, and o_running (which live in SRAM) are updated. After this step, o_running is $≈[3.8711,3.8711]\approx [3.8711, 3.8711]$ . The + sign with the purple dotted arrow signifies this update step.

Outer Loop - Iteration 2 (Our "Process Tile 2"):

Copy from K^T and V: The previous blocks of K^T and V are discarded. The second block (logits 6, 2, 1 and values v_4, v_5, v_6) is loaded into SRAM.
Compute Block on SRAM: The computation is repeated.
- A new, larger global maximum (m_new = 6) is found.
- This triggers the crucial rescaling of the existing d_running and o_running vectors.
- The local contributions are calculated and added. After this step, the final un-normalized o_running is $≈[4.3244,4.3244]\approx [4.3244, 4.3244]$ .

Outer Loop Finishes:

The loop is complete. The final normalization is performed in SRAM ( $orunning/drunningo_{\text{running}} / d_{\text{running}}$ ) to get the final output vector $≈[3.932,3.932]\approx [3.932, 3.932]$ .

Output to HBM:

The final, correct output vector is written from SRAM back to its designated row in the main output matrix in HBM.

Inner Loop Finishes:

If there are more rows in Q, the Inner Loop will continue to the next row and repeat the entire Outer Loop process.

Appendix A - How to derive activation matrices from weight matrices

For the purpose of focusing on the core concepts of FlashAttention, the walkthrough focused exclusively on the core attention calculation step: $softmax(QKT)V\text{softmax}(QK^T)V$ . This is the specific operation that FlashAttention optimizes.

This appendix provides the "prequel" step that was omitted to derive the activation matrices.

The Two Sets of Q, K, V

It's crucial to distinguish between:

The Weight Matrices (W_Q, W_K, W_V): These are the trainable parameters learned during the training of the Transformer model. They are part of a standard Linear layer. Their job is to project the input token embeddings into the query, key, and value spaces. They are the same for every input sequence.
The Activation Matrices (Q, K, V): These are the intermediate representations or activations. They are the result of multiplying the input embeddings by the weight matrices. These matrices are different for every input sequence and are the actual inputs to the attention calculation kernel that FlashAttention replaces. They are not trainable parameters themselves.

Think of it like a recipe:

The weight matrices (W_Q, W_K, W_V) are the instructions in the recipe book (fixed, learned).
The activation matrices (Q, K, V) are the actual ingredients you've prepared for one specific meal (changes every time you cook).

FlashAttention's innovation is in how to efficiently combine the prepared ingredients (Q, K, V), not in the initial preparation step itself.

The Omitted "Prequel" Step: Creating Q, K, and V

Here is the step that happens before our walkthrough begins.

Let's assume we have an input sequence of 6 tokens, and each token has an embedding dimension of 4. This is our input matrix X (not to be confused with the logit vector X from the walkthrough).

Input Embeddings X (size 6x4): $\begin{bmatrix} \text{embedding for token 1} \newline \text{embedding for token 2} \newline \vdots \newline \text{embedding for token 6} \end{bmatrix}$

Now, let's define our trainable weight matrices. Let's say the head dimension d is 2. The weight matrices must project from the embedding dimension (4) to the head dimension (2). So, they will all be size 4x2.

Weight Matrices W_Q, W_K, W_V (size 4x2, trainable): $WQ=[…],WK=[…],WV=[…]W_Q = \begin{bmatrix} \dots \end{bmatrix}, \quad W_K = \begin{bmatrix} \dots \end{bmatrix}, \quad W_V = \begin{bmatrix} \dots \end{bmatrix}$

The Q, K, and V activation matrices are created with standard matrix multiplication:

\begin{aligned} Q &= X \cdot W_Q \quad (\text{results in a } 6 \times 2 \text{ matrix}) \newline K &= X \cdot W_K \quad (\text{results in a } 6 \times 2 \text{ matrix}) \newline V &= X \cdot W_V \quad (\text{results in a } 6 \times 2 \text{ matrix}) \end{aligned}

The V matrix we get from this calculation is precisely the V matrix we used in the walkthrough:

\begin{bmatrix} 1 & 1 \newline 2 & 2 \newline 3 & 3 \newline 4 & 4 \newline 5 & 5 \newline 6 & 6 \end{bmatrix}

Connecting to the Logits

In our walkthrough, we focused on a single query vector q attending to all the keys. This corresponds to taking the first row of the Q matrix.

\quad (\text{the first row, size } 1 \times 2)

The logit vector X from the walkthrough would then be calculated as:

\text{Logits} = q \cdot K^T

This multiplication would result in the 1 x 6 vector we started with:

\text{Logits} = \begin{bmatrix} 1 & 2 & 3 & 6 & 2 & 1 \end{bmatrix}

Summary

So, the full, un-omitted process is:

Linear Projections (The Omitted Prequel):
- Start with input embeddings X.
- Compute $\cdot W_Q$ , $\cdot W_K$ , $\cdot W_V$ using the trainable weight matrices. This is done with standard, highly optimized matrix multiplication libraries (GEMM).
FlashAttention Calculation (The Walkthrough):
- Take the resulting Q, K, and V activation matrices as input.
- Efficiently compute $\text{softmax}(QK^T)V$ in a single kernel without materializing the full attention matrix.

Appendix B - The Paper's pseudocode vs. this walkthrough

The goal of my walkthrough was to make the calculation as intuitive as possible to follow by hand. To do this, I slightly rearranged the math to keep the intermediate numbers as simple as possible, while still being mathematically equivalent to the paper's algorithm.

1. The Paper's "Algorithm FlashAttention (Tiling)"

The algorithm in the PDF maintains a normalized output vector o' at every step. Let's look closely at the update rule for the output:

o^\prime_i \leftarrow o^\prime_{i-1}\frac{d^\prime_{i-1}e^{m_{i-1}-m_i}}{d^\prime_i} + \frac{\sum_{j=1}^b e^{x_i[j]-m_i}V[j + (i-1)b, :]}{d^\prime_i}

Notice that both terms are divided by $di′d^\prime_i$ , the new total denominator. This means that at the end of each iteration i, the vector o'_i is the correctly normalized attention output for all the data processed up to that tile.

2. My Walkthrough's Method

Doing this normalization (division) at every single step can be cumbersome for a manual example. It's easier to work with un-normalized sums and perform a single division at the very end.

So, this walkthrough maintains an un-normalized running output o_running. Our update rule was:

orunning=oold⋅emold−mnew+olocalo_{\text{running}} = o_{\text{old}} \cdot e^{m_{\text{old}} - m_{\text{new}}} + o_{\text{local}}

This is mathematically equivalent to the numerator of the paper's formula. You can see this if you multiply the paper's entire update rule by $di′d^\prime_i$ :

\begin{aligned} o^\prime_i \cdot d^\prime_i &= \left( o^\prime_{i-1}\frac{d^\prime_{i-1}e^{m_{i-1}-m_i}}{d^\prime_i} + \frac{\sum \dots V[\dots]}{d^\prime_i} \right) \cdot d^\prime_i \newline o^\prime_i \cdot d^\prime_i &= (o^\prime_{i-1} d^\prime_{i-1}) e^{m_{i-1}-m_i} + \sum \dots V[\dots] \end{aligned}

If we define the $orunningo_{\text{running}}$ as the paper's $oi′⋅di′o^\prime_i \cdot d^\prime_i$ , then this equation is exactly the one used in the walkthrough:

My $onewo_{\text{new}}$ is $oi′⋅di′o^\prime_i \cdot d^\prime_i$
My $ooldo_{\text{old}}$ is $oi−1′⋅di−1′o^\prime_{i-1} \cdot d^\prime_{i-1}$
My $olocalo_{\text{local}}$ is $]\sum \dots V[\dots]$

The note was an attempt to explain this simplification: we chose to track the un-normalized numerator throughout the process for clarity and then perform the final division $ofinal/dfinalo_{\text{final}} / d_{\text{final}}$ only once.

Why normalizing only at the last step works

The reason this works is that the update rules for the un-normalized numerator (o_running) and the denominator (d_running) are designed to maintain a consistent relationship. At every step i, the correctly normalized output is simply the ratio of the running numerator to the running denominator at that step. By deferring the division to the end, we arrive at the same final ratio. The proof by induction below demonstrates this formally.

The Two Methods

Let's formally define the two methods we are comparing. We will use the paper's notation where $oi′o^\prime_i$ is the normalized output after tile i, and we'll introduce $O_i$ (capital O) as our un-normalized running output numerator from the walkthrough.

Method 1: Normalize at Each Step (The Paper's Algorithm)
The state after tile i is defined by:

$di′=di−1′emi−1−mi+(∑j∈Tileiexj−mi)d^\prime_i = d^\prime_{i-1} e^{m_{i-1}-m_i} + \left( \sum_{j \in \text{Tile}_i} e^{x_j-m_i} \right)$
$oi′=oi−1′di−1′emi−1−midi′+∑j∈Tileiexj−miVjdi′o^\prime_i = o^\prime_{i-1}\frac{d^\prime_{i-1}e^{m_{i-1}-m_i}}{d^\prime_i} + \frac{\sum_{j \in \text{Tile}_i} e^{x_j-m_i}V_j}{d^\prime_i}$

The final result is $oN′o^\prime_{N}$ after the last tile N.

Method 2: Normalize Only at the End (The Walkthrough's Method)
The state after tile i is defined by an un-normalized numerator $O_i$ and the same denominator $di′d^\prime_i$ :

$di′=di−1′emi−1−mi+(∑j∈Tileiexj−mi)d^\prime_i = d^\prime_{i-1} e^{m_{i-1}-m_i} + \left( \sum_{j \in \text{Tile}_i} e^{x_j-m_i} \right)$
$Oi=Oi−1emi−1−mi+∑j∈Tileiexj−miVjO_i = O_{i-1} e^{m_{i-1}-m_i} + \sum_{j \in \text{Tile}_i} e^{x_j-m_i}V_j$

The final result is calculated as $ON/dN′O_{N} / d^\prime_{N}$ at the very end.

The Proof of Equivalence

We want to prove that $oN′=ON/dN′o^\prime_N = O_N / d^\prime_N$ . We can prove this by induction, showing that the relationship $oi′=Oi/di′o^\prime_i = O_i / d^\prime_i$ holds true for every step i.

1. Base Case (i=1)

Let's check the first tile. Both methods start with $d0′=0d^\prime_0 = 0$ , $o0′=0⃗o^\prime_0 = \vec{0}$ , and $O0=0⃗O_0 = \vec{0}$ .

Method 1:

$)⏟0+∑j∈T1exj−m1Vjd1′=∑j∈T1exj−m1Vjd1′o^\prime_1 = \underbrace{o^\prime_0 \cdot (\dots)}{0} + \frac{\sum{j \in T_1} e^{x_j-m_1}V_j}{d^\prime_1} = \frac{\sum_{j \in T_1} e^{x_j-m_1}V_j}{d^\prime_1}$
Method 2:

$)⏟0+∑j∈T1exj−m1Vj=∑j∈T1exj−m1VjO_1 = \underbrace{O_0 \cdot (\dots)}{0} + \sum{j \in T_1} e^{x_j-m_1}V_j = \sum_{j \in T_1} e^{x_j-m_1}V_j$ The final result would be $O1/d1′O_1 / d^\prime_1$ .

Comparing the two, we see they are identical. The base case holds.

2. Inductive Hypothesis

Assume that the relationship is true for step i-1. That is, assume:

o^\prime_{i-1} = \frac{O_{i-1}}{d^\prime_{i-1}} \quad \implies \quad O_{i-1} = o^\prime_{i-1} d^\prime_{i-1}

3. Inductive Step

Now we must prove that the relationship holds for step i. Let's start with the formula for $oi′o^\prime_i$ from Method 1 and show it equals $Oi/di′O_i / d^\prime_i$ .

Start with the definition of $oi′o^\prime_i$ :

o^\prime_i = o^\prime_{i-1}\frac{d^\prime_{i-1}e^{m_{i-1}-m_i}}{d^\prime_i} + \frac{\sum_{j \in T_i} e^{x_j-m_i}V_j}{d^\prime_i}

Let's combine the two fractions over the common denominator $di′d^\prime_i$ :

o^\prime_i = \frac{\left(o^\prime_{i-1} d^\prime_{i-1}\right) e^{m_{i-1}-m_i} + \sum_{j \in T_i} e^{x_j-m_i}V_j}{d^\prime_i}

Now, look at the term in the parentheses: $(oi−1′di−1′)(o^\prime_{i-1} d^\prime_{i-1})$ . According to our Inductive Hypothesis, this is exactly equal to $O_{i-1}$ . Let's substitute it in:

o^\prime_i = \frac{O_{i-1} e^{m_{i-1}-m_i} + \sum_{j \in T_i} e^{x_j-m_i}V_j}{d^\prime_i}

Now, look at the entire numerator: $Oi−1emi−1−mi+∑j∈Tiexj−miVjO_{i-1} e^{m_{i-1}-m_i} + \sum_{j \in T_i} e^{x_j-m_i}V_j$ . This is precisely the definition of $O_i$ from Method 2.

So, we can substitute $O_i$ for the numerator:

o^\prime_i = \frac{O_i}{d^\prime_i}

This completes the proof. We have shown that if the relationship holds for step i-1, it must hold for step i. Since it holds for the base case i=1, it holds for all steps.

Intuitive Analogy: Calculating a Weighted Average

Think about calculating the final grade in a class. You have different assignments with different weights (scores).

Method 1 (Normalize at Each Step): After you get your first grade, you calculate your current average. When you get your second grade, you update your average based on the new grade and its weight relative to the first. You keep re-calculating your "running average" after every single assignment.
Method 2 (Normalize at the End): You collect all your points scored for each assignment (score * weight) in one column. You collect the total possible points (sum of all weights) in another column. You do this for the whole semester. At the very end, you do one single division: (total points scored) / (total possible points).

Both methods give you the exact same final grade. The second method simply defers the division. We can apply the same principle in FlashAttention: accumulate the un-normalized numerator (o_running) and the un-normalized denominator (d_running) separately and effectively perform the division at the end.

Online softmax by hand

Lewis Won — Thu, 28 Aug 2025 07:00:53 +0000

Table of Contents

Why study online softmax? The key to flash attention
Conceptual overview: tiled "safe" softmax
Conceptual Overview: Online Softmax
Mathematical conversion from element-wise to tile-wise formula
Appendix A: A simple example of numerical stability
Appendix B: Proof of correctness of online softmax

Why study online softmax? The key to flash attention

Softmax normalization is critical in the self-attention mechanism of transformers because it helps to manage "extreme values and offers more favorable gradient properties during training" (source: Building LLM from Scratch, chapter 3.3.1). For numerical stability, a "safe" softmax is widely used where in order to compute the final probability for any single element in every row of a matrix, we need to calculate two "global" statistics from the entirety of each row:
a. The maximum value of the row, $m = \max(S_i)$ .
b. The normalization denominator, $\sum_j \exp(S_{ij} - m)$ , which is the sum over all elements in the row.

(Note: Please refer to Appendix A for an illustration of the importance of the absolute maximum value (max(x)) for numerical stability.)

With a naive implementation of "safe" softmax, the two global statistics could potentially force a "multi-pass" approach where for each row, the following needs to be computed sequentially:

Read the entire row to find the max.
Calculate the denominator.
Compute the softmax values.

With online softmax, steps 1 (find the max) and 2 (calculate the denominator) can be carried out within the same read of the matrix, instead of two sequential reads. Online softmax achieves this by calculating the necessary global statistics in a single pass over the data. This contrasts with a naive tiled approach that requires two passes to find the statistics, followed by a third pass for the final calculation. By reducing the number of memory passes from three to two, online softmax offers significant efficiency gains, especially since matrix operations in GPUs tend to be memory-bound.

This article provides an intuitive understanding of online softmax by demonstrating by hand how both the max value and denominator of the "safe" softmax can be calculated in a single pass. This walkthrough is based on the article Online normalizer calculation for softmax by Maxim Milakov and Natalia Gimelshein, 2018.

Online softmax also provides the basis for Flash Attention to fuse the entire attention calculation into a single GPU kernel, which I will discuss in my next article.

Conceptual overview: tiled "safe" softmax

The "safe" softmax is inherently a multi-pass algorithm. This is because calculating the final probability for any single element requires knowing two global properties of the entire input vector:

The absolute maximum value (max(x)) for numerical stability.
The sum of all exponentiated values, which serves as the normalization denominator.

Tiling is a technique used to break down large matrices or vectors into smaller blocks, or "tiles". This is crucial for efficiency, especially on hardware like GPUs, as it allows smaller chunks of data to be loaded into fast, local memory (SRAM) for processing, rather than repeatedly accessing slower global memory (DRAM).

When applying tiling to the normal softmax, we process the input tile-by-tile. However, because we need the global maximum and the global sum, we cannot complete the calculation for any tile in a single go. This results in a multi-pass approach over the tiled data. (Note: The algorithm below describes an element-wise process, I will show its relationship to a tile-based approach.)

Let's illustrate this with a simple example. Consider the following 1x6 input vector (which can be thought of as a single row in a larger matrix):

\begin{bmatrix} 1 & 2 & 3 & 6 & 2 & 1 \end{bmatrix}

We will process this vector with a tile size of 3. This breaks our input X into two tiles:

T_1 = \begin{bmatrix} 1 & 2 & 3 \end{bmatrix}

T_2 = \begin{bmatrix} 6 & 2 & 1 \end{bmatrix}

The process follows the logic of "Algorithm 2: Safe softmax" from the image, but adapted for tiles.

Step 1: First Pass — Find the Global Maximum

In the first pass, we iterate through each tile to find its local maximum. Then, we find the maximum among all local maximums to determine the global maximum.

1a. Find Local Maximums

For each tile, we compute the maximum value, which we'll call m_local.

For Tile 1 ( $T_1$ ):

$m_1 = \max(\begin{bmatrix} 1 & 2 & 3 \end{bmatrix}) = 3$
For Tile 2 ( $T_2$ ):

$m_2 = \max(\begin{bmatrix} 6 & 2 & 1 \end{bmatrix}) = 6$

1b. Find Global Maximum

Now, we find the maximum of the local maximums to get the global maximum, m_global.

m_{global} = \max(m_1, m_2) = \max(3, 6) = 6

This value, m_global, corresponds to m_V in the provided "Algorithm 2".

Link back to the safe softmax algorithm

Line 1: m_0 ← -∞

Explanation: This initializes the variable that will hold the maximum value. It's set to negative infinity to ensure the first element of the vector x_1 becomes the first maximum.
In the Example: Before "Pass 1," we start with a conceptual maximum of $−∞-\infty$ .

Lines 2-4: for k ← 1, V do ... m_k ← max(m_{k-1}, x_k)

Explanation: This is the first pass. The loop iterates through every element of the input vector X to find the single largest value. After the loop finishes, m_V holds the global maximum.
In the Example: This corresponds to our "Step 1: First Pass".
- We first found the local max for Tile 1 ( $m_1=3$ ) and Tile 2 ( $m_2=6$ ).
- We then found the maximum of these, yielding the global max m_global = 6. This is the final value of m_V after this loop completes.

Step 2: Second Pass — Calculate the Global Denominator

In the second pass, we again iterate through each tile. This time, we use the m_global to calculate a local sum of exponentials for each tile. These local sums are then aggregated to get the global denominator.

The formula for each element within a tile is $e^{(x_j - m_{global})}$ .

2a. Calculate Local Denominators

For Tile 1 ( $T1=[123]T_1 = \begin{bmatrix} 1 & 2 & 3 \end{bmatrix}$ ):

\begin{aligned} d_1 &= e^{(1 - 6)} + e^{(2 - 6)} + e^{(3 - 6)} \newline &= e^{-5} + e^{-4} + e^{-3} \newline &\approx 0.006738 + 0.018316 + 0.049787 \newline &= 0.074841 \newline \end{aligned}

For Tile 2 ( $T2=[621]T_2 = \begin{bmatrix} 6 & 2 & 1 \end{bmatrix}$ ):

\begin{aligned} d_2 &= e^{(6 - 6)} + e^{(2 - 6)} + e^{(1 - 6)} \newline &= e^{0} + e^{-4} + e^{-5} \newline &\approx 1 + 0.018316 + 0.006738 \newline &= 1.025054 \newline \end{aligned}

2b. Calculate Global Denominator

The global denominator, d_global, is the sum of all local denominators. This corresponds to d_V in the algorithm.

\begin{aligned} d_{global} &= d_1 + d_2 \newline &\approx 0.074841 + 1.025054 \newline &= 1.099895 \end{aligned}

Link back to the safe softmax algorithm

Line 5: d_0 ← 0

Explanation: This initializes the variable for the denominator (the sum of exponentials).
In the Example: Before "Pass 2," our sum is zero.

Lines 6-8: for j ← 1, V do ... d_j ← d_{j-1} + e^(x_j - m_V)

Explanation: This is the second pass. The loop iterates through every element x_j again. For each element, it subtracts the global maximum m_V (found in the first pass), exponentiates the result, and adds it to the running sum. After the loop, d_V holds the global denominator.
In the Example: This corresponds to our "Step 2: Second Pass".
- For Tile 1, we calculated $d1=e(1−6)+e(2−6)+e(3−6)≈0.0748d_1 = e^{(1-6)} + e^{(2-6)} + e^{(3-6)} \approx 0.0748$ .
- For Tile 2, we calculated $d2=e(6−6)+e(2−6)+e(1−6)≈1.0251d_2 = e^{(6-6)} + e^{(2-6)} + e^{(1-6)} \approx 1.0251$ .
- The loop's final result, d_V, is the sum $d1+d2≈1.0999d_1 + d_2 \approx 1.0999$ .

Step 3: Third Pass — Calculate Final Softmax Probabilities

In the final pass, we iterate through the tiles one last time. For each element x_i, we compute its exponentiated value (scaled by m_global) and then divide by the d_global to get the final softmax probability y_i.

The formula is:

y_i = \frac{e^{(x_i - m_{global})}}{d_{global}}

3a. Compute Softmax for Each Tile

For Tile 1 ( $T1=[123]T_1 = \begin{bmatrix} 1 & 2 & 3 \end{bmatrix}$ ):
- $y1=e(1−6)1.099895=e−51.099895≈0.0067381.099895≈0.006126y_1 = \frac{e^{(1 - 6)}}{1.099895} = \frac{e^{-5}}{1.099895} \approx \frac{0.006738}{1.099895} \approx 0.006126$
- $y2=e(2−6)1.099895=e−41.099895≈0.0183161.099895≈0.016652y_2 = \frac{e^{(2 - 6)}}{1.099895} = \frac{e^{-4}}{1.099895} \approx \frac{0.018316}{1.099895} \approx 0.016652$
- $y3=e(3−6)1.099895=e−31.099895≈0.0497871.099895≈0.045266y_3 = \frac{e^{(3 - 6)}}{1.099895} = \frac{e^{-3}}{1.099895} \approx \frac{0.049787}{1.099895} \approx 0.045266$
For Tile 2 ( $T2=[621]T_2 = \begin{bmatrix} 6 & 2 & 1 \end{bmatrix}$ ):
- $y4=e(6−6)1.099895=e01.099895≈11.099895≈0.909177y_4 = \frac{e^{(6 - 6)}}{1.099895} = \frac{e^{0}}{1.099895} \approx \frac{1}{1.099895} \approx 0.909177$
- $y5=e(2−6)1.099895=e−41.099895≈0.0183161.099895≈0.016652y_5 = \frac{e^{(2 - 6)}}{1.099895} = \frac{e^{-4}}{1.099895} \approx \frac{0.018316}{1.099895} \approx 0.016652$
- $y6=e(1−6)1.099895=e−51.099895≈0.0067381.099895≈0.006126y_6 = \frac{e^{(1 - 6)}}{1.099895} = \frac{e^{-5}}{1.099895} \approx \frac{0.006738}{1.099895} \approx 0.006126$

Link back to the safe softmax algorithm

Lines 9-11: for i ← 1, V do ... y_i ← e^(x_i - m_V) / d_V

Explanation: This is the third pass. This final loop computes the softmax probability y_i for each element. It re-calculates the exponentiated value for x_i (just as in the second pass) and then divides it by the global denominator d_V.
In the Example: This corresponds to our "Step 3: Third Pass".
- $y1=e(1−6)/1.0999≈0.0061y_1 = e^{(1-6)} / 1.0999 \approx 0.0061$
- $y2=e(2−6)/1.0999≈0.0167y_2 = e^{(2-6)} / 1.0999 \approx 0.0167$
- ...and so on for all six elements.

Final Result

Combining the results from all tiles gives the final softmax output vector Y:

\approx \begin{bmatrix} 0.0061 & 0.0167 & 0.0453 & 0.9092 & 0.0167 & 0.0061 \end{bmatrix}

The sum of these probabilities is approximately 1.0, as expected. This multi-pass process is the standard, "normal" way of computing softmax in a tiled fashion, and stands in contrast to the "online softmax" which is designed to reduce the number of passes over the data.

Conceptual Overview: Online Softmax

The "Online Softmax," as described in the paper by Milakov and Gimelshein, is a clever single-pass algorithm. Its key innovation is the ability to calculate the softmax probabilities tile-by-tile without needing to first compute the global maximum and global sum over the entire input. It achieves this by maintaining a running maximum and a running denominator. When a new tile is processed, these running statistics are updated.

The core of the method lies in how the running denominator is updated when a new, larger maximum value is found. It uses a "telescoping sum" property to rescale the previous sum, effectively correcting it to be consistent with the new maximum. This avoids the need for a second or third pass over the data, which is a significant advantage for memory-bound operations on hardware like GPUs. Proof that online softmax is correct is in Appendix B. Note that the original proof was found in Online normalizer calculation for softmax, pages 3 - 4, but I instead referenced Yi Wang's FlashAttention (Part 2): Online Softmax as I found his writing to be more intuitive. (Note: The algorithm below describes an element-wise process, I will show its relationship to a tile-based approach.)

Let's use the same example vector and tile size:

\begin{bmatrix} 1 & 2 & 3 & 6 & 2 & 1 \end{bmatrix}

With a tile size of 3, we have two tiles:

T_1 = \begin{bmatrix} 1 & 2 & 3 \end{bmatrix}

T_2 = \begin{bmatrix} 6 & 2 & 1 \end{bmatrix}

The process will follow the logic of "Algorithm 3: Safe softmax with online normalizer calculation". We will process the vector tile by tile, from left to right, maintaining and updating two running values:

$m_{running}$ : The running maximum found so far.
$d_{running}$ : The running denominator (the sum of exponentials) scaled by the current $m_{running}$ .

Initialization

Before processing the first tile, we initialize our running statistics as per the algorithm:

$m0=mrunning=−∞m_0 = m_{running} = -\infty$
$d_0 = d_{running} = 0$

These are represented in lines 1 and 2 of the online softmax algorithm.

Step 1: Process Tile 1 ( $T1=[123]T_1 = \begin{bmatrix} 1 & 2 & 3 \end{bmatrix}$ )

We now process the first tile. For this tile, we will compute a new maximum ( $m_{new}$ ) and a new denominator ( $d_{new}$ ).

1a. Find the New Maximum for the Tile

First, find the maximum value within the current tile and compare it to our running maximum.

Local max of $T_1$ :

m_{T_1} = \max(\begin{bmatrix} 1 & 2 & 3 \end{bmatrix}) = 3

The new overall maximum is:

m_{new} = \max(m_{running}, m_{T_1}) = \max(-\infty, 3) = 3

1b. Update the Running Denominator

Next, we update the running denominator. The formula from Algorithm 3 (line 5) combines the previous denominator with the contribution from the current tile.

d_{new} \leftarrow d_{old} \times e^{(m_{old} - m_{new})} + \sum_{x_j \in \text{current tile}} e^{(x_j - m_{new})}

Let's apply this:

$mold=mrunning=−∞m_{old} = m_{running} = -\infty$
$d_{old} = d_{running} = 0$
$m_{new} = 3$

\begin{aligned} d_{new} &= (0 \times e^{(-\infty - 3)}) + (e^{(1 - 3)} + e^{(2 - 3)} + e^{(3 - 3)}) \newline &= 0 + (e^{-2} + e^{-1} + e^{0}) \newline &\approx 0.1353 + 0.3679 + 1 = 1.5032 \end{aligned}

1c. Update Running Statistics

After processing the first tile, our running statistics are:

$m_{running} = 3$
$drunning≈1.5032d_{running} \approx 1.5032$

At this point, we could calculate intermediate (and incorrect) softmax values for the first tile using these stats, but the key insight of the online method is to wait until all tiles are processed. The final values will need to be rescaled.

Step 1: Process Tile 2 ( $T2=[621]T_2 = \begin{bmatrix} 6 & 2 & 1 \end{bmatrix}$ )

Now we move to the second tile and repeat the process, using the statistics from the previous step as our "old" values.

1d. Find the New Maximum

Local max of $T_2$ :

m_{T_2} = \max(\begin{bmatrix} 6 & 2 & 1 \end{bmatrix}) = 6

The new overall maximum is: $m_{new} = \max(m_{running}, m_{T_2}) = \max(3, 6) = 6$

We have found a new, larger maximum. This will trigger the rescaling part of the update formula.

1e. Update the Running Denominator

We use the same formula as before, but now our old values are the results from Tile 1.

$m_{old} = 3$
$dold≈1.5032d_{old} \approx 1.5032$
$m_{new} = 6$

\begin{aligned} d_{new} &= (d_{old} \times e^{(m_{old} - m_{new})}) + (\sum_{x_j \in T_2} e^{(x_j - m_{new})}) \newline &\approx (1.5032 \times e^{(3 - 6)}) + (e^{(6 - 6)} + e^{(2 - 6)} + e^{(1 - 6)}) \newline &\approx (1.5032 \times e^{-3}) + (e^{0} + e^{-4} + e^{-5}) \newline &\approx (1.5032 \times 0.04979) + (1 + 0.01832 + 0.00674) \newline &\approx 0.07484 + 1.02506 \newline &= 1.0999 \end{aligned}

The first term, $1.5032 \times e^{(3 - 6)}$ , is the crucial rescaling step. It takes the previous sum, which was calculated relative to the old maximum of 3, and correctly scales it to be relative to the new maximum of 6.

1f. Update Final Running Statistics

After processing all tiles, our final global statistics are:

$m_{final} = 6$
$dfinal≈1.0999d_{final} \approx 1.0999$

These values match the global maximum and denominator we found in the multi-pass "normal" softmax example.

Link back to the online softmax algorithm

Lines 3-6: for j ← 1, V do ...

Explanation: This is the main single pass of the algorithm. It iterates through the input vector element by element, updating the running max and denominator at each step.
In the Example: We processed this tile-by-tile, which is a block-wise version of this element-wise loop.

Line 4: $m_j ← max(m_{j-1}, x_j)$

Explanation: At each step j, this line updates the running maximum.
In the Example:
- When processing Tile 1 ( $x_1, x_2, x_3$ ), the running max m_j becomes 1, then 2, then finally 3. After Tile 1, our example had m_running = 3.
- When processing $x_4=6$ , the max max(3, 6) becomes 6. Our example showed this change when moving from Tile 1 to Tile 2, where m_new became 6.

Line 5: $d_j ← d_{j-1} × e^{(m_{j-1} - m_j)} + e^{(x_j - m_j)}$

Explanation: This is the core of the online update.
- The second term, $e^{(x_j - m_j)}$ , adds the contribution of the current element, scaled by the new running max.
- The first term, $d_{j-1} * e^{(m_{j-1} - m_j)}$ , is the crucial rescaling factor. If a new maximum is found ( $m_j > m_{j-1}$ ), this term correctly scales down the entire previous sum to be relative to the new, larger maximum. If the maximum doesn't change ( $m_j = m_{j-1}$ ), this term simply becomes $d_{j-1} × e^0 = d_{j-1}$ , so we just keep adding to the sum.
In the Example:
- For Tile 1, the max was always updating, but $d_{j-1}$ started at 0. The final sum for the tile was $d1=e(1−3)+e(2−3)+e(3−3)≈1.5032d_1 = e^{(1-3)} + e^{(2-3)} + e^{(3-3)} \approx 1.5032$ .
- The most important step was processing Tile 2. Our running denominator was d_old ≈ 1.5032 and m_old = 3. The new max became m_new = 6. The update $1.5032 × e^{(3 - 6)})$ is the direct application of $d_{j-1} × e^{(m_{j-1} - m_j)}$ , rescaling the entire sum from the first tile. We then added the contributions from Tile 2's elements.

Mathematical conversion from element-wise to tile-wise formula

You may notice that the online softmax pseudocode does not have a summation Σ which appears in the illustrative calculations. This is because the pseudocode describes the update rule on a per-element basis. The summation Σ appears in the example because we are performing a per-tile update. The summation is simply the result of applying the single-element rule to every element within the tile at the same time. It is a practical optimization that is mathematically equivalent.

Let's see how the tile-wise formula is derived directly from the element-wise formula.

The core element-wise update rule is:

d_j = d_{j-1} \cdot e^{m_{j-1} - m_j} + e^{x_j - m_j}

Imagine we have just finished processing Tile 1. Our running statistics are $m_{old}$ and $d_{old}$ . Now we want to process Tile 2, which contains the elements ${x_a, x_b, x_c}$ .

Instead of processing the whole tile at once, let's apply the element-wise rule three times in a row.

1. Processing the first element of the tile, $x_a$ :

m_{new_a} = \max(m_{old}, x_a)

d_{new_a} = d_{old} \cdot e^{m_{old} - m_{new_a}} + e^{x_a - m_{new_a}}

2. Processing the second element, $x_b$ (using the results from step 1):

m_{new_b} = \max(m_{new_a}, x_b)

d_{new_b} = d_{new_a} \cdot e^{m_{new_a} - m_{new_b}} + e^{x_b - m_{new_b}}

Let's substitute the expression for $d_{new_a}$ into this equation:

d_{new_b} = (d_{old} \cdot e^{m_{old} - m_{new_a}} + e^{x_a - m_{new_a}}) \cdot e^{m_{new_a} - m_{new_b}} + e^{x_b - m_{new_b}}

d_{new_b} = d_{old} \cdot e^{m_{old} - m_{new_b}} + e^{x_a - m_{new_b}} + e^{x_b - m_{new_b}}

Notice the pattern: the original $d_{old}$ is now scaled by the newest max, and the contributions from the elements processed so far ( $x_a, x_b$ ) are summed up, also scaled by the newest max.

3. Processing the third element, $x_c$ (using the results from step 2):

m_{new_c} = \max(m_{new_b}, x_c)

d_{new_c} = d_{new_b} \cdot e^{m_{new_b} - m_{new_c}} + e^{x_c - m_{new_c}}

Substituting the expression for $d_{new_b}$ :

d_{new_c} = (d_{old} \cdot e^{m_{old} - m_{new_b}} + e^{x_a - m_{new_b}} + e^{x_b - m_{new_b}}) \cdot e^{m_{new_b} - m_{new_c}} + e^{x_c - m_{new_c}}

d_{new_c} = d_{old} \cdot e^{m_{old} - m_{new_c}} + e^{x_a - m_{new_c}} + e^{x_b - m_{new_c}} + e^{x_c - m_{new_c}}

The Tile-Wise Optimization

In a practical implementation on parallel hardware (like a GPU), it's far more efficient to do the following:

Load the entire tile ${x_a, x_b, x_c}$ into fast local memory.
Find the local maximum for the tile: $m_{tile} = \max(x_a, x_b, x_c)$ .
Calculate the new overall maximum in one step: $m_{new} = \max(m_{old}, m_{tile})$ . Note that this final m_new is mathematically identical to the m_new_c we found after processing the last element in the step-by-step derivation above.
Apply the update rule in one go.

Looking at our final equation from the element-wise derivation:

d_{new_c} = d_{old} \cdot e^{m_{old} - m_{new_c}} + (e^{x_a - m_{new_c}} + e^{x_b - m_{new_c}} + e^{x_c - m_{new_c}})

We can rewrite the part in the parentheses using a summation:

\sum_{x_j \in \text{Tile}} e^{x_j - m_{new_c}}

If we substitute $m_{new_c}$ with our tile-wise $m_{new}$ , we get the exact formula used in the example:

d_{new} = d_{old} \cdot e^{m_{old} - m_{new}} + \sum_{x_j \in \text{Tile}} e^{x_j - m_{new}}

Step 2: Final Pass — Calculate Final Softmax Probabilities

Because the online algorithm only makes a single pass to find the statistics of max and sum, a final pass is still required to compute the output probabilities using the final global statistics.

The formula is the same as in the normal softmax: $yi=e(xi−mfinal)dfinaly_i = \frac{e^{(x_i - m_{final})}}{d_{final}}$

For Tile 1 ( $T1=[123]T_1 = \begin{bmatrix} 1 & 2 & 3 \end{bmatrix}$ ):
- $y1=e(1−6)1.0999=e−51.0999≈0.0061y_1 = \frac{e^{(1 - 6)}}{1.0999} = \frac{e^{-5}}{1.0999} \approx 0.0061$
- $y2=e(2−6)1.0999=e−41.0999≈0.0167y_2 = \frac{e^{(2 - 6)}}{1.0999} = \frac{e^{-4}}{1.0999} \approx 0.0167$
- $y3=e(3−6)1.0999=e−31.0999≈0.0453y_3 = \frac{e^{(3 - 6)}}{1.0999} = \frac{e^{-3}}{1.0999} \approx 0.0453$
For Tile 2 ( $T2=[621]T_2 = \begin{bmatrix} 6 & 2 & 1 \end{bmatrix}$ ):
- $y4=e(6−6)1.0999=e01.0999≈0.9092y_4 = \frac{e^{(6 - 6)}}{1.0999} = \frac{e^{0}}{1.0999} \approx 0.9092$
- $y5=e(2−6)1.0999=e−41.0999≈0.0167y_5 = \frac{e^{(2 - 6)}}{1.0999} = \frac{e^{-4}}{1.0999} \approx 0.0167$
- $y6=e(1−6)1.0999=e−51.0999≈0.0061y_6 = \frac{e^{(1 - 6)}}{1.0999} = \frac{e^{-5}}{1.0999} \approx 0.0061$

Link back to the online softmax algorithm

Lines 7-9: for i ← 1, V do ... y_i ← e^(x_i - m_V) / d_V

Explanation: This is a final pass to compute the results. After the main loop (lines 3-6) is complete, m_V and d_V hold the final, correct global statistics. This loop then uses those final values to calculate each y_i.
In the Example: This corresponds to our "Step 3: Final Pass". We took the final m_final = 6 and d_final ≈ 1.0999 and applied them to all original x_i values to get the final probabilities, which is exactly what this loop does.

Final Result

The final softmax vector is identical to the one produced by the normal softmax method:

\approx \begin{bmatrix} 0.0061 & 0.0167 & 0.0453 & 0.9092 & 0.0167 & 0.0061 \end{bmatrix}

The critical advantage is that the online method calculated the necessary global statistics ( $m_{final}$ and $d_{final}$ ) in a single pass over the data, whereas the normal tiled method required two passes just to get the statistics, followed by a third pass for the final calculation. This reduction in memory passes provides for efficiency gains with online softmax.

Appendix A: A simple example of numerical stability

Let's illustrate why this stabilization is crucial with a simple vector, especially when using lower precision floating-point numbers like float16, which is common in modern GPUs for accelerating deep learning workloads. The maximum value for a float16 is 65,504.

Consider the input vector:

\begin{bmatrix} 2 & 4 & 12 \end{bmatrix}

Case 1: The Naive (Unstable) Approach

Using the direct definition of softmax, we compute $exp(x_i)$ for each element:

$exp⁡(2)≈7.389\exp(2) \approx 7.389$
$exp⁡(4)≈54.598\exp(4) \approx 54.598$
$exp⁡(12)≈162,754.79\exp(12) \approx 162,754.79$

Now, let's analyze this from a float16 perspective. With $\begin{bmatrix} 1 & 2 & 12 \end{bmatrix}$ , we proceed to calculate the denominator:
$\sum_{j=1}^N \exp(x_j) \approx 7.389 + 54.598 + 162,754.79 = 162,816.777$

As the maximum finite positive value for a float16 is approximately 65,604, the sum 162,816.777 exceeds the maximum value, and hence causes a catastrophic overflow in float16, resulting in inf (infinity). As a result, the final probabilities would be computed as:

$y1=7.389∞=0y_1 = \frac{7.389}{\infty} = 0$
$y2=54.598∞=0y_2 = \frac{54.598}{\infty} = 0$
$y3=162,754.79∞=NaNy_3 = \frac{162,754.79}{\infty} = \text{NaN}$ (Not a Number, from $∞∞\frac{\infty}{\infty}$ )

The result is a vector of [0, 0, NaN], which is completely incorrect and would destroy any subsequent training or inference.

Case 2: The Safe (Stable) Approach

Now, let's use the numerically stable formula.

Step 1: Find the maximum value, m.

\max(\begin{bmatrix} 1 & 2 & 12 \end{bmatrix}) = 12

Step 2: Subtract m from each element.

\begin{bmatrix} 1-12 & 2-12 & 12-12 \end{bmatrix} = \begin{bmatrix} -11 & -10 & 0 \end{bmatrix}

Step 3: Compute the exponentials of the new values.

$exp⁡(−11)≈0.0000167\exp(-11) \approx 0.0000167$
$exp⁡(−10)≈0.0000454\exp(-10) \approx 0.0000454$
$exp⁡(0)=1\exp(0) = 1$

Notice that all these values are small, well-behaved numbers between 0 and 1. There is absolutely no risk of overflow.

Step 4: Calculate the denominator, d'.

\sum_{j=1}^N \exp(x_j - m) \approx 0.0000167 + 0.0000454 + 1 = 1.0000621

Step 5: Compute the final probabilities.

$y1=exp⁡(−11)1.0000621≈0.00001671.0000621≈0.0000167y_1 = \frac{\exp(-11)}{1.0000621} \approx \frac{0.0000167}{1.0000621} \approx 0.0000167$
$y2=exp⁡(−10)1.0000621≈0.00004541.0000621≈0.0000454y_2 = \frac{\exp(-10)}{1.0000621} \approx \frac{0.0000454}{1.0000621} \approx 0.0000454$
$y3=exp⁡(0)1.0000621=11.0000621≈0.9999379y_3 = \frac{\exp(0)}{1.0000621} = \frac{1}{1.0000621} \approx 0.9999379$

The resulting softmax vector is approximately $[0.0000170.0000450.999938]\begin{bmatrix} 0.000017 & 0.000045 & 0.999938 \end{bmatrix}$ . The sum is 1.0, and we have avoided any numerical errors. This demonstrates that subtracting the maximum value is not just a theoretical trick—it is an essential step for making softmax work in practice.

Appendix B: Proof of correctness of online softmax

I have reproduced the proof by Yi Wang in FlashAttention (Part 2): Online Softmax below.

The definition of softmax is as follows:

\left\{ \frac{\exp(x_i)}{\sum_{j=1}^N \exp(x_j)} \right\}_{i=1}^N

If any $x_i$ is large (e.g. $xi≥11x_i \geq 11$ ), $exp(x_i)$ exceeds the maximum value of float16. To address this numerical instability, we compute an alternative form which gives equivalent result but numerically stable:

\left\{ \frac{\exp(x_i - m)}{\sum_{j=1}^N \exp(x_j - m)} \right\}_{i=1}^N

where $m = \max_{j=1}^N x_j$ . This form is safe because $xi−m≤0x_i - m \le 0$ , ensuring that $\exp(x_i - m) \le 1$ .

As mentioned in this article, online softmax seeks to allow the calculation of the max value $m$ and the denominator $∑j=1Nexp⁡(xj−m)\sum_{j=1}^N \exp(x_j - m)$ in parallel. More concretely, for any integer $\leq i \leq N$ , we want to be able to:

Calculate an intermediate $δi=∑j=1iexp⁡(xj−mi)\delta_i = \sum_{j=1}^i \exp(x_j - m_i)$ so that $δN=∑j=1Nexp⁡(xj−mN)\delta_N = \sum_{j=1}^N \exp(x_j - m_N)$
Since $δi\delta_i$ is inductive, it should depend on $δi−1\delta_{i-1}$ .
To allow parallel execution, $δi\delta_i$ must not depend on future values such as $xi+1,…x_{i+1}, \dots$ or $mi+1,…m_{i+1}, \dots$ .

We begin by considering:

\delta_i = \sum_{j=1}^i \exp(x_j - m_i)

To ensure $δi\delta_i$ depends on $δi−1\delta_{i-1}$ , which is:

\delta_{i-1} = \sum_{j=1}^{i-1} \exp(x_j - m_{i-1})

we need to split $δi\delta_i$ into two parts: one involving $δi−1\delta_{i-1}$ (which should not depend on $x_i$ or $m_i$ ), and the remaining terms that depend on $x_i$ and $m_i$ . The first step is straightforward – we separate the last term in the summation:

\delta_i = \sum_{j=1}^{i-1} \exp(x_j - m_i) + \exp(x_i - m_i)

Now, $x_i$ only appears in the second term. However, $m_i$ still appears in the summation. Let's take the next step:

\delta_i = \sum_{j=1}^{i-1} \exp(x_j - m_{i-1} + m_{i-1} - m_i) + \exp(x_i - m_i) \ = \left[ \sum_{j=1}^{i-1} \exp(x_j - m_{i-1}) \right] \exp(m_{i-1} - m_i) + \exp(x_i - m_i)

The expression inside the square brackets is exactly $δi−1\delta_{i-1}$ . Therefore, we have:

\delta_i = \delta_{i-1} \exp(m_{i-1} - m_i) + \exp(x_i - m_i)

This allows us to compute $δi\delta_i$ inductively in parallel with $m_i$ .

Tensor parallelism by hand

Lewis Won — Sat, 23 Aug 2025 11:18:13 +0000

Table of Contents

Motivation
What is tensor parallelism
When do we implement tensor parallelism
Setup
Forward pass
Why split Matrix A by columns, Matrix B by rows
Backward pass (backpropagation)
Communication cost of tensor parallelism
Code

Motivation

Tensor parallelism is widely used in training and deploying massive neural networks such as LLMs across multiple GPUs. While libraries such as DeepSpeed provide out-of-the-box solutions for tensor parallelism, understanding how to implement tensor parallelism by hand may help with:

Building intuition: instead of just knowing tensor parallelism splits tensors, you will understand how the choice of splitting by "columns" in one layer and "rows" in the next is essential for the math to work out.
Understanding performance bottlenecks: You will gain an understanding why tensor parallelism is communication intensive and generally recommended to be done over high-speed interconnects such as NVLink.
Debugging: An intuitive mental model of tensor parallelism may be useful for debugging when distributed training jobs fail and the error messages are cryptic.

This article can also serve as a complementary study guide for Stanford's CS336 Lecture 7 on Parallelism. This article was written with assistance of Gemini 2.5 Pro.

What is tensor parallelism

When a model becomes too massive to fit into a single GPU's memory, tensor parallelism offers a way to split up the workload. Tensor parallelism accomplishes two goals simultaneously:

Distributes Memory Load: It enables the combining of the memory of several GPUs to store a single, massive layer that would be too large for any one device.
Parallelizes Computation: Each GPU performs a smaller portion of the matrix multiplication using its local shard of the weight matrix. This allows the GPUs to work on the same forward or backward pass of a single layer at the same time, speeding up the computation.

When do we implement tensor parallelism

Tensor parallelism comes at a significant cost of communication overhead. This article will demonstrate why tensor parallelism requires a high-volume, synchronous communication step (All-Reduce) twice for every single block in the model—once during the forward pass and once during the backward pass. This creates a fundamental trade-off:

Pipeline Parallelism: In pipeline parallelism, you split groups of layers (stages) across GPUs. Communication is less frequent, happening only at the boundaries between stages. However, pipeline parallelism introduces a "bubble" of idle time at the beginning and end of each batch as the pipeline fills and drains, reducing hardware utilization.
Tensor Parallelism: In tensor parallelism, you split the work within each layer. This helps to keep all GPUs perfectly synchronized and busy, completely eliminating the pipeline "bubble." But the price for this high utilization is constant, layer-by-layer communication across all participating GPUs.

The intense communication pattern of tensor parallelism over pipeline parallelism leads to a rule of thumb for designing large-scale training setups:

Use Tensor Parallelism for Intra-Node Scaling: When you have a group of GPUs (e.g., the 8 A100s in a DGX node) connected with a high-speed backplane, tensor parallelism is the ideal way to scale up to handle massive layers.
Use Pipeline and Data Parallelism for Inter-Node Scaling: To scale beyond a single node, you typically combine tensor parallelism with pipeline parallelism, where each stage of the pipeline is a tensor-parallel group of GPUs. This minimizes the slow communication across nodes to only what is necessary between pipeline stages.

Setup

We will walk through the forward pass of a two-layer Multi-Layer Perceptron (MLP) block found in a Transformer, which consists of two main operations:

Y = GeLU(X * A)
Z = Dropout(Y * B)

We'll assume we have two GPUs, GPU 1 and GPU 2.

Source: Stanford CS 336 Language Modeling from Scratch, Lecture 7: Parallelism 1, 58 min 12 sec

Forward pass

We first begin with the forward pass, and the next section will walk through the backpropagation.

The forward pass in a neural network is the process of taking an input and propagating it through the layers of the model to generate a final output. In the context of tensor parallelism, this process is modified because the model's large weight matrices are too big for a single GPU and are instead sharded, or split, across multiple GPUs.

To execute a matrix multiplication, each GPU performs a smaller calculation in parallel using its local shard of the weights. This often follows a specific pattern where one layer's weights are split by columns (requiring no communication after the multiplication) and the next layer's weights are split by rows. This row-parallel approach produces partial results on each GPU that must then be aggregated through a collective communication operation, such as an All-Reduce, to ensure the final activation tensor is mathematically equivalent to the result that would have been computed on a single, monolithic GPU.

Step 1: Initial Setup - Defining and Splitting Matrices

First, let's define our input matrix X and the weight matrices A and B for our two linear layers. We will use small, simple numbers for clarity.

Let our input X be a 2x4 matrix:

\begin{pmatrix} 1 & 2 & 3 & 4 \newline 5 & 6 & 7 & 8 \end{pmatrix}

Let the first weight matrix A be a 4x6 matrix:

\begin{pmatrix} 1 & 2 & 3 & 4 & 5 & 6 \newline 7 & 8 & 9 & 1 & 2 & 3 \newline 4 & 5 & 6 & 7 & 8 & 9 \newline 1 & 2 & 3 & 4 & 5 & 6 \end{pmatrix}

Let the second weight matrix B be a 6x4 matrix:

\begin{pmatrix} 6 & 5 & 4 & 3 \newline 2 & 1 & 7 & 8 \newline 9 & 8 & 6 & 5 \newline 4 & 3 & 2 & 1 \newline 1 & 2 & 3 & 4 \newline 5 & 6 & 7 & 8 \end{pmatrix}

Now, we split these matrices across our two GPUs.

Matrix A is split vertically (by columns). GPU 1 gets the first 3 columns (A1), and GPU 2 gets the last 3 columns (A2). This is known as column parallelism.

$A_1 = \begin{pmatrix} 1 & 2 & 3 \newline 7 & 8 & 9 \newline 4 & 5 & 6 \newline 1 & 2 & 3 \end{pmatrix}$ (on GPU 1)

$A_2 = \begin{pmatrix} 4 & 5 & 6 \newline 1 & 2 & 3 \newline 7 & 8 & 9 \newline 4 & 5 & 6 \end{pmatrix}$ (on GPU 2)

Matrix B is split horizontally (by rows). GPU 1 gets the first 3 rows (B1), and GPU 2 gets the last 3 rows (B2). This is known as row parallelism.

$B_1 = \begin{pmatrix} 6 & 5 & 4 & 3 \newline 2 & 1 & 7 & 8 \newline 9 & 8 & 6 & 5 \end{pmatrix}$ (on GPU 1)

$B_2 = \begin{pmatrix} 4 & 3 & 2 & 1 \newline 1 & 2 & 3 & 4 \newline 5 & 6 & 7 & 8 \end{pmatrix}$ (on GPU 2)

We split Matrix A by columns, and Matrix B by rows, so as to minimise communication costs incurred by the all-reduce operation. I will go into more detail at the next section on why we split Matrix A by columns, and Matrix B by rows.

Step 2: Forward Pass for `Y = GeLU(X * A)`

This operation is parallelized using the column-sharded matrix A. The input X is broadcast to both GPUs.

On GPU 1: Calculate `Y_intermediate_1 = X * A1`

We perform matrix multiplication. The output will be a 2x3 matrix.

Y_{intermediate_1} = \begin{pmatrix} 1 & 2 & 3 & 4 \newline 5 & 6 & 7 & 8 \end{pmatrix} \times \begin{pmatrix} 1 & 2 & 3 \newline 7 & 8 & 9 \newline 4 & 5 & 6 \newline 1 & 2 & 3 \end{pmatrix}

Calculating each element:

Y_intermediate_1 = (1*1) + (2*7) + (3*4) + (4*1) = 1 + 14 + 12 + 4 = 31
Y_intermediate_1 = (1*2) + (2*8) + (3*5) + (4*2) = 2 + 16 + 15 + 8 = 41
Y_intermediate_1 = (1*3) + (2*9) + (3*6) + (4*3) = 3 + 18 + 18 + 12 = 51
Y_intermediate_1 = (5*1) + (6*7) + (7*4) + (8*1) = 5 + 42 + 28 + 8 = 83
Y_intermediate_1 = (5*2) + (6*8) + (7*5) + (8*2) = 10 + 48 + 35 + 16 = 109
Y_intermediate_1 = (5*3) + (6*9) + (7*6) + (8*3) = 15 + 54 + 42 + 24 = 135

So, the result of the matrix multiplication on GPU 1 is:

Y_{intermediate_1} = \begin{pmatrix} 31 & 41 & 51 \newline 83 & 109 & 135 \end{pmatrix}

On GPU 2: Calculate `Y_intermediate_2 = X * A2`

Similarly, we perform the multiplication on GPU 2.

Y_{intermediate_2} = \begin{pmatrix} 1 & 2 & 3 & 4 \newline 5 & 6 & 7 & 8 \end{pmatrix} \times \begin{pmatrix} 4 & 5 & 6 \newline 1 & 2 & 3 \newline 7 & 8 & 9 \newline 4 & 5 & 6 \end{pmatrix}

Calculating each element:

Y_intermediate_2 = (1*4) + (2*1) + (3*7) + (4*4) = 4 + 2 + 21 + 16 = 43
Y_intermediate_2 = (1*5) + (2*2) + (3*8) + (4*5) = 5 + 4 + 24 + 20 = 53
Y_intermediate_2 = (1*6) + (2*3) + (3*9) + (4*6) = 6 + 6 + 27 + 24 = 63
Y_intermediate_2 = (5*4) + (6*1) + (7*7) + (8*4) = 20 + 6 + 49 + 32 = 107
Y_intermediate_2 = (5*5) + (6*2) + (7*8) + (8*5) = 25 + 12 + 56 + 40 = 133
Y_intermediate_2 = (5*6) + (6*3) + (7*9) + (8*6) = 30 + 18 + 63 + 48 = 159

So, the result of the matrix multiplication on GPU 2 is:

Y_{intermediate_2} = \begin{pmatrix} 43 & 53 & 63 \newline 107 & 133 & 159 \end{pmatrix}

Apply GeLU Activation Function

Next, the Gaussian Error Linear Unit (GeLU) activation function is applied element-wise to the intermediate results on each GPU. The GeLU function is defined as:

\cdot \Phi(x)

where $Φ(x)\Phi(x)$ is the Cumulative Distribution Function (CDF) for the standard normal distribution. A common and precise formula is:

\cdot x \cdot (1 + erf(x / \sqrt{2}))

For simplicity in this manual example, let's assume GeLU(x) ≈ x for large positive x and GeLU(x) ≈ 0 for large negative x. Since all our numbers are large and positive, we'll approximate GeLU(x) = x. In a real scenario, the exact formula would be applied.

(Note: The erf function produces an "S"-shaped sigmoid curve that passes through the origin. As the input x gets very large and positive, erf(x) gets very close to 1. As the input gets very large and negative, erf(x) gets very close to -1.)

So, the outputs on each GPU are:

$Y_1 = GeLU(Y_{intermediate_1}) \approx \begin{pmatrix} 31 & 41 & 51 \newline 83 & 109 & 135 \end{pmatrix}$ (on GPU 1)

$Y_2 = GeLU(Y_{intermediate_2}) \approx \begin{pmatrix} 43 & 53 & 63 \newline 107 & 133 & 159 \end{pmatrix}$ (on GPU 2)

These two matrices, Y1 and Y2, represent the sharded output of the first layer. They remain on their respective GPUs and become the input for the next step.

Step 3: Forward Pass for `Z = Dropout(Y * B)`

This operation is parallelized using the row-sharded matrix B. The inputs Y1 and Y2 are already on the correct GPUs.

On GPU 1: Calculate `Z_partial_1 = Y1 * B1`

We multiply the local activation Y1 with the local weight shard B1.

Z_{partial_1} = \begin{pmatrix} 31 & 41 & 51 \newline 83 & 109 & 135 \end{pmatrix} \times \begin{pmatrix} 6 & 5 & 4 & 3 \newline 2 & 1 & 7 & 8 \newline 9 & 8 & 6 & 5 \end{pmatrix}

Calculating each element:

Z_partial_1 = (31*6) + (41*2) + (51*9) = 186 + 82 + 459 = 727
Z_partial_1 = (31*5) + (41*1) + (51*8) = 155 + 41 + 408 = 604
Z_partial_1 = (31*4) + (41*7) + (51*6) = 124 + 287 + 306 = 717
Z_partial_1 = (31*3) + (41*8) + (51*5) = 93 + 328 + 255 = 676
Z_partial_1 = (83*6) + (109*2) + (135*9) = 498 + 218 + 1215 = 1931
Z_partial_1 = (83*5) + (109*1) + (135*8) = 415 + 109 + 1080 = 1604
Z_partial_1 = (83*4) + (109*7) + (135*6) = 332 + 763 + 810 = 1905
Z_partial_1 = (83*3) + (109*8) + (135*5) = 249 + 872 + 675 = 1796

Result on GPU 1:

Z_{partial_1} = \begin{pmatrix} 727 & 604 & 717 & 676 \newline 1931 & 1604 & 1905 & 1796 \end{pmatrix}

On GPU 2: Calculate `Z_partial_2 = Y2 * B2`

Simultaneously, on GPU 2:

Z_{partial_2} = \begin{pmatrix} 43 & 53 & 63 \newline 107 & 133 & 159 \end{pmatrix} \times \begin{pmatrix} 4 & 3 & 2 & 1 \newline 1 & 2 & 3 & 4 \newline 5 & 6 & 7 & 8 \end{pmatrix}

Calculating each element:

Z_partial_2 = (43*4) + (53*1) + (63*5) = 172 + 53 + 315 = 540
Z_partial_2 = (43*3) + (53*2) + (63*6) = 129 + 106 + 378 = 613
Z_partial_2 = (43*2) + (53*3) + (63*7) = 86 + 159 + 441 = 686
Z_partial_2 = (43*1) + (53*4) + (63*8) = 43 + 212 + 504 = 759
Z_partial_2 = (107*4) + (133*1) + (159*5) = 428 + 133 + 795 = 1356
Z_partial_2 = (107*3) + (133*2) + (159*6) = 321 + 266 + 954 = 1541
Z_partial_2 = (107*2) + (133*3) + (159*7) = 214 + 399 + 1113 = 1726
Z_partial_2 = (107*1) + (133*4) + (159*8) = 107 + 532 + 1272 = 1911

Result on GPU 2:

Z_{partial_2} = \begin{pmatrix} 540 & 613 & 686 & 759 \newline 1356 & 1541 & 1726 & 1911 \end{pmatrix}

All-Reduce Communication Step

The logic behind splitting B by rows is that the final output Z is the sum of the partial results from each GPU. This requires a communication step called an All-Reduce. Each GPU sends its partial result to the others and receives their partial results, summing them all together.

Z_{unactivated} = Z_{partial_1} + Z_{partial_2}

Z_{unactivated} = \begin{pmatrix} 727 & 604 & 717 & 676 \newline 1931 & 1604 & 1905 & 1796 \end{pmatrix} + \begin{pmatrix} 540 & 613 & 686 & 759 \newline 1356 & 1541 & 1726 & 1911 \end{pmatrix}

Adding each element:

Z_{unactivated} = \begin{pmatrix} 727+540 & 604+613 & 717+686 & 676+759 \newline 1931+1356 & 1604+1541 & 1905+1726 & 1796+1911 \end{pmatrix}

Z_{unactivated} = \begin{pmatrix} 1267 & 1217 & 1403 & 1435 \newline 3287 & 3145 & 3631 & 3707 \end{pmatrix}

After the All-Reduce, the complete Z_unactivated matrix is present on both GPUs.

Apply Dropout

Dropout is a regularization technique used during training to prevent overfitting. It randomly sets a fraction of the output elements to zero at each update step.

Let's assume a dropout rate of 50%. We create a "dropout mask" with the same dimensions as Z_unactivated, where each position is randomly 0 or 1.

Example dropout mask:

\begin{pmatrix} 1 & 0 & 0 & 1 \newline 0 & 1 & 1 & 0 \end{pmatrix}

We multiply our Z_unactivated by this mask:

Z_{masked} = Z_{unactivated} \odot Mask = \begin{pmatrix} 1267 \cdot 1 & 1217 \cdot 0 & 1403 \cdot 0 & 1435 \cdot 1 \newline 3287 \cdot 0 & 3145 \cdot 1 & 3631 \cdot 1 & 3707 \cdot 0 \end{pmatrix}

Z_{masked} = \begin{pmatrix} 1267 & 0 & 0 & 1435 \newline 0 & 3145 & 3631 & 0 \end{pmatrix}

Finally, to ensure the expected sum of outputs remains the same, the remaining non-zero elements are scaled up by a factor of 1 / (1 - dropout_rate). In our case, this is 1 / (1 - 0.5) = 2. This is called "inverted dropout".

Z_{masked} \times 2 = \begin{pmatrix} 2534 & 0 & 0 & 2870 \newline 0 & 6290 & 7262 & 0 \end{pmatrix}

This final matrix Z is the output of our tensor-parallel multi-layer perceptron block.

Summary of the forward pass

Start: We have an input X and two GPUs.
Layer 1 (Column Parallelism):
- The weight matrix A is split into columns (A1, A2).
- Each GPU calculates a part of the output: X * A1 on GPU 1 and X * A2 on GPU 2.
- An activation function (GeLU) is applied independently on each GPU to these partial results (Y1, Y2). No communication is needed.
Layer 2 (Row Parallelism):
- The weight matrix B is split into rows (B1, B2).
- Each GPU multiplies its local activation from the previous step with its local weight shard: Y1 * B1 on GPU 1 and Y2 * B2 on GPU 2.
- An All-Reduce operation sums the partial results from all GPUs. This is the critical communication step.
- The final operation (Dropout) is applied to the complete, summed tensor, resulting in the final output Z.

Why split Matrix A by columns, Matrix B by rows

The goal is to sequence the operations to ensure the expensive All-Reduce step (where GPUs must synchronize and sum their results) happens as infrequently as possible. To understand this, let's define the two operations and their communication costs.

1. Column Parallelism (What we do to `A`)

When we split a weight matrix A by columns, the input X is broadcast to all GPUs. Each GPU then calculates a vertical "slice" of the output.

Mathematics: Y = X * [A1, A2] = [X * A1, X * A2] = [Y1, Y2]
Result: The output Y is naturally sharded (split) across the GPUs. Y1 is on GPU 1 and Y2 is on GPU 2.
Communication Cost (Forward Pass): Zero. No communication is needed to produce the sharded output Y. Each GPU works independently. The output is left "as is" on each device, ready for the next step.

2. Row Parallelism (What we do to `B`)

When we split a weight matrix B by rows, the input must already be sharded in a corresponding way (which is exactly what Column Parallelism just produced!). Each GPU multiplies its local input shard with its local weight shard.

Mathematics: Z = Y * B = [Y1, Y2] * [B1; B2] = (Y1 * B1) + (Y2 * B2) (Note: ; denotes stacking rows)
Result: Each GPU calculates a partial result. To get the final, correct output Z, these partial results must be summed together.
Communication Cost (Forward Pass): One All-Reduce operation. This is the expensive step where Z_partial_1 and Z_partial_2 are added across all GPUs.

Analyzing the Standard `Column -> Row` Pattern

Now, let's see how these two pieces fit together in our MLP block:

First Layer: Y = GeLU(X * A) (Column Parallelism)
- We split A by columns.
- The multiplication X * A produces a sharded activation [Y1, Y2] without any need for communication.
- Forward Communication Cost: 0
Second Layer: Z = Dropout(Y * B) (Row Parallelism)
- We split B by rows.
- The sharded activation [Y1, Y2] from the previous step is the perfect input for this layer.
- The multiplication Y * B produces partial results Z_partial_1 and Z_partial_2.
- To get the final output Z, we must perform an All-Reduce to sum the partial results.
- Forward Communication Cost: 1 All-Reduce

By sequencing the operations this way (Column -> Row), we have successfully computed the output for a full two-layer block with only one communication step.

What Happens If We Switch? (Row -> Column)

Let's imagine we tried to do it the other way around.

First Layer: Y = GeLU(X * A) (Row Parallelism)
- We split A by rows (A1, A2).
- For the math X * A to work, we would also have to split the input X by columns (X1, X2).
- The operation becomes (X1 * A1) + (X2 * A2).
- This produces partial results which must be summed.
- Forward Communication Cost: 1 All-Reduce. The output Y is now a complete (un-sharded) tensor, identical on all GPUs.
Second Layer: Z = Dropout(Y * B) (Column Parallelism)
- We split B by columns (B1, B2).
- The input Y is a complete tensor from the previous step.
- The operation Y * [B1, B2] produces a sharded output [Z1, Z2] without communication.
- Forward Communication Cost: 0

Why column then row?

As you can see, both patterns (Column -> Row and Row -> Column) result in the exact same amount of communication: one All-Reduce per two-layer block.

So why is the Column -> Row pattern the standard convention used in frameworks like Megatron-LM?

The reason is modularity and efficiency in the computational graph. The standard Transformer block consists of an Attention block followed by an MLP block. The Column -> Row pattern creates a perfect compositional unit:

The MLP block takes a complete tensor as input (like X).
It performs Column Parallel -> Row Parallel.
It outputs a complete tensor (Z, after the All-Reduce).

This complete tensor Z can then be fed directly into the next Attention block, which will also start with a column-parallel operation. This creates a clean, predictable, and efficient computational flow where the "sharded" nature of the activations is kept internal to the block, and communication is neatly packaged.

In short, the Column -> Row sequence is a deliberate design pattern that cleverly delays the necessary communication to the very last step within a block, creating a modular and efficient structure for building large neural networks.

Backward pass (backpropagation)

The goal of the backward pass is to calculate the gradients of the loss function with respect to all the learnable parameters (the weights A and B) and the input (X). These gradients tell us how to adjust the weights to reduce the loss.

The process starts with the gradient of the loss with respect to the final output, which we'll call dL/dZ or simply dZ. This initial gradient propagates backward through the network. For this example, we'll assume the process has already begun and the initial gradient dZ is a 2x4 matrix of ones.

Initial Gradient (assumed):

\begin{pmatrix} 1 & 1 & 1 & 1 \newline 1 & 1 & 1 & 1 \end{pmatrix}

Step 1: Backward Pass through Dropout

The first step is to backpropagate through the Dropout layer. The gradient is only allowed to pass through the neurons that were not "dropped out" during the forward pass. This means we apply the same dropout mask and scaling factor.

Mathematics

The gradient with respect to the pre-dropout output (Z_unactivated) is calculated by element-wise multiplying the incoming gradient dZ with the dropout mask and scaling it by the same factor used in the forward pass.

∂L∂Zunactivated=∂L∂Z⊙Mask×11−dropout_rate \frac{\partial L}{\partial Z_{unactivated}} = \frac{\partial L}{\partial Z} \odot Mask \times \frac{1}{1 - \text{dropout\_rate}}

Calculation

From our forward pass, our mask and scaling factor were:
$\begin{pmatrix} 1 & 0 & 0 & 1 \newline 0 & 1 & 1 & 0 \end{pmatrix}$ , Scale = 2

The gradient dZ is passed then through this inverted dropout layer:

dZ_{masked} = dZ \odot Mask = \begin{pmatrix} 1 & 1 & 1 & 1 \newline 1 & 1 & 1 & 1 \end{pmatrix} \odot \begin{pmatrix} 1 & 0 & 0 & 1 \newline 0 & 1 & 1 & 0 \end{pmatrix} = \begin{pmatrix} 1 & 0 & 0 & 1 \newline 0 & 1 & 1 & 0 \end{pmatrix}

dZ_{unactivated} = dZ_{masked} \times 2 = \begin{pmatrix} 2 & 0 & 0 & 2 \newline 0 & 2 & 2 & 0 \end{pmatrix}

This gradient, dZ_unactivated, is now available on both GPUs, as the All-Reduce in the forward pass made the output Z_unactivated identical on both.

Step 2: Backward Pass for `Z = Dropout(Y * B)`

Now we backpropagate through the second linear layer. This involves calculating the gradients with respect to the weights (dB) and the activations from the previous layer (dY).

Mathematics for Gradients `dY` and `dB`

Given the operation Z_unactivated = Y * B, the gradients are:

Gradient w.r.t. Y: $∂L∂Y=∂L∂Zunactivated⋅BT\frac{\partial L}{\partial Y} = \frac{\partial L}{\partial Z_{unactivated}} \cdot B^T$
Gradient w.r.t. B: $∂L∂B=YT⋅∂L∂Zunactivated\frac{\partial L}{\partial B} = Y^T \cdot \frac{\partial L}{\partial Z_{unactivated}}$

In our tensor parallel setup (row parallelism), Y = [Y1, Y2] and B is split into B1 and B2. The forward pass involved local multiplications (Y1*B1, Y2*B2) followed by an All-Reduce. The backward pass reverses this.

For dY (Activation Gradients): The gradient dZ_unactivated is multiplied by the transpose of the local weight shard. This is a local computation on each GPU and requires no communication.
- On GPU 1: $dY1=dZunactivated⋅B1TdY_1 = dZ_{unactivated} \cdot B_1^T$
- On GPU 2: $dY2=dZunactivated⋅B2TdY_2 = dZ_{unactivated} \cdot B_2^T$
For dB (Weight Gradients): The transpose of the local activation Y is multiplied by the incoming gradient.
- On GPU 1: $dB1=Y1T⋅dZunactivateddB_1 = Y_1^T \cdot dZ_{unactivated}$
- On GPU 2: $dB2=Y2T⋅dZunactivateddB_2 = Y_2^T \cdot dZ_{unactivated}$

Calculation for `dY`

On GPU 1: Calculate dY1 = dZ_unactivated * B1^T

B_1^T = \begin{pmatrix} 6 & 2 & 9 \newline 5 & 1 & 8 \newline 4 & 7 & 6 \newline 3 & 8 & 5 \end{pmatrix}

dY_1 = \begin{pmatrix} 2 & 0 & 0 & 2 \newline 0 & 2 & 2 & 0 \end{pmatrix} \times \begin{pmatrix} 6 & 2 & 9 \newline 5 & 1 & 8 \newline 4 & 7 & 6 \newline 3 & 8 & 5 \end{pmatrix} = \begin{pmatrix} (2*6+2*3) & (2*2+2*8) & (2*9+2*5) \newline (2*5+2*4) & (2*1+2*7) & (2*8+2*6) \end{pmatrix} = \begin{pmatrix} 18 & 20 & 28 \newline 18 & 16 & 28 \end{pmatrix}

On GPU 2: Calculate dY2 = dZ_unactivated * B2^T

B_2^T = \begin{pmatrix} 4 & 1 & 5 \newline 3 & 2 & 6 \newline 2 & 3 & 7 \newline 1 & 4 & 8 \end{pmatrix}

dY_2 = \begin{pmatrix} 2 & 0 & 0 & 2 \newline 0 & 2 & 2 & 0 \end{pmatrix} \times \begin{pmatrix} 4 & 1 & 5 \newline 3 & 2 & 6 \newline 2 & 3 & 7 \newline 1 & 4 & 8 \end{pmatrix} = \begin{pmatrix} (2*4+2*1) & (2*1+2*4) & (2*5+2*8) \newline (2*3+2*2) & (2*2+2*3) & (2*6+2*7) \end{pmatrix} = \begin{pmatrix} 10 & 10 & 26 \newline 10 & 10 & 26 \end{pmatrix}

Calculation for `dB`

On GPU 1: Calculate dB1 = Y1^T * dZ_unactivated

Y_1^T = \begin{pmatrix} 31 & 83 \newline 41 & 109 \newline 51 & 135 \end{pmatrix}

dB_1 = \begin{pmatrix} 31 & 83 \newline 41 & 109 \newline 51 & 135 \end{pmatrix} \times \begin{pmatrix} 2 & 0 & 0 & 2 \newline 0 & 2 & 2 & 0 \end{pmatrix} = \begin{pmatrix} (31*2) & (83*2) & (83*2) & (31*2) \newline (41*2) & (109*2) & (109*2) & (41*2) \newline (51*2) & (135*2) & (135*2) & (51*2) \end{pmatrix} = \begin{pmatrix} 62 & 166 & 166 & 62 \newline 82 & 218 & 218 & 82 \newline 102 & 270 & 270 & 102 \end{pmatrix}

On GPU 2: Calculate dB2 = Y2^T * dZ_unactivated

Y_2^T = \begin{pmatrix} 43 & 107 \newline 53 & 133 \newline 63 & 159 \end{pmatrix}

dB_2 = \begin{pmatrix} 43 & 107 \newline 53 & 133 \newline 63 & 159 \end{pmatrix} \times \begin{pmatrix} 2 & 0 & 0 & 2 \newline 0 & 2 & 2 & 0 \end{pmatrix} = \begin{pmatrix} (43*2) & (107*2) & (107*2) & (43*2) \newline (53*2) & (133*2) & (133*2) & (53*2) \newline (63*2) & (159*2) & (159*2) & (63*2) \end{pmatrix} = \begin{pmatrix} 86 & 214 & 214 & 86 \newline 106 & 266 & 266 & 106 \newline 126 & 318 & 318 & 126 \end{pmatrix}

Step 3: Backward Pass through GeLU Activation

Next, the gradients dY1 and dY2 are passed backward through the GeLU activation function.

Mathematics

The chain rule dictates that we multiply the incoming gradient by the derivative of the activation function evaluated at its original input.

\frac{\partial L}{\partial Y_{intermediate}} = \frac{\partial L}{\partial Y} \odot GeLU'(Y_{intermediate})

The derivative of GeLU is $\cdot (1 + erf(x/\sqrt{2})) + \frac{x}{\sqrt{2\pi}}e^{-x^2/2}$ . As in the forward pass, for the large positive numbers in our example, this derivative is very close to 1. We will use this approximation for simplicity.

\approx 1

Calculation

We simply pass the gradients through, as multiplying by 1 does not change them.
On GPU 1:

dY_{intermediate_1} = dY_1 \odot 1 = \begin{pmatrix} 18 & 20 & 28 \newline 18 & 16 & 28 \end{pmatrix}

On GPU 2:

dY_{intermediate_2} = dY_2 \odot 1 = \begin{pmatrix} 10 & 10 & 26 \newline 10 & 10 & 26 \end{pmatrix}

Step 4: Backward Pass for `Y = GeLU(X * A)`

Finally, we backpropagate through the first linear layer. This computes the gradients dA (for the weights) and dX (for the input).

Mathematics for Gradients `dX` and `dA`

Given Y_intermediate = X * A, the gradients are:

Gradient w.r.t. X: $∂L∂X=∂L∂Yintermediate⋅AT\frac{\partial L}{\partial X} = \frac{\partial L}{\partial Y_{intermediate}} \cdot A^T$
Gradient w.r.t. A: $∂L∂A=XT⋅∂L∂Yintermediate\frac{\partial L}{\partial A} = X^T \cdot \frac{\partial L}{\partial Y_{intermediate}}$

In our tensor parallel setup (column parallelism), the input X was broadcast, and the output Y_intermediate was sharded (Y_intermediate_1, Y_intermediate_2).

For dA (Weight Gradients): The calculation is local.
- On GPU 1: $dA1=XT⋅dYintermediate1dA_1 = X^T \cdot dY_{intermediate_1}$
- On GPU 2: $dA2=XT⋅dYintermediate2dA_2 = X^T \cdot dY_{intermediate_2}$
For dX (Input Gradients): Each GPU calculates a partial gradient for dX. These partial gradients must be summed together using an All-Reduce operation.
- On GPU 1: $dXpartial1=dYintermediate1⋅A1TdX_{partial_1} = dY_{intermediate_1} \cdot A_1^T$
- On GPU 2: $dXpartial2=dYintermediate2⋅A2TdX_{partial_2} = dY_{intermediate_2} \cdot A_2^T$
- Final Gradient: $dX = dX_{partial_1} + dX_{partial_2}$

Calculation for `dA`

On GPU 1: Calculate dA1 = X^T * dY_intermediate_1

X^T = \begin{pmatrix} 1 & 5 \newline 2 & 6 \newline 3 & 7 \newline 4 & 8 \end{pmatrix}

dA_1 = \begin{pmatrix} 1 & 5 \newline 2 & 6 \newline 3 & 7 \newline 4 & 8 \end{pmatrix} \times \begin{pmatrix} 18 & 20 & 28 \newline 18 & 16 & 28 \end{pmatrix} = \begin{pmatrix} (18+90) & (20+80) & (28+140) \newline (36+108) & (40+96) & (56+168) \newline (54+126) & (60+112) & (84+196) \newline (72+144) & (80+128) & (112+224) \end{pmatrix} = \begin{pmatrix} 108 & 100 & 168 \newline 144 & 136 & 224 \newline 180 & 172 & 280 \newline 216 & 208 & 336 \end{pmatrix}

On GPU 2: Calculate dA2 = X^T * dY_intermediate_2

dA_2 = \begin{pmatrix} 1 & 5 \newline 2 & 6 \newline 3 & 7 \newline 4 & 8 \end{pmatrix} \times \begin{pmatrix} 10 & 10 & 26 \newline 10 & 10 & 26 \end{pmatrix} = \begin{pmatrix} (10+50) & (10+50) & (26+130) \newline (20+60) & (20+60) & (52+156) \newline (30+70) & (30+70) & (78+182) \newline (40+80) & (40+80) & (104+208) \end{pmatrix} = \begin{pmatrix} 60 & 60 & 156 \newline 80 & 80 & 208 \newline 100 & 100 & 260 \newline 120 & 120 & 312 \end{pmatrix}

Calculation for `dX` (with All-Reduce)

On GPU 1: Calculate dX_partial_1 = dY_intermediate_1 * A1^T

A_1^T = \begin{pmatrix} 1 & 7 & 4 & 1 \newline 2 & 8 & 5 & 2 \newline 3 & 9 & 6 & 3 \end{pmatrix}

dX_{partial_1} = \begin{pmatrix} 18 & 20 & 28 \newline 18 & 16 & 28 \end{pmatrix} \times \begin{pmatrix} 1 & 7 & 4 & 1 \newline 2 & 8 & 5 & 2 \newline 3 & 9 & 6 & 3 \end{pmatrix} = \begin{pmatrix} (18+40+84) & (126+160+252) & (72+100+168) & (18+40+84) \newline (18+32+84) & (126+128+252) & (72+80+168) & (18+32+84) \end{pmatrix} = \begin{pmatrix} 142 & 538 & 340 & 142 \newline 134 & 506 & 320 & 134 \end{pmatrix}

On GPU 2: Calculate dX_partial_2 = dY_intermediate_2 * A2^T

A_2^T = \begin{pmatrix} 4 & 1 & 7 & 4 \newline 5 & 2 & 8 & 5 \newline 6 & 3 & 9 & 6 \end{pmatrix}

dX_{partial_2} = \begin{pmatrix} 10 & 10 & 26 \newline 10 & 10 & 26 \end{pmatrix} \times \begin{pmatrix} 4 & 1 & 7 & 4 \newline 5 & 2 & 8 & 5 \newline 6 & 3 & 9 & 6 \end{pmatrix} = \begin{pmatrix} (40+50+156) & (10+20+78) & (70+80+234) & (40+50+156) \newline (40+50+156) & (10+20+78) & (70+80+234) & (40+50+156) \end{pmatrix} = \begin{pmatrix} 246 & 108 & 384 & 246 \newline 246 & 108 & 384 & 246 \end{pmatrix}

All-Reduce Communication Step
Finally, the partial gradients for X are summed across GPUs.

dX_{partial_1} + dX_{partial_2} = \begin{pmatrix} 142 & 538 & 340 & 142 \newline 134 & 506 & 320 & 134 \end{pmatrix} + \begin{pmatrix} 246 & 108 & 384 & 246 \newline 246 & 108 & 384 & 246 \end{pmatrix} = \begin{pmatrix} 388 & 646 & 724 & 388 \newline 380 & 614 & 704 & 380 \end{pmatrix}

This final dX is the gradient with respect to the input of the MLP block, ready to be passed further backward to the previous layer in the network. The weight gradients dA1, dA2, dB1, and dB2 are used by the optimizer to update the model's weights on each respective GPU.

The Duality of Communication in Forward vs. Backward Pass

A fundamental principle of backpropagation in tensor parallelism is that the communication patterns are reversed.

If a layer requires an All-Reduce in the forward pass, its corresponding backward pass for the input gradient requires no communication (Identity).
If a layer requires no communication (Identity) in the forward pass, its corresponding backward pass for the input gradient requires an All-Reduce.

Here is a simple table summarizing the communication for the input gradients (dX and dY):

Layer Type (Forward)	Forward Pass Communication (`f`)	Backward Pass Communication (`g`)
Column Parallelism	`Identity` (No communication)	`All-Reduce` (Sum partial results)
Row Parallelism	`All-Reduce` (Sum partial results)	`Identity` (No communication)

(Note: The gradients for the weights, dA and dB, are always local calculations and never require an All-Reduce, so we can focus on the activation gradients.)

Backpropagation Through the Standard `Column -> Row` Pattern

Now, let's trace the backward pass through our standard two-layer block, keeping the table above in mind. The process starts with dZ, the gradient from the layer above. Because the forward pass ended with an All-Reduce, Z was a complete tensor on all GPUs, which means dZ is also a complete tensor.

Backprop through Layer 2: Z = Dropout(Y * B) (Row Parallelism)
- Forward Pass required: All-Reduce.
- Backward Pass requires: Identity (No communication).
- Calculation (dY = dZ * B^T): The complete gradient dZ is multiplied by the local weight shard B^T on each GPU. This directly produces the correct sharded activation gradient dY ([dY1, dY2]) without any need for communication.
Backprop through Layer 1: Y = GeLU(X * A) (Column Parallelism)
- Forward Pass required: Identity.
- Backward Pass requires: All-Reduce.
- Calculation (dX = dY * A^T): The sharded gradient dY from the previous step is multiplied by the local weight shard A^T. This produces a partial gradient for the input, dX_partial. To get the final, complete gradient dX, these partial results must be summed across all GPUs. This is the All-Reduce step.

As you can see, the backward pass for our Column -> Row block has exactly one All-Reduce operation, just like the forward pass.

What if we used `Row -> Column`?

The result would be perfectly symmetric:

Backprop through Layer 2 (Column Parallelism): The incoming gradient dZ would be sharded. This step would require an All-Reduce to produce a complete gradient dY.
Backprop through Layer 1 (Row Parallelism): The complete gradient dY would be the input. This step would require no communication (Identity) to produce the sharded input gradient dX.

The total communication remains one All-Reduce in the forward pass and one in the backward pass.

The Preference is for Modularity, Not Cost

The total cost is the same. So why do we prefer Column -> Row?

The reason is that this pattern creates a perfectly self-contained, modular block.

Input: The block takes a complete tensor (X).
Internal State: It creates a sharded tensor (Y) internally.
Output: It performs an All-Reduce at the end and outputs a complete tensor (Z).

The backward pass mirrors this perfectly:

Input: It takes a complete gradient (dZ).
Internal State: It creates a sharded gradient (dY) internally.
Output: It performs an All-Reduce at the end and outputs a complete gradient (dX).

This means you can stack these Transformer blocks one after another without worrying about whether the tensor you're passing is complete or sharded. Each block handles its own internal sharding and communication, always presenting a clean, complete tensor at its boundaries. This makes the overall architecture of a large model much simpler to design, implement, and debug.

Communication cost of tensor parallelism

In our walkthrough, communication between GPUs did not happen after every mathematical operation. It only occurred at very specific points where partial results needed to be combined to form a complete tensor. Let's pinpoint them:

Forward Pass Communication: This happened at the end of Step 3. To calculate the final Z_unactivated, we needed to sum the partial results from each GPU.
$Z_{unactivated} = Z_{partial_1} + Z_{partial_2}$
This summation is an all-reduce operation. GPU 1 sends its Z_partial_1 and receives Z_partial_2; GPU 2 sends Z_partial_2 and receives Z_partial_1. Both then compute the sum.
Backward Pass Communication: This happened at the end of Step 4. To calculate the final input gradient dX, we needed to sum the partial gradients from each GPU.
$dX = dX_{partial_1} + dX_{partial_2}$
This summation is also an all-reduce operation.

For a single pass through this two-layer block, we perform two full all-reduce operations.

Analyzing the Communication Volume

The formula for the communication volume in tensor parallelism per layer:

\left( \frac{n_{\text{devices}} - 1}{n_{\text{devices}}} \right)

Source: Stanford CS 336 Language Modeling from Scratch, Lecture 7: Parallelism 1, 1 hour, 1 min 8 sec

b: This is the batch size. In our example, the input matrix X had 2 rows, so $b = 2$ .
s: This is the sequence length. In a simple MLP, we can consider this to be 1, so $s = 1$ .
h: This is the hidden size of the tensor being communicated. The tensor we communicated in the forward pass was Z_partial, which had a shape of (2, 4). The size of the last dimension is 4, so $h = 4$ .
n_devices: The number of GPUs, which is $ndevices=2n_{\text{devices}} = 2$ .

In short, bsh represents the total number of elements in the activation tensor being communicated.

Volume of data in one all-reduce:

In the forward pass all-reduce, the tensor being communicated is Z_partial. The number of elements in this tensor is:

Volume_{elements} = b \times s \times h = 2 \times 1 \times 4 = 8 \text{ elements}

The all-reduce operation is a collective communication where each GPU sends its 8 elements and receives 8 elements from its peer.

The term $(ndevices−1ndevices)\left( \frac{n_{\text{devices}} - 1}{n_{\text{devices}}} \right)$ in the formula characterizes the volume of a standard ring all-reduce algorithm. For our 2 GPUs, this factor is:

\left( \frac{2 - 1}{2} \right) = \frac{1}{2}

This signifies that each GPU sends and receives data chunks in a way that scales with the number of devices.

The factor of 8 in the formula often accounts for two things: 4 bytes for a standard 32-bit floating-point number and a factor of 2 because communication happens in both the forward and backward passes.

Why the communication cost of tensor parallelism is high

The communication cost of tensor parallelism is high due to the combination of volume and frequency:

High-Cost Operation: The all-reduce is a "collective" operation that is much more complex and latency-sensitive than a simple point-to-point send/receive. It requires synchronization across all GPUs in the group.
Large Data Volume: The amount of data communicated in each all-reduce is the full activation tensor (bsh), which can be very large in modern LLMs.
High Frequency: This expensive, high-volume operation is not performed once, but twice for every single transformer block in the model during a full training step (forward + backward pass).

Therefore, while tensor parallelism is excellent at keeping all GPUs utilized (eliminating the "bubble" of pipeline parallelism), it comes at the cost of constant, high-volume communication across all participating devices. This is why the general recommendation is to use tensor parallel whenever we have low-latency, high-bandwidth interconnects like NVIDIA's NVLink, which are specifically designed to handle this intense communication pattern efficiently.

Code

For those who prefer following the code, the walk through guide for both the forward and backward pass are replicated below using Python.

Setup

import torch
import torch.nn.functional as F

# --------------------------------------------------------------------------
# 1. SETUP: Define tensors from the manual example
# --------------------------------------------------------------------------
# Set requires_grad=True to track gradients
X = torch.tensor([[1, 2, 3, 4],
                  [5, 6, 7, 8]], dtype=torch.float32, requires_grad=True)

A = torch.tensor([[1, 2, 3, 4, 5, 6],
                  [7, 8, 9, 1, 2, 3],
                  [4, 5, 6, 7, 8, 9],
                  [1, 2, 3, 4, 5, 6]], dtype=torch.float32, requires_grad=True)

B = torch.tensor([[6, 5, 4, 3],
                  [2, 1, 7, 8],
                  [9, 8, 6, 5],
                  [4, 3, 2, 1],
                  [1, 2, 3, 4],
                  [5, 6, 7, 8]], dtype=torch.float32, requires_grad=True)

# NOTE: For verification, we use an identity function to match the manual
# calculation's approximation of GeLU(x) ≈ x. In a real model, you would
# use F.gelu(x).
gelu = lambda x: x

# Define the fixed dropout mask from the manual example
dropout_mask = torch.tensor([[1, 0, 0, 1],
                             [0, 1, 1, 0]], dtype=torch.float32)
dropout_rate = 0.5
dropout_scale = 1 / (1 - dropout_rate)

print("--- Initial Tensors ---")
print(f"X:\n{X}\n")
print(f"A:\n{A}\n")
print(f"B:\n{B}\n")

Output

--- Initial Tensors ---
X:
tensor([[1., 2., 3., 4.],
        [5., 6., 7., 8.]], requires_grad=True)

A:
tensor([[1., 2., 3., 4., 5., 6.],
        [7., 8., 9., 1., 2., 3.],
        [4., 5., 6., 7., 8., 9.],
        [1., 2., 3., 4., 5., 6.]], requires_grad=True)

B:
tensor([[6., 5., 4., 3.],
        [2., 1., 7., 8.],
        [9., 8., 6., 5.],
        [4., 3., 2., 1.],
        [1., 2., 3., 4.],
        [5., 6., 7., 8.]], requires_grad=True)

Forward pass

# --------------------------------------------------------------------------
# 2. FORWARD PASS
# --------------------------------------------------------------------------
print("--- FORWARD PASS ---")

# --- Simulate Tensor Parallelism on 2 GPUs ---

# Split A by columns (Column Parallelism)
# "GPU 1" gets A1, "GPU 2" gets A2
A1 = A[:, :3]
A2 = A[:, 3:]

# Split B by rows (Row Parallelism)
# "GPU 1" gets B1, "GPU 2" gets B2
B1 = B[:3, :]
B2 = B[3:, :]

# --- Layer 1: Y = GeLU(X * A) ---
# Operation is local to each GPU, no communication needed.

# "GPU 1" computes its part
Y_intermediate_1 = X @ A1
Y1 = gelu(Y_intermediate_1)

# "GPU 2" computes its part
Y_intermediate_2 = X @ A2
Y2 = gelu(Y_intermediate_2)

print(f"Y1 (on GPU 1):\n{Y1}\n")
print(f"Y2 (on GPU 2):\n{Y2}\n")

# --- Layer 2: Z = Dropout(Y * B) ---

# Local matrix multiplication on each GPU
# Note: Y is implicitly [Y1, Y2]
Z_partial_1 = Y1 @ B1
Z_partial_2 = Y2 @ B2

print(f"Z_partial_1 (on GPU 1):\n{Z_partial_1}\n")
print(f"Z_partial_2 (on GPU 2):\n{Z_partial_2}\n")

# *** All-Reduce Communication Step ***
# The partial results are summed across GPUs.
Z_unactivated = Z_partial_1 + Z_partial_2
print(f"Z_unactivated (after All-Reduce):\n{Z_unactivated}\n")

# Apply Dropout
Z_masked = Z_unactivated * dropout_mask
Z = Z_masked * dropout_scale

print(f"Final Output Z:\n{Z}\n")

Output

--- FORWARD PASS ---
Y1 (on GPU 1):
tensor([[ 31.,  41.,  51.],
        [ 83., 109., 135.]], grad_fn=<MmBackward0>)

Y2 (on GPU 2):
tensor([[ 43.,  53.,  63.],
        [107., 133., 159.]], grad_fn=<MmBackward0>)

Z_partial_1 (on GPU 1):
tensor([[ 727.,  604.,  717.,  676.],
        [1931., 1604., 1905., 1796.]], grad_fn=<MmBackward0>)

Z_partial_2 (on GPU 2):
tensor([[ 540.,  613.,  686.,  759.],
        [1356., 1541., 1726., 1911.]], grad_fn=<MmBackward0>)

Z_unactivated (after All-Reduce):
tensor([[1267., 1217., 1403., 1435.],
        [3287., 3145., 3631., 3707.]], grad_fn=<AddBackward0>)

Final Output Z:
tensor([[2534.,    0.,    0., 2870.],
        [   0., 6290., 7262.,    0.]], grad_fn=<MulBackward0>)

Backward pass

# --------------------------------------------------------------------------
# 3. BACKWARD PASS
# --------------------------------------------------------------------------
print("--- BACKWARD PASS ---")

# Define a dummy loss. Z.sum() creates an initial gradient dZ of all ones.
loss = Z.sum()
print(f"Loss (for backprop):\n{loss}\n")

# Automatically compute gradients for all tensors with requires_grad=True
loss.backward()

Output

--- BACKWARD PASS ---
Loss (for backprop):
18956.0

Verification

# --------------------------------------------------------------------------
# 4. VERIFICATION: Check gradients against manual calculations
# --------------------------------------------------------------------------
print("--- VERIFYING GRADIENTS ---\n")

# Gradient for X (dX)
# This required an All-Reduce in the backward pass
print(f"Gradient for X (dX):\n{X.grad}\n")

# Gradients for A (dA1 and dA2)
# A.grad contains the gradients for the entire A matrix. We can verify
# by slicing it into the parts corresponding to A1 and A2.
print(f"Gradient for A1 (dA1):\n{A.grad[:, :3]}\n")
print(f"Gradient for A2 (dA2):\n{A.grad[:, 3:]}\n")

# Gradients for B (dB1 and dB2)
# Similarly, B.grad contains the full gradient.
print(f"Gradient for B1 (dB1):\n{B.grad[:3, :]}\n")
print(f"Gradient for B2 (dB2):\n{B.grad[3:, :]}\n")

Output

--- VERIFYING GRADIENTS ---

Gradient for X (dX):
tensor([[388., 646., 724., 388.],
        [380., 614., 704., 380.]])

Gradient for A1 (dA1):
tensor([[108., 100., 168.],
        [144., 136., 224.],
        [180., 172., 280.],
        [216., 208., 336.]])

Gradient for A2 (dA2):
tensor([[ 60.,  60., 156.],
        [ 80.,  80., 208.],
        [100., 100., 260.],
        [120., 120., 312.]])

Gradient for B1 (dB1):
tensor([[ 62., 166., 166.,  62.],
        [ 82., 218., 218.,  82.],
        [102., 270., 270., 102.]])

Gradient for B2 (dB2):
tensor([[ 86., 214., 214.,  86.],
        [106., 266., 266., 106.],
        [126., 318., 318., 126.]])

Routing and balancing losses with Mixture of Experts

Lewis Won — Sat, 16 Aug 2025 14:11:16 +0000

Table of Contents

Motivation
What are mixture of experts
Why study mixture of experts
Routing taxonomy
Top-k routing strategy
Deeper dive into the math of "token chooses expert" strategy
Heuristic balancing losses

Motivation

This article explains by hand how to (i) route tokens and (ii balance losses in routing when training Mixture of Experts (MoEs) large language models.

For those following the Stanford CS 336 class, this could be used as a study guide that complements Lecture 4 on Mixture of Experts.

This article does not seek to carry out a comprehensive review of the entire MoEs literature. Hence, I will only focus on the most popular variant of routing, i.e. (i) tokens choose experts where k = 2; and (ii) heuristic balancing losses.

This article was written with the assistance of Google Gemini 2.5.

What are mixture of experts

MoEs is a machine learning technique that divides a complex problem among multiple specialized models, known as "experts". A "gating network" or "router" then selects the most suitable expert or combination of experts to handle a given input. (Fun fact: sparse expert models are a thirty-year old concept.)

The core idea behind MoEs is to replace dense feed-forward network (FFN) layers with sparse MoEs layers. In a traditional dense model, the entire network is activated to process any input. In contrast, MoE models use conditional computation, meaning only a select few "expert" sub-networks are activated for each input token.

In the context of large language models, the more popular implementation is to replace multi-layer perceptron (MLP) with mixture of expert layer. The diagram below compares a dense model with a mixture of expert model. The blue FFN layers are the MLP.

Source: A Review of Sparse Expert Models in Deep Learning by Fedus et. al., 2022.

Note that there are also implementations of LLMs that split the attention heads in the self-attention mechanism (the red layer) as experts. However these implementations are not common because they are difficult to train consistently.

Why study mixture of experts

MoEs are getting popular because:

Selective activation, or sparsity, allow models to have a significantly larger number of parameters without a proportional increase in computational cost during pre-training and inference.
Fedus et. al. showed that models with more experts have lower test loss and get better perplexity faster.

Source: Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity by Fedus et. al., 2022

MoEs also provide a natural way to parallelise training and inference at the experts level, i.e. by putting each FFN on a different device.

Routing taxonomy

Broadly, there are four major variants of routing tokens to experts:

Top-k: pick the top-k most highly activated experts to process the token. This strategy is used in models from Qwen, DeepSeek, Mixtral, Grok etc.

Hashing: Route token based on its hash value. It has low compute cost and is relatively easier to implement as compared to the other routing strategies.

Reinforcement learning to learn routes. Uncommon nowadays due to significantly higher compute costs, and contributes to instability of training.

Solve a matching problem, such as linear assignment problems.

Source of diagrams: A Review of Sparse Expert Models in Deep Learning by Fedus et. al., 2022.

I will focus only on the routing strategy top-k as it is the most popular amongst modern LLMs.

Top-k routing strategy

The "choose top-k" strategy has three main variants:

1. Token chooses expert
2. Expert chooses token
3. Global routing via optimization

Source: A Review of Sparse Expert Models in Deep Learning by Fedus et. al., 2022.

Illustration of routing by hand

The core of the routing process is a matrix of weights, often generated by a gating network, which represents the affinity between each token and each expert. Let us consider a scenario with 4 tokens (T1, T2, T3, T4) and 4 experts (E1, E2, E3, E4), and the routing mechanism $S$ contains the scores for each token-expert pair:

\begin{pmatrix} & T1 & T2 & T3 & T4 \newline E1 & 3.1 & 1.2 & 0.5 & 1.8 \newline E2 & 0.8 & 2.5 & 1.9 & 1.1 \newline E3 & 2.2 & 0.7 & 3.5 & 0.4 \newline E4 & 1.5 & 2.8 & 1.3 & 3.2 \end{pmatrix}

The value $S_{ij}$ represents the score for sending token $j$ to expert $i$ . A higher score indicates a better fit.

We now go through each of the 3 routing strategies, assuming top-k = 2 as it is also the more popular strategy for MoE LLMs.

1. Token Chooses Expert

In this method, each token independently selects the two experts with the highest affinity scores. This approach allows each token to be processed by multiple specialized experts.

Calculation:

For each token column, we identify the two experts (rows) with the highest scores.

T1: The top two scores are 3.1 (E1) and 2.2 (E3).
T2: The top two scores are 2.8 (E4) and 2.5 (E2).
T3: The top two scores are 3.5 (E3) and 1.9 (E2).
T4: The top two scores are 3.2 (E4) and 1.8 (E1).

Resulting Expert Load:

We tally the tokens assigned to each expert:

E1 processes: T1, T4
E2 processes: T2, T3
E3 processes: T1, T3
E4 processes: T2, T4

With k=2, each expert is assigned exactly two tokens, resulting in a balanced load for this specific example. However, token-choice routing can often lead to load imbalance where some experts receive a disproportionate number of tokens.

2. Expert Chooses Token

With expert-choice routing, the decision-making is inverted: each expert selects the top k tokens it is best suited to process. This method is designed to enforce a balanced load, as each expert is assigned a fixed number of tokens.

Calculation:

For each expert row, we identify the two tokens (columns) with the highest scores.

E1: The top two scores are 3.1 (T1) and 1.8 (T4).
E2: The top two scores are 2.5 (T2) and 1.9 (T3).
E3: The top two scores are 3.5 (T3) and 2.2 (T1).
E4: The top two scores are 3.2 (T4) and 2.8 (T2).

Resulting Token Assignment:

E1 chooses: T1, T4
E2 chooses: T2, T3
E3 chooses: T1, T3
E4 chooses: T2, T4

In this particular case, the outcome is identical to the "Token Chooses Expert" method. Each expert processes two tokens, and each token is processed by two experts. This symmetry is a feature of this specific matrix and not a general rule. In many real-world scenarios, the two methods would produce different assignments.

3. Global Routing via Optimization

Global routing seeks the most optimal assignment across the entire matrix to maximize the total affinity score. The constraints for this problem with k=2 are:

Each token must be assigned to exactly two experts.
Each expert must process exactly two tokens.

This is a balanced assignment problem that can be solved with algorithms like the Hungarian algorithm or by formulating it as a maximum weight bipartite matching problem.

Calculation:

The goal is to select eight scores from the matrix S such that every row and every column is chosen exactly twice, and the sum of these eight scores is maximized.

We can find this optimal assignment by selecting the eight highest scores in the matrix, as long as they satisfy the constraints. The eight highest scores are: 3.5, 3.2, 3.1, 2.8, 2.5, 2.2, 1.9, and 1.8.

Let's check if this selection meets the constraints:

Selected Pairs: (E3, T3), (E4, T4), (E1, T1), (E4, T2), (E2, T2), (E3, T1), (E2, T3), (E1, T4).
Token Assignments (2 per token):
- T1: assigned to E1, E3
- T2: assigned to E4, E2
- T3: assigned to E3, E2
- T4: assigned to E4, E1
Expert Assignments (2 per expert):
- E1: assigned T1, T4
- E2: assigned T2, T3
- E3: assigned T3, T1
- E4: assigned T4, T2

The constraints are perfectly met. The resulting assignment is identical to the previous two methods for this specific example. The total maximized score is the sum of these values:
$3.5 + 3.2 + 3.1 + 2.8 + 2.5 + 2.2 + 1.9 + 1.8 = 21.0$

While all three methods produced the same outcome here, this highlights a scenario of perfect alignment. In practice, especially with larger and more varied score matrices, these methods would yield different assignments, each with its own trade-offs between computational cost, load balancing, and model performance.

A note on the complexity of global routing

Global routing frames the assignment as an optimization problem: find the assignment of tokens to experts that maximizes the total affinity score, subject to a set of global constraints. For our example with k = 2, the constraints are:

1. Each token must be assigned to exactly two experts.
2. Each expert must be assigned exactly two tokens.

This is a balanced assignment problem, which is computationally more complex than the local token or expert choice methods. While algorithms like the Hungarian algorithm or other linear programming solvers can find the optimal solution, they are often too slow to be practical for routing in LLMs.

In our specific example matrix, it happens that simply selecting the eight highest affinity scores—3.5, 3.2, 3.1, 2.8, 2.5, 2.2, 1.9, and 1.8—coincidentally satisfies these constraints. However, this is not a general strategy and would fail on most other matrices.

Deeper dive into the math of "token chooses expert" strategy

A natural question is how are the numbers in routing by hand example derived? As the strategy of "token chooses expert" is the most common strategy in modern MoE LLMs, let us work out this strategy by hand with an example, assuming k = 2.

We will be working through the following equations specifically:

Source: Stanford CS336 Language Modeling from Scratch, Lecture 4 on Mixture of Experts, 24 min 52 sec

Scenario Setup

Let's assume we are at a specific layer l in the model. We will focus on routing a single token, t.

Number of Experts (N): Let's say we have N=4 experts.
Top-K: We want to route our token to the top K=2 experts.
Token Representation ( $u_t^l$ ): This is the input vector for our token. For simplicity, let's assume it's a 3-dimensional vector.

u_t^l = \begin{pmatrix} 0.5 \newline 1.0 \newline -0.2 \end{pmatrix}

Expert Embeddings ( $e_i^l$ ): Each expert has a learnable weight vector (embedding) of the same dimension, which is used by the router to calculate affinity.

$e_1^l = \begin{pmatrix} 2.0 \newline 0.8 \newline -0.5 \end{pmatrix}$ , $e_2^l = \begin{pmatrix} -1.0 \newline 0.1 \newline 0.9 \end{pmatrix}$ , $e_3^l = \begin{pmatrix} 1.5 \newline 1.2 \newline 0.3 \end{pmatrix}$ , $e_4^l = \begin{pmatrix} -0.4 \newline -1.1 \newline 1.4 \end{pmatrix}$

Step 1: Calculate Raw Affinity Scores (Logits)

The first step is to determine how well our token $u_t$ aligns with each expert $e_i$ . This is done by calculating the dot product between the token's vector and each expert's embedding vector. This corresponds to the inner part of the third equation, $u_t^{l^T} e_i^l)$ .

Logit for Expert 1: $(utl)Te1l=(0.5×2.0)+(1.0×0.8)+(−0.2×−0.5)=1.0+0.8+0.1=1.9(u_t^l)^T e_1^l = (0.5 \times 2.0) + (1.0 \times 0.8) + (-0.2 \times -0.5) = 1.0 + 0.8 + 0.1 = 1.9$
Logit for Expert 2: $(utl)Te2l=(0.5×−1.0)+(1.0×0.1)+(−0.2×0.9)=−0.5+0.1−0.18=−0.58(u_t^l)^T e_2^l = (0.5 \times -1.0) + (1.0 \times 0.1) + (-0.2 \times 0.9) = -0.5 + 0.1 - 0.18 = -0.58$
Logit for Expert 3: $(utl)Te3l=(0.5×1.5)+(1.0×1.2)+(−0.2×0.3)=0.75+1.2−0.06=1.89(u_t^l)^T e_3^l = (0.5 \times 1.5) + (1.0 \times 1.2) + (-0.2 \times 0.3) = 0.75 + 1.2 - 0.06 = 1.89$
Logit for Expert 4: $(utl)Te4l=(0.5×−0.4)+(1.0×−1.1)+(−0.2×1.4)=−0.2−1.1−0.28=−1.58(u_t^l)^T e_4^l = (0.5 \times -0.4) + (1.0 \times -1.1) + (-0.2 \times 1.4) = -0.2 - 1.1 - 0.28 = -1.58$

Our vector of raw logits is: $[1.9, - 0.58, 1.89, - 1.58]$

Step 2: Normalize Scores with Softmax

The next step is to normalize these logits into a probability distribution using the Softmax function. This gives us the scores $s_{i,t}$ , which represent the router's confidence for each expert.

Equation: $si,t=Softmaxi(utlTeil)s_{i,t} = \text{Softmax}_i(u_t^{l^T} e_i^l)$

The Softmax formula is: $si,t=exp⁡(logiti)∑j=1Nexp⁡(logitj)s_{i,t} = \frac{\exp(\text{logit}i)}{\sum{j=1}^{N} \exp(\text{logit}_j)}$

Exponentials:
- $exp⁡(1.9)≈6.686\exp(1.9) \approx 6.686$
- $exp⁡(−0.58)≈0.560\exp(-0.58) \approx 0.560$
- $exp⁡(1.89)≈6.619\exp(1.89) \approx 6.619$
- $exp⁡(−1.58)≈0.206\exp(-1.58) \approx 0.206$
Sum of exponentials:
$6.686 + 0.560 + 6.619 + 0.206 = 14.071$
Normalized scores $s_{i,t}$ :
- $s1,t=6.686/14.071≈0.475s_{1,t} = 6.686 / 14.071 \approx 0.475$
- $s2,t=0.560/14.071≈0.040s_{2,t} = 0.560 / 14.071 \approx 0.040$
- $s3,t=6.619/14.071≈0.470s_{3,t} = 6.619 / 14.071 \approx 0.470$
- $s4,t=0.206/14.071≈0.015s_{4,t} = 0.206 / 14.071 \approx 0.015$

Our vector of normalized scores is: $s_t = [0.475, 0.040, 0.470, 0.015]$ . Note that these values sum to 1.

Step 3: Select Top-K and Determine Gating Weights

Now we apply the Top-K logic to determine the final gating weights, $g_{i,t}$ . The gate value is equal to the score $s_{i,t}$ if it's in the Top-K, otherwise it's zero.

Equation: $si,t∈TopK(sj,t,K=2)0,otherwiseg_{i,t} = \begin{cases} s_{i,t}, & \text{if } s_{i,t} \in \text{TopK}({s_{j,t}}, K=2) \newline 0, & \text{otherwise} \end{cases}$

Our scores are [0.475, 0.040, 0.470, 0.015].
The two highest scores (K=2) are 0.475 (for Expert 1) and 0.470 (for Expert 3).
The final gating weights are:
- $g_{1,t} = s_{1,t} = 0.475$
- $g_{2,t} = 0$
- $g_{3,t} = s_{3,t} = 0.470$
- $g_{4,t} = 0$

This means only Expert 1 and Expert 3 will be activated for this token.

Step 4: Calculate Final Output

Finally, the output of the MoE layer ( $h_t^l$ ) is a weighted sum of the outputs from the selected experts, added to the original token representation (a residual connection).

Equation: $htl=∑i=1N(gi,tFFNi(utl))+utlh_t^l = \sum_{i=1}^{N} (g_{i,t} \text{FFN}_i(u_t^l)) + u_t^l$

Plugging in our gating values:

htl=(0.475⋅FFN1(utl))+(0⋅FFN2(utl))+(0.470⋅FFN3(utl))+(0⋅FFN4(utl))+utlh_t^l = (0.475 \cdot \text{FFN}_1(u_t^l)) + (0 \cdot \text{FFN}_2(u_t^l)) + (0.470 \cdot \text{FFN}_3(u_t^l)) + (0 \cdot \text{FFN}_4(u_t^l)) + u_t^l

This simplifies to:

htl=0.475⋅FFN1(utl)+0.470⋅FFN3(utl)+utlh_t^l = 0.475 \cdot \text{FFN}_1(u_t^l) + 0.470 \cdot \text{FFN}_3(u_t^l) + u_t^l

The final output is the combination of the outputs from the two chosen experts, weighted by their routing scores, plus the original input token vector.

A Note on Variants (Mixtral, DBRX)

Models like Mixtral and DBRX use a slightly different approach where they apply Softmax after selecting the Top-K.

Using our example, that would mean:

Select Top-K Logits: Take the raw logits [1.9, -0.58, 1.89, -1.58]. The Top-2 are 1.9 and 1.89.
Apply Softmax to only the Top-K:
- $g1,t=exp⁡(1.9)exp⁡(1.9)+exp⁡(1.89)=6.6866.686+6.619≈0.503g_{1,t} = \frac{\exp(1.9)}{\exp(1.9) + \exp(1.89)} = \frac{6.686}{6.686 + 6.619} \approx 0.503$
- $g3,t=exp⁡(1.89)exp⁡(1.9)+exp⁡(1.89)=6.6196.686+6.619≈0.497g_{3,t} = \frac{\exp(1.89)}{\exp(1.9) + \exp(1.89)} = \frac{6.619}{6.686 + 6.619} \approx 0.497$
- $g_{2,t} = 0$ , $g_{4,t} = 0$

In this variant, the final gating weights for the chosen experts are re-normalized to sum to 1.

Heuristic balancing losses

System efficiency requires that we use experts evenly, so that computation is evenly distributed across all experts. In the worst case where all tokens are routed to a single expert, it is as bad as having a single dense model while having significantly higher memory overhead from the unused experts.

To encourage a balanced load, an auxiliary loss is often added to the total model loss during training. The following section illustrates the load balancing loss introduced in the Switch Transformers paper. It is important to note that this specific formulation was designed for a top-1 routing scenario (k = 1), where each token is dispatched to the single expert with the highest score (argmax). We use a top-1 example here to faithfully explain the paper's original equations. See Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity.

1. Equation Explanations

Here is a breakdown of each equation:

Equation (4): The Auxiliary Loss

This equation defines the overall auxiliary loss that is added to the model's main training loss (e.g., cross-entropy).

\text{loss} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i \quad (4)

loss: The final calculated auxiliary loss value.
α (alpha): A small, constant hyperparameter (the text suggests 10^-2) used to scale this auxiliary loss. It ensures that load balancing doesn't overwhelm the primary goal of the main loss function.
N: The total number of experts in the MoE layer.
f_i: The fraction of tokens in the batch that are actually dispatched to expert i.
P_i: The average fraction of the router's probability (or "confidence") allocated to expert i across all tokens in the batch.

Meaning: This loss is essentially a scaled dot product between the vector of token fractions (f) and the vector of average router probabilities (P). This loss is minimized when the tokens are distributed uniformly, meaning each expert i receives 1/N of the tokens (f_i = 1/N) and the router assigns an average probability of 1/N to each expert (P_i = 1/N).

Equation (5): Fraction of Tokens Dispatched

This equation calculates f_i, the actual proportion of tokens assigned to each expert.

f_i = \frac{1}{T} \sum_{x \in B} 1\{\text{argmax } p(x) = i\} \quad (5)

T: The total number of tokens in the batch B.
p(x): The output of the router for a single token x. It's a probability vector of size N, where each element is the probability of sending the token to the corresponding expert.
argmax p(x): This finds the index of the expert that received the highest probability for token x.
1{...}: This is an indicator function. It returns 1 if the condition inside is true (i.e., if expert i was chosen for token x) and 0 otherwise.

Meaning: This formula counts how many tokens in the batch are assigned to expert i (based on the highest router probability) and then divides by the total number of tokens to get the fraction. As noted in the text, the use of argmax makes this function non-differentiable, which has implications for the backward pass.

Equation (6): Fraction of Router Probability

This equation calculates P_i, the average router probability assigned to each expert.

P_i = \frac{1}{T} \sum_{x \in B} p_i(x) \quad (6)

p_i(x): The specific probability assigned to expert i for token x. This is the i-th element of the vector p(x).

Meaning: This formula calculates the average "vote" or probability mass the router gives to expert i across all tokens in the batch. Unlike f_i, this calculation is based on the soft probabilities from the router and is fully differentiable.

2. Illustrative Example

Let's use a simple example to make these equations concrete.

Number of experts N: 3
Number of tokens in batch T: 4
Hyperparameter α: 0.01

Assume our router network processes the 4 tokens and outputs the following probability distributions (p(x)). Each row corresponds to a token, and each column corresponds to an expert.

	Expert 1	Expert 2	Expert 3
Token 1	0.7	0.2	0.1
Token 2	0.1	0.8	0.1
Token 3	0.6	0.3	0.1
Token 4	0.2	0.2	0.6

This matrix represents the p_i(x) values.

3. Forward Pass Demonstration

Using our example matrix, we will now calculate the auxiliary loss.

Step 1: Calculate `f` (Fraction of Tokens)

We apply Equation (5) by finding the argmax for each token (row):

Token 1: argmax([0.7, 0.2, 0.1]) = Expert 1
Token 2: argmax([0.1, 0.8, 0.1]) = Expert 2
Token 3: argmax([0.6, 0.3, 0.1]) = Expert 1
Token 4: argmax([0.2, 0.2, 0.6]) = Expert 3

Now, we count the assignments for each expert:

Expert 1 was chosen 2 times.
Expert 2 was chosen 1 time.
Expert 3 was chosen 1 time.

Finally, we calculate the fractions (f_i = count / T):

$f_1 = 2 / 4 = 0.5$
$f_2 = 1 / 4 = 0.25$
$f_3 = 1 / 4 = 0.25$

So, our f vector is [0.5, 0.25, 0.25].

Step 2: Calculate `P` (Fraction of Probability)

We apply Equation (6) by averaging the probabilities in each column of our matrix:

$P_1 = (0.7 + 0.1 + 0.6 + 0.2) / 4 = 1.6 / 4 = 0.4$
$P_2 = (0.2 + 0.8 + 0.3 + 0.2) / 4 = 1.5 / 4 = 0.375$
$P_3 = (0.1 + 0.1 + 0.1 + 0.6) / 4 = 0.9 / 4 = 0.225$

So, our P vector is [0.4, 0.375, 0.225].

Step 3: Calculate the Final Auxiliary Loss

We use Equation (4) with the f and P vectors we just calculated:

\text{loss} = \alpha \cdot N \cdot \sum_{i=1}^{N} f_i \cdot P_i

\cdot 3 \cdot ( (0.5 \cdot 0.4) + (0.25 \cdot 0.375) + (0.25 \cdot 0.225) )

\cdot ( 0.2 + 0.09375 + 0.05625 )

\cdot ( 0.35 ) = 0.0105

This value, 0.0105, is the auxiliary loss for this batch, which will be added to the main model loss before the backward pass.

4. Backward Pass Demonstration

During backpropagation, we need to calculate the gradient of the loss with respect to the model's parameters, specifically the weights of the router network. The gradient flows backward from the loss through the calculations we performed.

A key detail from the text is that the f vector is not differentiable due to the argmax function. In practice, this is handled by treating f as a constant during the backward pass. This is an application of a concept similar to a straight-through estimator (STE), where the gradient for the non-differentiable argmax operation is effectively passed through as if it were an identity function for the forward pass values. As a result, the gradient signal only flows back through the differentiable part of the loss equation, which is the P_i term (the average router probability).

Step 1: Gradient of the Loss w.r.t. `P_i`

We start by finding the derivative of the loss function with respect to each element of the P vector.

\frac{\partial \text{loss}}{\partial P_i} = \frac{\partial}{\partial P_i} \left( \alpha \cdot N \cdot \sum_{j=1}^{N} f_j \cdot P_j \right)

Since f is treated as a constant, the derivative is:

\frac{\partial \text{loss}}{\partial P_i} = \alpha \cdot N \cdot f_i

Let's calculate this for our example:

$∂loss∂P1=0.01⋅3⋅f1=0.03⋅0.5=0.015\frac{\partial \text{loss}}{\partial P_1} = 0.01 \cdot 3 \cdot f_1 = 0.03 \cdot 0.5 = 0.015$
$∂loss∂P2=0.01⋅3⋅f2=0.03⋅0.25=0.0075\frac{\partial \text{loss}}{\partial P_2} = 0.01 \cdot 3 \cdot f_2 = 0.03 \cdot 0.25 = 0.0075$
$∂loss∂P3=0.01⋅3⋅f3=0.03⋅0.25=0.0075\frac{\partial \text{loss}}{\partial P_3} = 0.01 \cdot 3 \cdot f_3 = 0.03 \cdot 0.25 = 0.0075$

Step 2: Gradient of `P_i` w.r.t. `p_i(x)`

Next, we need the gradient of P_i (the average probability) with respect to the individual router output p_i(x) for a specific token x.

P_i = \frac{1}{T} \sum_{x \in B} p_i(x) \implies \frac{\partial P_i}{\partial p_i(x)} = \frac{1}{T}

For our example, this is 1 / 4 = 0.25.

Step 3: Combine Gradients (Chain Rule)

Using the chain rule, we find the gradient of the loss with respect to each individual router probability p_i(x).

\frac{\partial \text{loss}}{\partial p_i(x)} = \frac{\partial \text{loss}}{\partial P_i} \cdot \frac{\partial P_i}{\partial p_i(x)} = (\alpha \cdot N \cdot f_i) \cdot \frac{1}{T}

This gradient value is the same for all tokens x in the batch. Let's calculate it:

$∂loss∂p1(x)=0.015⋅14=0.00375\frac{\partial \text{loss}}{\partial p_1(x)} = 0.015 \cdot \frac{1}{4} = 0.00375$
$∂loss∂p2(x)=0.0075⋅14=0.001875\frac{\partial \text{loss}}{\partial p_2(x)} = 0.0075 \cdot \frac{1}{4} = 0.001875$
$∂loss∂p3(x)=0.0075⋅14=0.001875\frac{\partial \text{loss}}{\partial p_3(x)} = 0.0075 \cdot \frac{1}{4} = 0.001875$

This gives us a gradient vector [0.00375, 0.001875, 0.001875]. This gradient is then backpropagated to the router's softmax layer for every token. The optimizer will use this gradient to update the router's weights, encouraging it to produce probabilities that lead to a more balanced load in the next iteration. In this case, since expert 1 was over-utilized (f_1 was high), it receives the largest gradient signal, which will push the router to assign it lower probabilities in the future.

Plain Guide to Einops

Lewis Won — Sat, 09 Aug 2025 09:47:24 +0000

Table of Contents

Why learn Einops?
Why read this guide instead of the official Einops tutorial
What's in the tutorial
Transposing tensors
Composition of axes
Decomposition of axis
Einops reduce
Repeating elements
Einsum
Summary of linear algebra operations using Einsum

Why learn `Einops`?

Einops provides a way for tensor manipulations to be self-documenting. Key advantages of Einops over traditional methods such as reshape include:

Reliable: In addition to doing what is expected, it also explicitly fails for wrong inputs
Readable: Codes with Einops are self-documenting and hence easier to read without requiring additional comments
Maintainable: readable codes are more maintainable

There are also other claimed benefits such as time and memory efficiency which is not touched upon in this article.

For example, given a tensor of shape (2, 3, 4) called encode_weights:

# encode_weights with shape (2, 3, 4)
encode_weights = np.array([
    [[0.1, 0.2, 0.3, 0.4],   # pos=0, dim1=0
     [0.5, 0.6, 0.7, 0.8],   # pos=0, dim1=1
     [0.9, 1.0, 1.1, 1.2]],  # pos=0, dim1=2
    [[2, 3, 4, 5],           # pos=1, dim1=0
     [6, 7, 8, 9],           # pos=1, dim1=1
     [1, 0, 1, 0]]           # pos=1, dim1=2
])

And we want to flatten it while preserving the first dimension into the following:

[[0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.  1.1 1.2]
 [2.  3.  4.  5.  6.  7.  8.  9.  1.  0.  1.  0. ]]

Previously without Einops, we could use reshape as below to achieve the intended result. Here's how to interpret the code:

Keep the first dimension unchanged: encode_weights.shape[0] preserves the size of the first axis.
Flatten all remaining dimensions: The -1 tells NumPy/PyTorch to automatically calculate the size needed to flatten all other dimensions into a single dimension.

encode_weights.reshape(encode_weights.shape[0], -1)

However, when manipulating tensors, we are primarily concerned with inputs and outputs, and not really how the operations are carried out to achieve the inputs and outputs. From the code, it is not clear what are the shapes of the input and output tensors, which will have to be separately documented with comments. This could present a problem in a large code base where code drifts could happen, and outdated comments could confuse instead of help.

In contrast, we could use the rearrange operation in Einops to have a self-documenting way of manipulating the tensor to achieve the same outcome. From the code, we can see that:

encode_weights has 3 dimensions, and is of shape (2, 3, 4)
The output has 2 dimensions, and is of shape (2, 3 * 4)
The code should fail if the dimensions specified are wrong, e.g. code drifts.

rearrange(encode_weights, "pos dim1 dim2 -> pos (dim1 dim2)", pos=2, dim1=3, dim2=4)

Why read this guide instead of the official `Einops` tutorial

Einops has a great multi-part tutorial covering all the key features, and this guide serves to complement the official Einops tutorials, not to replace these tutorials.

While the use of the e i n o p s tensors as examples may be useful for fields such as computer vision, I found that I could learn better by being able to trace numbers in tensors visually.

Being a complement, I will largely follow the structure of Einops tutorial part 1, but using small tensors with numbers instead of images.

What's in the tutorial

I will use two tensors, embeddings and encode_weights to demonstrate the following operations: rearrange, reduce, repeat, and einsum. I will also link the operations back to concepts in linear algebra.

import numpy as np
from einops import rearrange

# New embeddings with shape (1, 2, 3)
embeddings = np.array([
    [[1, 2, 3],   # pos=0
     [4, 5, 6]]   # pos=1
])

# encode_weights with shape (2, 3, 4)
encode_weights = np.array([
    [[0.1, 0.2, 0.3, 0.4],   # pos=0, dim1=0
     [0.5, 0.6, 0.7, 0.8],   # pos=0, dim1=1
     [0.9, 1.0, 1.1, 1.2]],  # pos=0, dim1=2
    [[2, 3, 4, 5],           # pos=1, dim1=0
     [6, 7, 8, 9],           # pos=1, dim1=1
     [1, 0, 1, 0]]           # pos=1, dim1=2
])

Transposing tensors

In linear algebra, the transpose of a matrix is an operator which flips a matrix over its diagonal, i.e. switching the row and column indices of matrix $A$ by producing another matrix, often denoted by $A^T$ . Wikipedia has a good animation to illustrate this operation:

Source: Wikipedia

Using the embedding tensor as an example, it has one batch, two pos, and dim1 of three, we can transpose it along pos and dim1 by switching the two axes using rearrange:

print("original:")
print(embeddings)
print("transposed:")
print(rearrange(embeddings, "batch pos dim1 -> batch dim1 pos", batch=1, pos=2, dim1=3))

We get the following output:

original:
[[[1 2 3]
  [4 5 6]]]
transposed:
[[[1 4]
  [2 5]
  [3 6]]]

We can also switch the axes batch and pos, so now we get two batch-es and one pos:

print("original:")
print(embeddings)
print("transposed:")
print(rearrange(embeddings, "batch pos dim1 -> pos batch dim1", batch=1, pos=2, dim1=3))

We get the following output:

original:
[[[1 2 3]
  [4 5 6]]]
transposed:
[[[1 2 3]]

 [[4 5 6]]]

We can verify that transpose of embeddings transposed returns us the original embeddings tensor:

embeddings_transposed = rearrange(embeddings, "batch pos dim1 -> batch dim1 pos", batch=1, pos=2, dim1=3)
embeddings_returned = rearrange(embeddings_transposed, "batch dim1 pos -> batch pos dim1", batch=1, pos=2, dim1=3)
print(embeddings_returned)

The output:

[[[1 2 3]
  [4 5 6]]]

Composition of axes

Axis composition means combining multiple axes into one. The order of composition matters, which I will also demonstrate.

We can use round brackets ( and ) to compose two axes into a single axis. Using the encode_weights tensor as an example, the following code manipulates the original batch of two, pos of three and dim1 of four to (batch * pos) of six, and dim1 of four.

print("original:")
print(encode_weights)
print("Composing batch and pos dimensions to a new height dimension:")
print(rearrange(encode_weights, "batch pos dim1 -> (batch pos) dim1", batch=2, pos=3, dim1=4))

The output is as below:

original:
[[[0.1 0.2 0.3 0.4]
  [0.5 0.6 0.7 0.8]
  [0.9 1.  1.1 1.2]]

 [[2.  3.  4.  5. ]
  [6.  7.  8.  9. ]
  [1.  0.  1.  0. ]]]
Composing batch and pos dimensions to a new height dimension:
[[0.1 0.2 0.3 0.4]
 [0.5 0.6 0.7 0.8]
 [0.9 1.  1.1 1.2]
 [2.  3.  4.  5. ]
 [6.  7.  8.  9. ]
 [1.  0.  1.  0. ]]

If we swap the new axis (batch * pos) with (pos * batch), we will instead interleave each pos with each batch. An intuitive way to understand this operation is to break up the operation into two steps:

Step 1: Transpose batch with pos
Step 2: Build the new composite axis (pos * batch)

print("original:")
print(encode_weights)
print("Step 1: transpose `batch` with `pos`")
encoded_weights_transposed = rearrange(encode_weights, "batch pos dim1 -> pos batch dim1", batch=2, pos=3, dim1=4)
print(encoded_weights_transposed)
print("Step 2: build the new composite axis (`pos` * `batch`)")
print(rearrange(encoded_weights_transposed, "pos batch dim1 -> (pos batch) dim1", batch=2, pos=3, dim1=4))

With the following as output:

original:
[[[0.1 0.2 0.3 0.4]
  [0.5 0.6 0.7 0.8]
  [0.9 1.  1.1 1.2]]

 [[2.  3.  4.  5. ]
  [6.  7.  8.  9. ]
  [1.  0.  1.  0. ]]]
Step 1: transpose `batch` with `pos`
[[[0.1 0.2 0.3 0.4]
  [2.  3.  4.  5. ]]

 [[0.5 0.6 0.7 0.8]
  [6.  7.  8.  9. ]]

 [[0.9 1.  1.1 1.2]
  [1.  0.  1.  0. ]]]
Step 2: build the new composite axis (`pos` * `batch`)
[[0.1 0.2 0.3 0.4]
 [2.  3.  4.  5. ]
 [0.5 0.6 0.7 0.8]
 [6.  7.  8.  9. ]
 [0.9 1.  1.1 1.2]
 [1.  0.  1.  0. ]]

We can also have a single composite axis across all three axes, i.e.:

print(rearrange(encode_weights, "batch pos dim1 -> (batch pos dim1)", batch=2, pos=3, dim1=4))

This returns a vector of dimension 1 X (2 * 3 * 4) or simply 1 X 24, as seen from the output below:

[0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.  1.1 1.2 2.  3.  4.  5.  6.  7.
 8.  9.  1.  0.  1.  0. ]

Decomposition of axis

As compared to axis composition which involves combining axes, axis decomposition breaks up a single axis into multiple axes.

The example below takes the original dim1, and breaks it up into dim1 and dim2. Note that we have reused the original dim1 of length 4 for this illustration, where dim1 is only of length 2. This is an illustration of how the names of dimensions in Einops are local to each Einops operation.

print("original:")
print(encode_weights)
print("Break-up dim into two axes")
print(rearrange(encode_weights, "batch pos (dim1 dim2) -> batch pos dim1 dim2", batch=2, pos=3, dim1=2, dim2=2))

Has the following output:

original:
[[[0.1 0.2 0.3 0.4]
  [0.5 0.6 0.7 0.8]
  [0.9 1.  1.1 1.2]]

 [[2.  3.  4.  5. ]
  [6.  7.  8.  9. ]
  [1.  0.  1.  0. ]]]
Break-up dim into two axes
[[[[0.1 0.2]
   [0.3 0.4]]

  [[0.5 0.6]
   [0.7 0.8]]

  [[0.9 1. ]
   [1.1 1.2]]]


 [[[2.  3. ]
   [4.  5. ]]

  [[6.  7. ]
   [8.  9. ]]

  [[1.  0. ]
   [1.  0. ]]]]

For completeness, it is illegal to decompose pos1=3 to pos1=1 and pos2=2. The reason is because the multiplication of the decomposed axes must equal the original axis, and hence we receive the error of 3 != 2.

print("original:")
print(encode_weights)
print("Break-up pos into two axes")
print(rearrange(encode_weights, "batch (pos1 pos2) dim1 -> batch pos1 pos2 dim1", batch=2, pos1=1, pos2=2, dim1=4))

Output is an error:

EinopsError:  Error while processing rearrange-reduction pattern "batch (pos1 pos2) dim1 -> batch pos1 pos2 dim1".
 Input tensor shape: (2, 3, 4). Additional info: {'batch': 2, 'pos1': 1, 'pos2': 2, 'dim1': 4}.
 Shape mismatch, 3 != 2

`Einops` reduce

A reduce operation is a computation that aggregates multiple values from a collection into a single result. These operations include min, max, sum, mean, prod, any, all. In Einops, the missing axes are reduced.

The following reduce operation of batch1 pos1 dim1 -> reduces the tensor into a single scalar, and the operation max returns the maximum value.

from einops import reduce

print("original:")
print(encode_weights)
print("Return the max value as a scalar")
print(reduce(encode_weights, "batch1 pos1 dim1 -> ", "max", batch1=2, pos1=3, dim1=4))

Which gives us the following output, i.e. 9.0:

original:
[[[0.1 0.2 0.3 0.4]
  [0.5 0.6 0.7 0.8]
  [0.9 1.  1.1 1.2]]

 [[2.  3.  4.  5. ]
  [6.  7.  8.  9. ]
  [1.  0.  1.  0. ]]]
Return the max value as a scalar
9.0

Note that instead of returning a scalar, we can retain the three dimensional tensor with either () or 1, which represents an axis of length 1. Both () and 1 are equivalent.

from einops import reduce

print("original:")
print(encode_weights)
print("Return the max value as a scalar:")
print(reduce(encode_weights, "batch1 pos1 dim1 -> 1 1 1", "max", batch1=2, pos1=3, dim1=4))
print("Assert that `()` and `1` are equivalent:")
print(np.array_equal(reduce(encode_weights, "batch1 pos1 dim1 -> 1 1 1", "max", batch1=2, pos1=3, dim1=4), reduce(encode_weights, "batch1 pos1 dim1 -> () () ()", "max", batch1=2, pos1=3, dim1=4)))

original:
[[[0.1 0.2 0.3 0.4]
  [0.5 0.6 0.7 0.8]
  [0.9 1.  1.1 1.2]]

 [[2.  3.  4.  5. ]
  [6.  7.  8.  9. ]
  [1.  0.  1.  0. ]]]
Return the max value as a scalar:
[[[9.]]]
Assert that `()` and `1` are equivalent:
True

Repeating elements

As the name suggests, repeat means repeating elements in the tensor. I will demonstrate the following variations:

Repetition along a new axis
Repetition along an existing axis
Interleaving of repetition
Representing repetition as an integer instead of variable name

Repetition along a new axis

In the following code, we introduced a new axis repeat, which transforms embeddings from shape of (1, 2, 3) to (1, 2, 3, 4). Note that the order the new axis repeat matters -- by placing repeat of size two right after dim1, we are as converting each element in dim1 into an array of length 2 two, where the element is repeated twice.

print("original:")
print(embeddings)
print("Repeat with new axis:")
repeat(embeddings, "batch pos dim1 -> batch pos dim1 repeat", batch=1, pos=2, dim1=3, repeat=2)

The output is below.

original:
[[[1 2 3]
  [4 5 6]]]
Repeat with new axis:
array([[[[1, 1],
         [2, 2],
         [3, 3]],

        [[4, 4],
         [5, 5],
         [6, 6]]]])

Repetition along an existing axis

We can repeat along the axis dim1 by introducing composition dim1 and repeat into a single axis (dim1 repeat).

print("Tensor from previous step:")
print(repeat(embeddings, "batch pos dim1 -> batch pos dim1 repeat", batch=1, pos=2, dim1=3, repeat=2))
print("Repeat with new axis:")
repeat(embeddings, "batch pos dim1 -> batch pos (dim1 repeat)", batch=1, pos=2, dim1=3, repeat=2)

The output is as follows:

Tensor from previous step:
[[[[1 1]
   [2 2]
   [3 3]]

  [[4 4]
   [5 5]
   [6 6]]]]
Repeat with new axis:
array([[[1, 1, 2, 2, 3, 3],
        [4, 4, 5, 5, 6, 6]]])

Interleaving of repetition

The order of composition matters, and by swapping the composite axis of (dim1 repeat) with (repeat dim1), we can interleave the the new axis repeat with dim1.

print("Tensor from previous step:")
print(repeat(embeddings, "batch pos dim1 -> batch pos (dim1 repeat)", batch=1, pos=2, dim1=3, repeat=2))
print("Repeat with new axis:")
repeat(embeddings, "batch pos dim1 -> batch pos (repeat dim1)", batch=1, pos=2, dim1=3, repeat=2)

The output is as follows:

Tensor from previous step:
[[[1 1 2 2 3 3]
  [4 4 5 5 6 6]]]
Repeat with new axis:
array([[[1, 2, 3, 1, 2, 3],
        [4, 5, 6, 4, 5, 6]]])

Representing repetition as an integer instead of variable name

Instead of representing the new axis as repeat=2, we can simply use the integer 2 to achieve the same effect:

print("Tensor from previous step:")
repeat(embeddings, "batch pos dim1 -> batch pos (repeat dim1)", batch=1, pos=2, dim1=3, repeat=2)
print("Repeat with integer 2:")
repeat(embeddings, "batch pos dim1 -> batch pos (2 dim1)", batch=1, pos=2, dim1=3)

The output shows that both operations give the same result:

Tensor from previous step:
array([[[1, 2, 3, 1, 2, 3],
        [4, 5, 6, 4, 5, 6]]])
Repeat with integer 2:
array([[[1, 2, 3, 1, 2, 3],
        [4, 5, 6, 4, 5, 6]]])

Einsum

Einsum is a function that implements the Einstein summation convention for performing multi-dimensional array operations. To understand Einsum, I begin with a brief introduction of Einstein Notation, which provides a shorthand way to express instructions for combining numbers.

Let us begin with an example

Imagine you have two vectors where a is [1,2,3] and b is [4,5,6]. You want to multiply the first number in a with the first in b, the second with the second, and so on and so forth, which can be expressed as:

(a_1 \times b_1) + (a_2 \times b_2) + (a_3 \times b_3)

This is also known as a dot product. We can also express the summation using the Sigma symbol $∑\sum$ :

\sum_{i=1}^{3} a_i b_i

This means: "For every number i from 1 to 3, take the i-th element of a and the i-th element of b, multiply them, and then add everything up."

Note that the summation almost always happens over an index that appears twice, which we can further shorten using Einstein Notation.

Two rules of Einstein Notation

Rule 1: The Repeating Index means "Sum"

If you see an index (a little letter) appears twice on one side of the equation, it means you should sum over all possible values of that index. This repeated index is called a "dummy index".

Let us first rewrite the vectors, with vector a as a row vector (with a subscript representing "covariant"), and vector b as a column vector (with a superscript representing "contravariant"). Note that in a linear space, vectors retain their values after such transformations, which may not necessarily be true in non-linear spaces.

a_i = \begin{pmatrix} 1 & 2 & 3 \end{pmatrix}

b^i = \begin{pmatrix} 4 \newline 5 \newline 6 \end{pmatrix}

So, our long sum from before:

\sum_{i=1}^{3} a_i b^i

Becomes in Einstein Notation:

c = a_i b^i

Because the index i appears twice (once with a and once with b), the summation is implied.

Example: Dot Product of two vectors

Let's use our vectors from before. We'll write the row vector with a subscript and the column vector with a superscript.

a_i = \begin{pmatrix} 1 & 2 & 3 \end{pmatrix}

b^i = \begin{pmatrix} 4 \newline 5 \newline 6 \end{pmatrix}

The Einstein Notation is:

c = a_i b^i

This tells you to sum over the index i:

c = a_1 b^1 + a_2 b^2 + a_3 b^3 = (1)(4) + (2)(5) + (3)(6) = 32

We represent the operation as python code below:

from einops import einsum

A = np.array([
    [1,2,3]
])
B = np.array([
    [4],
    [5],
    [6]
])

einsum(A, B, "i j, j i-> ")

The einsum operation returns a scalar of value 32, which is expected:

np.int64(32)

Rule 2: The "Free" Index informs the output's shape

If an index appears only once on one side, it's called a "free index". This index must also appear on the other side of the equation. It tells you the shape or dimension of the result.

Example: Matrix-Vector Multiplication

Let's take a matrix A and a vector v.

$\begin{pmatrix} 1 & 2 \newline 3 & 4 \end{pmatrix}$ , which we can write as $AjiA^i_j$ .

$\begin{pmatrix} 5 \newline 6 \end{pmatrix}$ , which we'll write with a subscript, $v_j$ .

The result of their multiplication is a new vector, let's call it u. The standard equation is:
$ui=∑jAjivju^i = \sum_{j} A^i_j v^j$

The Einstein Notation for this:
$ui=Ajivju^i = A^i_j v^j$

We see the index j is repeated (once up, once down). So, we sum over j.
We see the index i appears only once on the right side as a superscript. This means it's a free index, and the result, u, will be a column vector with the superscript i.

Let's calculate the first element of u, which is $u^1$ . This is where i=1:

u^1 = A^1_j v^j = A^1_1 v^1 + A^1_2 v^2 = (1)(5) + (2)(6) = 5 + 12 = 17

Now for the second element, $u^2$ , where i=2:

u^2 = A^2_j v^j = A^2_1 v^1 + A^2_2 v^2 = (3)(5) + (4)(6) = 15 + 24 = 39

So, the resulting column vector is $ui=(1739)u^i = \begin{pmatrix} 17 \newline 39 \end{pmatrix}$ .

So, the resulting vector is $ui=(1739)u^i = \begin{pmatrix} 17 \newline 39 \end{pmatrix}$ .

The python code representing the above operation is below:

from einops import einsum

A = np.array([
    [1,2],
    [3,4]
])
B = np.array([
    [5],
    [6]
])

einsum(A, B, "i j, j i-> i")

Which returns the resulting two-element vector as expected:

array([17, 39])

More Examples

Matrix-Matrix Multiplication

Let's multiply two matrices, A and B, to get a new matrix C.

$Aji=(1234)A^i_j = \begin{pmatrix} 1 & 2 \newline 3 & 4 \end{pmatrix}$ , $Bkj=(5678)B^j_k = \begin{pmatrix} 5 & 6 \newline 7 & 8 \end{pmatrix}$ .

The Einstein Notation is:

C^i_k = A^i_j B^j_k

The index j is repeated (summed over).
The indices i (upstairs) and k (downstairs) are free, which tells us the result C will be a matrix with a superscript i and subscript k.

Let's calculate one element, say $C^1_1$ (where i=1 and k=1):

C^1_1 = A^1_j B^j_1 = A^1_1 B^1_1 + A^1_2 B^2_1 = (1)(5) + (2)(7) = 5 + 14 = 19

How about $C12C^2_1$ (where i=2 and k=1):

C^2_1 = A^2_j B^j_1 = A^2_1 B^1_1 + A^2_2 B^2_1 = (3)(5) + (4)(7) = 15 + 28 = 43

If you calculate all the elements, you get the final matrix C.

The above operation can be represented as code below:

from einops import einsum

A = np.array([
    [1,2],
    [3,4]
])
B = np.array([
    [5, 6],
    [7, 8]
])

einsum(A, B, "i j, j k-> i k")

Trace of a Matrix (Sum of diagonal elements)

The trace of a matrix is the sum of its diagonal elements.
For our matrix $Aji=(1234)A^i_j = \begin{pmatrix} 1 & 2 \newline 3 & 4 \end{pmatrix}$ , the trace is $1 + 4 = 5$ .

The Einstein Notation is:

T = A^i_i

Here, the index i is repeated (once up, once down), so it means sum over i.

T = A^1_1 + A^2_2 = 1 + 4 = 5

By setting the superscript and the subscript to be the same dummy index, we are telling it to only look at the diagonal elements and sum them up.

We can implement trace using code:

from einops import einsum

A = np.array([
    [1,2],
    [3,4]
])

trace = einsum(A, "i i ->")

print(f"The matrix A is:\n{A}")
print(f"The trace of A is: {trace}")

Which returns the following output:

The matrix A is:
[[1 2]
 [3 4]]
The trace of A is: 5

How the code works:

einsum looks at the input matrix A and the string "i i ->".
It identifies the elements of A where the first and second indices are identical: A[0,0] and A[1,1].
These elements are 1 and 4.
Because the output is specified as a scalar (->), it sums these elements together: 1 + 4 = 5.

The result is 5, which is the correct trace of the matrix A.

Summary of linear algebra operations using `Einsum`

ajcr has a nice summary of how to replicate the more common linear algebra operations using Einsum, which I reproduce below. The original blog post is available here. Note that ajcr assumes the numpy implementation of einsum. If you are using the einops variant, remember to leave a space between each index, e.g. i,ij->i in numpy should be represented as i,i j->i in einops.

Let A and B be two 1D arrays of compatible shapes (meaning the lengths of the axes we pair together either equal, or one of them has length 1):

Now let A and B be two 2D arrays with compatible shapes:

Note that we can use the spread operator '...' to conveniently represent axes which we are not interested in, e.g. np.einsum('...ij,ji->...', a, b) will multiply just the last two axes of a with the 2D array b.

ZeRO by hand with a 4-parameter model

Lewis Won — Fri, 01 Aug 2025 13:29:38 +0000

Note as of 2 Aug 2025, GMT+8 1am: Updated article to include activation checkpointing.

Table of Contents

Motivation
Why study ZeRO
Setup
Why make a copy of the model weights for optimizer states
Why do we have to make copies of the momentum and variance and not just recompute them on the fly?
Stage 1
Stage 2
Stage 3

Motivation

This article seeks to implement by hand the Zero Redundancy Optimizer (ZeRO). The goal is to build an intuition for how to implement fully sharded data parallelism to enable training at scale.

This article has been written with the assistance of Google Gemini 2.5 Pro, and with reference to Stanford CS336 lecture 7 on parallelism.

Why study ZeRO

ZeRO optimizes memory and improves training speed while increasing the model size that can be efficiently trained (source - ZeRO: Memory Optimizations Toward Training Trillion Parameter Models). This is done by sharding the model's parameters, gradients, and optimizer states across multiple GPUs, with minimal memory redundancy in data-parallel training. ZeRO is implemented in three stages:

Stage 1: Partitions the optimizer states across the data-parallel workers.
Stage 2: In addition to partitioning the optimizer states, this stage also partitions the gradients.
Stage 3: In addition to gradients and optimizer states, this stage also partitions the model parameters themselves.

The following diagram is lifted from the original ZeRO paper itself. Assuming only a relatively small model (by modern standard) of 7.5 billion parameters, we can expect 120 GB of memory consumption on each GPU. This far exceeds the 80GB VRAM on a H100. We can bring down the memory consumption per device to 31.4GB, 16.6GB, and finally 1.9GB with ZeRO stage 1, 2 and 3 respectively.

(source - ZeRO: Memory Optimizations Toward Training Trillion Parameter Models)

Setup

We assume the following setup:

Hardware: 2 GPUs (GPU0, GPU1)
Model: a tiny model with just 4 weights.
- W_fp16 = [w1, w2, w3, w4]: The low-precision (fp16) weights used for the fast forward and backward passes.
Optimizer States:
- W_fp32 = [w1_fp32, w2_fp32, w3_fp32, w4_fp32]: The high-precision (fp32) master copy of the weights.
- Momentum M and Variance V: Momentum M = [m1, m2, m3, m4] and Variance V = [v1, v2, v3, v4]. These are also in fp32.

Why make a copy of the model weights for optimizer states

The original fp16 weights are upcasted to fp32 and replicated for optimizer states for reducing rounding errors during training.

Let's say a gradient for a specific weight is very small. For example, gradient = 0.00001. The learning rate might also be small, say learning_rate = 0.001.

The update to the weight would be learning_rate * gradient = 0.00000001.

In FP32 (32-bit): A 32-bit floating-point number has enough precision to represent this tiny update. You can add 0.00000001 to an existing weight and the value will actually change.
In FP16 (16-bit): A 16-bit float has much less precision. A number as small as 0.00000001 might be rounded down to zero. If you add zero to your FP16 weight, the weight does not change. The update is completely lost.

If this happens for many steps in a row, your model stops learning for that parameter.

Why do we have to make copies of the momentum and variance and not just recompute them on the fly?

The key innovation of optimizers like SGD with Momentum, RMSprop, and Adam is that they do not just look at the gradient from the current mini-batch. They maintain a memory of past gradients to make a more intelligent update. This memory helps them navigate the loss landscape more effectively, avoiding oscillations and speeding up convergence.

This article assumes Adam is used.

A. Why We Must Store Momentum (The First Moment)

Momentum is formally defined as an exponentially moving average (EMA) of the gradients.

The update rule is:
momentum_t = β₁ * momentum_{t-1} + (1 - β₁) * gradient_t

We can analyze the following:

momentum_t: The momentum we are calculating for the current step t.
momentum_{t-1}: This is the critical part. To calculate the new momentum, you absolutely need the value of the momentum from the previous step (t-1).
gradient_t: The gradient from the current mini-batch.
β₁: The decay rate (e.g., 0.9). It controls how much "memory" of past gradients is kept.

The "On-the-Fly" Problem: If you tried to calculate momentum "on the fly" without storing it, you would not have momentum_{t-1} as the value is gone from the last step. It is also computationally infeasible to re-calculate the momentum of the entire history of training at every step.

Analogy: Imagine you're calculating the average temperature for the last 30 days. On day 31, to update your average, you don't re-measure the temperature of all 30 previous days. You simply take the average you already calculated for the first 30 days and combine it with the new measurement. momentum_{t-1} is that stored average. It summarizes the entire history of gradients up to this point.

B. Why We Must Store Variance (The Second Moment)

The logic is identical for the variance term (the "v" in Adam). It is an EMA of the squared gradients.

The update rule is:
variance_t = β₂ * variance_{t-1} + (1 - β₂) * (gradient_t)²

Again, to calculate the variance for the current step t, you must have the stored value of the variance from the previous step, variance_{t-1}. Without storing this state between training steps, the algorithm simply breaks. It loses its "memory" of how much the gradients have been varying in the past, which is the key information it uses to adapt the learning rate for each parameter.

Stage 1

The core idea is to eliminate memory redundancy of the optimizer states. This is significant for popular optimizers such as Adam, where for every model parameter (weight), we store its optimizer states. For an optimizer like Adam, this consists of a 4-byte fp32 master copy of the weight, a 4-byte momentum, and a 4-byte variance. This totals 12 bytes of memory for optimizer states per parameter, compared to just 2 bytes for the fp16 parameter itself.

Stage 1 consists of the following four steps, which I will go through one at a time:

Step 1: Every rank computes a full gradient on their subset of the batch.
Step 2: Reduce-scatter the gradients - incur a communication cost proportional to the number of parameters (specifically, (N-1)/N * #params for a ring algorithm, where N is the number of GPUs).
Step 3: Each machine updates their param using their gradient + state.
Step 4: All gather the parameters - incur #params communication cost.

Source: Stanford CS336 Language Modeling from Scratch | Spring 2025 | Lecture 7: Parallelism 1, 21 min 50 sec

Initialization and partitioning

In ZeRO Stage 1, we start with a standard Data Parallel setup where the batch of training data is split among the GPUs. The key innovation is that while the model weights are replicated, the optimizer-related states are partitioned to save memory.

Model: A simple 2-layer neural network:
1. h = w1*x1 + w2*x2
2. a = ReLU(h) (ReLU(x) is max(0, x))
3. y_pred = w3*a + w4
Loss Function: Half Mean Squared Error: Loss = 0.5 * (y_pred - y_true)^2
Replicated Data (Present on BOTH GPU 0 and GPU 1):
- The fp16 model weights: W_fp16 = [w1=2.0, w2=-3.0, w3=1.0, w4=0.5]
Input Data (Different for each GPU):
- GPU 0: Input = [x1=1, x2=3], True Label (y_true) = 5
- GPU 1: Input = [x1=2, x2=1], True Label (y_true) = 7
Partitioned Optimizer States (these will be used in the parameter update step):
- GPU 0 holds the fp32 master weights, momentum, and variance for w1 and w2.
- GPU 1 holds the fp32 master weights, momentum, and variance for w3 and w4.

Memory Saving: Each GPU holds the full fp16 model (2 bytes/param) but only half of the fp32 model (4 bytes/param), half the momentum (4 bytes/param), and half the variance (4 bytes/param). We have eliminated the redundancy for all 12 bytes/parameter of optimizer states.

Step 1: Compute full gradient

Each GPU computes a full gradient on their subset of the batch.

On GPU0: Forward and backward pass

GPU 0 performs its calculations completely independently.

Forward Pass (GPU 0)

Calculate h: h = w1*x1 + w2*x2 = (2.0 * 1) + (-3.0 * 3) = 2 - 9 = -7.0
Apply Activation a: a = ReLU(h) = ReLU(-7.0) = 0
Calculate Prediction y_pred: y_pred = w3*a + w4 = (1.0 * 0) + 0.5 = 0.5
Calculate Loss: Loss = 0.5 * (0.5 - 5)^2 = 0.5 * (-4.5)^2 = 0.5 * 20.25 = 10.125

Note on Activation Checkpointing: During this forward pass, to save memory, we do not store the intermediate values of h and a. They will be recomputed during the backward pass.

Backward Pass (GPU 0)

Now we go backward using the chain rule to find the gradient of the Loss with respect to each weight.

Gradient wrt y_pred: dL/dy_pred = (y_pred - y_true) = 0.5 - 5 = -4.5
Recompute activation a for gradient calculation: Before calculating gradients for w3 and w4, we must first recompute the activation a which was discarded after the forward pass. This requires re-running the first part of the model.
- Recompute h: h = w1*x1 + w2*x2 = (2.0 * 1) + (-3.0 * 3) = -7.0
- Recompute a: a = ReLU(h) = ReLU(-7.0) = 0
Gradient wrt w4: dL/dw4 = (dL/dy_pred) * (dy_pred/dw4) The derivative of w3*a + w4 wrt w4 is 1. = -4.5 * 1 = -4.5
Gradient wrt w3: dL/dw3 = (dL/dy_pred) * (dy_pred/dw3) The derivative of w3*a + w4 wrt w3 is a. = -4.5 * a = -4.5 * 0 = 0
Gradient wrt a: dL/da = (dL/dy_pred) * (dy_pred/da) The derivative of w3*a + w4 wrt a is w3. = -4.5 * w3 = -4.5 * 1.0 = -4.5
Gradient wrt h: dL/dh = (dL/da) * (da/dh) The derivative of ReLU(h) is 0 if h<0 and 1 if h>0. Since h=-7.0, the derivative is 0. = -4.5 * 0 = 0
Gradient wrt w2 and w1: Since dL/dh is 0, any gradient flowing back from it will also be zero.
- dL/dw2 = (dL/dh) * (dh/dw2) = 0 * x2 = 0
- dL/dw1 = (dL/dh) * (dh/dw1) = 0 * x1 = 0

Result from GPU 0: Grad0 = [grad_w1=0, grad_w2=0, grad_w3=0, grad_w4=-4.5]

On GPU 1: Forward and Backward Pass

GPU 1 performs the same calculations with its own data.

Forward Pass (GPU 1)

Calculate h: h = w1*x1 + w2*x2 = (2.0 * 2) + (-3.0 * 1) = 4 - 3 = 1.0
Apply Activation a: a = ReLU(h) = ReLU(1.0) = 1.0
Calculate Prediction y_pred: y_pred = w3*a + w4 = (1.0 * 1.0) + 0.5 = 1.5
Calculate Loss: Loss = 0.5 * (1.5 - 7)^2 = 0.5 * (-5.5)^2 = 0.5 * 30.25 = 15.125

Note on Activation Checkpointing: Again, the values for h and a are discarded to save memory.

Backward Pass (GPU 1)

Gradient wrt y_pred: dL/dy_pred = (y_pred - y_true) = 1.5 - 7 = -5.5
Recompute activation a for gradient calculation:
- Recompute h: h = w1*x1 + w2*x2 = (2.0 * 2) + (-3.0 * 1) = 1.0
- Recompute a: a = ReLU(h) = ReLU(1.0) = 1.0
Gradient wrt w4: dL/dw4 = -5.5 * 1 = -5.5
Gradient wrt w3: dL/dw3 = -5.5 * a = -5.5 * 1.0 = -5.5
Gradient wrt a: dL/da = -5.5 * w3 = -5.5 * 1.0 = -5.5
Gradient wrt h: dL/dh = (dL/da) * (da/dh) Since h=1.0, the derivative of ReLU(h) is 1. = -5.5 * 1 = -5.5
Gradient wrt w2: dL/dw2 = (dL/dh) * (dh/dw2) The derivative of w1*x1 + w2*x2 wrt w2 is x2. = -5.5 * x2 = -5.5 * 1 = -5.5
Gradient wrt w1: dL/dw1 = (dL/dh) * (dh/dw1) The derivative of w1*x1 + w2*x2 wrt w1 is x1. = -5.5 * x1 = -5.5 * 2 = -11.0

Result from GPU 1: Grad1 = [grad_w1=-11.0, grad_w2=-5.5, grad_w3=-5.5, grad_w4=-5.5]

Step 2: Average and Scatter Gradients (Reduce-scatter)

The system performs a Reduce-scatter operation. This single operation accomplishes two things:

One, average the gradient tensors from all GPUs (i.e. reduce) and two, scatter the result:

g_avg1 = (0 + (-11.0)) / 2 = -5.5
g_avg2 = (0 + (-5.5)) / 2 = -2.75
g_avg3 = (0 + (-5.5)) / 2 = -2.75
g_avg4 = (-4.5 + (-5.5)) / 2 = -5.0

The synchronized gradient tensor on both GPUs is:
Grad_avg = [-5.5, -2.75, -2.75, -5.0]

Two, instead of creating a full tensor with these results on every GPU, the reduce-scatter operation immediately sends the relevant partition to its destination GPU. The full Grad_avg is a conceptual intermediate that is never fully materialized in every GPU's memory:

GPU 0 is responsible for w1 and w2, so it receives the first partition of the averaged gradient:

Grad_partition_0 = [-5.5, -2.75]

GPU 1 is responsible for w3 and w4, so it receives the second partition of the averaged gradient:

Grad_partition_1 = [-2.75, -5.0]

Step 3: Update partitioned parameters

Each GPU now updates it assigned partition of the parameters using its local partition of the gradient. The calculations happen in full fp32 precision.

Let us first recap the parameters on both GPUs:

On BOTH GPU 0 and GPU 1:
- The current fp16 model is: W_fp16 = [w1=2.0, w2=-3.0, w3=1.0, w4=0.5]
- The averaged gradient tensor is: Grad_avg = [g1=-5.5, g2=-2.75, g3=-2.75, g4=-5.0]
Partitioned Optimizer States:
- GPU 0 holds the fp32 master weights, momentum (m), and variance (v) for w1 and w2.
- GPU 1 holds the fp32 master weights, momentum (m), and variance (v) for w3 and w4.

For this example, let's define our optimizer hyperparameters. We'll use the Adam optimizer.

Learning Rate (lr): 0.1
beta1: 0.9
beta2: 0.999
epsilon: 1e-8
We will assume this is the very first training step (t=1), so all initial momentums and variances are 0.

Each GPU was able to compute the gradients for all four weights because the calculation only required two things it had locally:

The complete set of W_fp16 weights.
Its own unique mini-batch of data.

I will now show how GPU 0 uses [-5.5, -2.75] to update its partition (w1, w2), and GPU 1 uses [-2.75, -5.0] to update its partition (w3, w4). This is where each GPU works only on its assigned partition. The calculations happen in full fp32 precision.

On GPU 0 (Responsible for w1 and w2)

GPU 0 uses g1=-5.5 and g2=-2.75 from the Grad_avg tensor.

Updating w1:

Retrieve local states: w1_fp32 = 2.0, m0_1 = 0, v0_1 = 0.
Update momentum: m1_1 = beta1*m0_1 + (1-beta1)*g1 = 0.9*0 + 0.1*(-5.5) = -0.55.
Update variance: v1_1 = beta2*v0_1 + (1-beta2)*g1^2 = 0.999*0 + 0.001*(-5.5)^2 = 0.03025.
Bias Correction (since t=1):
- m_hat = m1_1 / (1 - beta1^t) = -0.55 / (1 - 0.9) = -5.5.
- v_hat = v1_1 / (1 - beta2^t) = 0.03025 / (1 - 0.999) = 30.25.
Calculate new w1_fp32: w1'_fp32 = w1_fp32 - lr * m_hat / (sqrt(v_hat) + epsilon) = 2.0 - 0.1 * (-5.5) / (sqrt(30.25) + 1e-8) = 2.0 - 0.1 * (-5.5) / 5.5 = 2.0 - 0.1 * (-1) = 2.1.
Cast to fp16 for the computational model: w1'_fp16 = 2.1.

Updating w2:

Retrieve local states: w2_fp32 = -3.0, m0_2 = 0, v0_2 = 0.
Update momentum: m1_2 = 0.9*0 + 0.1*(-2.75) = -0.275.
Update variance: v1_2 = 0.999*0 + 0.001*(-2.75)^2 = 0.0075625.
Bias Correction:
- m_hat = -0.275 / 0.1 = -2.75.
- v_hat = 0.0075625 / 0.001 = 7.5625.
Calculate new w2_fp32: w2'_fp32 = w2_fp32 - lr * m_hat / (sqrt(v_hat) + epsilon) = -3.0 - 0.1 * (-2.75) / (sqrt(7.5625) + 1e-8) = -3.0 - 0.1 * (-2.75) / 2.75 = -3.0 - 0.1 * (-1) = -2.9.
Cast to fp16: w2'_fp16 = -2.9.

State on GPU 0 after updates: Its local fp16 model is now [2.1, -2.9, 1.0, 0.5] (partially updated).

On GPU 1 (Responsible for w3 and w4)

GPU 1 uses g3=-2.75 and g4=-5.0 from the Grad_avg tensor.

Updating w3:

Retrieve local states: w3_fp32 = 1.0, m0_3 = 0, v0_3 = 0.
The calculations are identical to w2 since the gradient is the same (-2.75).
Calculate new w3_fp32: w3'_fp32 = 1.0 - 0.1 * (-1) = 1.1.
Cast to fp16: w3'_fp16 = 1.1.

Updating w4:

Retrieve local states: w4_fp32 = 0.5, m0_4 = 0, v0_4 = 0.
The calculations are similar to w1 since the gradient magnitude is close.
m1_4 = -0.5, v1_4 = 0.025.
m_hat = -5.0, v_hat = 25.
Calculate new w4_fp32: w4'_fp32 = 0.5 - 0.1 * (-5.0) / (sqrt(25) + 1e-8) = 0.5 - 0.1 * (-1) = 0.6.
Cast to fp16: w4'_fp16 = 0.6.

State on GPU 1 after updates: Its local fp16 model is now [2.0, -3.0, 1.1, 0.6] (partially updated).

At this point, the W_fp16 models on the two GPUs are inconsistent with each other. This must be fixed before the next forward pass.

Step 4: Synchronize the Model (All-Gather)

The final step is to make the W_fp16 computational model identical on all GPUs. This is done by having each GPU broadcast the new fp16 weights from its partition to everyone else.

Broadcast:
- GPU 0 sends its updated partition: [w1'=2.1, w2'=-2.9].
- GPU 1 sends its updated partition: [w3'=1.1, w4'=0.6].
Gather and Update:
- GPU 0 receives [1.1, 0.6] from GPU 1. It uses this to overwrite the old values for w3 and w4 in its local W_fp16 copy.
- GPU 1 receives [2.1, -2.9] from GPU 0. It uses this to overwrite the old values for w1 and w2 in its local W_fp16 copy.

Final State (End of Training Step)

After the all-gather, the state on both GPUs is now consistent and ready for the next iteration:

Final Synchronized W_fp16 on both GPUs: [2.1, -2.9, 1.1, 0.6]
Updated Optimizer States (Still Partitioned):
- On GPU 0: w1_fp32=2.1, m1_1=-0.55, v1_1=0.03025 and w2_fp32=-2.9, m1_2=-0.275, v1_2=0.0075625.
- On GPU 1: w3_fp32=1.1, m1_3=-0.275, v1_3=0.0075625 and w4_fp32=0.6, m1_4=-0.5, v1_4=0.025.

The loop is now complete. The system successfully updated the weights, saved a significant amount of memory by not replicating the optimizer states, and ended with a consistent model ready for the next forward pass. This completes ZeRO stage 1.

Why do we break up All-reduce into Reduce-scatter and All-gather?

We know from my previous article that Reduce-scatter + All-gather = All-reduce, in the sense that:

The outputs on both sides of the equation are equivalent
The communication cost on both sides of the equation are equivalent (assuming both implement the ring-based algorithms)

With a Naive Distributed Data Parallel (DPP) where each GPU computes gradients on its local data batch, implementing All-reduce in one step means each GPU sends its calculated gradients to all other GPUs, and they all compute the sum (or average) and end up with the identical, full gradient tensor.

By splitting up All-reduce into two steps:

Reduce-scatter: instead of every GPU getting the full averaged gradient tensor, this operation computes the average and then immediately scatters partitions of the averaged gradients to the GPUs responsible for them. Each GPU only receives the specific slice of the gradients it needs to update its portion of the optimizer states.
All-gather: After each GPU updates its local partition of the model's parameters, they need to be synchronized across all GPUs for the next forward pass. This is done with an All-gather operation, where each GPU broadcasts its updated parameter partition to all other GPUs.

Stage 2

In Stage 1, we partitioned the optimizer states, which provides the largest memory savings. However, a notable memory redundancy remains: during the backward pass, every GPU must store the full gradient tensor until the final reduction step. For very large models, this gradient tensor can itself become a memory bottleneck.

The core idea of Stage 2 is to eliminate this gradient redundancy by overlapping communication and computation. Instead of waiting for the entire backward pass to finish, Stage 2 begins reducing gradients as soon as they are calculated for a specific layer and then immediately discards them. This approach lowers the peak memory usage and can improve performance by hiding the network latency of communication behind the computation of the next layer.

Stage 2 consists of the following steps:

Step 1: Incremental backward pass and overlapped gradient reduction. With activation checkpointing, during the forward pass, intermediate activations are not stored. During the backward pass, they are recomputed layer-by-layer just before the gradients for that layer are calculated.
Step 2: Update partitioned parameters
Step 3: All-gather the parameters

Step 1: Incremental Backward Pass & Overlapped Gradient Reduction

This is the key innovation of Stage 2. The backward pass and gradient communication are interleaved. As we backpropagate through the model, layer by layer, we perform the following steps:

1a. Recompute Activations (due to checkpointing): Before calculating the gradients for a specific layer, the necessary input activations for that layer (which were discarded during the forward pass) must be recomputed.
1b. Incremental Reduction: Immediately after a layer's gradients are computed, a reduce operation is triggered for them. In practice, this is often a reduce-scatter that sums the gradients and sends the correct partition to the GPU that owns that parameter slice.
1c. Immediate Memory Deallocation: Once a layer's gradients have been sent to the appropriate worker, they are no longer needed on the local GPU and that memory is freed. This ensures the GPU never needs to hold the gradients for the entire model at once.

Let's trace this step-by-step with our running example.

A. Backpropagation Through the Final Layer (w3, w4)

The backward pass proceeds in reverse order, so it encounters the layer involving w3 and w4 first.

Recompute Activations: To calculate gradients for this layer, we need the activation a. Since it was not stored during the forward pass, each GPU must recompute it.
- GPU 0: Re-runs the first layer's computation: h = (2.0*1) + (-3.0*3) = -7.0, so the recomputed a = ReLU(-7.0) = 0.
- GPU 1: Re-runs the first layer's computation: h = (2.0*2) + (-3.0*1) = 1.0, so the recomputed a = ReLU(1.0) = 1.0.
Compute Local Gradients: From our work in Stage 1, we know the locally computed gradients for this layer on each GPU:
- GPU 0: Grad0_layer2 = [grad_w3=0, grad_w4=-4.5]
- GPU 1: Grad1_layer2 = [grad_w3=-5.5, grad_w4=-5.5]
Immediate Reduce-Scatter: The system does not wait. It immediately performs a reduce-scatter operation only on these gradients. The goal is to compute the average and deliver the result to GPU 1, which is responsible for w3 and w4.
- Average grad_w3: (0 + (-5.5)) / 2 = -2.75
- Average grad_w4: (-4.5 + (-5.5)) / 2 = -5.0
Deliver Partition: The resulting averaged gradient partition [-2.75, -5.0] is sent to GPU 1.
Free Memory:
- GPU 0 no longer needs Grad0_layer2. This memory is freed.
- GPU 1 has received its final partition, so it can free the temporary memory it used to hold Grad1_layer2.

B. Backpropagation Through the First Layer (w1, w2)

While the first communication may still be in flight, the backward pass continues to the next layer, which involves w1 and w2.

Recompute Activations: In this case, the input to the first layer is the raw input data (x1, x2), which is already in memory, so no recomputation is needed for the activations required by this specific layer's gradient calculation.
Compute Local Gradients: We again use the results from our Stage 1 calculation:
- GPU 0: Grad0_layer1 = [grad_w1=0, grad_w2=0]
- GPU 1: Grad1_layer1 = [grad_w1=-11.0, grad_w2=-5.5]
Immediate Reduce-Scatter: A second, independent reduce-scatter is performed on these gradients. The target is GPU 0, which owns w1 and w2.
- Average grad_w1: (0 + (-11.0)) / 2 = -5.5
- Average grad_w2: (0 + (-5.5)) / 2 = -2.75
Deliver Partition: The resulting averaged gradient partition [-5.5, -2.75] is sent to GPU 0.
Free Memory: Both GPUs can now free the memory used to store the local gradients for w1 and w2.

State After Incremental Reduction:
At the end of this overlapped process, the backward pass is complete. The GPUs did not store the full gradient tensor or the full set of activations from the forward pass. Instead, they each hold only the final, averaged partition that they need for the optimizer step.

GPU 0 Memory: Holds the gradient partition Grad_part0 = [-5.5, -2.75].
GPU 1 Memory: Holds the gradient partition Grad_part1 = [-2.75, -5.0].

Step 2: Update Partitioned Parameters

This step is now identical to Stage 1. Each GPU has the precise gradient information it needs to update its slice of the fp32 master parameters. We use the same Adam optimizer settings as before (lr=0.1, beta1=0.9, beta2=0.999, t=1).

On GPU 0 (Responsible for w1 and w2)

GPU 0 uses its local gradient partition [-5.5, -2.75].

Updating w1:

States: w1_fp32 = 2.0, m0_1 = 0, v0_1 = 0. Gradient g1 = -5.5.
Update: The calculation is identical to Stage 1, resulting in:
- m1_1 = -0.55
- v1_1 = 0.03025
- w1'_fp32 = 2.1
Cast to fp16: w1'_fp16 = 2.1.

Updating w2:

States: w2_fp32 = -3.0, m0_2 = 0, v0_2 = 0. Gradient g2 = -2.75.
Update: The calculation is identical to Stage 1, resulting in:
- m1_2 = -0.275
- v1_2 = 0.0075625
- w2'_fp32 = -2.9
Cast to fp16: w2'_fp16 = -2.9.

On GPU 1 (Responsible for w3 and w4)

GPU 1 uses its local gradient partition [-2.75, -5.0].

Updating w3:

States: w3_fp32 = 1.0, m0_3 = 0, v0_3 = 0. Gradient g3 = -2.75.
Update: The calculation is identical to Stage 1, resulting in:
- w3'_fp32 = 1.1
Cast to fp16: w3'_fp16 = 1.1.

Updating w4:

States: w4_fp32 = 0.5, m0_4 = 0, v0_4 = 0. Gradient g4 = -5.0.
Update: The calculation is identical to Stage 1, resulting in:
- w4'_fp32 = 0.6
Cast to fp16: w4'_fp16 = 0.6.

At this point, just like in Stage 1, the replicated W_fp16 models on the two GPUs are inconsistent.

Step 3: All-Gather the Parameters

This final step is also identical to Stage 1. An all-gather operation is required to ensure the replicated W_fp16 model is consistent on all GPUs before the next forward pass can begin.

Broadcast Partitions:
- GPU 0 sends its newly updated parameters: [2.1, -2.9].
- GPU 1 sends its newly updated parameters: [1.1, 0.6].
Gather and Reconstruct: Each GPU receives the partition from the other and assembles the full, synchronized model.

Final State (End of Training Step)

The training step is complete. The state on both GPUs is now consistent and ready for the next iteration:

Final Synchronized W_fp16 on both GPUs: [2.1, -2.9, 1.1, 0.6]
Updated Optimizer States (Still Partitioned):
- On GPU 0: w1_fp32=2.1, w2_fp32=-2.9, along with their new momentum and variance.
- On GPU 1: w3_fp32=1.1, w4_fp32=0.6, along with their new momentum and variance.

Stage 2 achieves the exact same numerical outcome as Stage 1 but does so with a lower peak memory footprint by overlapping communication and computation during the backward pass.

Stage 3

In Stage 2, we eliminated redundancies in optimizer states and gradients. However, one last major redundancy exists: the model parameters (W_fp16) themselves are still fully replicated on every GPU.

The purpose of Stage 3 is to eliminate this final parameter redundancy. It achieves this by partitioning the parameters (W_fp16) across all GPUs from the very beginning. Each GPU is now responsible for storing and updating only its small slice of the complete model.

This creates a new challenge: during the forward and backward pass, a GPU will often need access to parameters that it does not hold in its own memory. ZeRO solves this by dynamically fetching the required parameter partitions from other GPUs just-in-time for the computation of each layer, and discarding them immediately after use. This is accomplished through a sequence of all-gather, compute, and reduce-scatter operations that are tightly integrated into the training loop.

Stage 3 consists of the following steps:

Step 1: Initialization (different from stage 1)
Step 2: Forward pass (layer-by-layer, discarding weights and activations)
Step 3: Backward pass (recomputing activations and integrating gradient reduction)
Step 4: Optimizer update (perfectly local)
Step 5: Synchronization

Step-by-Step Example of Stage 3

Step 1: Initialization (The Key Change)

The memory layout is now fundamentally different. Everything is partitioned from the start.

GPU 0 Memory:
- Parameter Partition: W_part0 = [w1_fp16=2.0, w2_fp16=-3.0]
- Optimizer States (fp32 master weights, momentum, variance) for w1 and w2 only.
GPU 1 Memory:
- Parameter Partition: W_part1 = [w3_fp16=1.0, w4_fp16=0.5]
- Optimizer States (fp32 master weights, momentum, variance) for w3 and w4 only.
Input Data (Same as before):
- GPU 0: Input = [x1=1, x2=3], True Label (y_true) = 5
- GPU 1: Input = [x1=2, x2=1], True Label (y_true) = 7

Memory Saving: This provides a massive memory reduction. Neither GPU holds the full model.

Step 2: Forward Pass (Layer-by-Layer)

The forward pass proceeds one layer at a time, gathering the necessary weights for each computation.

1. Executing Layer 1 (h = w1*x1 + w2*x2)

Problem: To compute h, both GPUs need w1 and w2, but only GPU 0 holds them.
Communication (All-Gather): An all-gather operation is triggered for the first layer's parameters.
- GPU 0 sends its shard: [2.0, -3.0].
- GPU 1 has no parameters for this layer, so it contributes an empty tensor.
- After the operation, both GPUs now have a temporary, full copy of the needed weights: [w1=2.0, w2=-3.0].
Computation (Forward Local): Each GPU computes its local value for h. These are the same calculations as in Stage 1.
- GPU 0: h = (2.0 * 1) + (-3.0 * 3) = -7.0
- GPU 1: h = (2.0 * 2) + (-3.0 * 1) = 1.0
Memory Management (Free Full Weights and Activations): The temporary full weight tensor [w1, w2] is immediately discarded from memory on both GPUs. With activation checkpointing, they also discard the intermediate result h, only keeping the subsequent activation a = ReLU(h).
- GPU 0: a = ReLU(-7.0) = 0, then discards h.
- GPU 1: a = ReLU(1.0) = 1.0, then discards h.

2. Executing Layer 2 (y_pred = w3*a + w4)

Problem: Now the GPUs need w3 and w4, which are only on GPU 1.
Communication (All-Gather): A new all-gather is performed for the second layer's parameters.
- GPU 1 sends its shard: [1.0, 0.5].
- GPU 0 contributes an empty tensor.
- After the operation, both GPUs have a temporary copy of [w3=1.0, w4=0.5].
Computation (Forward Local): The final prediction and loss are calculated.
- GPU 0:
  - y_pred = (1.0 * 0) + 0.5 = 0.5
  - Loss = 0.5 * (0.5 - 5)^2 = 10.125
- GPU 1:
  - y_pred = (1.0 * 1.0) + 0.5 = 1.5
  - Loss = 0.5 * (1.5 - 7)^2 = 15.125
Memory Management (Free Full Weights): The weights [w3, w4] and activation a are discarded. At the end of the forward pass, no GPU holds a complete set of weights or intermediate activations.

Step 3: Backward Pass (Integrated Gradient Reduction and Activation Recomputation)

The backward pass mirrors the forward pass, but as gradients are computed for a layer, they are immediately reduced and scattered to their owning GPU. Before computing gradients for a layer, it must first recompute the necessary activations.

1. Backprop through Layer 2

Activation Recomputation: To calculate gradients with respect to w3 and w4, we first need the activation a. Since it was discarded, it must be recomputed by running a forward pass through Layer 1.
- Communication (All-Gather): An all-gather is performed for Layer 1's parameters, bringing [w1=2.0, w2=-3.0] to both GPUs.
- Re-Computation: Each GPU recomputes h and then a.
  - GPU 0: h = -7.0, recomputed a = 0.
  - GPU 1: h = 1.0, recomputed a = 1.0.
- Memory Management: The temporary [w1, w2] tensor used for recomputation is immediately discarded.
Communication (All-Gather): To calculate gradients for w3 and w4, we need the weights again. An all-gather brings [w3=1.0, w4=0.5] to both GPUs.
Computation (Backward Local): Each GPU computes the local gradients for w3 and w4 using its recomputed activation a.
- GPU 0: dL/dy_pred = 0.5 - 5 = -4.5.
  - dL/dw4 = -4.5 * 1 = -4.5.
  - dL/dw3 = -4.5 * a = -4.5 * 0 = 0.
  - Local Gradients (GPU 0): Grad_part_L2_0 = [grad_w3=0, grad_w4=-4.5].
- GPU 1: dL/dy_pred = 1.5 - 7 = -5.5.
  - dL/dw4 = -5.5 * 1 = -5.5.
  - dL/dw3 = -5.5 * a = -5.5 * 1.0 = -5.5.
  - Local Gradients (GPU 1): Grad_part_L2_1 = [grad_w3=-5.5, grad_w4=-5.5].
Communication (Reduce-Scatter): A reduce-scatter is performed only on these Layer 2 gradients.
- Reduce (Average):
  - g_avg3 = (0 + (-5.5)) / 2 = -2.75
  - g_avg4 = (-4.5 + (-5.5)) / 2 = -5.0
- Scatter: The result is sent to the owning GPU. Since GPU 1 owns w3 and w4, it receives the averaged gradient partition.
  - GPU 1 stores: Grad_final_1 = [-2.75, -5.0].
Memory Management (Free Full Weights and Activations): The [w3, w4] tensor and the recomputed activation a are discarded from both GPUs.

2. Backprop through Layer 1

Activation Recomputation: The input to Layer 1 is the original batch data (x1, x2), which is still available in memory. Therefore, no activation recomputation is needed for this step.
Communication (All-Gather): The weights [w1=2.0, w2=-3.0] are gathered again for the gradient calculation.
Computation (Backward Local): Each GPU computes local gradients for w1 and w2.
- GPU 0: The backpropagated gradient dL/dh was 0, so the local gradients are 0.
  - Local Gradients (GPU 0): Grad_part_L1_0 = [grad_w1=0, grad_w2=0].
- GPU 1: The backpropagated gradient dL/dh was -5.5.
  - dL/dw1 = -5.5 * x1 = -5.5 * 2 = -11.0.
  - dL/dw2 = -5.5 * x2 = -5.5 * 1 = -5.5.
  - Local Gradients (GPU 1): Grad_part_L1_1 = [grad_w1=-11.0, grad_w2=-5.5].
Communication (Reduce-Scatter): A reduce-scatter is performed on Layer 1 gradients.
- Reduce (Average):
  - g_avg1 = (0 + (-11.0)) / 2 = -5.5
  - g_avg2 = (0 + (-5.5)) / 2 = -2.75
- Scatter: GPU 0 owns w1 and w2, so it receives the final gradient partition.
  - GPU 0 stores: Grad_final_0 = [-5.5, -2.75].
Memory Management (Free Full Weights): The [w1, w2] tensor is discarded.

Step 4: Optimizer Update (Perfectly Local)

Each GPU now has the exact gradients it needs for its local parameter partition. The optimizer step requires no communication.

On GPU 0 (updates w1, w2):
- Uses its stored gradient Grad_final_0 = [-5.5, -2.75].
- The calculations are identical to Stage 1, resulting in:
  - New w1'_fp32 = 2.1.
  - New w2'_fp32 = -2.9.
- New Parameter Partition on GPU 0: W_part0 = [2.1, -2.9].
On GPU 1 (updates w3, w4):
- Uses its stored gradient Grad_final_1 = [-2.75, -5.0].
- The calculations are identical to Stage 1, resulting in:
  - New w3'_fp32 = 1.1.
  - New w4'_fp32 = 0.6.
- New Parameter Partition on GPU 1: W_part1 = [1.1, 0.6].

Step 5: Synchronization

There is no final all-gather step. The parameters are intended to live in their partitioned state. The "synchronization" was handled dynamically and incrementally throughout the forward and backward passes. The system is immediately ready for the next training iteration, starting again with the all-gather for the first layer's weights.

The final state of the system is a fully sharded model with updated weights:

GPU 0 holds: [w1=2.1, w2=-2.9] and their corresponding optimizer states.
GPU 1 holds: [w3=1.1, w4=0.6] and their corresponding optimizer states.

From Scatter to All-Reduce: A Plain-English Guide to Collective Operations

Lewis Won — Fri, 25 Jul 2025 10:40:41 +0000

Table of Contents

Motivation
What are collective operations
Communication load and time-space tradeoff
How does data sharding affect time-space-communication tradeoff
Rank vs node
Common collective operations
Broadcast
Scatter
Gather
Reduce
All-reduce
Reduce-scatter
All-gather
How is all-gather different from gather?
Why Reduce-Scatter + All-gather = All-reduce
Appendix A: Non-blocking all-gather

Motivation

Collective operations are fundamental to distributed computing, but it was not easy to read documentation such as NCCL to fully understand such operations.

This article will instead work through simple examples for each of the more common operations, which I believe is the quickest way to understand these operations.

As the intent of this article is to understand the operations at a high level, I will not be going into optimization patterns such as tree-based or ring-based algorithms to maximize parallelism and minimize network contention. If you would like to go into further details, you may refer to resources such as the MPI tutorial.

This article was written with the assistance of Google Gemini 2.5 Pro.

What are collective operations

In parallel and distributed computing, collective operations are specialized functions that involve communication among a group of processes working together on a task. These processes are identified by a unique integer ID called a rank. Instead of a single process sending a message to another single process (which is called point-to-point communication), collective operations are coordinated actions that all ranks in a defined group must participate in.

This group of processes is formally known as a communicator. While programs often start with a default communicator that includes all processes (like MPI_COMM_WORLD in MPI), it's possible to create smaller subgroups or communicators. This allows for collective operations to be performed on just a subset of the total available processes, which is a powerful feature for complex parallel algorithms. (Source: MPI_COMM_WORLD)

Collective operations are fundamental building blocks in parallel programming libraries like the Message Passing Interface (MPI) and NVIDIA's NCCL for GPU-based computing. These libraries provide optimized implementations of collective operations, so developers do not have to implement them from scratch.

Communication load and time-space tradeoff

In traditional algorithm analysis, time complexity quantifies the number of operations an algorithm performs, while space complexity measures the amount of memory it requires. However, in distributed systems, where data and computation are spread across multiple nodes, the cost of communication could become a dominant factor in overall performance.

Communication cost can be broken down into:

Latency (α): The time it takes to initiate a message transfer, regardless of its size. This is often influenced by factors like network protocol overhead and routing.
Bandwidth Cost (β): The cost per unit of data (e.g., per byte) to transmit the message over the network.

Other considerations to manage communication cost include:

Network Topology: The physical and logical arrangement of nodes (e.g., ring, mesh, tree, fat-tree) has a profound impact on the performance of collective operations. An algorithm that is optimal on a fully connected network may perform poorly on a network with high contention for certain links.
Contention: When multiple messages try to traverse the same network link simultaneously, contention occurs, leading to delays. This is a critical factor in real-world performance that is not captured by the simple α-β model.
Synchronization Overhead: Collective operations are blocking, meaning all processes in the group must wait for the operation to complete before proceeding. This synchronization can introduce significant idle time, especially if there are "straggler" nodes that are slower than others. (Note: Many libraries also provide non-blocking variants, which allow a process to initiate the collective operation and then immediately proceed with other tasks, checking for the operation's completion later. This technique is key to hiding network latency by overlapping communication with computation. See Appendix A on Non-blocking all-gather for an example.)
In-network Computation: Modern networking hardware can sometimes perform reduction operations directly within the network switches, which can dramatically reduce the amount of data that needs to be transmitted and lower the overall latency.

How does data sharding affect time-space-communication tradeoff

Time Complexity: By sharding the data, each node is responsible for processing a smaller portion of the overall dataset. This can significantly reduce the local computation time on each node. For example, if a computation has a time complexity of O(N) on a single machine, distributing the data across 'P' processors could ideally reduce the computation time on each processor to O(N/P).
Space Complexity: Similarly, sharding reduces the amount of memory required on each individual node, as each only needs to store its assigned shard of the data.
Communication Load: However, increasing the number of shards (and thus the number of participating nodes) often increases the communication overhead. Collective operations require coordination and data exchange between nodes. For instance, in an all-reduce operation, as the number of nodes increases, more messages may need to be sent across the network, and the complexity of coordinating these messages can grow.

Rank vs node

Think of a node as a physical, standalone computer. In the context of high-performance computing (HPC), a cluster is made up of many of these nodes connected by a network. Each node typically has its own:

Processors (CPUs) with multiple cores
Memory (RAM)
Operating System
Potentially one or more GPUs (as is common with NVIDIA systems)

A rank is a unique integer ID assigned to a specific process in a parallel computing job. When you launch a parallel application (for example, using a framework like MPI or a library like NCCL), you specify how many parallel processes you want to run. Each of these processes is assigned a rank, starting from 0. Within a communication group (e.g. MPI's MPI_COMM_WORLD), each process has a unique rank. This allows the user to send and receive messages to and from specific processes.

Multiple ranks could run on a single node, or may also be distributed across multiple nodes. For example, ranks 0 and 1 could be on NodeA, and ranks 2 and 3 on Node B.

What about communication cost difference between ranks and nodes?

While sharding inherently increases the need for communication, the actual overhead experienced by individual ranks is heavily influenced by their physical location within the computing cluster.

Intra-Node Communication: Multiple ranks can execute on a single node. Communication between these ranks is generally faster and more efficient as it occurs within the same machine, often over high-speed interconnects like NVLink for GPUs.
Inter-Node Communication: When ranks on different nodes need to communicate, the data must traverse the network, which introduces significantly higher latency and lower bandwidth compared to intra-node communication.

Therefore, how data sharding impacts communication load is not uniform across all ranks. The placement of ranks across nodes becomes a critical factor. A well-designed sharding strategy will aim to minimize the frequency and volume of inter-node communication by keeping data that needs to be processed together on the same node, thus allowing the corresponding ranks to communicate more efficiently.

Common collective operations

Broadcast

A broadcast operation sends the same piece of data from one designated "root" rank to all other ranks in the group.

Source: Nvidia NCCL documentation

The formula out[i] = in[i] is self-explanatory. Imagine Rank 2 has a single, important value, say the number 42, that all other ranks need for a calculation.

Initial State:
- Rank 0: has no data
- Rank 1: has no data
- Rank 2: has data 42
- Rank 3: has no data
Step-by-Step Process:
1. Initiation: Rank 2 is designated as the root and initiates the broadcast.
2. Transmission: Rank 2 sends the value 42 to all other ranks (Rank 0, Rank 1, and Rank 3). In practice, this might happen through a direct send to each or through a more optimized network pattern like a tree, where Rank 2 sends to Rank 0 and Rank 1, and they in turn forward it to others.
3. Reception: Each of the other ranks receives the value 42.
Final State:
- Rank 0: has data 42
- Rank 1: has data 42
- Rank 2: has data 42
- Rank 3: has data 42

Analogy: A manager (Rank 2) sending the same memo (42) to every employee (the other ranks) in a department.

Scatter

A scatter operation takes an array of data on a single root rank and distributes different chunks of that array to the other ranks in the group. The data distribution for a scatter operation can be described by the following formula:

outY[i] = in[Y*count+i]

This formula shows how each destination rank Y fills its output buffer outY with data from the root's single input buffer in.

Explanation of the Formula

This formula describes a direct data copy from a specific segment of the root's large buffer to the smaller buffer of each destination rank.

in[...]: This is the source of the data. It refers to the single, large input buffer that exists only on the root rank.
outY[i]: This is the destination for the data. It represents the i-th element of the output buffer on a specific destination rank Y. Each rank, including the root, will have its own outY buffer after the operation.
[Y*count+i]: This is the source index. It's the "selection logic" that determines which chunk of the root's in buffer is sent to which rank.
- Y: The rank of the destination process that will receive the data chunk.
- count: The number of elements in each distributed chunk (i.e., the size of each outY buffer).
- i: The local index within that chunk.

In essence, the root rank uses the formula to slice its large in buffer. The chunk corresponding to Y=0 goes to rank 0, the chunk for Y=1 goes to rank 1, and so on, until the entire input buffer has been distributed.

Concrete Example:

Imagine Rank 0 holds an array of data [10, 20, 30, 40], and it wants to give each rank one element to work on.

Initial State:
- Rank 0: has no data
- Rank 1: has no data
- Rank 2: has data [10, 20, 30, 40]
- Rank 3: has no data
Step-by-Step Process:
1. Initiation: Rank 2, as the root, prepares its array for distribution. It conceptually splits the array into chunks, one for each rank (including itself).
2. Distribution: Rank 2 sends a unique chunk to each corresponding rank. It sends 10 to Rank 0, 20 to Rank 1, and 40 to Rank 3. It keeps the third chunk, 30, for itself. Using Rank 0 as an example, Y = 0 because Rank 0. count = 1 because each rank gets 1 integer. i = 0 because we want the first integer. So the value that rank 0 gets out of data is data[0*1+0] = data[0] = 10.
3. Reception: Each rank receives its assigned piece of the data.
Final State:
- Rank 0: has data 10
- Rank 1: has data 20
- Rank 2: has data 30
- Rank 3: has data 40

Analogy: A teacher (Rank 2) with a stack of different exam questions [10, 20, 30, 40] giving one unique question to each student (the other ranks).

Gather

A gather operation is the inverse of a scatter. It collects individual data values from all ranks in the group and assembles them into an array on a single designated root rank.

The data placement for a gather operation can be described with the following formula, which specifies how the root rank constructs its output buffer:

out[Y*count+i] = inY[i] (This is valid on the root rank only)

This formula describes a direct data copy from each rank's input buffer to a specific location in the root rank's output buffer.

inY[i]: This is the source of the data. It represents the i-th element from the input buffer of a specific source rank Y.
out[...]: This is the destination for the data. It refers to the large output buffer that is being constructed only on the designated root rank. This is the key difference between Gather and All-Gather.
[Y*count+i]: This is the destination index within the root's out buffer. It acts as the "placement logic" to ensure data is assembled correctly.
- Y: The rank of the source process whose data is being placed.
- count: The number of elements each rank is contributing (the size of each inY buffer).
- i: The local index within that source chunk.

In essence, the formula states that the root rank iterates through each source rank Y, takes its input data inY, and copies it into the correct block of the final out buffer.

Concrete Example:

Following the scatter example, let's say each rank has performed a calculation on its data and now has a result. Rank 0 needs to collect all these results.

Initial State:
- Rank 0: has result A
- Rank 1: has result B
- Rank 2: has result C
- Rank 3: has result D
Step-by-Step Process:
1. Initiation: Rank 2 is designated as the root rank that will receive all the data.
2. Transmission: Each rank sends its individual data value to Rank 2. Rank 0 sends A, Rank 1 sends B, and Rank 3 sends D.
3. Collection: Rank 2 receives the data from all ranks and assembles it into an array in a specific order (usually based on rank ID).
Final State:
- Rank 0: still has its data A
- Rank 1: still has its data B
- Rank 2: has data [A, B, C, D]
- Rank 3: still has its data D

Analogy: A teacher (Rank 2) collecting the completed, unique exam answers (A, B, C, D) from every student (the other ranks) to grade them.

Reduce

Source: Nvidia NCCL documentation

A reduce operation is similar to a gather, but as it collects data from all ranks, it combines them into a single final value using a specified operation (like sum, max, min, or logical AND). This final result is stored on only one root rank.

The formula out[i] = sum(inX[i]) describes an element-wise reduction across all participating ranks.

inX[i]: This refers to the i-th element of the input buffer from a specific rank X.
- in0[i] is the i-th element from rank 0's input buffer.
- in1[i] is the i-th element from rank 1's input buffer.
- in2[i] is the i-th element from rank 2's input buffer.
- in3[i] is the i-th element from rank 3's input buffer. Imagine each in buffer is an array of numbers. i is the index into that array.
sum(...): This specifies the reduction operation to be performed. In this case, it's a summation. For each index i, the operation takes the corresponding elements from all the input buffers and adds them together. So, the calculation is: in0[i] + in1[i] + in2[i] + in3[i].
out[i]: This represents the i-th element of the final output buffer. As the diagram shows, this output buffer exists only on the designated root rank (in this case, rank 2).

Putting It All Together

The expression out[i] = sum(inX[i]) means:

For every position i in the data buffers, the value at out[i] on the root rank is calculated by summing the values at that same position i from the input buffers of all ranks (rank 0, rank 1, rank 2, and rank 3).

A Concrete Example

Let's assume each rank has an input buffer with three elements:

rank 0 in0: [10, 20, 5]
rank 1 in1: [2, 4, 6]
rank 2 in2: [1, 1, 1]
rank 3 in3: [5, 10, 15]

The root is rank 2.

After the Reduce operation, the out buffer on rank 2 will be calculated as follows:

out[0] = sum(inX[0]) = in0[0] + in1[0] + in2[0] + in3[0] = 10 + 2 + 1 + 5 = 18
out[1] = sum(inX[1]) = in0[1] + in1[1] + in2[1] + in3[1] = 20 + 4 + 1 + 10 = 35
out[2] = sum(inX[2]) = in0[2] + in1[2] + in2[2] + in3[2] = 5 + 6 + 1 + 15 = 27

Final State:

rank 0, 1, and 3: Their input buffers remain unchanged. They have no out buffer.
rank 2 (root): Now has the out buffer: [18, 35, 27].

Analogy: Team members (ranks) each count the number of widgets they made. They report their counts up a chain of command, with each manager summing the counts from their reports, until the final boss (Rank 0) has the grand total.

It is worth noting that many libraries provide an 'in-place' option for these operations. When operating in-place, the result is stored directly back into the input buffer, which can be more memory-efficient as it avoids the need to allocate a separate output buffer.

All-reduce

Source: Nvidia NCCL documentation

An all-reduce operation is a combination of a reduce and a broadcast. It performs a reduction (combining values from all ranks into a single result), and then it broadcasts that final result back to all the ranks.

The formula out[i] = sum(inX[i]) is identical to the one for Reduce, but its implication is different because of where the output buffer (out) is stored.

inX[i]: This refers to the i-th element of the input buffer from a specific rank X.
- in0[i] is the i-th element from rank 0.
- in1[i] is the i-th element from rank 1, and so on.
sum(...): This is the reduction operation. For each index i, it calculates the sum of the elements at that position from all ranks: in0[i] + in1[i] + in2[i] + in3[i].
out[i]: This is the crucial difference from a standard Reduce. out[i] represents the i-th element of the output buffer, and as the diagram clearly shows, this out buffer is present on every single rank (rank 0, rank 1, rank 2, and rank 3).

So, the expression out[i] = sum(inX[i]) for an All-Reduce means:

For every position i, calculate the sum of the values at that position from the input buffers of all ranks, and place this final sum into the out[i] position on every participating rank.

A Concrete Example

Let's use the same initial data as the Reduce example:

rank 0 in0: [10, 20, 5]
rank 1 in1: [2, 4, 6]
rank 2 in2: [1, 1, 1]
rank 3 in3: [5, 10, 15]

The intermediate Reduce calculation is the same:

sum(inX[0]) = 10 + 2 + 1 + 5 = 18
sum(inX[1]) = 20 + 4 + 1 + 10 = 35
sum(inX[2]) = 5 + 6 + 1 + 15 = 27

The final, reduced buffer is [18, 35, 27].

Now, the Broadcast phase ensures this result is distributed.

Final State:

After the All-Reduce operation, every single rank will have an identical out buffer:

rank 0 out: [18, 35, 27]
rank 1 out: [18, 35, 27]
rank 2 out: [18, 35, 27]
rank 3 out: [18, 35, 27]

Analogy: After all team members (ranks) report their individual widget counts up the chain to get a grand total, the final boss (who now knows the total) sends out a company-wide email announcing the total number of widgets produced so that every employee is aware of the final number.

Reduce-Scatter

Source: Nvidia NCCL documentation

A reduce-scatter operation first combines data from all ranks using a specified operation (like a reduce) and then scatters the resulting combined data chunks back to the ranks. So, each rank ends up with a piece of the final result.

The formula outY[i] = sum(inX[Y*count+i]) combines the reduce and scatter steps into a single mathematical expression.

outY[i]: This represents the i-th element of the output buffer on a specific destination rank Y.
- out0[i] is the i-th element of the output buffer on rank 0.
- out1[i] is the i-th element of the output buffer on rank 1.
- ...and so on. This matches the diagram where each rank Y gets its own outY buffer.
sum(...): This is the reduction operation (e.g., sum, max, min). It combines data from all the input ranks.
inX[...]: This refers to the input buffer on rank X, where the sum is performed over all possible values of X.
[Y*count+i]: The index into the large input buffers. It maps the local index i on the output rank Y to a global index in the input data.
- Y: The rank of the process that will receive this specific piece of data.
- count: The number of elements in each final output chunk (i.e., the size of out0, out1, etc.). In the diagram, if each out buffer has, for example, 10 elements, then count would be 10.
- i: The local index within a final output chunk. It ranges from 0 to count-1.

The expression Y*count+i is an offset calculation. For rank 0 (Y=0), it calculates the sum for the first count elements. For rank 1 (Y=1), it calculates the sum for the next count elements, and so on.

Putting It All Together

The formula outY[i] = sum(inX[Y*count+i]) means:

The value of the i-th element on the output buffer of rank Y is the sum of all input elements at the global index (Y * count + i).

A Concrete Example

Let's assume there are 4 ranks and the goal is for each rank to receive 1 final value (so count = 1). This means the total result has 4 elements, so each input buffer must also have 4 elements.

Initial State:

rank 0 in0: [10, 20, 5, 1]
rank 1 in1: [2, 4, 6, 2]
rank 2 in2: [1, 1, 1, 3]
rank 3 in3: [5, 10, 15, 4]

The Calculation:

For rank 0 (Y=0): It needs to calculate its output buffer, out0. Since count=1, i can only be 0.
- out0[0] = sum(inX[0*1+0]) = sum(inX[0])
- out0[0] = in0[0] + in1[0] + in2[0] + in3[0] = 10 + 2 + 1 + 5 = 18
For rank 1 (Y=1): It needs to calculate out1. Again, i=0.
- out1[0] = sum(inX[1*1+0]) = sum(inX[1])
- out1[0] = in0[1] + in1[1] + in2[1] + in3[1] = 20 + 4 + 1 + 10 = 35
For rank 2 (Y=2): It needs to calculate out2.
- out2[0] = sum(inX[2*1+0]) = sum(inX[2])
- out2[0] = in0[2] + in1[2] + in2[2] + in3[2] = 5 + 6 + 1 + 15 = 27
For rank 3 (Y=3): It needs to calculate out3.
- out3[0] = sum(inX[3*1+0]) = sum(inX[3])
- out3[0] = in0[3] + in1[3] + in2[3] + in3[3] = 1 + 2 + 3 + 4 = 10

Final State:

rank 0: out0 = [18]
rank 1: out1 = [35]
rank 2: out2 = [27]
rank 3: out3 = [10]

Analogy: A group of authors (ranks) are collaboratively writing a book. They first combine their individual draft sections for each chapter into a single, final version (the reduce step). Then, each author takes one of the final chapters to do a final proofread (the scatter step).

All-gather

Source: Nvidia NCCL documentation

An all-gather operation is where every rank collects all the individual data values from all other ranks. Unlike a standard gather where only one root rank receives the data, in an all-gather, every rank ends up with a complete set of the data from all ranks.

The formula out[Y*count+i] = inY[i] describes a direct data copy and placement, not a mathematical reduction like sum. It dictates how the smaller input buffers are assembled into the larger output buffer.

inY[i]: This is the source of the data. It refers to the i-th element of the input buffer from a specific source rank Y.
- in0[i] is the i-th element from rank 0's input.
- in1[i] is the i-th element from rank 1's input, and so on.
out[...]: This is the destination for the data. It refers to the large output buffer that is being constructed on every rank.
[Y*count+i]: This is the destination index. It's the "placement logic" that determines where the source data inY[i] should be copied to in the final out buffer.
- Y: The rank of the source process whose data is being placed.
- count: The number of elements in each individual input chunk (i.e., the size of in0, in1, etc.).
- i: The local index within that source chunk, ranging from 0 to count-1.

This indexing calculation ensures that the input buffers are concatenated in the correct order (in0 first, then in1, then in2, etc.).

Putting It All Together

The formula out[Y*count+i] = inY[i] means:

The data from the i-th position of the input buffer of rank Y is copied directly into the final output buffer at the global index (Y * count + i). This process is repeated for all source ranks Y, and the resulting out buffer is identical on every rank.

A Concrete Example

Let's assume each rank has an input buffer with two elements (so count = 2).

Initial State:

rank 0 in0: [10, 11]
rank 1 in1: [20, 21]
rank 2 in2: [30, 31]
rank 3 in3: [40, 41]

The All-Gather operation will construct the out buffer on all ranks as follows:

Placing data from rank 0 (Y=0):
- out[0*2+0] = in0[0] => out[0] = 10
- out[0*2+1] = in0[1] => out[1] = 11
Placing data from rank 1 (Y=1):
- out[1*2+0] = in1[0] => out[2] = 20
- out[1*2+1] = in1[1] => out[3] = 21
Placing data from rank 2 (Y=2):
- out[2*2+0] = in2[0] => out[4] = 30
- out[2*2+1] = in2[1] => out[5] = 31
Placing data from rank 3 (Y=3):
- out[3*2+0] = in3[0] => out[6] = 40
- out[3*2+1] = in3[1] => out[7] = 41

Final State:

After the operation, every single rank will have the identical, fully assembled out buffer:

rank 0 out: [10, 11, 20, 21, 30, 31, 40, 41]
rank 1 out: [10, 11, 20, 21, 30, 31, 40, 41]
rank 2 out: [10, 11, 20, 21, 30, 31, 40, 41]
rank 3 out: [10, 11, 20, 21, 30, 31, 40, 41]

Analogy: At a business meeting, each department head (the ranks) has their department's quarterly report (A, B, C, D). They don't just send their reports to the CEO. Instead, they each make copies of their report and pass them around the table until every department head has a complete binder containing the reports from all other departments.

How is All-gather Different from Gather?

The key difference between gather and all-gather lies in who receives the final, complete collection of data.

In a gather operation, only one designated "root" rank collects the data from all the other ranks. The other ranks only send their data; they do not receive the final assembled array.
In an all-gather operation, every rank involved in the operation receives the final, complete set of data from all other ranks. It's as if a gather operation happens for every single rank.

Why Reduce-Scatter + All-gather = All-reduce

The equation reduce-scatter + all-gather = all-reduce illustrates how two more fundamental operations can be combined to create a more complex one. Let's trace this with a concrete example. The goal is an All-reduce with the sum operation.

1. Initial State: Each rank has an array of data. The size of the array equals the number of ranks.

Rank 0 in: [10, 20, 5, 1]
Rank 1 in: [2, 4, 6, 2]
Rank 2 in: [1, 1, 1, 3]
Rank 3 in: [5, 10, 15, 4]

2. Perform a Reduce-Scatter: This operation performs an element-wise sum and then scatters the resulting array, so each rank gets one piece of the final sum.

The first element of the result is 10+2+1+5 = 18. This goes to Rank 0.
The second element of the result is 20+4+1+10 = 35. This goes to Rank 1.
The third element of the result is 5+6+1+15 = 27. This goes to Rank 2.
The fourth element of the result is 1+2+3+4 = 10. This goes to Rank 3.

3. Intermediate State (after Reduce-Scatter): Each rank now holds its portion of the reduced result.

Rank 0 has: [18]
Rank 1 has: [35]
Rank 2 has: [27]
Rank 3 has: [10]

4. Perform an All-gather: Now, each rank takes its piece and shares it with all other ranks, assembling the full array everywhere.

Final State (after All-gather):
Rank 0 has [18, 35, 27, 10]
Rank 1 has [18, 35, 27, 10]
Rank 2 has [18, 35, 27, 10]
Rank 3 has [18, 35, 27, 10]

This final state, where every rank holds the complete array of summed values [18, 35, 27, 10], is exactly the outcome of a single All-reduce operation.

Any communication cost difference between reduce-scatter + all-gather vs all-reduce?

At a high level, the theoretical communication cost of an efficient ring-based all-reduce is t he same as ring-based reduce-scatter followed by a ring-based all-gather. A ring-based all-reduce is, for all practical purposes, implemented as a fused reduce-scatter and all-gather operation. They are not two different competing algorithms in this context; rather, one is built directly from the other two. By breaking the problem down this way, data can be continuously streamed around the ring of nodes, ensuring that network links are always busy and that each rank is always either sending or receiving data. This pipelining approach is highly efficient because it avoids the distinct synchronization points that would exist if you performed a full Reduce to one node, which then had to perform a full Broadcast.

Appendix A: Non-blocking all-gather

Let us imagine a scenario where each rank needs to gather model weights from all other ranks and then use those weights to calculate a local gradient.

Scenario 1: Blocking All-gather

In this scenario, computation and communication happen in a strict sequence. The time spent waiting for the network is "wasted" time for the CPU.

// Each rank has its local weights in `local_weights_buffer`
// `all_weights_buffer` is the output buffer

1. print("Starting blocking All-gather...")
2. MPI_Allgather(local_weights_buffer, ..., all_weights_buffer, ...); // <-- PROGRAM PAUSES HERE
   // All ranks wait here until every rank has received all data.
   // The CPU is idle during this time.
3. print("All-gather complete. All weights are now in all_weights_buffer.")

4. print("Now, starting gradient calculation...")
5. calculate_gradients(all_weights_buffer); // <-- Can only start after step 2 is fully done.

Scenario 2: Non-blocking All-gather

By overlapping the independent computation (Step 4) with the communication (started in Step 2), you effectively "hide" the network latency. The total execution time can be significantly reduced.

// Buffers are the same, plus a request handle
MPI_Request request;

1. print("Initiating NON-BLOCKING All-gather...");
2. MPI_Iallgather(local_weights_buffer, ..., all_weights_buffer, ..., &request); // <-- RETURNS IMMEDIATELY

3. print("All-gather is running in the background. Now doing other work...");
4. perform_independent_computation(); // <-- THIS IS THE OVERLAP!
   // This could be anything that doesn't need `all_weights_buffer`.
   // e.g., preparing data for the *next* iteration.

5. print("Finished other work. Now I need the weights. Waiting for All-gather to complete...");
6. MPI_Wait(&request, ...); // <-- PROGRAM PAUSES HERE, ONLY IF IT'S NOT ALREADY DONE.
   // Guarantees that the operation is complete and the buffer is safe to read.

7. print("All-gather is confirmed complete.");
8. calculate_gradients(all_weights_buffer); // <-- Safe to run now.

DEV Community: Lewis Won

How does low-rank adaptation for large language models work

Table of Contents

Why study LoRA? The challenge of fine-tuning massive models

Conceptual overview: Full Fine-Tuning

A Simple Example

Link back to the Math Equation on Full Fine-tuning

Conceptual Overview: Low-Rank Adaptation (LoRA)

The Same Example with LoRA

Mathematical Walk-through of Forward Pass with LoRA

Mathematical Walk-through of the Backward Pass

1. The Setup

2. The Chain Rule in Action

3. Calculating the Gradient for B (grad_B)

4. Calculating the Gradient for A (grad_A)

5. Updating the Weights

Link Back to the Math Equation on LoRA

LoRA in the Self-Attention Module

Mathematical Walk-through of the Forward Pass

1. Define Inputs

2. Calculate the Original Path (Frozen)

3. Calculate the LoRA Path (Trainable)

4. Combine for the Final Query Vector

5. The Bigger Picture

Mathematical Walk-through of the Backward Pass

1. The Setup

2. Calculating the Gradient for B_q

3. Calculating the Gradient for A_q

4. Updating All LoRA Weights

Appendix A: The "No Additional Inference Latency" Trick

Appendix B: Why Does the Low-Rank Hypothesis Make Sense?

FlashAttention by hand

Introduction to FlashAttention

Vanilla self-attention mechanism without FlashAttention

Step 1: Calculate the Full Attention Score Matrix (A = softmax(X))

1a. Find the Global Maximum (m)

1b. Compute the Exponentials and the Denominator (d)

1c. Normalize to get the Final Attention Vector A

Step 2: Multiply Attention Scores by the Value Matrix ( O=A⋅VO = A \cdot VO=A⋅V )

Comparison and Key Differences

Conceptual Overview: Fused Tiled Attention

FlashAttention by hand

Link back to the FlashAttention algorithm

Step 1: Process Tile 1 (i=1)

1a. Find the New Maximum

Link back to the FlashAttention algorithm

1b. Calculate Local Denominator and Local Output

Link back to the FlashAttention algorithm

1c. Update Running Statistics

Link back to the FlashAttention algorithm

Step 2: Process Tile 2 (i=2)

2a. Find the New Maximum

Link back to the FlashAttention algorithm

2b. Calculate Local Denominator and Local Output

Link back to the FlashAttention algorithm

2c. Update Running Statistics

Link back to the FlashAttention algorithm

Final Result

Link back to the FlashAttention algorithm

Walkthrough of the FlashAttention diagram

High-Level Overview

Mapping the Diagram Components to Our Walkthrough

1. The Loops

2. The Memory Hierarchy (SRAM vs. HBM)

3. Putting it all Together: Tracing Our Walkthrough on the Diagram

Appendix A - How to derive activation matrices from weight matrices

The Two Sets of Q, K, V

The Omitted "Prequel" Step: Creating Q, K, and V

Connecting to the Logits

Summary

Appendix B - The Paper's pseudocode vs. this walkthrough

Why normalizing only at the last step works

The Two Methods

The Proof of Equivalence

1. Base Case (i=1)

2. Inductive Hypothesis

3. Inductive Step

Intuitive Analogy: Calculating a Weighted Average

Online softmax by hand

Why study online softmax? The key to flash attention

3. Calculating the Gradient for `B` (`grad_B`)

4. Calculating the Gradient for `A` (`grad_A`)

2. Calculating the Gradient for `B_q`

3. Calculating the Gradient for `A_q`

Step 1: Calculate the Full Attention Score Matrix (`A = softmax(X)`)

1a. Find the Global Maximum (`m`)

1b. Compute the Exponentials and the Denominator (`d`)

1c. Normalize to get the Final Attention Vector `A`

Step 2: Multiply Attention Scores by the Value Matrix ( $\cdot V$ )

Step 1: Process Tile 1 (`i=1`)

Step 2: Process Tile 2 (`i=2`)

Step 1: Process Tile 1 ( $T1=[123]T_1 = \begin{bmatrix} 1 & 2 & 3 \end{bmatrix}$ )

Step 1: Process Tile 2 ( $T2=[621]T_2 = \begin{bmatrix} 6 & 2 & 1 \end{bmatrix}$ )

Step 2: Forward Pass for `Y = GeLU(X * A)`

On GPU 1: Calculate `Y_intermediate_1 = X * A1`

On GPU 2: Calculate `Y_intermediate_2 = X * A2`

Step 3: Forward Pass for `Z = Dropout(Y * B)`

On GPU 1: Calculate `Z_partial_1 = Y1 * B1`

On GPU 2: Calculate `Z_partial_2 = Y2 * B2`

1. Column Parallelism (What we do to `A`)

2. Row Parallelism (What we do to `B`)

Analyzing the Standard `Column -> Row` Pattern

Step 2: Backward Pass for `Z = Dropout(Y * B)`

Mathematics for Gradients `dY` and `dB`

Calculation for `dY`

Calculation for `dB`

Step 4: Backward Pass for `Y = GeLU(X * A)`

Mathematics for Gradients `dX` and `dA`

Calculation for `dA`