Introduction

These are my notes on How to Scale Your Model, written in my own words. This document is only meant to be written, but not read.

1. Bounds

Terminology

There are three ways time passes in serving algorithms:

1. The computation itself, represented by:

   $T_{math} = \frac{\text{Computation FLOPS}}{\text{Accelerator FLOP/s}}$

2. The bandwidth within a chip from accelerator memory (high-bandwidth memory) to compute core.

3. The communication between chips:

   $T_{comms} = \frac{\text{Communication Bytes}}{\text{Memory Bandwidth Bytes/s}}$

To maximize training efficiency, we look at each of these metrics and determine the "limiting reactant". Does the bottleneck lie in computation or in communication?

Arithmetic intensity

One way to determine the bottleneck is via arithmetic intensity (AI):

$$\text{Arithmetic Intensity} = \frac{\text{Computation FLOPS}}{\text{Communication Bytes}}$$

In other words, for every byte moved, how much work are we performing on it? This is the arithmetic intensity of the algorithm. A high AI (not to be confused with artificial intelligence) indicates a bottleneck in compute, since we have a high load of computation per byte moved, while a low AI represents a bottleneck in communication. We find the critical arithmetic intensity by finding the arithmetic intensity of an algorithm that exceeds the accelerator's intensity, which is an indication of being compute bound.

A simple example for AI is finding the dot product between two vectors: $bf16[N] \cdot bf16[N] = bf16[1]$. There are $N$ product operations, $N-1$ summations, and each bf16 consists of 2N bytes: $AI = \frac{N + N - 1}{2N + 2N + 2} \approx \frac{1}{2}$.

Consider the additional bytes we use to write the result back into memory.

Roofline plots

Roofline plots are used to visualize computation and communication bounds. The x-axis represents AI and the y-axis represents peak FLOPs realized. An important computation not easily identifiable in the book is that the realized FLOPS is the product between bandwidth and AI. The ratio of FLOPS per byte moved multiplied by the number of bytes moved (bandwidth) gives us the total amount of FLOPS realized.

At lower AI, we are communication bound. As we increase the AI, we eventually become compute bound.

This inflection point can be determined by many variables such as dimension size, batch size and bandwidth. A higher bandwidth achieves the inflection point earlier since the compute bound ceiling is fixed for the hardware's peak FLOPs/s. Higher bandwidths scale more steeply in realized FLOPs, hitting this fixed ceiling sooner than lower bandwidths.

A quick note on sharding

Inflection points are dependent on batch size (B) when performing matmul on a single chip. For a [B, D] x [D, F] matmul, the [D, F] weight matrix is reused across all B inputs. The FLOPs done over the fixed bytes loaded from the weight matrix scales only with B. This is considering that we are performing large matmuls where B is comparatively small to D and output dimension. Under this regime, the arithmetic intensity approximately equals to batch size.

Conversely, when we shard the matmul, the inflection point is dependent on the dimension size (D). Batch size increases compute and communication equally, hence B is irrelevant. D, however, increases compute without increasing communication between chips. Importantly, D does not scale communication while scaling computation FLOPs because the reduction only occurs locally within the chips: the result maintains the same shape (X[:, :D//2] @ Y[:D//2, :]).

2. TPUs

Architecture

TPUs are highly specialized for matrix multiplications. They consist of a simple architecture consisting of a TensorCore that communicates with a high bandwidth memory (HBM). Within the TensorCore consist of the matrix multiply unit (MXU), vector unit (VPU), scalar unit and vector memory (VMEM).

The MXU performs large matmuls of size bf8[16, 128] @ bf[128, 128] for most generations. Weight matrices need to be padded to at least size of 128 for such a computation. The VPU performs general operations such as activations and vector operations. The scalar unit acts like a CPU communicating instructions between the VMEM and the MXU. The VMEM is a high bandwidth but lightweight storage that communicates between the HBM and the TensorCore.

If matrices can be stored inside the light storage of the VMEM, lower batch sizes can be used to achieve critical AI due to the higher bandwidth compared to the communication bottleneck of the HBM. The problem, though, is that VMEM has far less storage than HBM.

The aforementioned architecture describes the TensorCore. A TPU chip consists of 2-4 TensorCores attached to a shared HBM (generally). Four TPU chips are arranged in a set called a tray, connected to a CPU host via the PCIe network. Such chips are connected at an even larger scale to their four nearest neighbors via the ICI network in a pod.

Chips within a pod are connected via a structure that enables a maximum distance of $N/2$. Larger pods of size 16x16x16 are called super pods.

Systolic Arrays

I want to briefly mention the beautiful design of the systolic array utilized in the TPU's MXU to perform matmuls efficiently. The systolic array reuses values by "flowing" them through the array and tracking their accumulated values. fleetwood has an excellent animation.

For a traditional matmul of two $N \times N$ square matrices, we perform $N^3$ operations and read $2N^3$ values from memory (each $N^3$ operation requires two values read from memory, though this is the worst-case scenario with no reuse). With a systolic array, we still perform $N^3$ operations but only read each matrix from memory one time, totaling to $2N^2$, a factor of $N$ more efficient.

There are drawbacks to systolic memory, however, such as requiring the whole [128, 128] array to be populated. For a [64, 64] matmul, a significant portion of the array is empty with zero padding, wasting computational resources. Moreover, if we wish to multiply a larger matrix such as [512, 512], we tile the larger matmul into 128x128 chunks. Yet this espouses a memory bandwidth problem of communication between tiles.

3. Sharding

This section in particular was cognitively intense. The textbook is quite compressed and I wasn't satisfied with many of the explanations on a first pass read.

To begin, sharding is the breaking up of matrices across multiple accelerators. This is often due to large matrices being unable to fit into the high bandwidth memory of a single chip. Moreover, this also allows parallelism for lower latency.

Sharding notation

Sharding notation is as follows: $A[I_X, J_Y]$. We have multiple sets of axes. One set of axes refers to the axes of the matrix while another refers to the axes of the accelerator devices. In this case, we have the matrix axis $I$ sharded along the $X$ axis of the accelerators. The orientation and instantiation of devices in code is: Mesh(devices=((0, 1), (2, 3)), axis_names=('X', 'Y')). In this code, we have a 2x2 grid of accelerators with two axes. The accelerators within the grid communicate via the interchip connect (ICI).

Due to time, I won't be adding images, though the textbook has several visualizations I found incredibly valuable to test yourself through.

One note is that the $x,y$ axes span that of a traditional graph while the $i,j$ axes are flipped, where $i$ is the analog to the y-axis. I will go through two examples. In $A[I_X, J_Y]$, we first note that the axis going down $A$ is sharded by $X$, which spans the left-right axis. We split the $i$ axis in half and delegate it across the $X$ accelerators. We now do the same for the $j$ axis, splitting it in half, resulting in quadrants and delegating them across the $Y$ accelerators. In $A[I_{XY}, J]$, the values along the $I$ axis in the array $A$ are sharded across both accelerator axes $X$ and $Y$. The order matters, where we traverse sharding along the x-axis first before proceeding to the y-axis. By text, this likely does not mean much. It is important to go through these exercises actively.

Computation

We have a set of four computational regimes as well as a set of operations that navigate these regimes.

Neither multiplicand has a sharded contracting dimension

$A[I_X, J] \cdot B[J, K_Y] = C[I_X, K_Y]$. We can perform this operation without any difficulties, with an already sharded output.

One multiplicand has a sharded contracting dimension

$A[I, J_X] \cdot B[J, K_Y]$. There is a dimension mismatch. We use an AllGather function that gathers the shards along a dimension and reassembles the shards:

$$\text{AllGather}_X(A[I, J_X]) = A[I, J]$$

This sends the sharded data in each accelerator to all other accelerators within an axis via successive hops. AllGather can be done unidirectionally or bidirectionally. Considering a bidirectional AllGather, a single hop is defined by:

$$T_{hop} = \frac{V \cdot 2}{|X| \cdot W_{ICI}}$$

$V$ is the total bytes being moved, multiplied by two since it is bidirectional. $|X|$ is the number of accelerators along axis $X$ and $W_{ICI}$ is the bidirectional ICI bandwidth. The total time simplifies down to $T_{total} = \frac{V}{W_{ICI}}$ considering that we take $|X| / 2$ hops. The implication is that the total time is independent of the number of accelerators which must be gathered.

After some thought, I came up with a thought experiment: consider fixed memory and ICI bandwidth (precisely the equation above), but two regimes where we shard across a high number of accelerators and another regime sharded across fewer accelerators. The amount of memory required to saturate an accelerator is equal in both regimes, and each accelerator is connected to neighbors holding the same quantity. This localized view makes it clear that the accelerators ought to fill up in the same time.

Another note to make is that the idea of hops makes it misleadingly seem as if hops are a unit of time. It is true that the mesh with more accelerators requires more hops, but those hops do not need to take as much time as a mesh with fewer accelerators. Instead of hops, we ought to think of the data flow as continuous pipelines.

Both multiplicands have a contracting sharding dimension

$A[I, J_X] \cdot B[J_X, K] = C[I, K]$. This is a valid matmul but each device is left with a partial sum:

$$A[I, J_X] \cdot \text{LOCAL } B[J_X, K] = C[I, K]{U_X}$$

The local represents that we perform a partial sum but leave it unreduced. This ${U_X}$ indicates that the set of operations is unfinished and is pending the final sum.

The idea is equivalent to the sum of outer products. This is described by the AllReduce step, which takes twice the time of AllGather. It can be viewed as a sequence of a ReduceScatter followed by an AllGather step. ReduceScatter sums the partial sums in a manner similar to AllGather, yet leaves the result sharded. AllGather follows to copy the final result across all devices:

$$A[I, J_X] \cdot \text{LOCAL } B[J_X, K] = C[I, K]{U_X}$$

$$\text{AllReduce}_X(C[I, K]{U_X}) = C[I, K]$$

Both multiplicands have a non-contracting dimension sharded along the same axis

$A[I_X, J] \cdot B[J, K_X] = C[I_X, K_X]$. This is invalid math. One cannot shard both axes of a matrix across the same dimension. The solution is to AllGather either of the $X$ dimensions:

$$\text{AllGather}_X(A[I_X, J]) = A[I, J]$$

$$A[I, J] \cdot B[J, K_X] = C[I, K_X]$$

4. Transformers

I've formerly learned about the transformer architecture in class. The Illustrated Transformer is easily the most lucid explanation of the transformer.

Relating to scaling, however, we have batching and contracting dimenions in an array. When computing FLOPs, we multiply all of the dimensions where the contracting and batching dimensions are only counted once. Beyond this, there is not that much I'd find it useful to take away from this chapter.

5. Parallelization for Training

This chapter covers inter-chip parallelization for training. The central principle is that we aim to be compute bound. We ask whether the communication introduced by adding chips stays hidden behind the remaining compute.

Data Parallelism and FSDP

Data parallelism replicates the full model weights across chips and splits the batch. Each chip independently runs its shard, and an AllReduce after the backward pass synchronizes gradients so each replica updates identically. An important observation is that weight traffic does not depend on batch size — communication cost is a fixed AllReduce over parameters regardless of B.

Fully Sharded Data Parallelism (FSDP) goes further by sharding the weights rather than replicating them. Before each layer, an AllGather reconstructs the weights on each chip; a ReduceScatter handles gradients in place of the AllReduce. The tradeoff is more frequent communication for a lower per-chip memory footprint.

Compute-to-Communication Ratio

The compute-to-communication ratio grows linearly in per-device batch size. Staying compute bound requires:

$$\frac{B}{N} > \frac{C}{W}$$

where $B$ is the global batch size, $N$ is the number of chips, $C$ is bytes communicated, and $W$ is the ICI bandwidth. The per-device batch must generate enough compute to outpace the fixed communication cost.

Increasing parallelism (larger $N$) and shrinking batch size (smaller $B$) both reduce $B/N$, pushing us toward communication bound. Moreover, these often happen simultaneously in practice. Scaling is fundamentally the art of keeping enough work on each chip to bury the comms.

6. Inference

Inference splits into two fundamentally different jobs: prefill and generation. Training had one bottleneck (network communication) and one goal (stay compute bound). Inference is more complicated because prefill and generation sit on opposite sides of the roofline, and almost every design decision flows from which of the two we are optimizing and what is bottlenecking it.

Prefill and Generation

To begin, prefill is compute bound and generation is memory bound.

When processing a prompt, hundreds or thousands of tokens flow through the same weights simultaneously. The arithmetic intensity bar is cleared trivially: $B_{crit} \approx 240$ tokens on TPU v5e and approximately 280 on H100. Prefill behaves just like training, and we optimize it by maximizing FLOPs utilization.

Generation, however, produces one token at a time per request due to the sequential dependency. A single token cannot reuse weights across a batch dimension. To clear the same $B_{crit} \approx 240$ bar, we would need 240 concurrent requests batched together, which is genuinely difficult to achieve. Below that threshold, we are memory-bandwidth bound. Step time reduces to how long it takes to stream the parameters and KV caches out of HBM, and FLOPs sit idle.

The KV Cache

The KV cache is the new variable that changes the character of inference entirely. In training, memory is dominated by optimizer state and activation checkpoints. In inference there is no optimizer, no gradients, and activations are negligible. Every request, however, carries its own KV cache that grows with both batch size and sequence length.

This is why generation stays stubbornly memory bound. Weights are shared across the batch and are therefore amortizable, but KV caches are not. A larger batch means proportionally more KV bytes to load. Attention during generation consequently has low, constant arithmetic intensity: we are reading a large cache to perform a trivial amount of math, and no batching trick resolves this.

An important observation the textbook notes: for small-batch generation, per-step latency is lower-bounded by:

$$T_{step} \geq \frac{\text{params} + \text{KV cache (bytes)}}{\text{HBM bandwidth}}$$

Sharding and Disaggregation

The prefill and generation asymmetry changes how we shard.

Prefill can use any scheme training uses, including Megatron-style tensor parallelism up to the ICI bound and sequence parallelism beyond that. Generation is more restricted. FSDP is not viable since moving weights over ICI when already HBM bound only worsens the bottleneck. Data parallelism is pointless since separate replicas accomplish the same thing without the coordination overhead. The only viable option is model parallelism.

One note is that because generation is memory bound rather than compute bound, we can over-shard beyond the usual ICI limit to cut latency. FLOPs are not the constraint, so the typical over-sharding penalty does not apply.

The practical consequence is disaggregated serving. Prefill wants minimal sharding at batch size 1, while generation wants heavy sharding at large batch. The solution is to split them onto separate servers and ship KV caches over the network between them. This is the architecture behind real serving engines such as JetStream.

Conclusion

A recent life update has refocused my attentions elsewhere. I sped through the past two chapters in several hours and have also skipped Training LLaMa, Serving LLaMa, Profiling, and All About JAX. I plan on devoting more time to GPUs, however, when I have more time in the future.

In total, I spent around fifteen hours. The problems, though tedious, exposed many holes in my understanding that I would later return to patch if I cared enough to.

Most importantly, the purpose of this writing is to put my first technical blog post on the internet. Although these are simply my notes, I've always wanted to write and to show that to the world, though it would be remiss to ignore the exponential learning benefits of explaining topics in one's own words. I'll throw this on X. There will be more to come.