The wall a bigger cluster does not solve
Suppose an engineer at an Indian AI lab is handed a 64-GPU slice of a shared national compute cluster — H100s, 80 GB of high-bandwidth memory each, connected by NVLink within a node and InfiniBand across nodes — and told to pre-train a 7-billion-parameter Hindi-English foundation model from scratch. The instinct of anyone who has trained a ResNet on two GPUs is to reach for torch.nn.DataParallel-style thinking: put a full copy of the model on every GPU, split the batch across them, average the gradients. Sixty-four GPUs, sixty-four times the throughput, done.
It does not work, and the reason it does not work is not a throughput problem — it is a memory problem that no amount of extra GPUs fixes by itself. Pure data parallelism replicates the entire model, its gradients, and its optimizer state onto every single GPU. Adding the sixty-fifth GPU does not shrink what sits on GPU number one. If the model does not fit on one GPU, it does not fit on sixty-four GPUs run in pure data-parallel mode either — you have just built sixty-four expensive copies of the same out-of-memory error. This is the actual engineering content of "pre-training at scale": not just the scaling laws that tell you how many tokens and parameters to want, but the systems mechanics — ZeRO sharding, tensor parallelism, pipeline parallelism, activation recomputation — that make it physically possible to hold a model that large in GPU memory at all, and the communication topology that keeps thousands of GPUs from spending all their time waiting on each other. That is the ground this chapter covers.
Anatomy of the memory wall: what "16 bytes per parameter" means
Start with the arithmetic a training engineer actually does before writing a single line of distributed code. Modern pre-training uses mixed-precision training (Micikevicius et al., 2018): the forward and backward pass run in 16-bit floating point for speed, but Adam's optimizer state and a master copy of the weights are kept in 32-bit precision for numerical stability. Rajbhandari, Rasley, Ruwase, and He (2020), in the ZeRO paper, total this up per parameter Ψ:
- fp16 parameters: 2 bytes
- fp16 gradients: 2 bytes
- fp32 master-weight copy (Adam): 4 bytes
- fp32 first moment (momentum): 4 bytes
- fp32 second moment (variance): 4 bytes
Total: 16 bytes per parameter, before a single activation tensor is stored. Apply this to a real architecture — LLaMA-7B (Touvron et al., 2023), which has Ψ = 7×10⁹ parameters: 7×10⁹ × 16 = 1.12×10¹¹ bytes = 112 GB. A single H100 has 80 GB. The model's optimizer state alone — not its activations, not its data — cannot fit on one GPU. This is the wall. It exists whether you have one GPU or one thousand, because pure data parallelism does nothing to relieve it.
ZeRO: sharding the model state instead of replicating it
The fix, from the same ZeRO paper, is to stop replicating the 16Ψ bytes across every data-parallel rank and instead partition it, so each of the N GPUs in a data-parallel group holds only a 1/N slice. ZeRO does this in three stages of increasing aggressiveness:
- ZeRO-1 shards only the optimizer states (the 12Ψ bytes of master weights + momentum + variance). Parameters and gradients (4Ψ) stay replicated.
- ZeRO-2 additionally shards the gradients (2Ψ), so only the fp16 parameters (2Ψ) stay fully replicated.
- ZeRO-3 shards the parameters too. Every GPU permanently holds only 16Ψ/N bytes of model state, and gathers the full parameters for a given layer via an all-gather communication step just before that layer's forward or backward pass, then discards them again.
The trade-off is exactly what you would expect: each stage buys memory by spending communication. ZeRO-1 and ZeRO-2 add roughly the same communication volume as plain data parallelism (a single gradient all-reduce per step). ZeRO-3 adds a parameter all-gather on every layer, every forward pass and every backward pass — much more traffic, which is why ZeRO-3 is normally paired with fast intra-node interconnect. Trace the memory arithmetic for the 7B model on N = 64 GPUs:
def zero_stage_memory(psi, N):
fp16_params = 2 * psi
fp16_grads = 2 * psi
fp32_master = 4 * psi
fp32_m = 4 * psi # Adam first moment
fp32_v = 4 * psi # Adam second moment
opt_states = fp32_master + fp32_m + fp32_v # 12*psi
naive = fp16_params + fp16_grads + opt_states # 16*psi, on every GPU
stage1 = fp16_params + fp16_grads + opt_states / N
stage2 = fp16_params + (fp16_grads + opt_states) / N
stage3 = (fp16_params + fp16_grads + opt_states) / N
return naive, stage1, stage2, stage3
psi = 7e9
naive, s1, s2, s3 = zero_stage_memory(psi, 64)
print(f"naive (plain DP): {naive/1e9:.2f} GB/GPU")
print(f"ZeRO-1: {s1/1e9:.2f} GB/GPU")
print(f"ZeRO-2: {s2/1e9:.2f} GB/GPU")
print(f"ZeRO-3: {s3/1e9:.2f} GB/GPU")
# Output:
# naive (plain DP): 112.00 GB/GPU
# ZeRO-1: 29.31 GB/GPU
# ZeRO-2: 15.53 GB/GPU
# ZeRO-3: 1.75 GB/GPU
ZeRO-1 alone takes the model state from 112 GB — impossible on an 80 GB GPU — down to 29.31 GB, comfortably inside the budget with room left for activations. ZeRO-3 pushes it to 1.75 GB, at the cost of the extra all-gather traffic described above. This is why the 64-GPU cluster works and pure data parallelism did not: the GPUs are still running data-parallel training in the sense that each processes a different batch shard, but the model state is now partitioned across them rather than duplicated on each one.
The other half of the budget: activation memory
ZeRO only addresses parameters, gradients, and optimizer state. A transformer's forward pass also has to keep activation tensors in memory for use during the backward pass, and at 7B-model sequence lengths this is not a rounding error — it can exceed the model state itself. Korthikanti, Casper, Lym, McAfee, Andersch, Shoeybi, and Catanzaro (2022) derive the per-layer activation memory for a standard transformer block (self-attention plus MLP, mixed precision, standard non-fused attention) as:
activation_bytes_per_layer = s * b * h * (34 + 5 * a * s / h)
where s is sequence length, b is micro-batch size, h is hidden dimension, and a is the number of attention heads. Plug in LLaMA-7B's actual architecture (Touvron et al., 2023): h = 4096, a = 32, L = 32 layers, trained at sequence length s = 2048 with micro-batch b = 1:
def activation_memory_per_layer(s, b, h, a):
return s * b * h * (34 + 5 * a * s / h)
h, a, L, b = 4096, 32, 32, 1
for s in (2048, 4096):
per_layer = activation_memory_per_layer(s, b, h, a)
total = per_layer * L
print(f"s={s}: {per_layer/1e6:.1f} MB/layer, {total/1e9:.1f} GB total (32 layers)")
# Output:
# s=2048: 956.3 MB/layer, 30.6 GB total (32 layers)
# s=4096: 3254.8 MB/layer, 104.2 GB total (32 layers)
At the model's actual training sequence length, uncompressed activations alone cost 30.6 GB — comparable to the entire ZeRO-1 model state. Add that to the 29.31 GB from ZeRO-1 above and you are past 59 GB before counting the input embeddings, output logits, or any framework overhead; the margin on an 80 GB GPU is thinner than it looks. This is why activation checkpointing (Chen, Xu, Zhang, and Guestrin, 2016) is not an optional optimization at this scale but standard practice: instead of storing every intermediate activation, store only the input to each transformer layer and recompute the rest during the backward pass by re-running that layer's forward computation. The stored footprint drops to roughly 2 bytes (fp16) × s × b × h per layer:
for s in (2048, 4096):
ckpt_per_layer = 2 * s * 1 * 4096
ckpt_total = ckpt_per_layer * 32
print(f"s={s}: checkpointed total = {ckpt_total/1e6:.1f} MB (32 layers)")
# Output:
# s=2048: checkpointed total = 536.9 MB (32 layers)
# s=4096: checkpointed total = 1073.7 MB (32 layers)
Full checkpointing takes the 30.6 GB down to 0.54 GB — a 57× reduction — at the cost of re-running the forward pass a second time during backward, roughly a 30–33% increase in compute (FLOPs) for that saved memory. Production systems rarely checkpoint every layer; Korthikanti et al.'s selective recomputation instead checkpoints only the memory-heavy attention-softmax intermediates and keeps cheaper activations resident, recovering most of the memory savings for a much smaller compute tax. The general principle — trade recomputation FLOPs for memory — is the same one behind gradient checkpointing in any deep network; it simply becomes load-bearing rather than optional once sequence length and hidden size both scale.
When sharding alone is not enough: splitting the model itself
ZeRO-3 makes a model's parameters small on a per-GPU basis, but it does so by moving communication onto the critical path of every layer's forward and backward pass. Push ZeRO-3 to a 70B or 175B model over a cluster spanning many nodes and that all-gather traffic increasingly has to cross the slower inter-node network (InfiniBand, not NVLink), and step time degrades. At that scale, labs instead cut the model itself into pieces that live permanently on different GPUs — this is model parallelism, and it comes in two complementary flavors that Megatron-LM (Shoeybi, Patwary, Puri, LeGresley, Casper, and Catanzaro, 2019) and GPipe (Huang et al., 2019) each pioneered separately.
Tensor parallelism splits individual weight matrices — e.g., a transformer's attention and MLP projection matrices — column-wise or row-wise across GPUs, so each GPU computes a partial result and the pieces are combined with an all-reduce inside the layer. This requires an all-reduce on every single forward and backward pass through every split layer — very high communication frequency, so Megatron-LM deliberately keeps tensor-parallel groups confined to GPUs on the same NVLink-connected node (typically 8-way), where bandwidth is highest and latency lowest.
Pipeline parallelism instead splits the model by depth: GPU 0 holds layers 1–8, GPU 1 holds layers 9–16, and so on. Communication between pipeline stages is just the activation tensor flowing forward and the gradient tensor flowing backward at each stage boundary — infrequent, point-to-point, and far more tolerant of the slower inter-node network. This is why Narayanan et al. (2021) map tensor-parallel groups within a node and pipeline-parallel stages across nodes: each parallelism strategy is deliberately paired with the interconnect that matches its communication pattern.
Naive pipelining has its own cost, though: while GPU 3 waits for GPU 0's output on the very first micro-batch, GPUs 1 and 2 sit idle — the "pipeline bubble." With p pipeline stages and m micro-batches per training step, the fraction of step time lost to the bubble (GPipe's original schedule) is (p − 1) / m:
def bubble_fraction(p, m):
return (p - 1) / m
for m in (8, 32):
print(f"p=8, m={m}: bubble fraction = {bubble_fraction(8, m):.3f}")
# Output:
# p=8, m=8: bubble fraction = 0.875
# p=8, m=32: bubble fraction = 0.219
With only as many micro-batches as pipeline stages (m = p = 8), 87.5% of the step is wasted bubble — pipelining eight ways to get almost no speedup. Quadrupling the micro-batch count to 32 cuts the bubble to 21.9%. This is why production pipeline schedules always run many micro-batches per step, and why the 1F1B ("one-forward-one-backward") schedule from PipeDream (Harlap et al., 2018) improves on GPipe further by interleaving backward passes as soon as they become available instead of waiting for every micro-batch's forward pass to finish first, which reduces the peak activation memory that must be held in flight without changing the bubble fraction formula itself.
Correcting a common misconception
The misconception worth naming directly: "if I have enough GPUs, I can always train a bigger model by just adding more of them to a data-parallel run." This confuses two independent axes — throughput and memory. Data parallelism scales throughput: more replicas process more micro-batches per second. It does nothing for the memory ceiling, because every replica still needs the full model state resident on its own GPU. Consider GPT-3 at 175 billion parameters (Brown et al., 2020): 175×10⁹ × 16 bytes = 2.8×10¹² bytes = 2,800 GB of model state under naive mixed-precision Adam — thirty-five 80 GB GPUs' worth of memory, on a single replica, before a single activation is stored. No number of additional data-parallel replicas shrinks that 2,800 GB; you could have a million GPUs and each one would still need to hold the whole thing. The only ways to bring per-GPU memory below the hardware ceiling are to shard the model state (ZeRO) or to cut the model into pieces distributed across GPUs (tensor and pipeline parallelism) — memory-relieving strategies, which are then combined with data parallelism for throughput. They solve different problems and neither substitutes for the other.
3D parallelism: combining all three
Production pre-training runs at trillion-token scale combine all three axes at once — this is what Narayanan et al. (2021) call PTD-P (pipeline, tensor, data parallelism), commonly shortened to 3D parallelism. A GPU's identity is a coordinate in a three-dimensional grid: which pipeline stage it holds, which tensor-parallel shard within that stage, and which data-parallel replica of the whole arrangement it belongs to. The diagram below shows the minimal case — 2 pipeline stages × 2 tensor-parallel ranks × 2 data-parallel replicas = 8 GPUs — with the three communication patterns each drawn in the direction they actually travel.
Read the grid as a coordinate system: a GPU's total memory load is its 1/N slice of the ZeRO-sharded optimizer state (N = the data-parallel width, since sharding happens within a DP group), plus its slice of the model's layers (from the pipeline split), plus its slice of each layer's weight matrices (from the tensor split). The three communication types are deliberately mapped to three different network tiers by bandwidth requirement — tensor-parallel all-reduces are the most frequent and stay on NVLink within a node; pipeline sends are infrequent and can cross InfiniBand between nodes; data-parallel gradient all-reduces happen once per step and can even be overlapped with the backward pass computation itself to hide their latency almost entirely.
Token budgets and the trillion-token data pipeline
Everything so far has been about fitting a fixed-size model into GPU memory. The other half of "pre-training at scale" is deciding how many tokens that model should actually see, and what it takes to keep 64 or more GPUs fed with a corpus that large without ever loading the whole thing into RAM.
Kaplan et al. (2020) first showed that pre-training loss falls as a predictable power law in compute, model size, and dataset size, but Hoffmann et al. (2022) — the Chinchilla paper — corrected the implied allocation: for a fixed training-compute budget, loss is minimized when the number of training tokens D and the parameter count N are scaled together at roughly D ≈ 20N, not by growing N alone. For a 7B model, that compute-optimal point is:
def tokens_needed(psi, ratio=20):
return psi * ratio
psi = 7e9
chinchilla_optimal = tokens_needed(psi)
llama_actual = 1.0e12 # Touvron et al., 2023
print(f"Chinchilla-optimal tokens for 7B params: {chinchilla_optimal/1e9:.0f}B tokens")
print(f"LLaMA-7B actual training tokens: {llama_actual/1e12:.1f}T tokens ({llama_actual/psi:.0f} tokens/param)")
# Output:
# Chinchilla-optimal tokens for 7B params: 140B tokens
# LLaMA-7B actual training tokens: 1.0T tokens (143 tokens/param)
LLaMA-7B trains on roughly seven times the Chinchilla-optimal token count. This is deliberate, not wasteful: Touvron et al. (2023) optimize for inference cost at a target quality bar, not training-compute cost at a target quality bar. A smaller model trained far past its compute-optimal point costs more to train once but less to serve for the rest of its deployed life, and serving cost dominates total cost of ownership for any model queried at scale — which is exactly the regime a 7B Hindi-English foundation model deployed across an Indian AI lab's products would sit in.
Seeing a trillion tokens is a data-engineering problem as much as a compute one. LLaMA's 32,000-token vocabulary means each token ID fits in an unsigned 16-bit integer, so the tokenized corpus itself — not the raw text — occupies:
def corpus_size_bytes(n_tokens, vocab_size):
bytes_per_token = 2 if vocab_size < 65536 else 4 # uint16 vs uint32 token ids
return n_tokens * bytes_per_token
print(f"{corpus_size_bytes(1e12, 32000)/1e12:.1f} TB tokenized corpus on disk")
# Output:
# 2.0 TB tokenized corpus on disk
No single GPU's host machine needs that 2 TB resident in RAM at once. The standard approach is to store the tokenized corpus as one or more flat binary files of raw token IDs and access them through a memory-mapped array (e.g. NumPy's memmap): the operating system pages in only the byte ranges actually read, so a training process can stream sequentially through a multi-terabyte file while holding only a small working set in memory. The corpus is then partitioned once, up front, into as many contiguous shards as there are data-parallel replicas, and each replica's data-loading workers read only their own shard — with 64 GPUs running pure data parallelism, each rank streams roughly 2 TB / 64 ≈ 31.25 GB of the corpus, never the whole thing. Crucially, sharding happens only along the data-parallel axis: every tensor-parallel and pipeline-parallel rank within one data-parallel replica must consume the identical micro-batch, since they are jointly computing a single forward/backward pass over that data — sharding along the D dimension of the earlier 3D-parallelism grid, never along the T or P dimensions.
Finally, tokens per step translate directly into the number of optimizer steps a full pass over the corpus requires. LLaMA-7B trains with sequences of length 2048 and a global batch of 2048 sequences, for 2048 × 2048 ≈ 4.19 million tokens per optimizer step:
global_batch_tokens = 2048 * 2048
steps = 1e12 / global_batch_tokens
print(f"{global_batch_tokens:,} tokens/step -> {steps:,.0f} steps to see 1T tokens once")
# Output:
# 4,194,304 tokens/step -> 238,419 steps to see 1T tokens once
Almost a quarter-million optimizer steps, each one requiring every GPU in the 3D-parallel grid to complete its slice of the forward pass, backward pass, and the communication described in the sections above, before the next step's data is even loaded. The token budget and the memory budget are not separate problems — the number of tokens a lab chooses to train on is what determines how many steps the parallelism scheme above has to survive without falling over.
Active recall
Attempt each question before reading its answer.
- A 13B-parameter model trains under mixed-precision Adam on N = 128 GPUs. Compute the model-state memory per GPU under ZeRO-1, ZeRO-2, and ZeRO-3.
- Using the activation-memory formula from this chapter with LLaMA-7B's architecture (h = 4096, a = 32, L = 32, b = 1), compute total (non-checkpointed) activation memory across all 32 layers at sequence length s = 512 and at s = 1024. Does doubling s from 512 to 1024 exactly double the total, and why or why not?
- A colleague says: "We have 500 GPUs, let's just use pure data parallelism to train a 70B model." Using the 16-bytes-per-parameter rule, explain concretely why this fails regardless of GPU count, and name the two categories of fix.
- Why does ZeRO-3 require an all-gather of a layer's parameters immediately before that layer's forward pass, while ZeRO-1 does not need this step at all?
- Megatron-LM confines tensor-parallel groups to GPUs within a single NVLink node rather than spreading them across nodes. Give the communication-pattern argument for why.
- A pipeline has p = 8 stages. If the training step uses m = 16 micro-batches, compute the GPipe-schedule bubble fraction, and state one way (other than increasing m) that a real system reduces the practical cost of that bubble.
Worked answers
1. Ψ = 13×10⁹, N = 128. Naive: 13×10⁹ × 16 = 208 GB/GPU. ZeRO-1: 4Ψ + 12Ψ/N = 52 GB + (156×10⁹/128) = 52 GB + 1.22 GB ≈ 53.22 GB/GPU. ZeRO-2: 2Ψ + 14Ψ/N = 26 GB + (182×10⁹/128) = 26 GB + 1.42 GB ≈ 27.42 GB/GPU. ZeRO-3: 16Ψ/N = 208×10⁹/128 ≈ 1.625 GB/GPU.
2. At s = 512: per-layer = 512 × 1 × 4096 × (34 + 5×32×512/4096) = 512 × 4096 × 54 ≈ 113.2 MB/layer, × 32 layers ≈ 3.62 GB total. At s = 1024: per-layer = 1024 × 1 × 4096 × (34 + 5×32×1024/4096) = 1024 × 4096 × 74 ≈ 310.4 MB/layer, × 32 layers ≈ 9.93 GB total. That is a 2.74× increase, not 2×. The formula sbh(34 + 5as/h) has two terms with different scaling in s: the "34" constant scales only linearly with s through the leading sbh factor, but the "5as/h" term scales with s twice — once through the leading sbh and once through the s inside 5as/h — because it captures the attention score matrix, which is quadratic in sequence length (an s×s attention matrix per head). At s = 512 the bracket is 34+20=54, so the quadratic term is already more than a third of the total; doubling s to 1024 makes the bracket 34+40=74, so the quadratic term now dominates the bracket even more, which is why the total grows faster than linearly but has not yet reached the ~3.4× seen when doubling from 2048 to 4096 elsewhere in this chapter — the larger s gets, the closer the scaling approaches true s² as the quadratic term increasingly swamps the constant one.
3. 70×10⁹ × 16 = 1.12×10¹² bytes = 1,120 GB of model state per replica under naive mixed-precision Adam. An H100 has 80 GB. Pure data parallelism puts a full, undivided copy of that 1,120 GB on every one of the 500 GPUs — it does not fit on any single GPU regardless of how many GPUs exist, because data parallelism only splits the batch, never the model state. The two categories of fix: (a) shard the model state itself across GPUs within a group (ZeRO stages 1–3), or (b) partition the model architecturally across GPUs (tensor parallelism splitting weight matrices, pipeline parallelism splitting layers). Data parallelism is then layered back on top of either fix for throughput, not as a substitute for it.
4. Under ZeRO-3, no GPU permanently holds a complete copy of any parameter — each holds only a 1/N shard of every layer's weights. To actually run the matrix multiplications of a given layer's forward pass, the full weight matrix for that layer must exist somewhere in memory momentarily, so every GPU in the group all-gathers the missing shards from its peers immediately before using that layer, then can discard the reassembled full weights right after. ZeRO-1 never shards the parameters themselves (only the optimizer states, which are only needed during the optimizer's update step, not during forward/backward), so every GPU already holds the complete parameters it needs to compute — no gather required.
5. Tensor parallelism requires an all-reduce on every forward pass and every backward pass through every split layer — many times per training step, each one blocking further computation until it completes. That communication pattern is bandwidth-hungry and latency-sensitive, so it needs the fastest possible interconnect, which is NVLink between GPUs on the same physical node (hundreds of GB/s), not InfiniBand between nodes (tens of GB/s, higher latency). Pipeline parallelism, by contrast, only exchanges one activation tensor and one gradient tensor per micro-batch at each stage boundary — far less frequent and more tolerant of inter-node latency — so it is the parallelism dimension assigned to cross the slower network between nodes.
6. Bubble fraction = (p − 1) / m = 7 / 16 = 0.4375, so 43.75% of the step time is bubble under the naive GPipe schedule. Beyond simply increasing m, the practical fix used in production (PipeDream's 1F1B schedule, Harlap et al., 2018, and Narayanan et al.'s interleaved variant, 2021) is to reorder the schedule so each GPU alternates forward and backward micro-batches instead of running every forward before any backward — this does not change the bubble-fraction formula itself, but it caps the number of activation buffers a GPU must hold in flight at once, which lets a system afford a larger m (and thus a smaller realized bubble) within the same memory budget.
Think About It
Think about this: How would you explain llm pre-training at scale: from theory to trillion tokens to a friend who has never seen a computer? What real-world analogy would you use? Imagine you had to build a system using these concepts — what would be your first step? Try this: before moving on, write down three things you learned and one question you still have.
Key Takeaways — Summary and Recap
Let us recap what we covered: the core ideas behind llm pre-training at scale: from theory to trillion tokens, how they connect to real-world applications, and why they matter for your journey in computer science. Remember these key points as you move forward. For competitive exam preparation (CBSE, JEE, BITSAT), focus on understanding the WHY behind each concept, not just the WHAT.