11.2 Distributed Training: Data & Model Parallelism
Key Takeaways
- Data parallelism copies the full model to each device, splits each batch across devices, and synchronizes gradients, usually with all-reduce.
- Model parallelism splits a model that doesn't fit on one device, either within layers (tensor parallelism) or across layers (pipeline parallelism).
- Sharded data parallelism such as FSDP or ZeRO divides parameters, gradients, and optimizer states across workers to train models larger than one device's memory.
- Reduction Server is an all-reduce algorithm on Agent Platform that speeds multi-node GPU data-parallel training using NCCL; it runs in worker pool 2.
- When scaling the global batch size across more devices, the learning rate usually needs adjusting, often with warmup, to keep training stable.
The exam guide asks you to understand the options for distributed training on GPUs and TPUs using data and model parallelism strategies. The core question is always: does the model fit on one device? If it does and you need speed, use data parallelism. If it doesn't, add model parallelism or sharding.
Data Parallelism
Every device holds a full copy of the model. Each global batch is split into per-device micro-batches, each device computes gradients, and the gradients are combined before the weights update.
| Variant | How gradients combine | Pros | Cons |
|---|---|---|---|
| Synchronous all-reduce | All workers average gradients each step (ring or tree all-reduce over NCCL) | Consistent model and predictable convergence | The slowest worker sets the pace. Communication overhead grows with model size |
| Asynchronous parameter server | Workers push gradients to parameter servers that hold the weights. Workers don't wait for each other | Tolerates slow or preempted workers. Scales to many CPU workers | Stale gradients can hurt convergence |
Framework strategies
| Framework | Single machine, multiple GPUs | Multiple machines |
|---|---|---|
| TensorFlow | MirroredStrategy | MultiWorkerMirroredStrategy (sync), ParameterServerStrategy (async), TPUStrategy |
| PyTorch | DistributedDataParallel (DDP) | DDP across nodes (launched with torchrun), FSDP |
| JAX | pmap / sharding APIs | Multi-host sharding across TPU or GPU slices |
Scaling the batch
With N devices at the same per-device batch, the global batch grows N times. Larger batches usually need a higher learning rate (a common starting heuristic is to scale it proportionally) plus warmup, and sometimes a different optimizer (LAMB or LARS for very large batches). If accuracy drops as you add devices, check the learning rate schedule before blaming the hardware.
Model Parallelism
| Technique | What gets split | When to use |
|---|---|---|
| Tensor (intra-layer) parallelism | Individual weight matrices across devices | Very wide layers, such as large transformer attention and MLP blocks. Needs fast interconnect |
| Pipeline (inter-layer) parallelism | Consecutive layers placed on different devices, with micro-batches flowing through stages | Very deep models. Watch for idle "bubbles" between stages |
| Sharded data parallelism (PyTorch FSDP, DeepSpeed ZeRO) | Parameters, gradients, and optimizer states sharded across data-parallel workers | Models too large for one GPU's memory, with simpler code than full tensor or pipeline parallelism |
| Expert parallelism | Mixture-of-experts experts on different devices | MoE architectures |
Large-model training often combines them. For example: tensor parallelism within a node over fast GPU links, pipeline parallelism across nodes, and data parallelism across replicas.
Distributed Training on Agent Platform
Cluster structure
| Worker pool | Role |
|---|---|
workerPoolSpecs[0] | Primary replica (exactly 1) |
workerPoolSpecs[1] | Workers |
workerPoolSpecs[2] | Parameter servers or Reduction Server |
workerPoolSpecs[3] | Evaluators |
Code reads CLUSTER_SPEC (general) or TF_CONFIG (TensorFlow) to find its role and peers. Use the same container image in every pool to avoid version mismatches.
Reduction Server
Gradient communication between nodes can dominate step time. Reduction Server is an all-reduce algorithm that Agent Platform provides as a container image for worker pool 2. It can raise throughput and cut latency for multi-node GPU data-parallel training. Requirements:
- GPU workers, with TensorFlow or PyTorch configured for multi-host data-parallel training using NCCL all-reduce.
- Primary and worker containers that support it: prebuilt TensorFlow 2.3+ or PyTorch 1.4+ training containers, or a custom container with NCCL 2.7+ and the
google-reduction-serverpackage.
Interconnects
Machine families for H100, H200, B200, and GB200 include high-bandwidth GPU networking (GPUDirect-TCPXO on H100 Mega, GPUDirect-RDMA on H200, B200, and GB200). Tensor parallelism across nodes without fast interconnects scales poorly.
Other orchestration choices
- Managed Training Clusters for reserved large-scale clusters, including Slurm-based workflows.
- Ray on Agent Platform with Ray Train for Python-native distributed training.
- Kubeflow Trainer on GKE for Kubernetes-native jobs (Chapter 9).
Fault Tolerance at Scale
The more machines a job uses, the more likely one fails during a long run.
- Checkpoint frequently to Cloud Storage, and make restarts resume from the latest checkpoint.
- Synchronous jobs stop when any worker fails, so enable automatic restarts and make startup idempotent.
- Asynchronous parameter-server training tolerates losing individual workers better, at the cost of gradient staleness.
- For multi-week runs, prefer reserved capacity (Managed Training Clusters or reservations) over Spot VMs.
Choosing a Strategy
| Situation | Strategy |
|---|---|
| Model fits on one GPU, and training on one GPU takes too long | Synchronous data parallelism (DDP or MirroredStrategy), multiple GPUs in one node first |
| Scaling across many GPU nodes, with communication as the bottleneck | Multi-node data parallelism + Reduction Server |
| Many CPU workers, some preemptible, sparse models | Asynchronous parameter server strategy |
| Model doesn't fit on one GPU (for example, full fine-tuning at 7-13B) | FSDP/ZeRO sharding |
| Very large transformer (tens to hundreds of billions of parameters) | Tensor + pipeline + data parallelism on fast-interconnect clusters or TPU slices |
| Large matrix-heavy model on TPUs | TPU slice with data parallelism, and model sharding via JAX or PyTorch/XLA |
Worked Scenario
A company trains a 1-billion-parameter vision model. One A100 takes 10 days. The goal is under 2 days.
- The model fits in 80 GB, so start with data parallelism: 8 A100s in one node with DDP, which is roughly 7× faster after communication overhead.
- Scale the learning rate with the larger global batch, and add warmup.
- To go further, use 2 nodes × 8 GPUs, and add Reduction Server in worker pool 2 to cut the gradient-synchronization bottleneck.
- Checkpoint to Cloud Storage and use Dynamic Workload Scheduler so all 16 GPUs start together.
A 400-million-parameter model fits comfortably on one A100, but training takes 6 days and the team wants it done in under 1 day. What is the most appropriate first strategy?
A team's multi-node PyTorch DDP training on GPUs spends most of each step waiting on gradient communication. They use prebuilt PyTorch training containers on Agent Platform. What should they add?
A 13-billion-parameter model must be fully fine-tuned, but its weights, gradients, and optimizer states don't fit in one GPU's memory. The team wants the least code complexity. Which approach fits best?