Training Infrastructure

Watch a single GPU crawl through a trillion-token dataset, then scale up to thousands of GPUs across many servers, coordinating gradients and parameters until a fully trained model emerges. ChatGPT is only a familiar entry point: the dataset size and GPU counts shown are illustrative, not published details of how any particular model was trained.

What infrastructure is required to train a large AI model? AI training infrastructure is the data pipeline, accelerators, high-speed network, storage, job scheduler and monitoring that together keep a training run fed, synchronised and recoverable. A large model is trained by repeating one loop millions of times: read a batch of tokenised text, run it forward through the network to get a loss, run backward to get a gradient for every parameter, combine the gradients from every worker, let an optimiser update the parameters, and occasionally write a checkpoint. Distribution is what you do when that loop does not fit on one accelerator, either because the model's training state exceeds device memory or because one device would need years to process the data. The model is then split, replicated or sharded across many accelerators, and the work shifts from raw arithmetic to keeping them fed with data, synchronised over a fast network, and recoverable when a component fails.

That is why the chapter is about systems and not only about GPUs. Four things must hold at once for a run to be productive: the input pipeline delivers tokens at least as fast as the accelerators consume them, the training state fits in memory with room for activations, the communication that combines gradients is mostly hidden behind computation, and the job can resume from a recent checkpoint after the failure that will eventually happen. Fail any one and an expensive cluster idles.

We will follow one running example: a 7-billion-parameter language model pretrained on 64 accelerators (8 nodes of 8) with a global batch of 1,048,576 tokens. The hardware is invented and round so the arithmetic is checkable. Hardware internals are covered in AI Accelerators and the parallelism mathematics in Parallel Processing; here we assemble them into a training system.

Everything else in this chapter is an answer to a problem the single step creates, so start there. A step turns a batch of text into a small change in the parameters. The data loader assembles a batch of token sequences (for example 512 sequences of 2,048 tokens). The forward pass pushes them through every layer, storing intermediate results called activations, and ends with a loss: a number that measures how badly the model predicted each next token (see next-token prediction). The backward pass walks the layers in reverse, using the stored activations to compute the gradient of the loss with respect to every parameter. In a distributed job each worker has seen only part of the batch, so gradients are synchronised across workers. Finally the optimiser reads the gradients and its own running statistics and writes new parameter values. Every so often the whole state is saved as a checkpoint.

Three details turn this from a diagram into an engineering problem. First, the forward pass must keep activations until the backward pass consumes them, so memory use peaks in the middle of the step, not at the start. Second, the gradient synchronisation sits on the critical path: no worker can update until it has the combined gradient, so the slowest worker or the slowest link sets the pace of everybody. Third, the optimiser is stateful. It remembers running averages of past gradients, so an interrupted run cannot be resumed from the weights alone. The state that must survive a restart is much larger than the model you eventually ship. These three facts motivate the memory accounting, the communication patterns and the checkpoint design that follow.

A training dataset is a manufactured artefact, and its quality limits the model more than most infrastructure choices do. The pipeline runs in a fixed order. Collection begins with rights and provenance: which sources you may use, under what licence or consent, with what personal data removed. Cleaning strips boilerplate, markup, encoding damage and low-quality text. Deduplication removes exact and near-duplicate documents, which otherwise waste compute, bias the model toward repeated text, and, critically, leak into evaluation sets. Tokenisation converts text to integer token ids with a fixed vocabulary. Sharding packs the ids into large files, and streaming plus shuffling feeds them to workers in an order that approximates random sampling.

Two properties of this pipeline matter to infrastructure. First, token data is small relative to its cost to produce: our global batch of 1,048,576 tokens stored as 4-byte integer ids is only 4.2 MB per step, so bandwidth for text is rarely the problem. The troubles are access patterns: millions of tiny files, random reads that defeat the storage system's caching, decompression and on-the-fly tokenisation burning CPU, or a loader worker that blocks on a slow remote read. (Image and video training have much larger batches and much more genuine bandwidth demand.) Second, the dataset must be a versioned artefact. A run should record that it consumed "shards v7, manifest hash X", so that a later question such as "which documents did this model see?" or "did this benchmark leak?" can be answered. Shuffling needs a design too: a global shuffle of the sample index with a seed is reproducible and resumable; a local shuffle buffer on each worker is cheap but depends on arrival order.

Data starvation, worked through

Return to the 64-accelerator run. Its step takes 4.30 s of compute (derived below), and each of the 8 nodes must supply 64 sequences per step: 64 / 4.30 = 14.9 sequences per second per node. Suppose (illustrative) the nodes read cold shards from remote storage and each loader sustains only 10.7 sequences per second. Producing a node's batch then takes 64 / 10.7 = 6.0 s, longer than the 4.3 s of compute. With prefetching the loader works on batch n+1 while the accelerators run step n, but it can only hide the shorter of the two times; the step settles at max(4.3, 6.0) = 6.0 s. Accelerator busy fraction falls to 4.3 / 6.0 = 71.7 percent, and throughput from 243,810 to 174,763 tokens per second. Without prefetch the step is load then compute: 6.0 + 4.3 = 10.3 s, 41.7 percent busy. Nothing on the accelerators looks wrong; compute kernels are as fast as ever. The loss shows up as gaps.

The fixes are all about the loader, not the accelerators: larger sequential shards, more loader workers per node, pre-tokenising offline instead of in the training loop, a local NVMe cache in front of object storage, and deeper prefetch queues. The diagnostic is equally specific: measure the time each rank spends waiting on the input iterator at the start of the step. If that wait is more than a few percent of step time, no amount of tuning of the model code will help.

Before choosing how to distribute a model, count what one training replica must hold. Training holds five kinds of memory, and only the first is what people usually mean by "the model".

1. 1

   Weights (parameters)

   The model itself: parameter count times bytes per parameter in the compute precision, such as 2 bytes in BF16.
2. 2

   Gradients

   One gradient value per parameter, usually the same precision as the weights, alive from the backward pass until the optimiser consumes them.
3. 3

   Optimiser state

   Running statistics the optimiser keeps per parameter. Adam and AdamW keep two moment tensors (first and second moments), commonly in FP32 for numerical stability. Plain SGD keeps none, SGD with momentum keeps one.
4. 4

   Activations

   Intermediate outputs of every layer, retained for the backward pass. These scale with batch size, sequence length and depth, not with parameter count alone, and can be traded for recomputation (activation checkpointing).
5. 5

   Buffers and fragmentation

   Communication buckets, temporary workspaces for kernels, the memory allocator's fragmentation and framework overhead. Small per item, and the reason a computed total is always a lower bound.

Weights, gradients and optimiser state are called the model state; they scale with parameters. Activations scale with the batch and are the part you can tune without changing the model.

An illustrative 1-billion-parameter budget

Take 1 billion parameters (1 × 10⁹) and decimal gigabytes (1 GB = 10⁹ bytes; GiB would be 2³⁰ bytes and 7 percent larger). In BF16, a parameter is 2 bytes, so weights are 1e9 × 2 B = 2 GB and gradients another 2 GB. Adam keeps two moment tensors, each FP32 at 4 bytes per parameter: 2 × 4 B × 1e9 = 8 GB. The four-line total is 2 + 2 + 8 = 12 GB. Many mixed-precision recipes also keep an FP32 master copy of the weights, because small updates would be lost when added to BF16 numbers; that adds 4 B × 1e9 = 4 GB, for 16 GB, the commonly quoted 16 bytes per parameter for Adam in mixed precision. Activations, temporary buffers and fragmentation come on top.

Scaling to the running 7-billion-parameter model gives weights 14 GB, gradients 14 GB, Adam moments 56 GB and master weights 28 GB: 112 GB of model state, which exceeds the memory of an 80 GB accelerator (an illustrative size) before a single activation is stored. Two conclusions follow directly. A plain replicated copy of this model cannot be trained on one such device, which is the reason the techniques in the next sections exist. And inference on the same model needs only the 14 GB of weights plus its KV cache, which is why a model you cannot train on a device may still be servable on it (see Inference Infrastructure).

Three variants are worth distinguishing. Full-parameter training updates every weight and pays the whole bill above. Parameter-efficient fine-tuning (PEFT, section on fine-tuning below) freezes the base weights and trains small adapters, so gradients and optimiser state exist only for the adapters. Inference holds weights and per-request caches, no gradients or optimiser state. Activation memory depends on none of these choices in a simple way: it tracks batch and sequence length, so the same model on the same device can fit at micro-batch 2 and fail at micro-batch 8.

The simplest way to use many devices when the model fits on one, and the usual starting point for multi-GPU training, is data parallelism. Every worker holds a full replica of the model, each processes a different slice of the batch, and the gradients are averaged so that all replicas apply the same update and stay identical. Because the loss is an average over examples, the average of the per-worker gradients equals the gradient of the whole batch; that equality is the correctness argument for the entire technique.

The averaging is an all-reduce: every worker contributes a vector and every worker receives the sum. A ring all-reduce over N workers moves about 2(N-1)/N times the gradient size through each worker's link, nearly independent of N, which is why data parallelism scales so far. It is organised in buckets: the framework starts reducing the gradients of late layers while the backward pass is still computing earlier ones, hiding much of the communication behind compute. For the 1B model at 2 GB of BF16 gradients on 8 workers, each worker transfers 2 × (7/8) × 2 GB = 3.5 GB per step; at an illustrative effective 400 GB/s per-worker bandwidth that is 8.75 ms if it were not overlapped at all. For the 7B model the same arithmetic gives 14 GB × 1.75 = 24.5 GB, about 61 ms at the same bandwidth, and slower links or multi-node paths make it larger. See Parallel Processing for the full cost derivation.

A two-worker step with real numbers

Take the smallest possible model: one weight w = 1.0 predicting y from x, loss = mean((w·x − y)²), with gradient mean(2x(wx − y)). Eight examples are split across two workers, four each, and the learning rate is 0.01. Worker 0 computes a local gradient of −13.5 on its four examples, worker 1 computes −9.5 on its four. The all-reduce sums them to −23.0 and each divides by the number of workers (2) to get −11.5. A single worker on all eight examples computes exactly −11.5. Both replicas then apply w ← 1.0 − 0.01 × (−11.5) = 1.115, so they start the next step identical.

data = [(1.0, 2.0), (2.0, 3.0), (3.0, 7.0), (4.0, 7.0),
        (1.5, 2.0), (2.5, 6.0), (3.5, 6.0), (0.5, 2.0)]
w, lr = 1.0, 0.01

def grad(w, shard):
    # d/dw of mean((w*x - y)^2) = mean(2*x*(w*x - y))
    return sum(2 * x * (w * x - y) for x, y in shard) / len(shard)

shards = [data[:4], data[4:]]            # one shard per worker
local = [grad(w, s) for s in shards]     # [-13.5, -9.5]
avg = sum(local) / len(local)            # all-reduce(sum) / world_size
print(avg, grad(w, data))                # -11.5 -11.5
print(w - lr * avg)                      # 1.115

Two bookkeeping terms need separating. The micro-batch is what one worker pushes through the model in one forward/backward pass, limited by activation memory. The global batch is micro-batch × data-parallel workers × gradient-accumulation steps. In our 7B run, 64 workers with 8 sequences each is 512 sequences, so the global batch is 512 × 2,048 = 1,048,576 tokens with no accumulation. If memory forced micro-batch 4, accumulating over 2 micro-batches gives the same global batch: gradients are summed locally for the first micro-batch and the all-reduce runs only after the last, so communication is paid once per optimiser step. The important point is that the optimiser sees one global batch, so changing the number of workers without changing the global batch should not change training behaviour, while changing the global batch usually requires retuning the learning rate and warm-up.

Data parallelism needs a full replica per device. When even the model state or a single layer's activations exceed device memory, the model itself must be divided. There are two classic ways to cut a network, and real systems use both together with data parallelism.

Tensor (intra-layer) parallelism splits the matrices inside each layer across devices. A large matrix multiply is divided by columns or rows, every device computes its part, and results are combined with an all-reduce or all-gather. In a transformer block this happens a small, fixed number of times per layer in each direction. That makes communication frequent, latency sensitive and on the critical path of every layer, so tensor parallelism is normally confined to devices joined by the fastest links, typically within one node, and used with small degrees such as 2, 4 or 8.

Pipeline (inter-layer) parallelism assigns consecutive groups of layers to different devices, called stages. A batch is split into microbatches that flow through the stages like items on an assembly line, and only activations at stage boundaries cross the network. That is far less traffic than tensor parallelism, so pipelines are the usual way to span nodes. The cost is the pipeline bubble: at the start of a step the later stages have nothing to do, and at the end the earlier stages sit idle.

With p stages and m microbatches in this simple fill-and-drain schedule, the idle fraction is (p − 1) / (m + p − 1). For p = 4 and m = 4 that is 3/7, or 42.9 percent, as the figure shows. Using m = 32 microbatches reduces it to 3/35 = 8.6 percent; more microbatches means smaller ones, which has its own efficiency cost, and cleverer schedules interleave forward and backward passes to reduce activation memory and bubble together. The point to carry away is the trade: pipeline parallelism buys memory fit across nodes with modest bandwidth and pays with idle time that you control through the microbatch count.

Large runs combine the three. A hypothetical 64-device job could use tensor parallelism 8 within a node, pipeline parallelism 2 across two nodes, and data parallelism 4 across the resulting 16-device model replicas: 8 × 2 × 4 = 64. Each dimension is placed on the network tier that matches its communication pattern: tensor inside the fastest domain, pipeline across nodes, data parallel gradient all-reduce across everything else. The derivations of these costs, and why each pattern needs the bandwidth it does, are in Parallel Processing; what matters for infrastructure is that the job's parallel layout is a mapping onto physical topology. A scheduler that places ranks on arbitrary free GPUs can turn a well-designed layout into a slow one.

There is a middle path between replicating everything and rewriting the model structure. Plain data parallelism is wasteful: all 64 replicas hold identical copies of the 56 GB of Adam moments and 28 GB of master weights. Sharded data parallelism, introduced for deep learning in the ZeRO paper (Rajbhandari et al., 2019, "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models") and implemented in several frameworks under names such as ZeRO and FSDP (fully sharded data parallel), keeps the data-parallel programming model but partitions the redundant state across the workers. The ZeRO stages shard progressively more:

1. 1

   Stage 1: shard optimiser state

   Each worker owns 1/N of the moments and master weights and updates only its slice, then the updated parameters are shared. Gradients and weights are still replicated.
2. 2

   Stage 2: also shard gradients

   Instead of an all-reduce, gradients are reduce-scattered so each worker receives only the averaged gradient for the slice it owns.
3. 3

   Stage 3: also shard the weights

   No worker holds the full model at rest. Before each layer runs, its weights are all-gathered from peers and freed afterwards; this is the idea behind FSDP-style training.

Apply it to the earlier budgets (decimal GB, model state only, activations extra). For 1B parameters on 8 workers with 16 bytes per parameter: replicated 16 GB; stage 1 is 2 + 2 + 12/8 = 5.5 GB; stage 2 is 2 + 14/8 = 3.75 GB; stage 3 is 16/8 = 2 GB. For the 7B run on 64 workers: replicated 112 GB; stage 1 is 14 + 14 + 84/64 = 29.3 GB; stage 2 is 14 + 98/64 = 15.5 GB; stage 3 is 112/64 = 1.75 GB. The 112 GB that could not fit an 80 GB device becomes comfortably small, leaving the memory for activations.

Memory is not free here; it is traded for communication and complexity. Replicated data parallelism moves roughly 2 × the gradient size per step (the all-reduce). Stage 3 must all-gather the weights for the forward pass, again for the backward pass, and reduce-scatter the gradients: about 3 × the parameter size, around 1.5 times the traffic, and the gathers are on the critical path unless they are prefetched and overlapped with the previous layer's compute. Stages 1 and 2 add little traffic. Which stage is worth it depends on whether memory or the network is your limiting resource: if the model fits comfortably, stage 1 or 2 captures most of the benefit for little cost, and stage 3 is for when it does not. Implementations differ in granularity (per layer or per flat parameter group), in how they overlap, in what they offload to host memory or NVMe, and in how they checkpoint, so treat the stage numbers as a conceptual map, not a guarantee of identical behaviour across frameworks.

A training job is one tenant of a larger system whose parts have different responsibilities. It helps to separate the control plane, which decides who runs where and watches health, from the data plane, which carries the tokens, gradients and checkpoints that make up the actual training work.

The scheduler is the first piece that has no analogue in a single-machine script. Training needs gang scheduling: all 64 ranks must start together, because a collective with a missing member never completes, so a scheduler that starts 63 and waits holds 63 accelerators idle. It also applies quotas and priorities between teams, prefers placements that respect topology (all ranks of a tensor-parallel group on one node), and decides what to do on failure. Orchestration then launches containers with the right image, environment, ranks and secrets. Isolation between tenants covers more than containers: separate storage prefixes, credentials scoped to the dataset a team is allowed to read, and network policy keep one team's data and checkpoints away from another. Monitoring watches accelerators, hosts, network links and the job's own step metrics. It matters most in distributed runs because a failure is often detected as a collective that hangs, not as a clean exception.

Two physical facts shape the data plane. The fabric is hierarchical: links inside a node are much faster than links between nodes, and a collective can only go as fast as its slowest hop, so topology awareness matters. And storage is shared and bursty: all ranks read at once at startup, and all write at once at each checkpoint, so the storage system is sized for the peak of both, not the average. None of this is visible in the model code, which is why a training platform team exists in practice.

A checkpoint is a snapshot from which the run can continue as though it had never been interrupted. Saving the weights alone is enough to serve the model but not to resume training, and the difference is the most common source of subtly wrong restarts. A resumable checkpoint contains the model weights (including the FP32 master copy if the recipe uses one), the optimiser state (Adam moments and step count), the learning-rate scheduler position, the random-number-generator state for each rank, and the data-loader progress: which samples have been consumed. Without the optimiser state, the first updates after resume are taken with reset moments and the loss typically spikes. Without the scheduler, the learning rate restarts from warm-up. Without the data position, the run repeats or skips data.

Size follows from the memory accounting. For the 7B run, a resumable checkpoint stores FP32 master weights (4 B), Adam moments (8 B) per parameter: 12 B × 7e9 = 84 GB, decimal. The BF16 weights (14 GB) can be regenerated from the master copy, and a separate inference-only export of 14 GB is what you hand onwards to serving. With 64 workers sharing the state, each rank writes its own shard of 84 / 64 = 1.31 GB. If shared storage sustains an illustrative 10 GB/s aggregate write, the checkpoint takes 84 / 10 = 8.4 s, during which the naive implementation stalls training. Sharded (distributed) checkpointing writes each rank's slice in parallel instead of gathering everything onto one process, which both avoids a memory spike and uses the aggregate bandwidth; some implementations can also reload under a different number of workers by resharding.

The consistency rules follow from the diagram. A checkpoint is valid only when every shard is present and verified, so the manifest is the commit record and is written last; the directory is renamed, or a pointer file switched, only afterwards, using whatever atomic primitive the storage system offers. Object stores differ from file systems here: some offer no atomic directory rename, in which case a manifest-last convention with a verified checksum list does the job. Restart code must scan back from the newest checkpoint to the newest valid one, and retention should keep several, since the newest may be the one being written when a node died.

The data-loader position deserves care. The robust design makes sample order a pure function of (seed, epoch, global sample index), so the only state to save is the global index of the next sample. Restart then needs no replay of the loader and works even if the worker count changes. A design that stores each loader worker's local buffer position is faster to write but ties the checkpoint to the exact original layout.

Return to the 7B run. Checkpoints are written every 500 steps, and the last committed one is step 8,000. At step 8,400 one node loses an accelerator. What happens, and what does it cost? First the step time: 6 × parameters × tokens per step is the standard approximation for dense-transformer training FLOPs (6 × 7e9 × 1,048,576 = 4.40e16 FLOPs per step, ignoring attention and recomputation). If each accelerator's peak is an illustrative 4.0e14 FLOP/s and we achieve 40 percent of it, delivered compute is 64 × 4.0e14 × 0.40 = 1.024e16 FLOP/s, so a step takes 4.40e16 / 1.024e16 = 4.30 s and the run processes about 243,800 tokens per second. This is the step time used throughout the chapter.

1. 1

   Detection (about 60 s, illustrative)

   The other 63 ranks block inside the collective that waits for the dead rank. A watchdog timeout, or the orchestrator noticing the lost node, declares the job failed. Detection time is a tunable: shorter timeouts recover faster but may kill a job that is merely slow.
2. 2

   Reschedule and relaunch (about 300 s)

   The orchestrator tears down the remaining processes, finds a healthy replacement node (or waits for the failed one to be repaired or swapped), and relaunches all 64 ranks as a gang.
3. 3

   Load the checkpoint (about 120 s)

   Ranks scan for the newest valid checkpoint (step 8,000), read their shards (84 GB total, about 8.4 s at 10 GB/s plus initialisation, framework start-up and cold caches), and restore weights, optimiser state, scheduler, RNG and sample index.
4. 4

   Replay steps 8,000 to 8,400

   The 400 steps already computed are computed again. Nothing from the failed attempt is reused: the updates from steps 8,001 to 8,400 existed only in the memory of the lost workers.

The cost: 400 steps × 4.30 s = 1,720 s of work repeated, plus 480 s of recovery, 2,200 s in total, which at 64 accelerators is 64 × 2,200 / 3,600 = 39.1 accelerator-hours for a single failure. Had the checkpoint interval been 100 steps, the repeated work would have been at most 430 s. Had it been 2,000 steps, up to 8,600 s.

Duplicate and skipped data on restart

Replay raises an honest question: does the model see steps 8,000 to 8,400 twice? Not in any harmful sense. The restored model state is the one from step 8,000, so the updates from the failed attempt are discarded, and the replayed steps are the first time those updates are applied to this state. Duplicate exposure is a real risk only when the pieces of the checkpoint disagree. If the data position is stored separately and was last written at step 8,400, restoring weights from 8,000 and data from 8,400 skips 400 × 512 = 204,800 sequences. If the data position is stale in the other direction, the same sequences are fed to a model that already trained on them. Likewise, resuming with a new shuffle seed after a restart, or with a different worker count that reassigns shards inconsistently, can silently duplicate or drop data. The remedy is to treat weights, optimiser, scheduler, RNG and data position as one atomically committed unit.

Choosing the checkpoint interval

More frequent checkpoints reduce lost work but cost write time. A first-order model with checkpoint cost C seconds, interval τ seconds, mean time between failures M and restart cost R gives expected waste per unit of time of roughly C/τ + τ/(2M) + R/M. The first term is checkpoint overhead, the second is the average half-interval of redone work, the third is recovery. Minimising over τ gives the classic Young/Daly approximation τ ≈ √(2CM). Suppose (illustrative) each accelerator has a mean time between failures of 384 hours, so a cluster of 64 fails on average every 384 / 64 = 6 hours (21,600 s); failures are assumed independent and exponentially distributed, which real hardware only roughly satisfies. A synchronous checkpoint costs C = 8.4 s:

M = 21,600 s        (cluster MTBF = 384 h / 64 devices = 6 h)
R = 480 s           (detect 60 + reschedule 300 + load 120)

every 500 steps  (tau = 2,150 s, C = 8.4 s):
  C/tau = 0.39%   tau/2M = 4.98%   R/M = 2.22%   total = 7.6%

Young/Daly: tau = sqrt(2 * 8.4 * 21,600) = 602 s  (about 140 steps)
  C/tau = 1.40%   tau/2M = 1.39%   R/M = 2.22%   total = 5.0%

Asynchronous save, training stalls only C = 1 s:
  tau = sqrt(2 * 1 * 21,600) = 208 s  (about 48 steps)
  C/tau = 0.48%   tau/2M = 0.48%   R/M = 2.22%   total = 3.2%

Scale changes the picture sharply. Cluster MTBF is device MTBF divided by device count under the independence assumption: 384 h / 4,096 devices is about 5.6 minutes, a failure interval shorter than a single naive checkpoint. This is why large runs use asynchronous checkpointing (copy state to host memory, then write to storage in the background while training continues), keep the latest checkpoint in node-local or in-memory tiers, and invest heavily in fast detection and automated restart.

The explorer below uses the same model: try 64 devices with 8.4 s and then 1 s checkpoints, and then 4,096 devices.

The headline metrics are tokens per second (the number that translates to calendar time for a given token budget) and steps per second (which depends on batch size and is only comparable at the same global batch). For efficiency you also need a notion of how much of the hardware's potential you are using. Model FLOPs utilisation (MFU) is defined here as model FLOPs delivered per second divided by the aggregate peak FLOP/s of the accelerators: (6 × N × tokens per second) / (devices × peak). For our run, 6 × 7e9 × 243,810 = 1.024e16 FLOP/s divided by 64 × 4.0e14 = 2.56e16 gives 0.40. Peak is a theoretical number from the datasheet; delivered is what the job achieves, and MFU is the ratio. It is a coarse figure by design: it counts only useful model FLOPs (not recomputation) and the 6N approximation omits attention, so use it for comparing runs of the same architecture, not for declaring a number good in the abstract. The figure of 40 percent here is illustrative.

MFU says how far from peak you are, not why. Decompose the step instead. Profile one rank across a step and attribute time to categories:

Two numbers in this profile look inconsistent and teach something. Accelerators are running kernels for 3.60 of 4.30 seconds, 83.7 percent, yet MFU is 40 percent. These measure different things. A device's utilisation counter reports that some kernel was executing, not that arithmetic units were saturated; kernels limited by memory bandwidth, small matrices, and imperfect fusion deliver a fraction of peak FLOP/s while the device is "busy". Closing the gap between 83.7 and 40 is a kernel and precision problem (see AI Accelerators); closing the other 16 percent is a systems problem.

The systems part has a clear order of attack. Exposed communication (0.45 s) is gradient reduction not hidden behind backward compute; it responds to bucket sizes, overlap, topology-aware placement and, for ZeRO-3, prefetching of the next layer's weights. Input wait (0.15 s) points to the loader, as in the starvation example. The slowest rank also matters: step time is a maximum over ranks, so a single thermally throttled device or a node with a bad link drags all 64. Comparing per-rank step times and communication-wait times usually reveals the outlier immediately. A repeatable method: measure tokens per second at 1 node and at 8 nodes at fixed per-device batch; find where the curve bends; profile at that point.

Scaling efficiency formalises that comparison. Suppose (illustrative) 8 accelerators deliver 31,200 tokens per second. Perfect weak scaling to 64 would give 8 × 31,200 = 249,600; the measured 243,810 is 97.7 percent efficient. A drop in efficiency as scale grows is a signal that communication or stragglers, not compute, is becoming the limit, which is the moment to look at topology, bucket sizes and sharding stage rather than at kernels.

The simplest honest cost model starts with accelerator-hours and then adds everything that is not that. For the 7B run, suppose (illustrative) a budget of 140 billion tokens. At 243,810 tokens per second the run needs 140e9 / 243,810 = 574,200 s = 159.5 hours of wall-clock time; on 64 accelerators that is 159.5 × 64 = about 10,200 accelerator-hours. This is the floor, the cost if nothing ever goes wrong.

The floor is multiplied by reality. Failures and restarts: with a 6-hour cluster MTBF, a 160-hour run expects about 26.6 failures. If each costs on average half a checkpoint interval (250 steps = 1,075 s) of redone work plus 480 s of recovery, 1,555 s, the expected loss is 26.6 × 1,555 s = 11.5 hours, 7.2 percent of the run. Everything before the run counts too: experiments to choose the learning rate, data mixture and parallel layout, often several times the cost of the final run for a new recipe; aborted and diverged runs; evaluation jobs; and the data pipeline itself (processing, deduplicating and storing the corpus). The final run is rarely the most expensive line.

accelerator_cost   = accelerator_hours * rate_per_accelerator_hour
                     (rate = rental rate, or amortised purchase + hosting)
power_and_hosting  = energy_kwh * rate_per_kwh + facility share      (if not in rate)
storage            = dataset copies + checkpoints retained * size * duration
interconnect       = fabric share of the cluster; egress if data crosses regions
experiments        = sum of ablations, sweeps and aborted runs (accelerator-hours)
failures           = (expected_failures * (lost_work + recovery)) * devices
people             = platform, data and research engineering time

total ~ accelerator_cost * (1 + failure_overhead) + experiments + storage + interconnect + people

Two lessons about economics are specific to infrastructure. Efficiency compounds: raising delivered utilisation from 40 to 48 percent cuts the run from 159.5 to 132.9 hours, saving about 26.6 hours on every one of the 64 accelerators, so profiling is an investment with a quantifiable return. And reliability is a cost line: automated restart, fast checkpointing and early detection are not operational polish, because each minute of an idle 64-accelerator cluster is paid for. The question of whether to train at all, versus fine-tuning or calling a hosted model, is a cost decision and is revisited in Cost Optimization.

A run that finishes is not a model that is good. Evaluation during and after training answers whether it learned, and the answer has several layers. Validation loss on held-out data, tracked at regular intervals, is the continuous signal: it falls smoothly when training is healthy, spikes or diverges when something is wrong, and flattens when you have run out of useful data or capacity. Task benchmarks measure what the loss does not, such as instruction following, reasoning, coding or the domain-specific questions your product cares about; they are run on checkpoints, usually less often because they are slow. For a support assistant the most relevant evaluation is your own: a held-out set of real policy questions with reference answers, scored for correctness and groundedness against retrieved documents.

Data leakage is the failure that makes all of these numbers untrustworthy. If benchmark questions, or near-duplicates of them, appear in the training data, scores measure memorisation. The defence is in the data pipeline: deduplicate against evaluation sets before training, keep evaluation data in a separate, access-controlled location, and version both so you can prove which documents a model could have seen. A leaked benchmark is a data-versioning failure first.

Reproducibility means that someone can answer, a year later, how a model was produced. That needs a run record: the git commit of the training code, the container image digest, the library and driver versions, the job configuration, seeds, the dataset manifest hash, the parallel layout, and the sequence of checkpoints with the evaluation results attached to each. Exact bitwise reproduction is usually out of reach (see the replay note above), so the realistic standard is that a rerun lands within the noise of a documented seed variance. Release gating then turns evaluation into a decision: a candidate checkpoint is promoted to a registry only if it passes thresholds on validation loss, task benchmarks, safety checks, and a no-regression comparison against the current production model, and the promotion records its lineage. That registry entry is the handoff point to serving; how the artefact is then deployed, routed and versioned is the subject of Model Serving.

Most teams that train at all fine-tune an existing model, and the infrastructure picture changes substantially. A fine-tuning job starts from pretrained weights, uses a dataset thousands to millions of times smaller, and runs for hours not weeks. Whether it fits on one accelerator is the memory question again, and parameter-efficient fine-tuning changes its answer.

With LoRA-style adapters, the base weights are frozen and small low-rank matrices are trained alongside them. Suppose (illustrative) the adapters total 40 million parameters on the 7B model. The frozen BF16 base weights take 14 GB and need no gradients or optimiser state. The adapters need, at 16 bytes per parameter with FP32 master and Adam, 40e6 × 16 B = 0.64 GB. The total is about 14.64 GB plus activations, well within a single 80 GB device, compared with 112 GB of model state for full fine-tuning.

The consequences are practical. A job that fits on one device has no gradient all-reduce, no pipeline bubble, no gang scheduling and a trivially small checkpoint (the adapter alone is tens to hundreds of megabytes), so most of this chapter's machinery disappears, and so does most of the cost. When one device is enough, use one; distribution adds failure modes. The reasons to scale out anyway are specific: the base model plus activations does not fit on one device even with adapters (larger models, long sequences), you need full-parameter updates, or the dataset is large enough that wall-clock time matters. Then the same tools apply: data parallelism over adapter jobs, sharding for large bases, and careful placement.

Fine-tuning also has a boundary worth stating plainly for our assistant. It is good at consistent tone, output format and task behaviour. It is a poor way to teach the assistant this quarter's refund policy, which will change next quarter and which different tenants may or may not be allowed to see; that belongs in an authorised, updatable retrieval layer. And whichever route you take, the output is an artefact whose resource needs at serving time differ from its needs during training, which is where Inference Infrastructure picks up.

The pieces above come together in two artefacts: a declarative job specification that the scheduler and the training code both read, and a resume routine that makes restarts safe. Both are shown as illustrative pseudocode, not a specific framework's API; real launchers and checkpoint libraries differ in names and options. The values match the running 7B example.

job:
  name: policy-lm-7b-pretrain
  image: registry.example/train@sha256:<digest>      # pinned, not a tag
  gang: { nodes: 8, accelerators_per_node: 8 }       # all-or-nothing scheduling
  priority: batch-high
  max_restarts: 20
  hang_timeout_seconds: 60                            # collective watchdog

data:
  manifest: shards/v7/manifest.json                   # versioned, hashed
  manifest_sha256: <hash>
  sequence_length: 2048
  global_batch_sequences: 512                         # = 1,048,576 tokens
  shuffle: { seed: 1234, mode: global_index }         # order = f(seed, index)
  loader: { workers_per_node: 8, prefetch_batches: 4, local_cache: nvme }

parallelism:
  data_parallel: 64
  sharding: zero3_style                               # weights, grads, optimiser state
  microbatch_sequences: 8
  grad_accumulation: 1
  activation_recompute: selective

optimiser:
  name: adamw
  lr_schedule: { warmup_steps: 2000, peak: <lr>, decay: cosine }
  moment_dtype: fp32
  master_weights: fp32
  compute_dtype: bf16

checkpoint:
  every_steps: 500
  async_write: true
  format: sharded
  keep_last: 3
  save: [weights, optimiser, scheduler, rng_per_rank, next_sample_index]
  verify_on_commit: true

evaluation:
  every_steps: 1000
  validation_set: eval/heldout-v3                    # deduplicated against train
  gate: { max_val_loss: <threshold>, regress_vs: prod-model }

def save_checkpoint(step, state, rank, world, store):
    tmp = f"ckpt/step_{step}.partial"
    store.write(f"{tmp}/shard_{rank}.bin", state.shard(rank))   # parallel, per rank
    barrier()                                                    # every shard flushed
    if rank == 0:
        manifest = {
            "step": step,
            "world_size": world,
            "next_sample_index": state.next_sample_index,        # data progress
            "scheduler": state.scheduler_state,
            "shards": store.checksums(tmp),
            "data_manifest_sha256": state.data_manifest_hash,
        }
        store.write(f"{tmp}/manifest.json", manifest)
        store.rename(tmp, f"ckpt/step_{step}")                   # commit point
        store.write("ckpt/latest", f"step_{step}")
        prune_old(store, keep=3)
    barrier()

def resume(rank, world, store):
    for name in newest_first(store.list("ckpt/step_*")):         # ignores *.partial
        manifest = store.read_json(f"ckpt/{name}/manifest.json")
        if manifest and store.verify(f"ckpt/{name}", manifest["shards"]):
            state = load_sharded(store, name, rank, world)       # may reshard
            assert manifest["data_manifest_sha256"] == current_data_hash()
            set_rng_state(state.rng[rank])
            return state, manifest["step"], manifest["next_sample_index"]
    return init_fresh_state(), 0, 0                              # nothing valid: start over

state, step, sample_index = resume(rank, world, store)
while step < total_steps:
    batch = loader.batch_at(sample_index, seed)                  # pure function of index
    loss = train_step(state, batch)
    step += 1
    sample_index += global_batch_sequences
    if step % ckpt_every == 0:
        save_checkpoint(step, state, rank, world, store)

Read the design for what it prevents rather than what it does. Writing under a .partial name and committing by rename means a crash mid-save cannot produce a checkpoint that looks valid and is not. Storing the sample index inside the same manifest as the weights eliminates the skipped-or-duplicated data case from the failure trace. Verifying checksums on resume converts silent corruption into an explicit fallback to the previous checkpoint. And the loader being a pure function of seed and index means a restart on a different number of workers still reads the same global sequence. What this blueprint omits is just as instructive: automatic detection of a hung collective, the restart policy in the orchestrator (max_restarts above), and a way to mark a node bad so the replacement is not the same faulty hardware. Those belong to the cluster layer described earlier.

Once the run passes its release gate, its exported artefact leaves training. From there the objectives change entirely: not maximum tokens per second over a week, but low latency and steady cost for many concurrent users, which needs different batching, memory management and scheduling.

Training produces one durable asset, a set of weights. Turning those weights into a service that answers many requests at once is a different systems problem, with its own bottlenecks in memory bandwidth, KV cache growth and scheduling, and it is taken up in the next lesson.

1. 1

   A run needs four things at once: a data pipeline that keeps the accelerators fed, a model state that fits in memory, communication overlapped with compute, and a job that can resume after failure.

2. 2

   Count memory before choosing a parallelism strategy: with Adam-style mixed precision the model state is about 16 bytes per parameter (the 7B example needs 112 GB before activations), so data parallelism alone needs a full replica per device.

3. 3

   Reach for tensor and pipeline parallelism when a replica does not fit, and for ZeRO or FSDP sharding to remove redundant state while keeping the data-parallel programming model; each trades memory for communication and complexity.

4. 4

   Diagnose throughput by decomposing the step: a starved loader, exposed communication and idle kernels have different fixes. MFU says how far from peak you are, not why.

5. 5

   A checkpoint is valid only when every shard of model, optimiser, RNG and data-loader position is present and committed; make sample order a pure function of (seed, epoch, step) so a resumed run replays correctly.

6. 6

   Pick the checkpoint interval from the failure rate: cluster MTBF shrinks with device count, and the Young/Daly estimate τ ≈ √(2CM) gives the first-order balance between checkpoint cost and lost work.

7. 7

   Evaluate with leakage checks and keep a run record (code, data, config, seeds) so someone can answer a year later how the model was produced.