Kahibaro
Discord Login Register

14.5 Distributed Training Concepts Overview

Why Distributed Training Exists

Distributed training is the umbrella term for training one model using more than one compute device, usually multiple GPUs, sometimes across multiple machines. The motivation is not only speed. It is also about fitting larger models or larger effective batch sizes than a single device can handle, and about keeping training wall clock time reasonable as datasets grow. In PyTorch, distributed training typically means replicating some part of your work across processes and coordinating gradients, parameters, or data.

A useful mental model is that training has three main costs, compute, memory, and communication. With one GPU, communication is mostly internal. With multiple GPUs, communication becomes a first class constraint because devices must exchange information, most commonly gradients during backpropagation.

Distributed training is not automatically faster. If communication overhead dominates, your multi GPU run can be slower than single GPU training.

Core Approaches: Data Parallel, Model Parallel, and Pipeline Parallel

Data parallelism is the most common starting point. Each GPU holds a full copy of the model. Each GPU processes a different mini batch, computes gradients locally, then gradients are aggregated so that all model copies stay in sync. In deep learning practice, when people say distributed training, they often mean data parallel training.

Model parallelism splits the model across devices. This helps when the model is too large to fit on one GPU. It requires sending activations forward and gradients backward between devices, which increases communication and complexity. PyTorch can support this with manual placement of submodules on different devices, and there are higher level libraries, but it is not the beginner default.

Pipeline parallelism is a structured form of model parallelism where layers are partitioned into stages, and micro batches flow through the stages like an assembly line. This improves utilization but adds scheduling complexity and can change optimization dynamics due to micro batch boundaries.

For an absolute beginner, the key distinction is that data parallelism keeps the code closest to your single GPU training loop, while model and pipeline parallelism often require architectural changes and more careful debugging.

Processes, Ranks, and World Size

PyTorch distributed training is process based. Instead of one Python process controlling multiple GPUs, the common pattern is one process per GPU. Each process is assigned a rank, which is its unique integer ID in the distributed job. The world size is the total number of processes participating.

You will also encounter the idea of a process group, which is a set of processes that can communicate using collectives such as all reduce. In most basic setups, you have one global process group that contains all ranks.

In most PyTorch distributed setups, one process controls one GPU. Do not try to drive multiple GPUs from one process unless you are intentionally using a different parallelism strategy.

What Gradient Synchronization Means

In data parallel training, each rank computes gradients on its local mini batch. To keep model replicas identical, the ranks perform an all reduce operation over gradients. Conceptually, if each rank computes a gradient $g_i$, the synchronized gradient is often the average

$$
g = \frac{1}{N}\sum_{i=1}^{N} g_i
$$

where $N$ is the world size. After this, each rank applies the optimizer step using the same gradient, so parameters remain identical across ranks.

This averaging interacts with learning rate choices. If you keep the per GPU batch size fixed and increase the number of GPUs, the global batch size increases. Many training recipes adjust learning rate with global batch size, but the exact rule depends on model, optimizer, and task. The important concept here is that distributed training changes the effective batch size unless you compensate.

DDP vs DataParallel: The Practical Default

Two PyTorch options are commonly mentioned. torch.nn.DataParallel is older and uses one process with multiple threads, gathering outputs to a main device. It is convenient but usually slower and less scalable.

torch.nn.parallel.DistributedDataParallel, commonly called DDP, is the standard approach. It uses one process per GPU and is designed for performance, scalability, and correctness. Even on a single machine with multiple GPUs, DDP is typically recommended over DataParallel.

You do not need the full API details in this chapter, but you should remember that most real multi GPU training in PyTorch is done with DDP, launched as multiple processes.

Prefer DistributedDataParallel over DataParallel for serious training. DataParallel can become a bottleneck and behaves differently in edge cases.

Feeding Data Correctly: Distributed Sampling

Data parallel training only works well if each rank sees different data. If every rank iterates through the same DataLoader without coordination, you waste compute and you effectively train on duplicated batches.

The usual fix is a distributed sampler that partitions the dataset across ranks. Each epoch, the sampler shuffles in a coordinated way so that different ranks get different indices. In practice, this is often implemented with torch.utils.data.distributed.DistributedSampler, and you must also ensure the sampler is told the epoch number so that shuffling changes each epoch in a synchronized manner.

The key idea is simple: in distributed data parallel training, the dataset is split across ranks, but the model parameters are kept synchronized.

Communication Backends and Hardware Topology

Collective communication relies on a backend. On GPUs, the common backend is NCCL, designed for high throughput GPU to GPU collectives. On CPU or some environments, Gloo is common.

Hardware topology matters. GPUs connected by fast links communicate more efficiently than GPUs that must communicate through slower paths. Across multiple machines, network bandwidth and latency can dominate performance.

You do not need to memorize backend details, but you should understand that when multi GPU scaling is disappointing, the cause is often communication, not computation.

Batch Size, Global Batch Size, and Training Dynamics

Distributed training changes the notion of batch size. If each GPU processes a batch of size $b$ and you have $N$ GPUs, then the global batch size is $B = N b$. Your gradient is now based on more samples per optimizer step.

Larger global batch sizes can improve throughput, but may hurt generalization or require different learning rate schedules and warmup. There is no universal rule that bigger is always better. Distributed training is as much about systems tradeoffs as it is about modeling.

A practical takeaway is to track and report both per device batch size and global batch size in experiments so your results remain interpretable.

Synchronization Points and Performance Pitfalls

Distributed training introduces synchronization points where ranks must wait for each other. If one rank is slower due to data loading, CPU bottlenecks, or uneven work, all ranks can stall.

Common causes include slow DataLoader workers, expensive preprocessing on CPU, data stored on slow disks, and having different amounts of work per batch. In DDP, another source of overhead is gradient synchronization, especially with many small tensors or frequent synchronization triggers.

A general performance principle is that you want more computation per synchronization event. This is one reason why very small models may not scale well across many GPUs.

The job runs at the speed of the slowest rank. If one GPU is data starved or slower, all GPUs wait during synchronization.

Single Node vs Multi Node at a Conceptual Level

Single node multi GPU training runs multiple ranks on one machine. Multi node training extends this to multiple machines, adding network communication and more failure modes. Conceptually, the same primitives are used. Practically, multi node requires correct initialization, addressing, and often a job scheduler environment.

From a learning perspective, it is helpful to first become comfortable with single node DDP, then generalize to multi node once you understand ranks, world size, and distributed sampling.

How to Think About Debugging Distributed Training

Distributed bugs can look strange because you may only see logs from one process, and failures can happen only on certain ranks. A helpful mindset is to separate correctness issues from scaling issues.

Correctness issues include data duplication across ranks, inconsistent model states, non deterministic behavior that differs by rank, and silent mismatches in batch sizes. Scaling issues include poor GPU utilization, high communication overhead, and DataLoader bottlenecks.

A very practical habit is to make sure you can reproduce the same training behavior on one GPU first, then move to multiple GPUs while changing as little as possible.

Summary of Key Concepts to Carry Forward

Distributed training in PyTorch is usually data parallel training with DDP, meaning one process per GPU, replicated model parameters, and synchronized gradients via all reduce. The dataset is partitioned across ranks using distributed sampling so each GPU does unique work. Performance depends on balancing compute with communication and avoiding stalls caused by the slowest rank. Distributed training changes global batch size, which can change optimization behavior, so batch size and learning rate choices must be tracked carefully.

Views: 61

Comments

Please login to add a comment.

Don't have an account? Register now!