About this series

Transformers power modern language models, but explanations often jump from a colorful attention diagram to dense matrix notation. This series takes the middle route: clear intuition, enough mathematics to verify the mechanism, and practical implications for model architecture and GPU inference. It assumes familiarity with programming, arrays and basic algebra, not previous experience building Transformers.

Reading path. Parts 1–4 build one attention calculation. Parts 5–8 assemble a Transformer block, including positional information and feed-forward networks (FFNs). Parts 9–13 connect the math to memory-efficient GPU execution and text generation. Parts 14–15 close with training, implementation and evidence-based performance analysis.

Editorial conventions.  $T$ denotes token positions, $d_{model}$ the number of hidden features per token, $d_k$ Query/Key width, $B$ batch size and $H$ attention-head count. Dimensions are usually written as rows × columns. Approximate costs are analytic estimates, not measurements of a particular GPU.

1. Part 1 – Why a Word Needs Its Neighbors
2. Part 2 – Q, K and V: The Matching Engine
3. Part 3 – Why Divide Attention Scores by $\sqrt{d_k}$?
4. Part 4 – Softmax, Causal Masks and a Three-Token Walkthrough
5. Part 5 – Multiple Heads and Positional Information
6. Part 6 – Normalization and Residual Connections
7. Part 7 – Trace One Complete Transformer Block
8. Part 8 – The Feed-Forward Network Is Not Just Another Matrix
9. Part 9 – The GPU Cost of Attention
10. Part 10 – FlashAttention Without the Mystery
11. Part 11 – Why Predicting the Next Token Is a Learning Problem
12. Part 12 – Prefill, Decode and the First Generated Token
13. Part 13 – KV Cache, Grouped-Query Attention and PagedAttention
14. Part 14 – What Gradients and Distributed GPUs Really Change ← You are here
15. Part 15 – Build, Verify and Profile Your Own Attention

Overview

So far we have described what a Transformer computes during a forward pass. Training adds a reverse journey: a loss is measured, and backpropagation calculates how sensitive that loss is to each trainable parameter. At scale, the model and its intermediate tensors may also be distributed across many GPUs.

Those are different challenges. Gradients explain how the weights learn; distributed execution explains how the required state and computation are partitioned. Neither should be confused with a colorful attention heatmap.

Trace one gradient through a two-Key attention example

Take one Query attending to two Keys, with scaled scores $s=\left[0,ln3\right]$. Softmax gives attention weights $a=\left[0.25,0.75\right]$. Suppose their scalar Values are $v=\left[1,3\right]$. The attention output is

Choose a deliberately simple squared-error teaching target $y=2$ and loss $\mathcal{L}=\frac{1}{2}\left(o-y\right)^2=0.125$. This is not the usual vocabulary cross-entropy training objective; it isolates the attention derivative.

The chain rule gives the sensitivity of loss to each score as

Substituting the numbers produces $\partial\mathcal{L}/\partial s_1=-0.1875$ and $\partial\mathcal{L}/\partial s_2=+0.1875$. They sum to zero. That makes sense: adding the same constant to both scores does not change their softmax distribution, so the loss is unchanged by a shared score shift.

If a trainable Query parameter affects the scores, another chain-rule step propagates these score sensitivities to that parameter. The same principle applies through the projections, FFN, residuals and earlier layers.

What a high attention weight cannot prove

A high attention weight tells us that one head assigned a large mixing coefficient to one Key position. It does not establish that the corresponding token caused the final answer. Its Value vector may have small or transformed coordinates; subsequent output projections and FFNs can change its contribution; other heads and layers also participate.

Attention visualization can help identify routing patterns. Strong causal interpretations require more evidence, such as controlled interventions or attribution techniques with their own limitations. It is a mistake to treat the heatmap as a fully faithful explanation of model decisions.

Why one GPU is not always enough

A large training job may exceed the memory capacity or compute budget of one device. There are several ways to partition it, each with a different communication cost.

Data parallelism runs different examples on model replicas and combines gradient information. Tensor parallelism partitions work within a layer, such as projection features or attention heads. Pipeline parallelism places different groups of consecutive layers on different devices and moves activations between stages. Context parallelism distributes token positions across devices, requiring suitable exchange or streaming of Keys/Values for attention across the sequence.

Fully Sharded Data Parallel (FSDP) shards model-training state, including parameters and usually gradients and optimizer-related state under its chosen configuration. Its main purpose is to reduce replication of training state; it is not by itself an answer to quadratic token-pair work.

Partition the correct dimension

Think of four different questions: Is the batch too large? Are layer matrices too large? Is the depth too great for one device? Is the sequence too long to keep and process on one accelerator? Different answers point to different parallelization strategies.

Communication matters. Tensor-parallel fragments may require all-reduce or gather operations to combine partial results. Pipeline stages exchange activations. Context-parallel attention moves or aggregates data associated with remote token positions. The bandwidth and latency of GPU interconnects such as NVLink/NVSwitch or network RDMA-capable fabrics can influence whether the partition pays off.

Training memory is not inference cache memory

During training, backpropagation may require saving or recomputing activations and attention-related intermediate values. A full-sequence score tensor may have quadratic token dimensions. The number of learned projection weights, however, does not grow merely because the sequence grows longer.

During inference decode, the long-lived KV cache is a different category of state. Keeping weights, optimizer state, saved training activations, temporary attention intermediates and serving KV cache separate prevents misleading capacity calculations.

Takeaway. Gradients explain parameter updates, but attention maps do not explain final outputs on their own. Distributed training partitions batches, features, layers, sequences or training state, and every partition introduces communication trade-offs that must be measured.

Next: Part 15 turns the concepts into repeatable numerical checks and disciplined profiling.