Chapter 27 Mapping and Estimating Large-Scale Training Systems

27.1 A model fits, but its training step does not

Suppose one million weights occupy two million bytes. A device has ten million bytes available. It is tempting to conclude that training fits with room to spare. Chapter 19’s update, however, needed gradients, and its optimizer retained information between steps. Chapter 20 also saved intermediate values so that the backward pass could follow the forward computation. None of those objects is included in the two-million-byte weight file.

We will build a resource estimate for an assigned small model, then distribute the work. This is a planning exercise that runs on a CPU with scalar arithmetic. It does not require a GPU or claim that the book operates an industrial training cluster. MG-2 already supplies a real training and recovery example; its float64 implementation has a different memory ledger from the mixed-precision scenario below.

27.2 Persistent arrays and the values between them

Example 27.1 — A one-million-parameter ledger. Assign $N=1{,}000{,}000$ trainable parameters. Working weights and accumulated gradients each use two bytes per parameter. A separate high-precision master copy uses four bytes per parameter. Adam has two four-byte moment arrays: one tracks a moving average of gradients, the other of squared gradients. These are declared storage assumptions, not properties of every optimizer or framework.

Table 27.1 — Derived model-state storage

Object	Why it remains available	Bytes per parameter	Total bytes
Working weights	Forward/backward computation	2	2,000,000
Gradients	Contributions before the update	2	2,000,000
Master weights	Accumulate finer updates	4	4,000,000
First and second moments	Adam’s update history	4 + 4	8,000,000
Total	All rows coexist in this estimate	16	16,000,000

The master copy and moments together account for 12 bytes per parameter. Some implementations organize precision or copies differently; counting a second copy that does not exist is as wrong as omitting one that does. Here, the ten-million-byte device already fails before activation storage is added.

Now assign a local microbatch of two sequences, 128 token positions per sequence, width 64, and four layers. For a deliberately simple activation estimate, suppose each layer retains ten arrays shaped like batch × positions × width, at two bytes per value. The count is

$$M_{\mathrm{act}}=2\times128\times64\times4\times10\times2=1{,}310{,}720\text{ bytes}.$$

The factors mean sequences, positions, coordinates, layers, saved arrays per layer, and bytes per coordinate, respectively. An activation is a value produced by a forward computation; it is not another trainable parameter. One array for one layer costs $2\times128\times64\times2=32{,}768$ bytes. Forty such arrays give the total, or 1.25 MiB.

Together, our state and saved activations require 17,310,720 bytes. Reserving an assigned 25% margin gives $17{,}310{,}720\times1.25=21{,}638{,}400$ bytes. The margin is a planning allowance, not a measured allocator peak. Attention-score arrays can grow with the square of sequence length in some implementations; embeddings, logits, temporary workspace, communication buffers and fragmentation need their own accounting. The ten-array assumption does not magically include them. A real allocation trace must replace this approximation before making a capacity promise.

Storage allocated for model state

1. Stored model numbersWeights 2 MB + master 4 MB
2. A list of numbersGradients 2 MB + moments 8 MB
3. An operationShard master/moments over 4 workers

Optimizer-only sharding reduces this storage from 16 MB to 7 MB per worker. Gradients are reset at update boundaries.

Within a step

1. A list of numbers40 saved arrays × 32,768 bytes
2. An operationKeep 16 arrays; recompute the rest
3. A list of numbers1,310,720 → 524,288 bytes

Recomputation saves activation storage and adds forward work.

Across a stop

1. Outside the modelConsistent weights + optimizer state
2. Outside the modelStep + RNG + data cursor + versions

Reload all required state at the same update boundary; a weight file alone cannot reproduce the continuation.

Figure 27.1 — Storage objects have different lifetimes. Original book diagram · CC BY-SA 4.0.

Figure 27.1 distinguishes state that persists across updates, activations that can be discarded or recomputed, and information needed to recover a run. All are storage, but their lifetimes differ. That difference creates the opportunities below.

27.3 Smaller numbers need a numerical plan

Why keep master weights if working weights are small? Consider an analogy with decimal rounding. A working weight stored to two decimal places is 1.00. Subtracting 0.003 and immediately rounding gives 1.00 again. Repeating that operation with only the rounded copy loses every update. If a finer master copy retains 0.997 and then 0.994, rounding the latter for computation gives 0.99. The small updates accumulated before becoming visible.

Actual FP16 and BF16 use binary floating-point formats, not this decimal rule. They both occupy two bytes but divide precision and exponent range differently. “Half the bytes” therefore does not imply identical overflow, underflow or rounding behavior. Mixed precision chooses which operations and stored arrays use which representation; it is not merely renaming every tensor.

A gradient can also become too small to represent. Loss scaling multiplies the loss before backward computation, producing correspondingly larger gradients, then removes that scale before updating. If a scale is 1,024, a gradient of 0.000001 becomes 0.001024 during the scaled calculation. Dividing by 1,024 restores its intended magnitude in sufficient precision. This cannot repair a value already lost before scaling, and too large a scale can overflow. A training loop must check nonfinite values and handle a skipped update consistently with its schedule. MG-2’s existing float64 run demonstrates the update sequence, not an FP16 loss-scaling implementation.

27.4 Several small batches can contribute to one update

Reducing a microbatch reduces its activation storage in our linear estimate. But changing the batch also changes which evidence one update averages. Gradient accumulation separates those choices: process several microbatches at the same parameter values, accumulate their contributions, then update once.

Example 27.2 — The denominator is prediction events. Take two microbatches with 2 and 6 valid target-token events. For one parameter, their mean gradients are 1 and 3. The sums are therefore $2\times1=2$ and $6\times3=18$. Across all eight events, the mean is $(2+18)/8=2.5$. Averaging the two reported means gives $(1+3)/2=2$, which gives the short microbatch too much influence.

For microbatch i with $n_i$ valid events and mean gradient $g_i$, the combined mean is

$$\bar g=\frac{\sum_i n_i g_i}{\sum_i n_i}.$$

The denominator counts events that actually contribute to the loss, excluding masked padding. The numerator reconstructs their gradient sum. With equal valid-event counts, dividing each mean microbatch loss by the number of microbatches is sufficient; with unequal counts, weights must reflect those event counts. Gradient clipping and the optimizer step belong at the declared combined-update boundary. Updating weights after the first microbatch would change the second gradient’s evaluation point, making this a different procedure.

An effective batch of eight equal-length sequences can use a local batch of two and four accumulation steps on one worker. With two data-parallel workers, local batch two and two accumulation steps also give eight sequences. For unequal lengths, publish effective valid-token counts as well. Accumulation saves activation space but still performs the work; it does not supply additional compute devices.

27.5 Splitting data, matrices and layers

A worker here is one cooperating training process. Data parallelism gives workers different data while they start from consistent model parameters. Each computes gradient contributions, and a collective operation combines them so that updates agree. If losses are token averages, the combination must honor the same denominator as Example 27.2.

Table 27.2 — A division of labor also creates traffic

Division	What is split	What must cross the boundary
Data parallelism	Training examples	Gradient sums/counts and consistent updated state
Tensor parallelism	Pieces of one matrix or layer	Partial results needed to finish layer operations
Pipeline parallelism	Consecutive layers	Activations forward and corresponding gradients backward
State sharding	Copies of optimizer, gradient or weight state	Reduced gradient shards or gathered weights when needed

Tensor parallelism might give one worker some output columns of a matrix and another the remaining columns; the following operation determines whether results must be concatenated or summed. Pipeline parallelism gives earlier and later layers to different workers. Microbatches let those stages overlap, but a stage can sit idle while the pipeline fills or drains. The schedule must also specify when weights change. A picture with four boxes does not prove fourfold speed.

Example 27.3 — Shard the largest copies. Use four workers with Example 27.1’s model. First shard only the 12-million-byte master/moment state evenly. Each worker keeps 3 million bytes of that state, plus its full 2-million-byte working weights and 2-million-byte gradients: 7 million bytes per worker. Adding the assigned local activations gives 8,310,720 bytes before buffers and margin.

If gradient storage is also ideally sharded, model state becomes $2+14/4=5.5$ MB per worker. If all 16 MB of model state is ideally sharded, the persistent share is 4 MB. These are progressively different arrangements, not numbers to substitute interchangeably. Layer execution may temporarily gather weights; peak memory exceeds a bare persistent share. Communication now supplies information that a worker no longer keeps locally.

Suppose a transfer must move an assigned one million bytes across a link delivering 100 million bytes per second. The payload alone needs $1{,}000{,}000/100{,}000{,}000=0.01$ seconds, or 10 ms, before latency and contention. If useful computation takes 2 ms and the transfer cannot overlap it, the iteration needs at least 12 ms. More workers can move the bottleneck from arithmetic to communication.

27.6 Discard now, reconstruct later, or recover after failure

Activation checkpointing keeps selected intermediate activations and recomputes missing portions during backward computation. In the assigned ledger, reducing retained arrays per layer from ten to four reduces saved activations to 524,288 bytes. Extra forward work reconstructs the omitted values. This trades computation for space; it does not reduce the parameter and moment rows in Table 27.1.

A recovery checkpoint serves another purpose: resume after stopping. Chapter 21’s MG-2 run already saves model and optimizer state, completed step, sampling/random state and configuration. A data cursor or batch-generator state is essential: restoring weights while selecting different subsequent examples is not the same continuation. Sharded training also needs shard identities, a consistent update boundary, tokenizer/data versions and evidence that all pieces belong together. A partially written file should not be treated as a completed checkpoint.

A practical small acceptance check is to compare uninterrupted updates with a save–reload–continue run under the same declared CPU stack. This is already implemented for MG-2; Part V does not claim a new distributed recovery experiment. Exact bitwise continuation across hardware or software changes requires additional evidence. A loadable weight file is enough for some model-running uses and insufficient for this training-recovery contract.

Our resource worksheet can now say what it estimates, what it omits, which array is split or recomputed, and what traffic follows. It cannot certify industrial training competence. Before changing a model again, another question becomes possible: what evidence would show that an internal component is responsible for a behavior?

27.7 Summary

Training memory includes working weights, gradients, optimizer/master state and activations with different lifetimes. Precision, accumulation, sharding, parallelism and recomputation address different terms. Each must preserve the intended update or explicitly change it. Recovery requires the state of the process, not just the latest weights.

27.8 Exercises

1. Keep Example 27.1’s parameter state but double the microbatch to four. Calculate saved activation bytes and the total before margin under the stated ten-array assumption.
2. Two microbatches have 3 and 1 valid events and mean gradients 2 and −2. Find the combined gradient. Explain the error in averaging the two means equally.
3. On four workers, someone reports 4 MB from fully sharding state as the complete peak memory. Identify three omitted terms.
4. A decimal working weight starts at 1.00 and each update subtracts 0.003. Compare two updates with immediate two-decimal rounding against keeping a finer master copy.
5. A resumed run loads weights and moments but reinitializes its data sampler. What can still be claimed, and what recovery test has it failed to establish?
6. Activation checkpointing halves saved activations but adds 30% to arithmetic time. Under what capacity condition can that still be useful, and why is it not evidence of a faster iteration?

27.9 Reference answers

1. Saved activations double to 2,621,440 bytes. Adding the unchanged 16,000,000-byte model state gives 18,621,440 bytes. Real attention or workspace costs require an additional model; this answer follows the supplied assumption only.
2. The event sum is $3\times2+1\times(-2)=4$, divided by four events, giving 1. An equal average of means gives 0 and assigns the single event as much weight as three events.
3. Local saved activations, temporary gathered weights, and communication/kernel buffers are three examples. Fragmentation and a safety margin may add more. Persistent divided storage is a lower-level term, not a measured peak.
4. Immediate rounding leaves 1.00 after both updates. The master copy retains 0.997, then 0.994; its rounded working copy becomes 0.99 after the second update. This decimal analogy shows why losing each tiny change differs from accumulating changes.
5. It can continue training from a saved numerical model/optimizer state. It has not established continuation of the same data sequence or equality with an uninterrupted run. Save the sampler state and compare the subsequent updates under the declared environment.
6. If the original run exceeds available memory and the reduced retained activations bring the real peak within capacity, recomputation can make training possible. More arithmetic may make each iteration slower; space feasibility and speed are different criteria.

27.10 References

• Micikevicius et al., Mixed Precision Training, 2018. Sections 3.1–3.2 explain master weights and loss scaling under FP16 conditions.
• Rajbhandari et al., ZeRO, 2020. Sections 3 and 5 separate replicated state categories, sharding and additional memory terms.
• Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, 2021. Section 2 describes parallel dimensions and pipeline scheduling; no reported cluster performance is used as our measurement.