Distributed Training: Scaling Machine Learning Across GPUs and Clusters

Introduction

The past decade of machine learning has been defined by scaling. From BERT's 340 million parameters in 2018 to GPT-4's estimated 1.8 trillion parameters in 2023, model sizes have grown by roughly 10x every two years. This growth is not arbitrary — it follows the scaling laws established by Kaplan et al. (2020) and Hoffmann et al. (2022), which show that model performance improves predictably with increased parameters, data, and compute when scaled together [1][2].

But scaling laws have a hardware bottleneck. A single H100 GPU has 80 GB of HBM3 memory and delivers 1979 TFLOPS of FP16 compute. A 70-billion-parameter model in FP16 requires 140 GB of parameter memory alone — 1.75x the capacity of a single GPU — plus memory for gradients, optimizer states, and activations. Training such a model on one GPU is physically impossible. Training a 500-billion-parameter model requires dozens of GPUs for memory capacity alone, even before considering throughput.

Distributed training solves this problem by spreading the computational and memory load across multiple devices. The core challenge is not merely parallelism — it is efficient parallelism. Amdahl's law dictates that the serial fraction of any workload bounds the maximum speedup. In distributed training, the serial fraction is dominated by communication: synchronizing gradients, transferring activations, and coordinating parameter updates across devices. The art of distributed training is maximizing compute utilization while minimizing communication overhead.

This guide covers the five major parallelism strategies — data, model, pipeline, tensor, and fully sharded data parallelism — along with their hybrid combinations. We examine the infrastructure required to support each strategy and provide a cost analysis framework for planning large-scale training runs. Understanding these strategies is essential for anyone training models larger than a single GPU can hold, or anyone who wants to train models faster by using multiple GPUs in parallel.

For background on the hardware that makes distributed training possible, see our guide on parallel processing and GPU architecture.

Data Parallelism

Data parallelism is the simplest and most widely used distributed training strategy. Each GPU holds a complete copy of the model and processes a different subset of the training batch. After each GPU computes its forward and backward pass independently, the gradients are synchronized across all GPUs via an all-reduce operation. Every GPU then applies the averaged gradients to its local model copy, ensuring all replicas remain identical.

Synchronous vs asynchronous.

In synchronous data parallelism, all GPUs finish their backward pass before any GPU applies the gradient update. This produces deterministic training — the effective batch size is worker_count x per_worker_batch_size — and maintains the same convergence properties as single-GPU training. In asynchronous data parallelism (also called parameter server architecture), each GPU updates the model independently as soon as its gradients are ready. This eliminates idle time but introduces stale gradients — a GPU might apply its update based on weights that have already been updated by other GPUs. Asynchronous training often converges to a worse solution or requires careful learning rate tuning [3].

Synchronous data parallelism with Distributed Data Parallel (DDP) in PyTorch is the production standard. The all-reduce operation that synchronizes gradients is the critical path. In the naive implementation, all GPUs communicate with all other GPUs in an O(P^2) pattern, which does not scale. Modern implementations use ring all-reduce, where GPUs are arranged in a logical ring and each GPU communicates only with its two neighbours. The total data transferred per GPU is 2 x N x (P-1) / P, where N is the model size in bytes and P is the number of GPUs. For large P, this approaches 2N — each GPU sends and receives roughly two copies of the full gradient tensor [4].

Gradient accumulation.

Gradient accumulation simulates larger batch sizes without increasing per-GPU memory. Each GPU accumulates gradients over multiple forward-backward passes before performing the all-reduce and optimizer step. This is useful when the per-GPU batch size is limited by memory — for example, training with batch size 1 per GPU but accumulating over 32 steps to achieve an effective batch size of 32 x num_GPUs. The trade-off is that the model parameters are not updated between micro-batches, which can reduce batch normalization accuracy and affect convergence dynamics.

Ring all-reduce is implemented by NVIDIA's NCCL library, which automatically selects the optimal algorithm based on message size and hardware topology. For small messages, NCCL uses ring all-reduce. For large messages (above approximately 256 MB), it switches to tree all-reduce, which reduces the number of communication steps from P-1 to log2(P) at the cost of higher bandwidth requirements on the root nodes.

# Ring all-reduce: each GPU communicates with two neighbours
# Bandwidth optimal: O(P) messages of size N/P
# For P GPUs and N elements per GPU:
#   Total data transferred per GPU = 2 * N * (P-1) / P
#   Time ≈ 2 * N / bandwidth + P * latency

# NCCL automatically selects the optimal algorithm:
import torch.distributed as dist

# All-reduce: sum gradients across all GPUs
# NCCL uses ring all-reduce for small messages (< 256 MB)
# and tree all-reduce for large messages
dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

# NCCL environment tuning for multi-node:
# export NCCL_IB_HCA=mlx5_0:1,mlx5_1:1
# export NCCL_ALGO=Ring
# export NCCL_PROTO=Simple
# export NCCL_NTHREADS=128
# export NCCL_NSOCKS_Per_THREAD=2

Data parallelism has two important limitations. First, every GPU must have enough memory to hold the full model — a 70B model in BF16 requires 140 GB for parameters alone, exceeding the 80 GB of an H100. Second, communication overhead scales with model size, not batch size — a larger model means more gradient data to synchronize per step, regardless of how much data each GPU processes. These limitations motivate the more advanced strategies discussed below.

Model Parallelism

Model parallelism splits the model across GPUs by layers. GPU 0 holds layers 1 through L, GPU 1 holds layers L+1 through 2L, and so on. During the forward pass, GPU 0 processes its layers, sends the activations to GPU 1, which processes its layers, and so on through the pipeline. The backward pass reverses the flow.

The defining characteristic of naive model parallelism is its sequential nature. At any given moment, only one GPU is active — the rest are idle, waiting for activations (forward) or gradients (backward) from their neighbours. This produces GPU utilization of approximately 1/P for P GPUs, which is abysmally inefficient. For this reason, naive model parallelism is almost never used in practice. It has been superseded by pipeline parallelism (which adds micro-batching to overlap computation) and tensor parallelism (which splits individual layers rather than layer sequences).

However, model parallelism in the broader sense — splitting a model horizontally across devices — is the foundation for all advanced parallelism strategies. The key insight is that model parallelism is required whenever a single GPU cannot hold the entire model in memory, regardless of batch size. For models exceeding 100B parameters, some form of model-level splitting is unavoidable.

For practical guidance on when and how to apply these techniques to language model training, see our LLM fine-tuning guide.

Pipeline Parallelism

Pipeline parallelism improves on naive model parallelism by dividing each batch into micro-batches and pipelining their execution through the model stages. While GPU 1 processes micro-batch 2 through its layers, GPU 0 can simultaneously work on micro-batch 3. This overlaps computation across GPUs and dramatically improves utilization.

GPipe and PipeDream.

GPipe (Huang et al., 2019) divides the batch into M micro-batches and pipelines them through the model stages [5]. All micro-batches complete their forward pass before any micro-batch begins its backward pass. This creates a pipeline bubble: at the start and end of each batch, the first and last stages are idle. The bubble overhead is proportional to (P-1) / (M+P-1), where P is the number of pipeline stages and M is the number of micro-batches. With M >> P, the bubble approaches zero, but memory usage scales linearly with M because all micro-batch activations must be stored for the backward pass.

PipeDream (Narayanan et al., 2019) uses a 1F1B (one-forward-one-backward) schedule [6]. Instead of completing all forward passes before any backward pass, each micro-batch alternates forward and backward steps in a staggered pattern. This reduces memory usage because activations are freed sooner, and it maintains the same pipeline bubble as GPipe. PipeDream also supports asynchronous gradient updates across stages, which can improve throughput at the cost of slight convergence differences.

Balanced partitioning.

The efficiency of pipeline parallelism depends critically on balanced partitioning — each pipeline stage should contain approximately the same amount of compute work. For transformer models, this is straightforward because all layers are identical: partition by equal layer counts. For heterogeneous models (mixture-of-experts, multi-modal architectures with different processing paths), automatic partitioning tools like Piper or the partitioner in DeepSpeed are essential to avoid straggler stages.

Gradient accumulation across stages adds another dimension. When using gradient accumulation, each micro-batch within an accumulation step follows the pipeline schedule, and the optimizer step is applied after all micro-batches complete both forward and backward passes. This is the standard configuration for large-scale training with pipeline parallelism.

Tensor Parallelism

Tensor parallelism splits individual operations — not layers — across GPUs. In a transformer model, each attention head and each feed-forward network layer can be partitioned across multiple GPUs. The key observation is that the independent heads in multi-head attention are trivially parallelizable: with T GPUs, each GPU computes N/T attention heads independently, and the results are concatenated.

The Megatron-LM approach (Shoeybi et al., 2019) defines two fundamental tensor-parallel patterns for transformer blocks [7]. For the feed-forward network with a weight matrix W of shape [d_model, d_ff], the matrix is split column-wise: GPU i holds W[:, i * d_ff/T : (i+1) * d_ff/T]. Each GPU computes its portion of the forward pass independently, and an all-reduce across the T GPUs combines the partial results. For the attention block, the QKV projection weight is split row-wise, and the output projection is split column-wise.

Communication requirements.

Tensor parallelism is communication-intensive. Every transformer layer requires two all-reduce operations: one after the attention computation and one after the feed-forward computation. For a model with L layers, the total communication per training step is 4 x L x model_dimension bytes per GPU (two all-reduces, each requiring 2x the data volume in bidirectional communication). This is significantly more than data parallelism for large models, which is why tensor parallelism requires high-bandwidth intra-node interconnects (NVLink or NVSwitch) with at least 600 GB/s per GPU.

For models that fit on a single node, tensor parallelism with TP=8 (matching the 8 GPUs in a DGX node) is the standard configuration. Going beyond TP=8 requires inter-node tensor parallelism, which is only feasible with InfiniBand or comparable high-bandwidth, low-latency interconnects. In practice, most training runs cap tensor parallelism at node level and use pipeline parallelism to scale across nodes.

Fully Sharded Data Parallelism (FSDP)

FSDP, introduced by Zhao et al. (2023) and implemented in PyTorch, unifies data parallelism with model-level memory savings [8]. It is based on the ZeRO (Zero Redundancy Optimizer) approach from Rajbhandari et al. (2020), which eliminates memory redundancy by sharding optimizer states, gradients, and parameters across data-parallel workers [9].

ZeRO stages.

ZeRO defines three stages of memory optimization. ZeRO-1 shards only the optimizer states across GPUs. For AdamW, which stores two state values per parameter (momentum and variance), this means each GPU holds 2/P of the total optimizer state. ZeRO-2 adds gradient sharding — each GPU stores only its assigned portion of the gradients. ZeRO-3 shards the model parameters themselves. At ZeRO-3, each GPU holds 1/P of the parameters, 1/P of the gradients, and 1/P of the optimizer states.

The memory savings are dramatic. A 70B model trained with full precision (FP32 for optimizer states, FP16 for parameters and gradients) normally requires 140 GB (parameters) + 140 GB (gradients) + 560 GB (optimizer states for AdamW) = 840 GB total. With ZeRO-3 across 8 GPUs, each GPU holds 17.5 GB of parameters, 17.5 GB of gradients, and 70 GB of optimizer states — 105 GB per GPU, which fits comfortably in an H100's 80 GB when using BF16 for parameters and gradients (reducing to 52.5 GB per GPU before activations).

Communication patterns.

FSDP uses all-gather to collect the full parameter set before each forward pass and reduce-scatter to distribute gradients after each backward pass. The all-gather operation requires each GPU to send (P-1)/P of the parameter data and receive (P-1)/P of the parameter data, totalling 2 x (P-1)/P x model_size bytes per forward pass. The reduce-scatter is symmetric. Compared to DDP, which performs a single all-reduce equal to the model size per step, FSDP performs N_layer all-gather/reduce-scatter pairs (one per wrapped layer), increasing total communication volume by roughly 1.5-2x depending on the wrapping policy.

CPU offloading.

ZeRO-3 can offload optimizer states and parameters to CPU memory, further reducing GPU memory usage at the cost of CPU-GPU transfer latency. The DeepSpeed library provides ZeRO-Infinity, which tiers memory across GPU HBM, CPU DRAM, and NVMe storage [10]. This enables training models with up to 100x more parameters than GPU memory would otherwise allow, though at significantly reduced throughput (typically 20-40% of the non-offloaded baseline).

# FSDP configuration for training a 70B parameter model
import torch
import torch.distributed as dist
from torch.distributed.fsdp import (
    FullyShardedDataParallel as FSDP,
    MixedPrecision,
    ShardingStrategy,
    BackwardPrefetch,
)
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
    checkpoint_wrapper,
    CheckpointImpl,
)

dist.init_process_group("nccl")
torch.cuda.set_device(local_rank)

mixed_precision = MixedPrecision(
    param_dtype=torch.bfloat16,
    reduce_dtype=torch.bfloat16,
    buffer_dtype=torch.bfloat16,
)

model = FSDP(
    model,
    sharding_strategy=ShardingStrategy.FULL_SHARD,
    mixed_precision=mixed_precision,
    auto_wrap_policy=transformer_auto_wrap_policy,
    backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
    device_id=local_rank,
    limit_all_gathers=True,
    use_orig_params=True,
)

# Enable activation checkpointing (selective)
for layer in model.layers:
    layer.checkpoint = checkpoint_wrapper(
        layer, checkpoint_impl=CheckpointImpl.NO_REENTRANT
    )

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, fused=True)

for batch in dataloader:
    loss = model(batch)
    loss.backward()
    for p in model.parameters():
        p.grad.div_(dist.get_world_size())
    optimizer.step()
    optimizer.zero_grad()

Mixed precision with FSDP.

FSDP integrates directly with mixed-precision training. The standard configuration uses BF16 for forward and backward computations (parameters, activations, gradients) while maintaining an FP32 master copy of weights for optimizer updates. This reduces memory usage by approximately 40% compared to full FP32 training while maintaining numerical stability for most models. The MixedPrecision wrapper in FSDP handles the precision casts automatically — parameters are stored in BF16, upcast to FP32 for the optimizer step, and re-sharded in BF16 after the update.

For more detail on how precision reduction affects model quality, see our guide on LLM quantization and compression.

Hybrid Parallelism

For models exceeding 100 billion parameters, no single parallelism strategy is sufficient. Production training runs combine all four strategies into a unified approach called hybrid or 3D/4D parallelism. The standard configuration for training a 500B+ parameter model uses tensor parallelism within each node (typically TP=8 for an 8-GPU node), pipeline parallelism across nodes (PP=8 to 64), and data parallelism across the remaining dimension (DP = total_GPUs / (TP x PP)).

Meta's Llama 3 405B training run provides a concrete example: 16,384 H100 GPUs organized as 2,048 nodes, with TP=8 within each node, PP=8 across nodes, and DP=128 across the remaining dimension. This configuration achieves approximately 45% model flops utilization (MFU), meaning the GPUs spend 45% of their time computing and 55% waiting for communication or synchronization [11]. While 45% MFU appears low, it is state of the art for multi-thousand-GPU training and represents years of engineering optimization.

4D parallelism.

A fourth dimension — sequence parallelism — is increasingly used for training with very long sequences (128K+ tokens). Sequence parallelism splits the sequence dimension of the attention computation across GPUs, leveraging the fact that the per-token compute of attention scales quadratically with sequence length. By distributing sequence positions across GPUs, each GPU processes a shorter sequence, reducing both memory and compute per device. Ring attention and DeepSpeed-Ulysses are the two leading implementations of sequence parallelism.

Cluster topology considerations.

The optimal hybrid configuration depends on cluster topology. In a cluster where each node has 8 GPUs connected via NVSwitch (7.2 TB/s bisection bandwidth), tensor parallelism within the node is essentially free in terms of communication overhead. If inter-node connectivity uses 400 Gbps InfiniBand (approximately 50 GB/s per link), pipeline parallelism (which communicates only between adjacent stages) wastes less bandwidth than data parallelism (which requires all-to-all communication). The general rule is: place the highest-communication strategy (tensor parallelism) within the fastest interconnect (NVLink/NVSwitch), and the lowest-communication strategy (data parallelism) across the slowest interconnect.

For a deeper analysis of cluster design and networking, see our complete guide to AI infrastructure.

Infrastructure Considerations

Distributed training is as much an infrastructure problem as a software problem. The choice of interconnect, storage, orchestration, and fault tolerance strategy directly determines training throughput and reliability.

Interconnect hierarchy.

The three tiers of GPU interconnect have dramatically different performance profiles. NVLink 4.0 provides 900 GB/s bidirectional bandwidth per GPU (450 GB/s each direction) within a node. NVLink 5.0, introduced with the B200, doubles this to 1.8 TB/s. InfiniBand NDR 400 provides 400 Gbps (50 GB/s) per link, typically deployed with 8 links per node for 400 GB/s aggregate. Ethernet (RoCE v2) provides 100-400 Gbps per link with higher latency and CPU overhead than InfiniBand. The practical impact: a tensor-parallel all-reduce that takes 10 microseconds on NVLink might take 200 microseconds on InfiniBand and 500 microseconds on Ethernet.

NCCL tuning.

NVIDIA's NCCL library is the communication layer used by PyTorch, TensorFlow, JAX, and most training frameworks. NCCL auto-detects the cluster topology and selects optimal communication algorithms (ring, tree, or NVLink direct). For multi-node training, key environment variables include NCCL_IB_HCA (InfiniBand HCA selection), NCCL_ALGO (algorithm override), NCCL_PROTO (Simple or LL protocol for large messages), and NCCL_DEBUG (info or trace for diagnosing communication bottlenecks). A common performance pitfall is running multi-node training without properly setting NCCL_IB_HCA, which causes NCCL to fall back to TCP/IP over the management network instead of using InfiniBand.

Checkpointing strategies.

For training runs spanning days or weeks, checkpointing is critical for fault tolerance. The standard approach is asynchronous checkpointing: the training loop periodically writes model state (parameters, optimizer state, learning rate, data loader state) to a distributed file system (Lustre, GPFS, or cloud object store) while continuing training in the background. The checkpoint frequency must balance storage cost against compute loss — a checkpoint every hour loses at most one hour of work on failure, but writing a 140 GB checkpoint (for a 70B model) every hour requires 3.4 TB/day of storage bandwidth. Efficient checkpointing uses distributed I/O where each GPU writes its sharded portion of the state independently, avoiding a single-writer bottleneck.

Fault tolerance.

At thousands of GPUs, hardware failures are a daily occurrence. The mean time between failures (MTBF) for an individual GPU in a large cluster is approximately 10,000 hours. For a 16,384-GPU cluster, the expected time between any GPU failure is roughly 10,000 / 16,384 = 0.61 hours (37 minutes). Production training frameworks must handle GPU failures gracefully through elastic training — when a GPU fails, the remaining GPUs redistribute the sharded state and continue from the last checkpoint. Technologies like TorchElastic and Kubernetes-based fault tolerance have made elastic training practical, though the overhead of re-sharding and resuming can cost 5-15 minutes per failure.

Containerization with Kubernetes.

Most large-scale training runs use Kubernetes for orchestration with GPU-specific operators (NVIDIA GPU Operator, Kubeflow, Volcano). A typical training pod contains 8 GPUs (matching a single node), and the scheduler ensures GPU-affinity through node selectors and taints. Dynamic GPU sharing — where multiple training jobs share a GPU cluster through time-slicing or MIG (Multi-Instance GPU) — is increasingly common for teams running multiple concurrent experiments. The trade-off is that sharing introduces variable performance due to GPU memory contention and thermal throttling.

Cost Analysis

Training large models is expensive. The following analysis compares the cost of training a 70B-parameter model using different parallelism configurations and deployment options. Estimates assume a training run of 500 billion tokens (approximately 2.5x the Chinchilla-optimal compute for a 70B model) and use BF16 mixed precision.

Note that total GPU-hours remain roughly constant up to ~128 GPUs because linear scaling is achieved — doubling the GPUs halves the training time without significant communication overhead. Beyond 128 GPUs, the communication overhead of data parallelism across nodes reduces scaling efficiency from 100% to approximately 80-90% (the GPU-hour total increases by ~5-10% when going from 128 to 1,024 GPUs). The cost per GPU-hour assumes on-demand H100 pricing at approximately $11.30/hr (average across major cloud providers).

Cloud vs on-premise.

On-premise GPU clusters involve a large capital expenditure but lower marginal cost. An 8-GPU H100 node costs approximately $300,000. A 1,024-GPU cluster requires 128 nodes at $38.4M capital cost, plus networking ($3-5M), storage ($1-2M), facility ($10-15M for power and cooling), and annual operating costs ($5-8M for power at $0.10/kWh). At $11.30/GPU/hr on-demand, 1,024 GPUs cost $11,571/hr. Break-even against on-premise requires approximately 2,000-3,000 hours of utilization per year, assuming a 3-year depreciation period.

Spot and reserved instances.

Spot/preemptible GPU instances reduce cost by 60-80% compared to on-demand. For training runs with robust checkpointing, spot instances are the most cost-effective option — a 70B model that costs $280K on on-demand GPUs can be trained for $56-84K on spot instances, assuming a 10-20% failure-and-restart overhead. Reserved instances (1-year commitment) provide 30-50% savings over on-demand, making them the preferred option for teams with continuous training workloads. Checkpoint storage adds $5-15K depending on frequency and retention policy, using cloud object storage at approximately $0.023/GB/month.

For a practical framework for managing training infrastructure costs, see our MLOps production guide.

Parallelism Strategy Comparison

The following table summarizes the five parallelism strategies across key dimensions. No single strategy is universally optimal — the right choice depends on model size, cluster topology, and throughput requirements.

Combining strategies: a decision framework.

For a model that fits on a single GPU, use data parallelism (DDP). For a model that exceeds single-GPU memory but fits on 8 GPUs, use FSDP with FULL_SHARD. For a model exceeding 8-GPU memory, add tensor parallelism (TP=8 within the node) and pipeline parallelism across nodes. For models above 500B parameters, add sequence parallelism and consider ZeRO-3 offloading for memory-constrained configurations. The optimal configuration is model-specific and requires profiling — the MFU achieved by different strategies varies by up to 2x depending on model architecture, sequence length, and batch size.