Distributed Training

Data parallelism replicates the full model on each device and splits the training data across them; gradients are synchronized (all-reduced) after each backward pass. Model parallelism splits the model itself across devices when it is too large for a single GPU, assigning different layers or parameter shards to different devices. Data parallelism scales batch size linearly with GPU count. Model parallelism is necessary when individual model parameters exceed single-GPU memory.

DDP uses a gradient bucketing scheme: gradients are grouped into buckets as they become ready during backward (from last to first layer due to backward-propagation order). When a bucket is full, DDP triggers an asynchronous all-reduce over NCCL, overlapping communication with the remaining backward computation. This pipelining hides most of the gradient communication latency. By default, bucket_cap_mb=25 MB; larger buckets reduce overhead but delay communication.

model = torch.nn.parallel.DistributedDataParallel(
    model, device_ids=[local_rank], bucket_cap_mb=25)

ZeRO (Zero Redundancy Optimizer) eliminates memory redundancy across data-parallel ranks. Stage 1 shards optimizer states (e.g., Adam m/v tensors) across GPUs. Stage 2 additionally shards gradients. Stage 3 also shards model parameters, so each GPU holds only 1/N of parameters, gradients, and optimizer states. Memory reduction is roughly 8x at Stage 3 for large models. ZeRO powers DeepSpeed and is also available via PyTorch FSDP (which implements ZeRO-3 semantics).

NCCL (NVIDIA Collective Communications Library) is a highly optimized library for GPU-to-GPU collective operations (all-reduce, broadcast, all-gather, reduce-scatter) that leverages NVLink, NVSwitch, and RDMA (GPUDirect) for maximum bandwidth. Gloo is a CPU-based communication library that is slower for GPU tensors because it requires copying data to CPU before communication. For GPU training, NCCL delivers near-peak interconnect bandwidth; Gloo is used as a fallback or for CPU tensors.

dist.init_process_group(backend='nccl')

All-reduce combines tensors from all ranks (typically by summing) and distributes the result to every rank. The ring algorithm divides the operation into two phases: reduce-scatter (each rank sends/receives chunks around the ring, accumulating partial sums) and all-gather (each rank distributes its fully-reduced chunk). Total data exchanged is 2(N-1)/N times the tensor size per rank, making it bandwidth-optimal. NCCL implements ring all-reduce (and tree variants for latency-sensitive small tensors).

dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

Pipeline parallelism splits model layers into stages across GPUs. Without scheduling, GPUs are idle waiting for the pipeline to fill/drain (bubble). The 1F1B (one-forward-one-backward) schedule interleaves forward and backward micro-batches: as soon as a stage completes a forward micro-batch, it starts backward on an earlier micro-batch. This keeps all stages busy and reduces bubble fraction from O(p-1)/m to O(p-1)/(m+p-1) where p is pipeline stages and m is micro-batch count.

Tensor parallelism splits individual weight matrices across GPUs along a chosen dimension. For a transformer MLP with two linear layers (A and B), column-parallel partitions A so each GPU computes part of the hidden expansion; row-parallel partitions B so each GPU computes part of the output projection. An all-reduce after B synchronizes results. This is the Megatron-LM approach. It requires tight synchronization on every MLP/attention block forward pass.

DistributedSampler divides the dataset indices among world_size processes by assigning each rank a strided subset of shuffled indices. Each process receives total_size // world_size samples per epoch; if the dataset is not evenly divisible, it pads to the nearest multiple. The sampler must be seeded with epoch number (sampler.set_epoch(epoch)) each epoch so different shuffles occur while maintaining cross-rank non-overlap. Without set_epoch, all epochs use the same ordering.

sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
for epoch in range(epochs):
    sampler.set_epoch(epoch)  # re-shuffle with epoch seed

Gradient compression reduces the volume of gradient data communicated during all-reduce by applying sparsification (sending only the top-k% largest gradients) or quantization (encoding gradients in fewer bits, e.g., 1-bit SGD, PowerSGD). This helps when network bandwidth is the bottleneck (e.g., cloud instances without NVLink). Error feedback (accumulating compression errors locally) is critical to maintain convergence. PowerSGD (low-rank approximation) achieves 10-100x compression with minor accuracy loss.

In synchronous SGD (standard DDP), all workers compute gradients on their local mini-batch, then block at an all-reduce barrier until every worker's gradient is aggregated — all workers perform the same parameter update. In asynchronous SGD (e.g., parameter server), workers push gradients and pull parameters independently without synchronization, achieving higher throughput but suffering from stale gradients that can harm convergence. Synchronous SGD is standard for modern GPU clusters with fast interconnects.

torchrun implements elastic training via a rendezvous mechanism: when a worker fails or is added, remaining workers rendezvous at a shared endpoint (e.g., etcd, c10d), agree on a new world_size, and restart training from the last checkpoint. It supports a min/max worker count range, allowing training to continue with fewer GPUs if some fail. The application must implement periodic checkpointing for restart. This is essential for large-scale training on cloud preemptible instances.

# Launch with elastic support:
# torchrun --nnodes=1:4 --nproc_per_node=8 train.py

DataParallel uses a single process with multiple threads (Python GIL limits parallelism) and scatters inputs to GPUs, running forward on each, then gathers outputs on GPU 0 which becomes a memory bottleneck. DDP uses separate processes (one per GPU), avoids the GIL, and synchronizes gradients via all-reduce without a designated gather GPU. DDP scales linearly, supports multi-node, and is significantly faster. DataParallel is simpler but deprecated for serious training.

3D parallelism combines data parallelism (DP), tensor parallelism (TP), and pipeline parallelism (PP) simultaneously to scale model training across thousands of GPUs. TP handles weight matrix splitting within a node (fast NVLink), PP splits layers across nodes, and DP replicates the 3D-parallel configuration across data-parallel replicas. Megatron-LM and DeepSpeed jointly developed this approach for GPT-style LLM training. Microsoft's Megatron-DeepSpeed implements all three.

In pipeline parallelism, each stage must recompute activations during backward passes for the micro-batches flowing through it. Gradient checkpointing (recompute-activations) is naturally compatible — each stage stores only a few activation checkpoints and recomputes the rest. This dramatically reduces activation memory per stage, allowing more micro-batches in flight (higher arithmetic intensity). Megatron-LM enables selective checkpointing of attention ops (most memory-intensive) while caching cheaper ops.

A parameter server (PS) is a centralized or sharded server that stores model parameters. Workers pull parameters, compute gradients on local data, and push gradients back to the PS which applies updates. The PS architecture is asynchronous-friendly and scales to large clusters but introduces staleness and potential bottlenecks at the server. Modern GPU training with fast interconnects (NVLink, InfiniBand) has largely moved to all-reduce architectures (DDP/FSDP) that are decentralized and more bandwidth-efficient.

Reduce-scatter combines a reduce (summing tensors across ranks) and a scatter (distributing disjoint chunks to each rank) in one communication operation. Each rank ends up with the fully-reduced value for its assigned chunk. FSDP uses reduce-scatter after backward: instead of all-reducing the full gradient tensor, FSDP reduce-scatters so each rank receives and stores only its parameter shard's gradient. This is followed by a local optimizer step on the shard, then an all-gather for the next forward.

dist.reduce_scatter_tensor(output, input, op=dist.ReduceOp.SUM)

Linear scaling rule (Goyal et al., 2017): multiply the learning rate by the number of GPUs (world_size) when the effective batch size scales proportionally. Rationale: gradient noise decreases with larger batches, so a larger step size maintains similar update magnitude. A linear warmup is essential because large LRs early in training are destabilizing. For very large scale (>8k GPUs), learning rate may saturate and square-root scaling or adaptive methods are used.

lr = base_lr * world_size  # linear scaling rule

All-gather collects a tensor chunk from every rank and concatenates them, so all ranks end up with the full combined tensor. In distributed inference with FSDP or tensor parallelism, all-gather is used to reconstruct full parameter shards before each forward computation. Each rank holds 1/N of a parameter; all-gather reconstructs the full weight matrix needed for the linear layer computation, then the result is optionally discarded to free memory after the layer.

dist.all_gather(tensor_list, local_tensor)

In tensor parallelism, activations (not just weights) are replicated across TP ranks, causing activation memory to not scale down with TP degree. Sequence parallelism (Megatron-LM v3) shards activations along the sequence dimension across TP ranks for operations that don't need full sequence context (like LayerNorm, Dropout). An all-gather reconstructs the sequence before attention; a reduce-scatter splits it after. Together, TP+SP reduces both weight and activation memory proportionally with the TP degree.

In MoE models, expert parallelism assigns different expert networks to different GPUs. The router sends tokens to the appropriate expert GPU via all-to-all communication, each expert processes its assigned tokens locally, and results are sent back via another all-to-all. This enables models with hundreds of experts to scale across many GPUs while each GPU only computes for a subset of experts. Expert parallelism is combined with TP and DP in frameworks like Tutel and MegaBlocks.

Communication-computation overlap runs gradient all-reduce communication in parallel with the remaining backward computation, hiding latency. DDP achieves this by attaching backward hooks to each parameter: as soon as a parameter's gradient is ready, its bucket's all-reduce starts asynchronously while backward continues computing gradients for earlier layers. This means by the time backward finishes, most gradient synchronization is complete, making overall training time closer to compute-only speed.

FP8 (8-bit floating point) uses two formats: E4M3 (4-bit exponent, 3-bit mantissa, good for activations/weights) and E5M2 (5-bit exponent, 2-bit mantissa, wider range for gradients). FP8 halves memory vs FP16 and doubles throughput on NVIDIA Hopper GPUs (H100), which have FP8 Tensor Cores. The Transformer Engine library from NVIDIA implements FP8 training with automatic scaling and cast-back-to-FP16 for non-quantization-friendly ops. AMD MI300X also supports FP8.

A process group is a subset of all distributed processes that participate in collective operations together. The default group includes all ranks, but custom groups created via dist.new_group([ranks]) allow collective operations (all-reduce, broadcast) to be performed only among a subset. This is essential for hybrid parallelism: you might all-reduce gradients within a DP group, all-gather parameters within a TP group, and send/receive activations within a PP group — each requiring different process groups.

tp_group = dist.new_group(ranks=[0, 1, 2, 3])  # tensor parallel group
dist.all_reduce(tensor, group=tp_group)

bfloat16 and float16 both use 16 bits, but bfloat16 allocates 8 bits to the exponent (matching float32's range) and only 7 bits to the mantissa. float16 has 5 exponent bits and 10 mantissa bits — higher precision but much smaller dynamic range, prone to overflow/underflow. In distributed training, gradient all-reduces can produce large values that overflow float16 but fit comfortably in bfloat16. On NVIDIA Ampere/Hopper and Google TPUs, bfloat16 is preferred because it eliminates the need for loss scaling (GradScaler).

model = model.to(torch.bfloat16)  # no GradScaler needed

DDP normally all-reduces gradients after every backward pass. With gradient accumulation (accumulating over N micro-batches before stepping), you want to skip all-reduce for the intermediate backward calls and only synchronize on the final one. model.no_sync() suppresses gradient synchronization inside the context. Calling backward N-1 times under no_sync() then once outside accumulates gradients locally and only communicates on the final step, reducing communication overhead by N-1x.

for i, (x, y) in enumerate(loader):
    with model.no_sync() if (i+1) % accum != 0 else contextlib.nullcontext():
        loss = criterion(model(x), y) / accum
        loss.backward()

dist.barrier() forces all processes to wait at the call site until every process reaches it, ensuring global synchronization. It is necessary before operations that must see consistent state across ranks (e.g., before loading a checkpoint that rank 0 just saved). However, overusing it introduces unnecessary synchronization bubbles and can deadlock if a subset of ranks never reaches the call. It should be called sparingly — only when true synchronization points are required.

if rank == 0:
    save_checkpoint(model)
dist.barrier()  # all ranks wait until rank 0 saves
model.load_state_dict(torch.load('ckpt.pt'))

Mixed-precision training stores activations and gradients in float16/bfloat16, while keeping a float32 master copy of parameters for the optimizer step. This halves activation memory and gradient memory compared to pure FP32. In distributed settings (DDP/FSDP), gradient all-reduces are performed in FP16 (smaller buffers, higher bandwidth utilization), and the float32 master weights are used only locally for optimization. FSDP can shard FP32 optimizer states across GPUs to further reduce memory.

Chinchilla (Hoffmann et al., 2022) showed that for a given compute budget, the optimal model size scales with the number of tokens, and optimal tokens per parameter is approximately 20. Earlier models (GPT-3) were undertrained — too large for their compute. For distributed training, this implies allocating compute to longer training runs on moderate-size models rather than simply scaling model size. It also shifts the bottleneck from model parallelism (fewer needed) to data throughput.

Activation checkpointing reduces activation memory from O(L) to O(sqrt(L)) by recomputing activations during backward instead of storing them. In distributed settings, this reduces per-GPU peak memory, enabling larger batch sizes per GPU (higher GPU utilization) and fitting larger models. The tradeoff is approximately 33% extra FLOPs for recomputation. In pipeline parallelism, it allows more micro-batches to be in flight simultaneously, reducing the pipeline bubble and improving efficiency despite the recompute cost.

from torch.utils.checkpoint import checkpoint_sequential
out = checkpoint_sequential(layers, segments=4, input=x)

With FSDP, each rank holds only a shard of the parameters. For saving, FSDP offers StateDictType options: FULL_STATE_DICT (rank 0 gathers and saves the complete model; memory-intensive), SHARDED_STATE_DICT (each rank saves its shard; more scalable), and LOCAL_STATE_DICT (raw local shards). Loading requires using the matching context manager. SHARDED_STATE_DICT with DistributedStateDictContext is recommended for large models to avoid OOM on rank 0.

from torch.distributed.fsdp import StateDictType
with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
    state = model.state_dict()

Intra-node communication (between GPUs on the same server) uses NVLink or PCIe, achieving 600 GB/s (NVLink 4.0, H100) or 32 GB/s (PCIe 4.0). Inter-node communication travels over InfiniBand (200–400 Gb/s) or Ethernet, which is 2–3 orders of magnitude lower bandwidth. This asymmetry dictates parallelism strategy: tensor parallelism (high-bandwidth, many collective ops) should stay intra-node; pipeline parallelism (fewer, larger point-to-point messages) crosses nodes; data parallelism gradient all-reduce also benefits from being intra-node first.

MFU is the fraction of a GPU's theoretical peak FLOPs that are actually used for model computation: MFU = (observed_throughput × model_FLOPs_per_forward_backward) / (GPU_peak_TFLOPS). Model FLOPs for a transformer are approximately 6ND per token (6 × parameters × tokens, accounting for forward+backward). MFU of 30–60% is typical for well-optimized large-model training. It is a hardware-agnostic metric for comparing training efficiency across systems.

Flash Attention (Dao et al., 2022) rewrites the scaled dot-product attention to be I/O-aware: it tiles Q, K, V into SRAM blocks, computes attention incrementally without materializing the full N×N attention matrix in GPU HBM. Memory usage drops from O(N^2) to O(N), and HBM read/write bandwidth is drastically reduced, making attention up to 3-5x faster on A100. In distributed training, this allows larger sequence lengths and batch sizes per GPU, directly improving MFU.

out = torch.nn.functional.scaled_dot_product_attention(
    q, k, v, is_causal=True)  # uses Flash Attention under the hood

Ring all-reduce arranges N GPUs in a logical ring. In the reduce-scatter phase, each GPU sends a 1/N chunk to its neighbor and accumulates. After N-1 steps, each GPU holds the fully-reduced value for its chunk. In the all-gather phase, each GPU sends its reduced chunk around the ring. Total data sent per GPU is 2(N-1)/N × tensor_size, which approaches 2× tensor_size for large N — bandwidth-optimal because adding more GPUs does not increase per-GPU communication. Latency scales O(N) but this is secondary at scale.

ZeRO-Infinity (Stage 4 effectively) extends parameter and optimizer state sharding to CPU RAM and NVMe SSDs, enabling training of models with hundreds of billions of parameters on clusters where GPU memory would otherwise be a hard limit. Parameters reside on NVMe/CPU and are streamed to GPU just-in-time for computation. CPU Adam optimizer updates are performed on CPU. NVMe bandwidth (~3 GB/s) is the bottleneck, requiring careful prefetching and overlapping to maintain GPU utilization.

find_unused_parameters=True tells DDP to scan the output of backward to identify parameters that did not receive gradients (unused in that forward pass) and exclude them from the gradient synchronization bucket. This is necessary for models with dynamic computation graphs (e.g., mixture-of-experts, conditional branches) where not all parameters are used every forward. Setting it to True adds overhead (a Python-level graph traversal); it should be False for standard models where all parameters are always used.

model = DDP(model, find_unused_parameters=True)  # for conditional architectures

Context parallelism (CP) shards the full attention context (keys and values) across GPUs, unlike sequence parallelism which only shards non-attention ops. In CP, each GPU holds a segment of Q, K, V. A ring-style attention kernel sends K and V chunks around the ring, computing local Q against each K,V chunk as it passes. This enables sequence lengths beyond what any single GPU's SRAM can handle and scales sub-quadratically with GPU count. NVIDIA Megatron-Core and Ulysses implement CP.

dist.broadcast(tensor, src=0) copies a tensor from a source rank to all other ranks. A critical use case is model initialization: after creating a model, rank 0 broadcasts its initial parameters to all other ranks to ensure every process starts from an identical model (avoiding divergence from random seed differences). DDP handles this automatically at construction, but when using FSDP with custom initialization, an explicit broadcast may be needed.

dist.broadcast(model.weight, src=0)

In distributed training, gradients are already synchronized (all-reduced) across ranks before clipping. The correct order is: (1) backward(), (2) unscale gradients if using GradScaler, (3) clip_grad_norm_() on local gradients (all ranks have identical gradients after all-reduce, so clipping is consistent), (4) optimizer.step(). If clipping is applied before all-reduce (in the no_sync context), each rank clips local gradients independently, which gives inconsistent results.

loss.backward()
scaler.unscale_(optimizer)  # before clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)

Activation offloading moves forward-pass activations (needed for backward) to CPU RAM instead of keeping them on GPU. When activation memory is the bottleneck (e.g., long sequences), this frees GPU VRAM for larger batch sizes or sequence lengths. The GPU-to-CPU PCIe transfer is the bottleneck (32 GB/s for PCIe 4.0), so activations must be moved asynchronously during computation. FSDP's cpu_offload parameter enables this. The tradeoff is data transfer latency vs freed VRAM.

from torch.distributed.fsdp import CPUOffload
model = FSDP(model, cpu_offload=CPUOffload(offload_params=True))

NVIDIA Nsight Systems profiles timeline of GPU kernels, NCCL communication ops, CPU activity, and CUDA streams. In distributed training, it reveals: (1) NCCL all-reduce duration and overlap with compute, (2) kernel launch gaps (CPU-bound overhead), (3) synchronization stalls, (4) load imbalance across pipeline stages. By identifying whether training is compute-bound or communication-bound, engineers decide whether to increase batch size, enable overlap, switch to reduce-scatter, or optimize kernel launch patterns.

DDP's static_graph=True flag signals that the same set of parameters will be used in every forward pass and that backward ops are identical across iterations. This allows DDP to disable the overhead of tracking which parameters participated in backward (find_unused_parameters traversal) and to pre-build gradient buckets once. It enables additional optimizations like gradient bucket fusion. For standard static architectures (ResNets, ViTs, GPT), enabling this can improve throughput by 5–10%.

model = DDP(model, static_graph=True)

Interleaved scheduling (Megatron-LM) assigns each GPU multiple non-consecutive stage chunks (virtual stages) instead of one contiguous range. With v virtual stages per GPU, the pipeline bubble fraction reduces from (p-1)/(m+p-1) to (p-1)/(vm+p-1) where m is micro-batches and p is pipeline stages. The tradeoff is more all-reduce communications between model chunks. For large m (many micro-batches), both approaches converge in bubble reduction, but interleaved is better for smaller m.

all_to_all sends different chunks of a tensor from each rank to every other rank — it is a personalized exchange where rank i sends data[j] to rank j and receives data[i] from rank j. Total data moved is world_size × chunk_size per rank. It is the key collective for expert parallelism in MoE: tokens are routed to specific expert GPUs (all_to_all dispatch), processed by their expert, then returned (all_to_all combine). It is also used in sequence-to-data-parallel transitions.

dist.all_to_all(output_list, input_list)

dist.reduce(tensor, dst=0) sums (or applies another op) the tensors across all ranks and stores the result only on the destination rank (rank 0 by default). dist.all_reduce(tensor) performs the same reduction but stores the result on ALL ranks. All-reduce is used in DDP (every rank needs updated gradients for local parameter update). Reduce is used when only one rank needs the aggregated result (e.g., accumulating validation metrics on rank 0 for logging).

dist.reduce(val_loss, dst=0)  # only rank 0 gets result
dist.all_reduce(grad, op=dist.ReduceOp.SUM)  # all ranks get result

In asynchronous distributed SGD, a worker computes gradients based on parameter version W_t but may apply them to W_{t+k} because other workers updated parameters in between. The k-step-old gradient is stale, potentially pointing in the wrong direction. Bounded staleness (SSP — Stale Synchronous Parallel) sets a maximum staleness threshold τ: fast workers block if they are more than τ steps ahead of the slowest. This bounds staleness while still allowing fast workers to proceed without full synchronization.

Megatron-LM uses a hierarchical strategy: TP degree is set to fit one model replica's layers on NVLink-connected GPUs (typically 8 on DGX nodes) due to frequent all-reduces. PP degree is set to divide model layers across nodes with minimal bubble. DP fills the remaining GPU count. The total GPU count = TP × PP × DP. Practitioners use empirical heuristics: TP=8 within-node, PP=1–8 across nodes, DP fills remaining. Profiling with different configs validates the choice.

# Total GPUs = TP x PP x DP
# E.g., 512 GPUs: TP=8, PP=8, DP=8

Data parallelism (DDP/FSDP) communicates gradients via all-reduce after backward — communication volume is proportional to parameter count, and it occurs once per training step. Pipeline parallelism communicates activations and gradients as point-to-point messages between adjacent stages — communication volume per micro-batch is proportional to the activation tensor at the stage boundary (batch × hidden_dim), and it occurs at every micro-batch boundary. DP is communication-intensive in gradient volume; PP is latency-sensitive with smaller per-message volume at higher frequency.

Checkpoint frequency determines how much training progress is lost when a job fails. With thousands of GPUs, hardware failures (GPU crash, network partition) occur frequently — mean time between failures for a 1000-GPU run can be hours. Checkpointing every N steps limits lost work to N steps × step_time. Trade-off: frequent checkpointing consumes I/O bandwidth and storage. Asynchronous checkpointing (saving to CPU RAM while training continues, then writing to disk) reduces I/O pause. Distributed checkpointing libraries (PyTorch Distributed Checkpoint, Nebula) parallelize writes across ranks.

Larger effective batch sizes (world_size × local_batch_size) reduce gradient noise by averaging over more samples — the gradient estimate has lower variance (scales as 1/sqrt(B)). This allows larger learning rates (linear scaling rule) and fewer steps to convergence, but also reaches a critical batch size beyond which additional parallelism yields diminishing returns. The linear scaling rule breaks down in the noise-dominated regime (small B) and the curvature-dominated regime (very large B). Measuring gradient signal-to-noise ratio (GSNR) helps diagnose when batch size is too large.