Distributed Model Parallelism

Suppose a transformer training process needs 96 GiB for model weights, gradients, and optimizer state, but one graphics processing unit (GPU) has only 80 GiB of high-bandwidth memory (HBM). Adding a second GPU does not make the model fit automatically. Software must decide which bytes each GPU stores, which matrix operations each GPU performs, and when the GPUs exchange partial results.

Distributed training frameworks normally run one central processing unit (CPU) process per GPU. The process receives an integer identifier called a rank. Rank 0 is a process identity, not a claim that its GPU is faster or physically central. A process group names the ranks that participate in a particular sequence of communication operations. PyTorch logs, profiler events, and NCCL messages use these rank and group identifiers to relate local activity to the complete job.

A memory report collected after model construction shows persistent allocations such as parameters and optimizer tensors. A profiler trace collected for one step shows temporary parameter gathers, gradient reductions, activation transfers, and idle intervals. Local tensor shapes printed at selected module boundaries show how each general matrix multiplication (GEMM) changed on every rank. GPU communication kernels can also occupy streaming multiprocessors (SMs) while model kernels wait. Together, the observations reveal whether a parallel method solved the original capacity problem and what new work the method introduced.

This chapter derives data, fully sharded, tensor, pipeline, expert, and combined parallelism from those concrete objects. Every section names the owner of each tensor, the exact local shape, the required communication, and the point in forward or backward execution where another rank's data becomes necessary.