Distributed training and scaling

All-reduce aggregates (typically sums) the locally computed gradients across all workers and distributes the identical averaged result back to every worker, so each replica applies the same parameter update. This keeps model copies bit-wise consistent across steps. Ring all-reduce makes per-worker communication cost independent of the number of workers (bandwidth-optimal).

ZeRO-1 shards optimizer states, ZeRO-2 additionally shards gradients, and ZeRO-3 additionally shards the model parameters themselves. Each stage reduces per-GPU memory at the cost of more communication, with ZeRO-3 gathering parameters on demand during forward/backward. This lets data-parallel training scale to models far larger than a single GPU's memory.

Tensor parallelism splits the computation of a single layer (e.g., the weight matrices of an MLP or attention) across devices, requiring all-reduce/all-gather within each layer's forward and backward. Pipeline parallelism assigns different contiguous layers to different devices and streams micro-batches through them. Tensor parallelism is communication-heavy and best kept within a high-bandwidth node, while pipeline parallelism communicates only activations at stage boundaries.

The bubble is the idle time on pipeline stages while the pipeline fills and drains, since early stages wait for later ones (and vice versa). With P stages and M micro-batches, the bubble fraction is approximately (P-1)/(M+P-1), so increasing M reduces the relative bubble overhead. This is why pipeline schedules like GPipe and 1F1B use many micro-batches.

Gradient accumulation runs several forward/backward passes on smaller micro-batches, summing gradients before a single optimizer step, thereby simulating a larger effective batch size without the memory of holding it all at once. In distributed setups it also lets you reduce all-reduce frequency (skip the sync with no_sync until the accumulation boundary), improving the compute-to-communication ratio. The learning rate / loss scaling must account for the accumulation count.

for i, batch in enumerate(loader):
    loss = model(batch) / accum_steps
    loss.backward()
    if (i + 1) % accum_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

FSDP shards parameters, gradients, and optimizer states across data-parallel ranks (conceptually like ZeRO-3) instead of replicating the full model on every GPU as DDP does. During forward/backward it all-gathers a layer's full parameters just-in-time, computes, then frees the non-owned shards, and reduce-scatters gradients. This trades extra communication for a large reduction in per-GPU memory.

Activation checkpointing reduces memory by discarding intermediate activations during the forward pass and recomputing them during the backward pass, instead of storing all of them. This roughly trades a ~30% increase in compute (one extra forward) for a large drop in activation memory, enabling larger models or batch sizes. It is essential for training deep transformers under memory constraints.

from torch.utils.checkpoint import checkpoint
out = checkpoint(self.block, x, use_reentrant=False)

The linear scaling rule says multiply the base learning rate by k when the batch size is multiplied by k, since the gradient estimate's variance drops proportionally. The key caveat is that this breaks down at very large batch sizes and requires a gradual learning-rate warmup to avoid instability early in training. Optimizers like LARS/LAMB were designed to push this rule to extreme batch sizes.

In ring all-reduce, N workers form a logical ring and perform a reduce-scatter followed by an all-gather, each in N-1 steps. Each worker sends and receives approximately 2*(N-1)/N times the size of the gradient tensor, which approaches a constant (~2x) as N grows, making per-worker communication essentially independent of cluster size. This avoids the bottleneck of a central parameter server.

All-gather concatenates each rank's local shard so that every rank ends up holding the full set of all shards. Reduce-scatter performs an element-wise reduction (e.g., sum) across ranks and then scatters disjoint slices of the result, so each rank holds only a reduced portion. An all-reduce is exactly a reduce-scatter followed by an all-gather.

MoE layers route each token to a subset of experts living on different devices, requiring two all-to-all collectives per MoE layer: one to dispatch tokens to their assigned experts and one to combine the expert outputs back. This all-to-all traffic, plus load imbalance across experts, is the main scaling bottleneck. Capacity factors and auxiliary load-balancing losses are used to keep expert utilization even.

3D parallelism combines data parallelism, tensor (intra-layer) parallelism, and pipeline (inter-layer) parallelism simultaneously to scale to thousands of GPUs. Tensor parallelism is placed within a node (high NVLink bandwidth), pipeline parallelism spans nodes, and data parallelism replicates across pipeline groups. This composition balances memory, compute, and communication for trillion-parameter models.

Standard BatchNorm computes statistics only over each GPU's local mini-batch, so with small per-GPU batches the statistics are noisy and differ from the intended global batch, hurting accuracy. SyncBatchNorm fixes this by all-reducing the mean and variance across all GPUs so normalization uses the true global batch statistics. Many large-scale recipes instead switch to GroupNorm/LayerNorm to avoid the cross-device sync entirely.

Gradients in fp16 can underflow to zero because fp16 has a narrow dynamic range, losing small gradient values. Loss scaling multiplies the loss by a large factor before backprop so gradients land in fp16's representable range, then unscales them before the optimizer step. Dynamic loss scaling automatically raises the factor when no overflow occurs and lowers it (skipping the step) when inf/NaN is detected.

scaler = torch.amp.GradScaler('cuda')
with torch.autocast('cuda', dtype=torch.float16):
    loss = model(x)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

bf16 has the same 8-bit exponent as fp32, giving it the same wide dynamic range, while sacrificing mantissa precision. This means bf16 rarely overflows/underflows and typically does not require loss scaling, simplifying stable training of large models. fp16 has more mantissa bits (higher precision) but a narrow range, making it more prone to numerical issues.

Synchronous SGD waits for all workers to compute gradients and aggregates them before each update, guaranteeing all replicas stay identical but being limited by the slowest worker (straggler). Asynchronous SGD lets workers update shared parameters independently without waiting, improving throughput but introducing stale gradients that can degrade convergence. Modern large-scale training overwhelmingly favors synchronous SGD for its stable, reproducible convergence.

Workers compute gradients on data shards and push them to parameter servers, which hold the master copy of parameters, apply updates, and serve the latest weights back to workers. A key drawback is that the parameter servers become a communication and bandwidth bottleneck (and a scaling/fault-tolerance concern) as the number of workers grows. This is why all-reduce based approaches have largely replaced it for dense synchronous training.

Sequence parallelism splits activations along the sequence-length dimension across devices for the operations not covered by tensor parallelism (notably LayerNorm and dropout), reducing the activation memory that tensor parallelism leaves replicated. It pairs with tensor parallelism: regions use all-gather to reconstruct sequences for attention/MLP and reduce-scatter to redistribute afterward. Together they cut activation memory while keeping total communication volume comparable to TP alone.

GPipe runs the forward pass for all micro-batches before any backward, so it must store activations for every in-flight micro-batch simultaneously, giving memory proportional to the number of micro-batches. 1F1B interleaves a backward pass right after each stage's forward becomes available, so each stage only keeps activations for about the pipeline-depth number of micro-batches in flight. This bounds activation memory by pipeline depth rather than micro-batch count.

NCCL (NVIDIA Collective Communications Library) provides optimized GPU collective primitives (all-reduce, all-gather, reduce-scatter, broadcast, all-to-all) over NVLink, PCIe, and InfiniBand. It is topology-aware, automatically building rings/trees that exploit the fastest interconnects (e.g., NVLink within a node, RDMA across nodes) to maximize bandwidth and minimize latency. Choosing the right collective and process placement to match topology is critical for scaling efficiency.

DDP groups gradients into contiguous buckets and launches an asynchronous all-reduce on a bucket as soon as all its gradients are ready, overlapping communication with the ongoing backward computation. This avoids one tiny all-reduce per parameter (high launch overhead) and hides communication latency behind compute. Bucket size is a tunable trade-off between overlap granularity and per-call efficiency.

LAMB applies layer-wise learning-rate adaptation: it normalizes the Adam-style update per layer and rescales it by the ratio of the weight norm to the update norm (a trust ratio). This keeps each layer's effective step size well-conditioned even at extreme batch sizes, preventing the instability that naive linear LR scaling causes. It famously trained BERT in 76 minutes using batch sizes up to ~32k.

Strong scaling fixes the total problem size (e.g., global batch) and adds workers to finish faster, where efficiency falls as communication and the unparallelizable fraction dominate (Amdahl's law). Weak scaling grows the problem size proportionally with the number of workers (e.g., constant per-GPU batch), aiming to keep time-per-step roughly constant. Distributed training usually targets weak scaling, since adding GPUs increases the global batch.

If all ranks (or rank 0) serialize the full sharded model and optimizer state to a single shared filesystem at once, it causes huge synchronous I/O stalls and memory spikes, scaling poorly to thousands of GPUs. Distributed/sharded checkpointing has each rank write only its own shard in parallel, often asynchronously, dramatically cutting checkpoint time. Tools like PyTorch Distributed Checkpoint (DCP) and asynchronous staging address this.

With Megatron-style tensor parallelism, each transformer layer requires two all-reduces in the forward pass (one after the attention output projection, one after the MLP output projection) and two corresponding all-reduces in the backward pass. So roughly four all-reduce operations per layer per step, each over the hidden-dimension activations. This per-layer cost is why tensor parallelism is confined to high-bandwidth intra-node links.

Gradient compression reduces the number of bits or elements sent during all-reduce, lowering communication cost in bandwidth-limited clusters and improving the compute-to-communication ratio. The main risk is that lossy compression introduces gradient error/bias that can slow or destabilize convergence, often requiring error-feedback (accumulating residuals) to recover accuracy. It is most useful when network bandwidth, not compute, is the bottleneck.

Because parameters are sharded, each layer must be all-gathered before it can compute, and doing this strictly on demand serializes communication with computation, leaving GPUs idle. Prefetching issues the all-gather for the next layer while the current layer is still computing, overlapping communication with compute to hide latency. Both forward and backward prefetching are tuned to maximize overlap without exceeding memory.

Offloading moves optimizer states, gradients, and/or parameters from GPU memory to CPU RAM (ZeRO-Offload) or NVMe disk (ZeRO-Infinity), letting you train models that otherwise won't fit in GPU memory. The trade-off is reduced throughput due to slower PCIe/NVMe transfer bandwidth and CPU-side optimizer compute. It is most attractive when GPU count is limited and fitting the model matters more than peak speed.

Place tensor parallelism within a single node to exploit fast NVLink/NVSwitch (since it is the most communication-intensive per step), use pipeline parallelism across nodes where only activations cross the slower interconnect, and apply data parallelism as the outermost dimension replicating across pipeline groups. This ordering minimizes traffic on the slowest links. Getting this mapping wrong is a common cause of poor scaling efficiency.

A distributed sampler partitions the dataset so each rank sees a disjoint, non-overlapping subset of samples per epoch, preventing redundant training on the same data and ensuring the global batch covers distinct examples. It must be re-seeded per epoch (set_epoch) so shuffling differs across epochs while staying consistent across ranks. Without it, all ranks would process identical data, wasting the parallelism.

sampler = DistributedSampler(dataset)
for epoch in range(E):
    sampler.set_epoch(epoch)
    for batch in DataLoader(dataset, sampler=sampler):
        ...

In synchronous training every step waits for the slowest worker, so a single straggler (due to thermal throttling, network jitter, or data-loading stalls) inflates total step time and idles all other GPUs. Mitigations include backup workers / redundant replicas (drop the slowest), better load balancing of data and compute, and overlapping communication with computation to absorb variance. Elastic frameworks can also reconfigure around persistently slow or failed nodes.

Expert parallelism places different experts on different devices and routes tokens to them via all-to-all, so each device holds a few complete experts. Tensor parallelism instead splits each individual expert's (or dense layer's) weight matrices across devices, communicating within the layer. They are often combined: experts are distributed via expert parallelism, and each large expert can additionally be tensor-parallelized.

During the backward pass, gradients for later layers become available first, so their all-reduce can be launched and run on the network while earlier layers' gradients are still being computed, hiding communication latency behind compute. Maximizing this overlap keeps GPUs busy and is one of the largest levers for high scaling efficiency, especially as cluster size grows and communication share rises. Frameworks achieve it via gradient bucketing and asynchronous collectives.

The Chinchilla finding is that for a fixed compute budget, model parameters and training tokens should be scaled roughly in equal proportion (about 20 tokens per parameter), rather than making the model as large as possible with relatively few tokens. Many earlier large models were undertrained because they over-prioritized parameter count. This reshaped how teams allocate compute between bigger models and more data. (Note: for inference-cost-aware deployment, models are often trained well past this point on far more tokens.)