Part IV: Parallel Deep Learning and Large Models
Chapter 15: Data-Parallel Deep Learning

Data-Parallel Deep Learning

Part III trained models where the data lived, on shards the scheduler scattered or on phones you were not allowed to touch. Part IV changes the question. The data is yours again, co-located on a cluster you own, and the new pressure is the model: a deep network whose single pass over a batch is so expensive that one accelerator cannot finish an epoch in any reasonable time. This chapter takes the simplest and most widely deployed answer, replicate the whole model on every accelerator, split the batch across them, and sum the gradients, and turns it from a one-line idea into a system that scales from one GPU to a multi-node cluster without leaving the accelerators idle on every step.

Conceptual illustration for Chapter 15: Data-Parallel Deep Learning

"Eight of us hold the exact same weights and never speak about anything except gradients. Once per step we shout our sums into the ring, agree on the average, and pretend we each trained the whole batch alone."

A Replica Who Has Made Peace With All-Reduce

Chapter Overview

This is the first chapter of Part IV, and it marks a turn in the book's argument. Part III asked how to learn from data spread across machines you did not fully control; Part IV asks how to train one very large model fast on a cluster you do. The binding constraint flips. There the difficulty was that the data could not move, was not identically distributed, and had to stay private. Here the data is yours to scatter and reshuffle freely, and the difficulty is that a single forward and backward pass over a deep network is too slow on one accelerator, so the model, the batch, or both must be spread across many. Data parallelism is the first and simplest way to spread the work, and it is the foundation the rest of the part builds on: model and pipeline parallelism in Chapter 16 enter only when the model itself no longer fits on one device, and the sharded optimizers of that chapter are data parallelism with the redundant state removed.

Read in order, the nine sections take you from "one GPU cannot finish the epoch" to "a multi-node cluster trains the same model at near-linear speedup, with the communication tucked invisibly behind the computation." The thread to watch is the exact-gradient identity of Part I made practical: data parallelism is correct because the summed gradients are exact, but it is fast only because the summation is overlapped, compressed, and scheduled with care. Every later parallelism strategy in the part, pipeline, tensor, expert, and sharded, is measured against the simplicity and the limits of the data-parallel baseline this chapter establishes.

Prerequisites

This chapter assumes the communication and optimization background built earlier in the book. From Chapter 4: Communication Primitives for Distributed Training you carry the single most important tool, all-reduce, and the ring algorithm that makes its bandwidth cost independent of the number of workers, because gradient synchronization in Sections 15.4 and 15.5 is an all-reduce on every step and the whole chapter turns on hiding its cost. From Chapter 10: Distributed Optimization you carry the data-parallel gradient itself, the identity that the average of per-worker gradients equals the full-batch gradient, the synchronous-versus-asynchronous tradeoff, and the large-batch convergence difficulties that return in Section 15.9 as the reason scaling efficiency is not free. The chapter assumes comfortable Python and PyTorch, a working understanding of mini-batch SGD and of training a neural network through forward and backward passes, and a basic picture of GPU memory and floating-point formats, since mixed precision in Section 15.8 is a statement about numerical range and storage. No prior multi-GPU or cluster experience is required; Section 15.2 builds the hardware picture from a single device upward.

Learning Objectives

Chapter Roadmap

What's Next?

This chapter establishes the data-parallel baseline: every accelerator holds a full copy of the model, and the only thing crossing the wire is gradients. That baseline has one breaking point, which the next chapter is built around. Data parallelism assumes the model fits on a single device, and the largest models do not. Chapter 16: Model, Pipeline, and Sharded Parallelism takes up training when the model itself must be split: tensor parallelism that shards individual layers across devices, pipeline parallelism that assigns consecutive layers to different devices and streams micro-batches through them, and the sharded data parallelism of ZeRO and FSDP that keeps the data-parallel structure of this chapter but partitions the redundant optimizer state, gradients, and parameters across the replicas rather than copying them. The all-reduce you mastered here reappears there decomposed into the reduce-scatter and all-gather of Chapter 4, and the gradient synchronization you learned to hide behind the backward pass becomes one collective among several that must be scheduled together. Read it next, and watch the single full replica of this chapter fracture into shards that each hold only their slice of a model too large for any one device.

Bibliography & Further Reading

Foundations of Data-Parallel Training

Goyal, P., Dollar, P., Girshick, R., Noordhuis, P., Wesolowski, L., Kyrola, A., Tulloch, A., Jia, Y., He, K. "Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour." arXiv:1706.02677, 2017. arxiv.org/abs/1706.02677

The work that showed data-parallel training scales to large batches with a linear learning-rate scaling rule and warmup, the empirical foundation for the scaling-efficiency discussion of Section 15.9.

📄 Paper

Li, S., Zhao, Y., Varma, R., Salpekar, O., Noordhuis, P., Li, T., Paszke, A., Smith, J., Vaughan, B., Damania, P., Chintala, S. "PyTorch Distributed: Experiences on Accelerating Data Parallel Training." Proc. VLDB Endowment, arXiv:2006.15704, 2020. arxiv.org/abs/2006.15704

The system paper behind PyTorch DistributedDataParallel, detailing gradient bucketing and computation/communication overlap, the direct reference for Sections 15.5 and 15.6.

📄 Paper

Sergeev, A., Del Balso, M. "Horovod: Fast and Easy Distributed Deep Learning in TensorFlow." arXiv:1802.05799, 2018. arxiv.org/abs/1802.05799

The paper introducing Horovod's ring-all-reduce-based, framework-agnostic approach to distributed training, the basis of Section 15.7.

📄 Paper

Large-Batch Optimization

You, Y., Gitman, I., Ginsburg, B. "Large Batch Training of Convolutional Networks (LARS)." arXiv:1708.03888, 2017. arxiv.org/abs/1708.03888

The layer-wise adaptive rate scaling optimizer that keeps large-batch data-parallel training stable, one answer to the convergence difficulties of Section 15.9.

📄 Paper

You, Y., Li, J., Reddi, S., Hseu, J., Kumar, S., Bhojanapalli, S., Song, X., Demmel, J., Keutzer, K., Hsieh, C.-J. "Large Batch Optimization for Deep Learning: Training BERT in 76 Minutes (LAMB)." arXiv:1904.00962, 2019. arxiv.org/abs/1904.00962

The adaptive large-batch optimizer that extended LARS to transformer training, the method behind the very-large-batch scaling that Section 15.9 treats as the frontier of data parallelism.

📄 Paper

Mixed Precision

Micikevicius, P., Narang, S., Alben, J., Diamos, G., Elsen, E., Garcia, D., Ginsburg, B., Houston, M., Kuchaiev, O., Venkatesh, G., Wu, H. "Mixed Precision Training." arXiv:1710.03740, 2017. arxiv.org/abs/1710.03740

The paper that introduced FP16 training with loss scaling and an FP32 master copy of the weights, the technical core of the mixed-precision enabler in Section 15.8.

📄 Paper

PyTorch. "Automatic Mixed Precision Package (torch.amp)." Official documentation. pytorch.org/docs/stable/amp.html

The reference for autocast and GradScaler, the production API that implements the mixed-precision training of Section 15.8 in a few lines.

🛠️ Tool

Frameworks and Tools

PyTorch. "Distributed Data Parallel (DDP) Notes." Official documentation. pytorch.org/docs/stable/notes/ddp.html

The reference describing DDP internals, gradient bucketing, and the reducer, the implementation companion to Sections 15.5 and 15.6.

🛠️ Tool

NVIDIA. "NCCL: NVIDIA Collective Communications Library." Official documentation. docs.nvidia.com/deeplearning/nccl

The GPU collective-communication library that implements the ring all-reduce underneath PyTorch DDP and Horovod, the substrate behind the gradient synchronization of Section 15.4.

🛠️ Tool

Hugging Face. "Accelerate: Run your PyTorch training across any distributed configuration." Official documentation. huggingface.co/docs/accelerate

The library that wraps DDP, mixed precision, and launch configuration behind a thin API, an example of the ecosystem tooling surveyed in Section 15.7.

🛠️ Tool