Chapter 14: Distributed Model Parallelism

Data parallelism and gradient synchronization

Derive the synchronized global-batch gradient and explain replicated memory, bucket readiness, and slowest-rank behavior in distributed data parallel training

Before you begin

Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.

Begin with one transformer that fits on one GPU and a batch of 64 training sequences. One GPU needs 800 milliseconds for forward computation, backward computation, and the optimizer update. Four GPUs could process four groups of 16 sequences concurrently, but four independent updates would no longer equal the original batch-64 update unless the four gradient contributions were combined first.

A distributed training launcher starts one process for each GPU and assigns each process a rank. Every rank stores the same initial model parameters and processes a different local microbatch of 16 sequences. After backward computation, rank r owns a local gradient tensor g_r. An all-reduce sums corresponding gradient elements across the four ranks and returns the sum to every rank before the optimizer changes any replica.

Observe the step with PyTorch Profiler on every worker. The profiler uses CPU instrumentation and the CUDA Profiling Tools Interface for GPU activity to write timestamped operator and kernel events. Gradient all-reduce events should begin as backward produces ready gradient buckets. A memory snapshot should show one full parameter, gradient, and optimizer-state allocation on every rank, because classic data parallelism replicates those categories.

Data parallelism names this replicated-model, partitioned-input arrangement. The method can increase sample throughput when local computation is large enough to cover or amortize gradient communication. The method cannot make an oversized parameter set fit, and a comparison that increases global batch answers a different optimization question from a fixed-global-batch speedup test.

Derive the synchronized gradient

Let rank r process a local microbatch and compute gradient g_r for the same parameter vector. For a global-batch mean gradient over P equal-size local batches, every rank needs g = (1/P) sum_r g_r before applying an equivalent optimizer update. An all-reduce can sum the local tensors and return the result to all ranks, followed by division or an equivalent scaling convention.

If local batch sizes differ, a simple average of rank means is not the sample mean. Weight each local contribution by its sample or token count according to the training objective. Variable sequence workloads must also define whether loss is normalized per sequence, token, or another unit.

Account for replicated memory

Classic data parallelism stores a full parameter copy on every rank. Each rank usually also stores full gradients and full optimizer state. Eight ranks therefore provide eight copies; the ranks do not divide these allocations by eight.

Activations depend on local microbatch and sequence shapes. Increasing data-parallel degree while holding global batch fixed reduces local examples and can reduce activation memory. Parameter-related capacity remains unchanged, which is why data parallelism alone cannot fit a model whose weights and optimizer state exceed one device.

Calculate communication per optimizer step

Each parameter usually contributes one gradient value. If a model has Pm parameters and BF16 gradients, logical gradient payload is approximately 2Pm bytes. A ring all-reduce sends about 2(D−1)/D times that payload per rank for data-parallel degree D, before protocol and alignment effects.

The payload is largely independent of local microbatch size. More local tokens add forward and backward compute between synchronizations, improving the compute-to-communication ratio. They also change global batch unless accumulation or data-parallel degree adjusts.

Understand gradient readiness

Backward traverses dependencies in roughly reverse forward order. Gradients for later layers often become ready first. A DDP implementation groups parameter gradients into buckets and launches a reduction when every gradient in the next required bucket is ready.

Parameter source order does not always equal readiness order. Shared modules, branches, activation checkpointing, or compiled graphs can change when hooks fire. Measure actual bucket readiness rather than assuming it from the source file.

Trade bucket startup against overlap

A small bucket can launch early while substantial backward compute remains. The small-bucket schedule pays more startup costs and can operate below the best bandwidth regime. A large bucket amortizes startup but waits for more gradients, reducing the interval in which the transfer can overlap backward.

The final bucket is especially important. Gradients for early forward layers become ready near the end of backward. Their reduction often remains after compute ends. Reordering bucket membership or reducing the late bucket can remove more step time than increasing peak aggregate bandwidth.

Preserve global-batch semantics

Global batch equals local microbatch × data-parallel degree × gradient accumulation steps when each factor is uniform. Doubling ranks while holding local work fixed doubles global batch. The increased-global-work experiment is weak scaling and can change convergence, data order, and learning-rate requirements.

A strong-scaling comparison holds global useful work fixed. Local batches become smaller as rank count grows, which can reduce kernel efficiency and expose more communication. Report throughput and time to quality so a systems improvement does not depend on an unexamined optimization change.

Include the slowest rank

Every synchronous reduction waits for its required participants. A slow input shard, longer sequence batch, throttled GPU, CPU delay, or earlier network fault can make one rank enter late. Early ranks accumulate wait inside the collective.

Balance work according to cost, not only sample count. For attention, equal token counts can still have unequal quadratic sequence work. Use per-rank distributions and last-arrival analysis when scaling falls.

Worked example: Choosing a gradient bucket size

A model has 8 GiB of gradients. With 1 MiB buckets, backward launches thousands of all-reduces. With 256 MiB buckets, only 32 all-reduces run, but the first one begins late.

Capture gradient-ready timestamps for a representative step. Backward lasts 180 milliseconds. Later-layer gradients are available from 25 milliseconds onward, while the final early-layer gradients become ready at 176 milliseconds.
Fit alpha–beta curves for the data-parallel group. The 1 MiB configuration pays several thousand startups and achieves poor message bandwidth. The 256 MiB configuration transfers efficiently but each bucket requires many gradients before it can launch.
Construct candidate buckets from the actual readiness order. For 16, 32, 64, 128, and 256 MiB, calculate ready time, predicted duration, remaining backward compute, and final exposed tail.
Measure each schedule. A 32 MiB bucket starts early but lengthens some backward kernels through HBM and SM contention. A 128 MiB bucket starts later but produces the shortest combined backward-plus-tail interval.
Inspect the final bucket. Move compatible late-ready gradients into two smaller buckets so part of the final communication can begin before the last layer completes. Include the extra startup in the model.
Verify identical global batch, loss normalization, reduction scaling, optimizer result, and long-run quality. Re-run the analysis if compilation or activation checkpointing changes gradient readiness.

Conclusion: The bucket is a scheduling unit connecting autograd readiness and collective economics. Its best size is defined by overlap, not network throughput alone.

Common errors

Claiming strong scaling after increasing global batch. More work per step can hide communication without reducing time for the original global work.
Dividing model-state memory by data-parallel degree. Classic data parallelism replicates parameter-related state.
Summing all-reduce durations as exposed overhead. Many buckets can overlap useful backward compute.
Assuming parameter source order is gradient-ready order. The executed autograd graph determines readiness.

Reference summary

Data parallelism gives every rank the same parameters and assigns different input examples to each rank. After backward, the gradients must represent the combined global batch before every replica applies the same optimizer update. Distributed data-parallel implementations group gradients into buckets and launch each bucket’s all-reduce when all of its gradients are ready.

Each rank normally retains full parameters, gradients, and optimizer state. Adding data-parallel ranks does not divide those allocations.
Gradient communication is approximately proportional to gradient bytes per optimizer step. Increasing local batch adds computation between synchronizations but does not directly shrink the gradient tensor.
Global batch equals local microbatch × data-parallel degree × gradient-accumulation steps. A scaling comparison must keep optimization and quality semantics explicit.
Small buckets become ready early but pay more startup operations. Large buckets transfer efficiently but can delay the first communication and leave a final exposed tail.
Synchronous progress follows the slowest rank at each rendezvous. Data preparation, sequence length, throttling, and compute imbalance affect global step time.

Knowledge check

Why does doubling data-parallel ranks not halve memory required for model parameters?

Show the answer guide

Classic data parallelism replicates the complete parameter vector on every rank so each local forward and backward pass can execute the same model.
Each rank normally also retains full gradients and optimizer state; all-reduce synchronizes values but does not change persistent ownership into shards.
Increasing data-parallel degree can reduce activation memory under a fixed global batch, but it does not divide parameter-related storage unless a separate sharding method is used.

Practice

Audit data-parallel memory and gradients

Run or model a two- or four-rank data-parallel step for one declared network. Record local batch, global batch, parameter, gradient, optimizer, and activation bytes on every rank and verify the gradient reduction convention.

Evidence to keep: Per-rank memory table, global-batch equation, local and reduced gradient values for a small check case, bucket readiness trace, and measured or modeled communication bytes.

When full replicated state no longer fits, ZeRO and fully sharded methods partition parameter-related state while reconstructing it around computation.

References

Sources and further reading

The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.