Appendix 2: Distributed AI Training Clusters

By Chander Kant Updated July 31, 2026

Chapter 1 divided clusters into throughput and capability categories. A distributed training job belongs on the capability side because its workers cooperate to update one model. The same installation can also act as a throughput cluster when it runs independent experiments or data-preparation jobs. Within one training job, data parallelism increases the rate of example processing, while model parallelism combines accelerators to provide the memory and computation the model needs. Together, those workers form one synchronized program spread across many failure domains.

Training also differs from the inference service described in Appendix 1. An inference cluster receives requests and tries to meet latency or throughput objectives while using fixed model parameters. A training job repeatedly changes those parameters, often for hours or days, and a slow or failed worker can stop every other worker in its group. The same Linux cluster may run both workloads, but they place different pressure on memory, communication, storage, scheduling, and recovery.

We will follow a training step from input data to a parameter update, then examine the ways that step can be divided across a cluster. The examples use current deep-learning terminology, but the underlying questions are familiar HPC questions: what state must fit, what data must move, which operations synchronize, where time is lost, and whether the result is correct. Answering them gives us the evidence needed to decide whether the complete system is useful.

Following one training step

A training run begins with a model, an objective, and a defined set of examples. A data pipeline selects examples, reads and transforms them, and assembles a batch. The model performs a forward pass to produce predictions. The training program compares those predictions with the desired result and computes a loss. A backward pass follows the chain of operations in reverse to calculate a gradient for each trainable parameter. The optimizer uses those gradients and its retained state to update the parameters. The next batch then enters a model that is slightly different from the one that processed the previous batch.

Figure A2.1: A simplified synchronous training step. Parallelism changes where each operation runs and how its state is divided, but the dependencies still join the workers into one repeated calculation.

On one device, these operations are already constrained by memory capacity, memory bandwidth, arithmetic rate, and the time needed to supply data. Distribution adds dependencies between workers. Replicated workers must agree on gradients before they apply an equivalent update. A layer divided among several accelerators cannot continue until the required pieces of its input or output arrive. Pipeline stages exchange activations in the forward direction and gradients in the reverse direction. A mixture-of-experts layer routes tokens to the devices that hold the selected experts and then brings the results back.

These synchronization points explain why the slowest participant matters. If one worker receives data late, encounters a corrected hardware error, shares a congested link, or performs more expert work than its peers, other workers can reach the next dependency and wait. The useful unit is therefore the completed training step, and ultimately a completed run that reaches the required model quality; peak accelerator throughput says little about time spent waiting for a peer.

What has to fit in memory

The word model is often used as if it described one block of memory. Training keeps several kinds of state. Parameters hold the values being learned. Gradients hold the direction of the current update. An optimizer may retain master parameter copies, momentum, variance estimates, or other history. The forward pass creates activations that the backward pass needs. Communication libraries reserve buffers, compiled kernels use workspaces, and the runtime needs enough free memory to operate without repeated allocation failures.

Model size alone therefore does not determine whether training fits. The number of bytes per parameter depends on the numerical format used for computation, gradients, master values, and optimizer state. The optimizer matters as much as the parameter count. Activation memory follows a different set of variables: microbatch size, sequence or image dimensions, layer shapes, attention implementation, pipeline schedule, and which intermediate results are retained. A configuration that fits a short sequence can fail immediately when the context grows.

Precision, microbatches, and recomputation

Mixed-precision training performs selected operations and stores selected tensors in lower-precision formats while retaining enough higher-precision state to keep optimization stable. It can reduce memory traffic and use faster matrix hardware, but the gain depends on numerical choices. Loss scaling, accumulation precision, reduction precision, overflow handling, and model-specific sensitivity all affect the result. A throughput comparison is incomplete unless it states these choices and confirms equivalent quality.

A large logical batch can be divided into microbatches. Gradient accumulation runs several microbatches before applying an optimizer update. For uniform data-parallel training, the relationship is approximately:

global batch = microbatch per worker x accumulation steps x data-parallel replicas

Tensor, pipeline, or expert parallel workers cooperating on the same examples do not each multiply the global batch. This distinction becomes important when a large job combines several parallel dimensions. Increasing accumulation can reduce the frequency of data-parallel gradient communication, but it also changes optimizer-step frequency, activation scheduling, and the time between parameter updates. If the global batch changes, the learning-rate schedule or optimization behavior may need to change as well.

Activation checkpointing, also called activation recomputation or rematerialization, saves memory by retaining only selected forward results and recomputing the omitted portions during backward. The term is easy to confuse with the persistent checkpoints used to restart a failed job. Activation checkpointing trades additional computation for device memory inside a step. A restart checkpoint trades storage traffic and pause time for less lost work across steps.

Parameters and optimizer state can also be offloaded to host memory or local NVMe. This extends capacity through a slower tier, but every required transfer joins the step's critical path unless it is successfully prefetched and overlapped. Host memory capacity, memory channels, PCIe topology, pinned-memory limits, and local storage endurance become part of the training design. Fitting a larger model is useful only when the resulting step time remains acceptable.

Ways to divide the work

Distributed training is a collection of parallel operations, and each method divides a different dimension of the calculation. Some methods increase throughput after the model already fits. Others are needed because parameters, activations, or optimizer state do not fit on one accelerator. Most large jobs combine several methods. A rank is the numerical identity assigned to one participating process, and a process group is the set of ranks that communicate for a parallel dimension. The plan must therefore describe these groups rather than only state a GPU count.

Each parallelism strategy creates a different memory and communication pattern
Strategy	What is divided	Typical communication	Primary reason to use it
Data parallelism	The batch; each worker normally holds a model replica	Gradient all-reduce, or reduce-scatter with sharded state	Process more examples per unit of time
Fully sharded data parallelism	Parameters, gradients, optimizer state, or a selected subset	Parameter all-gather and gradient reduce-scatter	Reduce repeated model state on each worker
Tensor parallelism	Rows, columns, heads, or other dimensions within a layer	Frequent all-reduce, all-gather, or reduce-scatter within the layer	Fit and execute layers that are too large for one device
Pipeline parallelism	Groups of layers; microbatches move through stages	Point-to-point activation and gradient transfers	Divide model depth and activation state
Sequence or context parallelism	Sequence positions and associated activations	Gather/scatter, ring, or point-to-point exchanges around attention and other operations	Train with longer sequences or reduce activation memory
Expert parallelism	Experts in a sparse mixture-of-experts layer	Token dispatch and return, commonly all-to-all	Spread a large expert set while activating only selected experts per token

Data parallelism and state sharding

In conventional synchronous data parallelism, every worker starts with an equivalent model, processes a different part of the batch, and calculates local gradients. An all-reduce combines those gradients so every replica applies an equivalent update. This works well when the complete training state fits on each worker and the gradient exchange can be overlapped with backward computation.

Repeating all model state on every worker eventually becomes wasteful. ZeRO and fully sharded data-parallel designs divide optimizer state first, then gradients, and finally parameters. A fully sharded worker gathers the parameters needed for a portion of the forward or backward pass, releases or reshards them according to policy, and uses reduce-scatter to leave each worker with its gradient partition. The memory saving is purchased with more frequent communication and stricter ordering. The useful sharding unit has to be large enough for efficient collectives but small enough to control peak memory.

Tensor, pipeline, sequence, and expert groups

Tensor parallelism divides work within a layer. One accelerator may hold a group of matrix columns and another a different group, with a collective joining their partial results. These exchanges can occur several times per layer, so tensor-parallel ranks are usually placed on the fastest available scale-up links. Extending a fine-grained tensor group across a slower or oversubscribed network can make communication dominate the layer.

Pipeline parallelism places consecutive groups of layers on different stages. Splitting a batch into microbatches lets the stages work concurrently, but the pipeline takes time to fill and drain. An unbalanced stage creates a bubble in which faster stages wait. More microbatches can reduce the fraction of time lost to the bubble while increasing scheduling complexity and the number of in-flight activations.

Sequence and context-parallel methods divide a long sequence and the activation state associated with it. The exact exchanges depend on the attention algorithm and framework. Their purpose is broader than adding devices to a batch: they make a longer context fit and distribute work whose memory grows with sequence length. A plan should name the method and its communication, because the terms are not used identically by every framework.

Expert parallelism distributes the experts of a sparse model. A router selects experts for each token, the tokens travel to the devices holding those experts, and the outputs return to their original sequence positions. This creates an all-to-all pattern and a load-balancing problem. If many tokens select the same expert, the device hosting it becomes a straggler while other expert devices are underused. Model quality, routing policy, capacity limits, token dropping or padding, and network behavior are coupled decisions.

The dimensions can be composed. A job with four-way tensor parallelism, eight pipeline stages, and sixteen data-parallel replicas uses 4 x 8 x 16 = 512 workers before any separate context or expert dimension is added. Each worker belongs to several groups with different traffic. The framework may build those groups automatically, but the operator still needs the map in order to allocate, place, diagnose, and reproduce the job.

The parallelism plan is a traffic map

The best logical division can perform poorly when it is placed without regard to the machine. Intra-node scale-up links, PCIe switches, CPU sockets, network adapters, and inter-node rails do not provide one uniform pool of bandwidth. Chapter 2 described the paths inside a node, and Chapter 3 followed them into the cluster fabric. Training assigns a repeated collective or point-to-point transfer to each of those paths.

Figure A2.2: One possible mapping of four-way tensor parallelism and two-way data parallelism. It keeps frequent tensor exchanges inside each node and sends gradient traffic across the scale-out fabric. Other parallel combinations require different maps.

Collective operations

A small vocabulary describes much of the traffic. A broadcast sends data from one rank to the group. An all-gather gives every rank all of the partitions. A reduce-scatter combines values and leaves each rank with a partition of the result. An all-reduce combines values and returns the complete result to every rank; it can often be understood as a reduce-scatter followed by an all-gather. An all-to-all lets every rank send a distinct portion to every other rank. Point-to-point sends connect particular pipeline stages or other neighbors.

The operation name is not enough to predict performance. Message size, frequency, group size, topology, concurrent collectives, and the amount of computation available for overlap all matter. Small transfers are sensitive to latency and software overhead. Large transfers approach a bandwidth problem. All-to-all traffic can pressure many switch paths at once. A collective library may choose rings, trees, direct paths, channels, and protocols according to the discovered topology, but it cannot create bisection bandwidth that the fabric does not have.

NCCL on NVIDIA systems and RCCL on AMD systems provide topology-aware GPU collectives and point-to-point operations. The training framework remains responsible for deciding when a collective is required and which ranks participate. The job launcher establishes processes and rank information, while Slurm, Kubernetes, or another cluster manager allocates the hosts and devices. Keeping these layers distinct makes failures easier to place.

Placement, rails, and direct paths

Communication and computation can overlap only when their dependencies permit it and when they do not contend for the same accelerator engines, memory bandwidth, CPU threads, or network resources. Gradient bucketization can start reducing earlier layers while backward computation continues through later layers. Prefetching can gather the next parameter shard while the current layer computes. A trace should confirm the overlap; a configuration flag proves only that the runtime attempted it.

GPU-aware RDMA can let a network adapter transfer data to or from accelerator memory without staging the payload through ordinary host memory. This removes a copy, but it still depends on Linux drivers, device support, PCIe routing, memory registration, IOMMU mode, and the relationship between the accelerator and network adapter. A worker bound to a distant CPU socket or NIC can cross internal links before it even reaches the fabric. Inspect the node topology and place ranks, CPU affinity, memory, and NIC selection together.

Multiple network adapters or rails can increase aggregate bandwidth and provide more paths, but only when the collective library, routing, cabling, and job placement use them as intended. Other jobs can inject traffic onto the same links. Storage and training traffic can collide if they share a fabric. A scaling run should record the allocated nodes and switch path so a topology change is not mistaken for random application variance.

Keeping the workers supplied with data

The training dataset may begin in object storage, a parallel file system, a scale-out file service, or an archive prepared by another pipeline. Supplying the bytes is only the first step. Software must locate and read each record; it may then decompress, decode, tokenize, filter, or augment the contents. Once the examples are ready, the loader assembles a batch and may copy it through pinned host memory before transferring it to device memory. A delay at any one of these stages can starve the training step.

Storage capacity and sequential bandwidth are only part of the requirement. Millions of small files create metadata and request pressure. A few compressed archives can shift the bottleneck to CPU decompression and make parallel access coarse. Remote object access adds request latency and service limits. A format optimized for archival or exchange may be a poor format for repeated shuffled reads. Measure the complete input pipeline with the same preprocessing and worker count used by training.

Distributed workers must also agree on which examples they process. A data-parallel sampler normally gives each replica a distinct partition and changes the shuffle deterministically between epochs. An iterable input source needs equivalent partition logic or workers may repeat and omit data. The restart state must identify the dataset version, epoch or stream position, shuffle seed, and any preprocessing state needed to continue without silently changing the sample sequence.

Local NVMe can stage a dataset or cache frequently used shards. The useful policy depends on dataset size, reuse, node churn, and the time allowed before a job begins. Copying the full dataset to every node can turn startup into the bottleneck and waste capacity. Reading everything remotely can create synchronized storage demand at each epoch boundary. Sharding data among nodes reduces replication but may add network transfers when placement changes. A cache needs an identity and eviction policy so stale data is not accepted as the current dataset.

Direct storage-to-GPU paths, such as GPUDirect Storage on compatible NVIDIA systems, can avoid a payload bounce through host memory. The CPU still arranges the operation, and preprocessing, filesystem, network, and compatibility limits remain. We should therefore verify the direct path as part of an end-to-end input measurement.

The Linux training stack

Linux joins the hardware and the training program. Firmware exposes the processors, accelerators, network adapters, and management interfaces. The kernel and device drivers establish memory, DMA, networking, process isolation, and error reporting. User-space accelerator runtimes, compilers, numerical kernels, and collective libraries expose those facilities to a framework. The framework builds the model graph, automatic differentiation, optimizer, parallel groups, and checkpoints. A container may package much of user space, but it still depends on the host kernel, drivers, devices, and network configuration.

PyTorch commonly runs one worker process per accelerator and forms distributed process groups for data, tensor, pipeline, or other dimensions. JAX can run one or more controller processes per host and represent arrays sharded across a global device mesh. MPI may launch processes or support components of another runtime, but the deep-learning framework usually owns the model operations and their collectives. These are examples of the same layered system rather than interchangeable names for it.

Every process needs a rank, local device assignment, world size, and a way to discover its peers. A rendezvous service or job launcher supplies this information. Initialization is a collective event: a wrong address, stale process, firewall rule, duplicate rank, or inconsistent world size can leave every worker waiting. Log hostnames, local and global ranks, device identities, process-group membership, and the first error from each host. Without that context, the last process to time out is often mistaken for the process that caused the failure.

Compatibility has to be managed as a matrix. The Linux kernel, accelerator driver, firmware, runtime, compiler, framework, collective library, network provider, container, and model extensions all participate. A newer package is not automatically compatible with the rest of the stack. Record a tested combination and change one boundary at a time. Chapter 6 develops this software environment in detail.

Reconstructing a run

A reproducible training record includes more than a container tag. It should identify four groups of evidence:

Training definition: source revision, model configuration, optimizer and schedule, numerical formats, batch construction, parallel dimensions, and random seeds.
Data: dataset manifest plus the preprocessing and tokenizer versions.
Recovery: checkpoint schema and the state needed to resume the same run.
Execution environment: framework and library versions, driver and firmware versions, environment settings, and allocated topology.

Container images should be identified by an immutable digest. Externally mounted code or libraries are also part of the environment and must be recorded.

Bit-for-bit repetition may still be unavailable when parallel floating-point reductions change order or kernels use nondeterministic algorithms. State the intended reproducibility boundary. A scientifically useful requirement may be equivalent convergence and evaluation quality within a defined range rather than identical final bits. Performance repetition has its own boundary: the same software can run differently after a topology, power, cooling, or competing-traffic change.

Scheduling and failure

A synchronous training job normally needs all of its requested workers before it can begin. This is a gang-scheduling requirement even when the scheduler uses different terminology. Starting half of a tensor or pipeline group provides no useful partial job. The request therefore needs the node count, accelerator count and type, memory, local storage, network or topology constraints, wall time, licenses if any, and a placement policy that keeps the parallel groups on suitable links.

Large fixed-shape requests can wait while enough compatible nodes become free. They can also leave fragments that other jobs cannot use. These are Chapter 7 policy questions, but the training plan supplies their cost: changing the data-parallel dimension may preserve the model layout, while changing a tensor or pipeline dimension can require a different checkpoint layout, batch, or performance tuning. Preemption is expensive when it discards work or requires a long checkpoint before resources can be released.

Synchronous progress makes a straggler visible to the whole group. The cause may begin in the application, with uneven input work, expert imbalance, or an imbalanced pipeline stage. It may instead be a system problem such as throttling, a failing link, retransmission, memory pressure, or an interfering background service. Placement adds another possibility when a rank uses the wrong NIC or NUMA node. Compare per-rank timelines rather than looking only at an average step time. A rank waiting in a collective may be healthy; the missing rank may still be loading data or computing an earlier operation.

A process, accelerator, node, or link failure often invalidates the current communicator. Some current collective libraries can abort, shrink, or rebuild communicators, and elastic launchers can form a new worker group. That does not guarantee that the model and algorithm can continue at a different world size. PyTorch's elastic launcher, for example, restarts surviving workers when membership changes and warns that ranks are not stable across restarts. The application still needs a valid checkpoint, data position, and parallel plan for the new group.

A checkpoint is a recovery contract

A restartable checkpoint normally includes model parameters, optimizer state, learning-rate schedule, mixed-precision scaler where used, step and epoch, random-number state, data position, and enough parallel-layout metadata to reconstruct the state. Saving only the weights may be sufficient for later inference, but it is not necessarily sufficient to continue training with the same optimization trajectory.

At scale, gathering all state onto one rank creates a memory and I/O bottleneck. Distributed checkpoint systems let ranks write shards in parallel. The checkpoint needs a manifest or commit record that distinguishes a complete generation from partial files left by a failed save. Retention should protect at least one known-good generation while a newer one is being written. A checkpoint that has never been restored is only an untested collection of files.

Checkpoint frequency balances pause or staging cost against expected lost work. A longer interval reduces write traffic but increases the amount of training repeated after failure. The calculation should include failure detection, scheduler requeue, node replacement, environment startup, checkpoint read, data-pipeline restoration, graph compilation or kernel warmup, and the work repeated since the save. Restore time can be more important than write time during an incident.

Asynchronous checkpointing can copy state to host memory and let training resume while storage writes continue. That shortens the visible pause only if the copy, CPU memory, PCIe traffic, and background write do not interfere with later steps. PyTorch's distributed asynchronous checkpoint guidance, for example, calls out extra CPU and pinned-memory requirements and recommends controlling concurrent requests. Another design stages shards to local storage and drains them to durable shared storage, which requires a clear rule for when the checkpoint is safe against node loss.

Measuring useful training

Samples per second or tokens per second describe the rate of work only after the model, sequence, batch, precision, and parallel plan are fixed. Step time should be reported as a distribution, not only a best iteration, because compilation, input stalls, checkpoints, congestion, and intermittent stragglers affect a long run. Separate warmup from the steady interval and retain the outliers for diagnosis.

Strong scaling gives the same total problem to more workers and asks whether elapsed time falls. If a baseline uses one group and a test uses p times as many workers, parallel efficiency is:

parallel efficiency = speedup / p

Weak scaling increases the work with the worker count and asks whether throughput per worker remains stable. Training studies often blur these definitions by changing global batch, sequence length, model size, or optimization settings while adding accelerators. Those may be useful configurations, but they do not isolate scaling. State what changed and whether the new run follows the same learning curve.

Accelerator utilization is also easy to overstate. A busy-device percentage does not show that useful model operations filled the device. Model FLOPs utilization estimates the model arithmetic completed per unit of time relative to an assumed peak, but both the model-operation count and the peak depend on definitions and numerical format. Use it with the measured step breakdown: compute kernels, collective communication, pipeline bubbles, input waits, checkpointing, and host overhead.

The final metric must include quality. A faster run that converges to a worse model has not completed the same job. MLPerf Training uses time to reach a defined quality target for this reason. An internal validation can use loss, accuracy, reward, scientific error, or another domain measure, but the target and evaluation data must be fixed before systems are compared. Time, energy, and cost per acceptable model are stronger system measures than peak operations alone.

A distributed training validation plan

The plan should begin with a model and dataset that the team can run repeatedly. It does not have to be the largest intended production model, but it must exercise the same software layers, collective patterns, input path, precision, and checkpoint mechanism. Record every result against a versioned test definition.

Prove one worker. Verify the model, data transforms, loss, optimizer, numerical settings, and checkpoint save/restore on one accelerator. Preserve a short correctness trace and expected quality progression.
Use every accelerator in one node. Confirm process-to-device mapping, CPU and memory affinity, peer paths, local collectives, power behavior, and step-time balance. Compare with the one-worker baseline using a clearly defined strong or weak scaling test.
Measure the communication path. Run collective tests across the exact nodes and rails intended for training. Cover the message sizes and operations used by the model rather than relying on one large-bandwidth number.
Cross the first node boundary. Examine rank traces for data loading, forward and backward compute, collective time, and waits. Confirm that fine-grained groups stayed on the intended links and that the scale-out groups use the intended adapters.
Scale in controlled steps. Hold the test definition fixed while increasing nodes. Record throughput, step-time distribution, scaling efficiency, communication fraction, input stalls, memory headroom, and power at each size.
Stress the data path. Start with cold caches, repeat with warm caches, and cross an epoch or shard boundary. Verify sample uniqueness, shuffle behavior, preprocessing capacity, metadata load, and the effect on accelerator waits.
Exercise checkpoint and restart. Measure the visible save pause, background interference, storage burst, completion marker, restart time, and repeated work. Interrupt a controlled run and verify that model, optimizer, data position, and quality progression continue as intended.
Create one approved failure. Under a safe test procedure, terminate a worker or remove a test node. Confirm error propagation, communicator cleanup, scheduler state, log context, retry policy, and return to useful training.
Run long enough to see the system. Include compilation, steady training, checkpoints, data transitions, and ordinary correctable events. A brief clean window does not reveal the tail behavior of a multi-day job.
Report the complete result. Publish the model and dataset identity, quality target, numerical settings, software stack, hardware, topology, parallel dimensions, batch construction, measurements, failures, and exclusions needed for another engineer to interpret the result.

Where training fits in the book

In Chapter 1's terms, a distributed training job is a capability workload: several accelerators cooperate on one model, and a delayed worker can hold up the group. Data-parallel replicas increase the number of examples processed per unit of time, borrowing the shape of throughput scaling, but their gradient synchronization means that they are not independent jobs. The broader training service also becomes a throughput cluster when it schedules many independent training runs, evaluations, or data-preparation tasks. In either case, the design should follow the work being divided and the communication required to join it.

The rest of this book develops the mechanisms that the appendix has used together. Chapter 2 covers the accelerator node and its memory and I/O paths. Chapter 3 covers the fabric. Chapter 4 covers training data and checkpoints as storage workloads. Chapter 5 carries the sustained power and heat into the facility. Chapters 6 through 8 cover the software, scheduler, and operation of the resulting Linux system. A training run exercises all of these layers at once.

Distributed AI Training Clusters