Data Parallelism › Limitations and Trade-offs

Despite its simplicity and effectiveness, pure data parallelism has fundamental constraints that become relevant as models and cluster sizes grow. Understanding these limitations shapes how you combine DDP with other distributed training strategies.

Data parallelism's most significant cost is memory: every GPU must hold a complete copy of the model, gradients, and optimizer states. For a model with $P$ parameters trained with AdamW in mixed precision, the memory cost per GPU is approximately:

$$\text{Memory per GPU} = P \cdot (2 + 2 + 4 + 8) = 16P \text{ bytes}$$

where the terms represent: 2 bytes for BF16 parameters, 2 bytes for BF16 gradients, 4 bytes for FP32 master weights (used by the optimizer for numerical stability), and 8 bytes for Adam's first and second moment estimates stored in FP32 (4 bytes each).

For a 7B parameter model, this is roughly 112 GB per GPU, far exceeding the 80 GB capacity of an A100 or H100. Pure data parallelism cannot train such models without additional techniques. ZeRO optimization (Chapter 6 of this part) addresses this by partitioning the optimizer states and gradients across GPUs. FSDP goes further and also partitions the parameters. Tensor parallelism and pipeline parallelism (Chapters 4 and 5) partition the model computation itself across GPUs.

This is why real-world large-scale training combines multiple parallelism strategies. Data parallelism is typically the outermost dimension: you have $D$ data-parallel replicas, and each replica itself uses tensor parallelism across $T$ GPUs within a node. The total GPU count is $D \times T$. The gradient all-reduce in the DDP dimension only needs to communicate between the $D$ data-parallel groups, not across all $D \times T$ GPUs.

DDP requires all ranks to reach the barrier at the end of each step before any rank can proceed to the next step. This is enforced by the all-reduce: every rank must contribute its gradients, and no rank receives the averaged result until all ranks have submitted. If one GPU is slower than the others (due to hardware variation, variable batch complexities, thermal throttling, or OS preemption), all other GPUs must wait at the barrier. This straggler effect directly reduces effective GPU utilization.

The straggler problem grows worse as the number of GPUs increases. With $N = 8$ GPUs, the slowest GPU on average is somewhat faster than with $N = 512$ GPUs, where the probability of having a significantly slow GPU in any given step is much higher. In large clusters, per-step variance in GPU timing is substantial, and straggler management becomes a first-class concern.

Practical straggler mitigation strategies include:

• Using homogeneous hardware (all GPUs the same model and generation)
• Setting appropriate timeouts on collective operations to detect and recover from hangs
• Using padding and masking for variable-length sequences rather than variable-length batches (which would cause load imbalance across ranks)
• Monitoring per-rank step times to identify consistently slow ranks (often indicating hardware faults)

PyTorch's distributed training infrastructure includes a watchdog mechanism that detects collective operation timeouts and raises informative errors, preventing silent hangs that would otherwise waste compute indefinitely.

Scaling to more GPUs means a larger global batch size, which can negatively affect model quality. Research has shown that there is a sweet spot for batch size in language model training. Below this "critical batch size," doubling the batch size roughly halves the number of training steps needed to reach a given loss level, so the total compute is constant. Above the critical batch size, further increases in batch size no longer produce proportional reductions in step count: you need more total compute (more tokens seen) to reach the same quality.

The critical batch size for language models is roughly related to the signal-to-noise ratio of the gradient. As the batch size increases, gradient variance decreases (noise goes down), but the useful signal (the true gradient direction) also becomes "used up" faster per update. The optimal regime balances these effects. Empirically, the critical batch size for large language model training is typically in the range of a few hundred thousand to a few million tokens per step.

When you add GPUs and correspondingly increase the global batch size, you may exceed this critical batch size. Training with an above-critical batch size is not wrong: the model will still converge, and you are not wasting compute in an absolute sense. But you may need to train for more total tokens than you would have with a smaller batch size to reach the same final quality. The linear scaling rule adjusts for this in the optimizer step size, but it cannot eliminate the fundamental relationship between batch size and gradient quality.

Grasping this trade-off is essential for practical distributed training decisions. Adding GPUs is not free in terms of model quality per training token. The optimal configuration depends on your hardware budget, time budget, and quality requirements. For many practical training scenarios, a combination of data parallelism, gradient accumulation, and careful batch size tuning achieves the best balance.

With many GPUs processing data rapidly, the data loading pipeline can become a bottleneck. Each rank needs its data on time for every step. With insufficient I/O throughput, GPUs stall waiting for data rather than doing computation. For large transformer models that process data quickly, this is a real risk, particularly when training on slow network filesystems or when tokenization is performed on-the-fly.

Solutions include:

• Multiple DataLoader workers: Use num_workers=8 or more in DataLoader to parallelize data fetching and preprocessing across CPU threads.
• Pinned memory: Use pin_memory=True in DataLoader to store batches in page-locked host memory, enabling faster CPU-to-GPU memory transfers via DMA.
• Pre-tokenized data: Store tokenized data in memory-mapped binary files (for example, NumPy .npy, HuggingFace datasets with Arrow format, or webdataset tar files) to avoid per-epoch tokenization overhead.
• Streaming datasets: For very large datasets (trillions of tokens), stream data from disk rather than loading the whole dataset into RAM. Libraries like webdataset and HuggingFace's streaming datasets support this pattern.
• Local SSD caching: Copy frequently accessed data from slow network storage to fast local NVMe SSDs on training nodes to reduce I/O latency.

Pure DDP has no built-in fault tolerance. If any rank fails (due to a hardware error, preemption, or network partition), the training job fails completely. The all-reduce collective requires all participants; any absent rank causes all others to hang indefinitely at the barrier.

For short training runs on reliable hardware, this is acceptable. For multi-week training runs on thousands of GPUs, hardware failures are not rare events but expected occurrences. GPU failures, network transients, and node preemptions happen daily on large clusters. Without fault tolerance, a single failure wastes all progress since the last checkpoint.

Solutions include:

• Frequent checkpointing: Save model and optimizer state every few hundred steps. Combined with elastic training frameworks that can restart on fewer GPUs, checkpoints bound the amount of wasted compute after a failure.
• Elastic training: PyTorch's torchrun and frameworks like PyTorch Lightning support elastic training, where the number of workers can change dynamically and training resumes with a different number of GPUs after a failure.
• Redundant training: Some large-scale training runs checkpoint to distributed storage systems (like GCS or S3) that are themselves fault-tolerant, making sure checkpoints survive even if the training nodes fail.