Distributed Training & Inference: From CPUs and GPUs to a Cluster

You can run a small model on a laptop, train a larger one on a GPU server, and spread an enormous one across a cluster. The difficult step is understanding what changes between those setups. Adding GPUs gives you more arithmetic capacity and more memory, but your program must decide how to use both—and how to pay for moving data between them.

Start with one question: what are we dividing? Different training examples, different layers, pieces of one matrix, or the stored state used to update the model? Those choices explain data parallelism, pipeline parallelism, tensor parallelism, and sharding. CUDA, NCCL, FSDP, DeepSpeed, and Ray then have distinct jobs within the resulting system.

1. CPU, GPU, and the role of CUDA

A CPU is a flexible processor. It runs the operating system, handles application logic, prepares data, and performs numerical computation. Modern CPUs have multiple cores and vector instructions, so “CPU” does not mean “one calculation at a time.” They are particularly useful when work involves varied instructions, branches, or a modest amount of arithmetic.

A GPU devotes more of its design to doing similar arithmetic across many values. A neural network repeatedly multiplies matrices: the same pattern of multiplication and addition applies to many rows, columns, and examples. That gives the GPU a large amount of parallel work. Small operations, irregular access, or frequent transfers can leave it underused; a GPU does not accelerate arbitrary code simply because it is installed.

In a typical discrete-GPU server, the CPU uses host RAM and each GPU has its own device memory, often called VRAM. The CPU launches a kernel—a function executed on the GPU by many threads. Data must be accessible to that device; weights and intermediate results can remain there across many operations. Some systems instead use unified memory architectures, so separate physical memory pools are not universal. NVIDIA’s programming model.

CUDA is NVIDIA’s parallel-computing platform and programming stack. It provides the runtime, programming tools, and libraries through which frameworks can execute NVIDIA GPU operations. In a usual PyTorch CUDA setup, your Python matrix multiplication dispatches to GPU kernels or libraries such as cuBLAS; you do not need to write a CUDA kernel yourself. A compatible NVIDIA driver and CUDA-enabled framework are needed for that path. Installing a compiler toolkit separately is not always necessary when the framework package already supplies its runtime dependencies. The CUDA platform.

2. One GPU, several GPUs, several nodes

A node is one server: CPUs, RAM, network interfaces, and possibly several GPUs. A single-node, multi-GPU job uses multiple devices inside that server. A multi-node, multi-GPU job crosses server boundaries as well. A common arrangement runs one training process per GPU. Each process receives a rank, its identifier within the distributed job; the total process count is the world size.

Inside a node, devices communicate over PCIe and, on supported systems, faster GPU fabrics such as NVLink/NVSwitch. Between nodes, traffic travels through network adapters and a network such as Ethernet or InfiniBand. The relevant questions are bandwidth, latency, contention, and actual connectivity—not just the number of GPUs. The fabric names alone do not establish the speed of a particular deployment.

Four 24 GB GPUs do not automatically behave like one 96 GB GPU. If your program creates a complete 40 GB model allocation on every device, it still fails. You need a method that explicitly partitions the model or its state. Fast links help move those partitions; they do not remove the need to partition them.

The hybrid view uses eight GPUs: two tensor shards per stage, two pipeline stages per model replica, and two data-parallel replicas. Each node contains one complete replica spread across four GPUs. This placement keeps the frequent tensor exchanges local and synchronizes corresponding training shards across nodes. Other mappings are possible; choose one for the actual communication pattern and hardware.

Here D is the data-parallel degree, T the tensor-parallel degree, and P the pipeline degree. This is the illustrated three-axis layout, not a universal formula covering every expert-, context-, or hybrid-sharding arrangement.

3. Data parallelism: same model, different examples

Imagine four people solving different pages of practice problems from the same textbook. Each has a complete copy of the model. Each processes different examples and computes gradients: the suggested changes to its parameters. Before the optimizer updates those parameters, the workers combine their suggestions.

In PyTorch DistributedDataParallel (DDP), gradient communication uses all-reduce and can overlap with backpropagation. With equally sized local batches and mean losses, averaging the local gradients gives the mean gradient over their combined batch. Matching initial weights and optimizer states then let each worker take the same update. DDP does not divide your dataset automatically; the input pipeline must give workers the intended distinct samples. PyTorch DDP tutorial.

Here b is the number of examples per microbatch on each data-parallel replica, D is the number of those replicas, and A is the number of microbatches accumulated before an optimizer update. With b = 2, D = 4, and A = 8, the effective batch is 64 examples. Tensor and pipeline workers cooperate on the same examples, so their counts do not multiply this batch size. For variable-length language-model data, also track tokens and the loss-normalization convention.

Data parallelism is useful when the model and training state already fit and you need to process more data per second. Replicated DDP keeps a full model and optimizer state on each GPU; it does not solve their memory footprint. Doubling the global batch can also change optimization behavior, so compare time to a target quality as well as raw step time.

4. Tensor parallelism: divide the work inside a layer

“Weight parallelism” is ambiguous. It can mean splitting a matrix calculation across devices, which is tensor parallelism (TP), or storing parameter shards between calculations, which we discuss under FSDP. The distinction is what each GPU computes after the split.

Take a linear layer y = Wx, with x represented as a column vector. If GPU 0 gets the top rows of W and GPU 1 gets the bottom rows, both read x and produce different coordinates of y. Joining those output slices recovers the complete result:

Alternatively, split W into left and right column blocks, and split x accordingly. Now both devices contribute to every output coordinate. Their results must be added:

Joining slices may use all-gather; summing contributions may use all-reduce or reduce-scatter. An optimized model can keep intermediate outputs sharded, so it does not necessarily gather after every layer. Library names such as “row-wise” and “column-wise” also depend on weight-storage conventions; the formulas above define the axes used in this example.

TP spreads both a layer’s parameters and its arithmetic across GPUs. It can make a large layer fit, but requires repeated communication through the network of layers. That is why fast connections matter so much, particularly when generating one token at a time. PyTorch tensor parallelism.

5. Pipeline parallelism: divide the sequence of layers

Pipeline parallelism (PP) assigns successive groups of layers to different devices. GPU 0 computes early layers, passes their activations to GPU 1, and so on. Backpropagation sends activation gradients in the reverse direction. Each stage owns its assigned layers rather than a complete model.

A single microbatch leaves most stages waiting. Split a batch into multiple microbatches and a stage can start the next one while its neighbor works on the previous one—like different stations on an assembly line. Idle periods while the pipeline fills or drains are called bubbles.

In this deliberately simple forward-only schedule, P equal-duration stages process M microbatches in M + P − 1 ticks. Four stages and one microbatch occupy 25% of the stage-time slots; four microbatches occupy 4/7, or about 57.1%. This compares two schedules with different amounts of work—it is not a measured speedup.

Training adds backward passes and saved activations. Schedules such as GPipe and one-forward-one-backward (1F1B) organize those passes differently. Uneven stage workloads, communication, and activation storage change real utilization. More microbatches can reduce bubbles, but tiny microbatches may use each GPU poorly. PyTorch pipeline schedules.

6. FSDP and DeepSpeed: stop duplicating training state

Training stores more than the weights. It also needs gradients, optimizer state such as Adam’s running moments, activations for backward, and temporary buffers. Sharding means each worker owns only a portion of some of that state.

DeepSpeed is a training and inference optimization library. ZeRO, its Zero Redundancy Optimizer approach, progressively removes duplicated training state from data-parallel workers. The stage number describes what is partitioned; it is not a ranking in which the highest number always runs fastest. DeepSpeed’s ZeRO guide.

These illustrative sizes correspond to BF16 weights and gradients plus two FP32 Adam moments, with no separate FP32 master weights. If your implementation maintains a master copy, include it in the optimizer-related state before dividing that state across workers. Actual dtypes and allocations depend on configuration.

Fully Sharded Data Parallel (FSDP) is PyTorch’s mechanism for sharding parameters, gradients, and optimizer state. In a common fully sharded configuration, workers gather the weights for the next wrapped unit, compute it on their different local examples, and release the unsharded weights afterward. Backward computation gathers what it needs again; reduce-scatter combines gradients and leaves each worker with its owned portion for the optimizer update. PyTorch FSDP2 tutorial.

That is the key difference from TP: FSDP gathers a unit so each data-parallel worker can compute that unit on its own examples; TP divides the unit’s computation across workers. Fully sharded FSDP and ZeRO-3 share the broad memory-saving idea, but their implementations, interfaces, and available policies are not identical.

A 3 GiB state subtotal does not promise a 3 GiB peak. Wrapping an entire model as one gather unit can create a much larger temporary allocation than gathering layer by layer. Prefetching may keep additional units live. Activation checkpointing trades recomputation for less saved activation memory; CPU or NVMe offload trades device memory for transfers and potentially slower steps. These are complementary choices, not extra ZeRO stages.

7. Ray starts work; NCCL exchanges GPU tensors

Ray is a distributed-computing framework. A Ray task is a function scheduled for remote execution; an actor is a stateful worker whose methods can be called remotely. This is useful when a workload combines GPU learning, CPU data preparation, environment simulations, and services with different resource needs. Ray Core concepts.

Ray Train packages the training function, worker count, and resource requirements into a distributed job. It starts workers and sets up the framework’s distributed environment. Your training code still chooses DDP, FSDP, or another supported strategy. Requesting four GPUs does not automatically split an arbitrary model across them. Ray Train overview.

NCCL, the NVIDIA Collective Communications Library, handles common GPU communication operations. All-reduce combines contributions and gives every participant the result. All-gather assembles everyone’s slices on each participant. Reduce-scatter combines contributions but leaves each participant with only its result slice. DDP and sharded training use these building blocks differently. NCCL collective operations.

Think in layers: the cluster manager allocates machines; a launcher or Ray starts workers; PyTorch and the chosen strategy define their computation; CUDA executes local NVIDIA GPU work; NCCL exchanges tensors between those workers. A Ray dashboard showing four available GPUs proves resource discovery. Successful, synchronized optimizer steps are separate evidence that training works.

8. Inference changes the objective

Ordinary inference uses fixed weights to produce outputs. It has no training backward pass, gradients, or optimizer updates. Autoregressive language-model serving instead stores a KV cache: attention keys and values from previous tokens, reused when generating the next token.

During prefill, the model processes the prompt, providing substantial parallel work across prompt tokens. During decode, each sequence advances one generated token at a time. Batching multiple sequences improves the amount of useful work per step, while longer contexts and more concurrent requests increase cache pressure. Request latency, total tokens per second, and maximum concurrency are different goals.

If a model fits comfortably on one device, independent replicas can serve different requests. They do not average gradients. If one model replica needs multiple GPUs, TP and PP divide it so those GPUs cooperate on requests. You can then replicate that entire group for more serving capacity. Splitting onto more devices can increase communication enough to hurt latency even when memory usage per device falls. vLLM parallelism and scaling.

FSDP’s repeated gather-and-release pattern is designed around training-state savings; it is usually not the default layout for low-latency token serving. An inference engine has its own weight placement, kernel, batching, and KV-cache decisions. Separating prefill and decode onto different workers is another serving option, but transferring the cache introduces an additional cost.

9. Choose the split that addresses the bottleneck

Specialized models add other axes. Expert parallelism places different mixture-of-experts networks on different devices and routes tokens to their selected experts. Sequence parallelism partitions selected activation operations along the sequence dimension; context parallelismdistributes long-context attention work and exchanges the information that attention still needs. These are additional tools for particular model structures and activation bottlenecks, rather than synonyms for data parallelism. Megatron Core parallelism guide.

For a useful comparison, hold the model, precision, sequence-length distribution, effective batch, and quality target constant. Separate initialization and compilation from steady-state execution. Record peak memory, useful tokens per second, time waiting for data or collectives, and end-to-end cost. For inference, also record time to first token and latency between generated tokens under the same request load.

The practical progression is to make one device correct, measure what limits it, and introduce the split that removes that limit. A cluster succeeds when workers spend enough time doing useful computation to repay the coordination and communication they introduce.