Training vs. Inference: GPU Efficiency and Parallelism

Why Use Multiple GPUs?

Training usually consumes more resources than inference because it computes gradients, and large datasets also reduce efficiency as the number of batches grows. Inference is not simple either: recent LLMs and VLMs often do not fit on one server GPU. Hugging Face's device_map="auto" and vLLM's tensor_parallel_size=4 are familiar ways to use several GPUs.

These can look similar to Distributed Data Parallel training, but their purposes are different. Training parallelism synchronizes gradients so several GPUs can train the same model. Inference parallelism distributes parameters and activations so that an oversized model fits in aggregate GPU memory or runs faster.

What Gets Parallelized?

Training Parallelism: DDP Is Data Parallelism

DDP replicates the entire model on every GPU. Each rank processes a different mini-batch, and gradients are averaged with all-reduce during backward propagation so all replicas perform the same update.

• GPU0: batch0 → forward → loss0
• GPU1: batch1 → forward → loss1
• GPU2: batch2 → forward → loss2
• GPU3: batch3 → forward → loss3

The losses do not need to be collected and averaged. What matters is the gradient derived from each local loss.

loss = criterion(model(x), y)
loss.backward()
optimizer.step()

Backward propagation walks through the computation graph layer by layer. When every process has computed a gradient for a parameter, DDP performs all-reduce(sum) and divides by world_size. If the local gradients are $g0=d(loss0)/d\theta$, $g1=d(loss1)/d\theta$, $g2=d(loss2)/d\theta$, and $g3=d(loss3)/d\theta$, every replica updates with $g=(g0+g1+g2+g3)/4$.

Inference Parallelism Has Several Forms

1. Hugging Face device_map="auto": a relay race. Layers, rather than individual layer weights, are placed on different GPUs. GPU0 runs its layers and hands the activation to GPU1, which continues the sequence. This solves model placement but does not make every GPU compute the same layer simultaneously.

2. Tensor Parallelism in Megatron or vLLM: sharing matrix multiplication. Several GPUs jointly compute one large Linear or Attention layer. For $Y=XW$, column-wise partitioning uses

$W=\begin{bmatrix}W_1&W_2&W_3&W_4\end{bmatrix}$,

$Y_1=XW_1,~Y_2=XW_2,~Y_3=XW_3,~Y_4=XW_4$, and $Y=\begin{bmatrix}Y_1&Y_2&Y_3&Y_4\end{bmatrix}$.

Every GPU receives $X$, and partial output channels are joined by all-gather. In row-wise partitioning,

$W=\begin{bmatrix}W_1 \newline W_2 \newline W_3 \newline W_4\end{bmatrix}$ and $X=\begin{bmatrix}X_1&X_2&X_3&X_4\end{bmatrix}$.

Each GPU computes $Z_i=X_iW_i$, followed by $Y=Z_1+Z_2+Z_3+Z_4$ through all-reduce. The first is like giving everyone the same ingredients to make different parts; the second gives everyone different ingredients whose partial results are combined.

3. Data-parallel inference: no merging. Every GPU holds a full model and receives different requests. It does not reduce parameter memory, but it is often the strongest option for batch throughput when the model fits on one GPU.

How Do We Train a Model That Does Not Fit on One GPU?

Pure DDP works for relatively small models or memory-light fine-tuning because every GPU keeps a full copy. Larger models combine it with other techniques.

ZeRO and FSDP

DDP replicates parameters, gradients, and optimizer state. ZeRO and FSDP shard them across GPUs:

• ZeRO-1 shards optimizer state.
• ZeRO-2 shards optimizer state and gradients.
• ZeRO-3 shards optimizer state, gradients, and parameters.

PyTorch FSDP provides behavior similar to ZeRO-3. Immediately before a layer's forward pass, its full weights are temporarily all-gathered. After computation they are resharded; backward repeats the gather-and-reshard process, and optimizer state remains sharded during the update.

Tensor Parallelism and Pipeline Parallelism

TP divides expensive operations inside a layer and works in training as well as inference. Pipeline Parallelism instead assigns blocks of layers to different stages. Unlike simple device_map="auto", a training pipeline splits a batch into micro-batches to reduce idle “pipeline bubbles.” While micro-batch 1 advances to the next GPU, micro-batch 2 begins on GPU0.

Combining Several Dimensions

All GPUs normally participate in both computation and sharding; there is no dedicated “storage GPU.” A 128-GPU job might use TP=8, PP=2, and DP=8, giving $8\times2\times8=128$. TP ranks hold different weight slices, PP ranks own different stages, and the same TP/PP model arrangement is replicated across DP groups.

Offloading is the exception: ZeRO-Offload can keep optimizer state in CPU memory or NVMe. This moves storage outside the GPU pool rather than assigning a GPU solely to storage.

In short:

• DDP replicates models, splits data, and all-reduces gradients.
• ZeRO/FSDP shards parameters, gradients, and optimizer state.
• TP splits operations within Linear and Attention layers.
• PP splits layer blocks into stages and moves activations between them.
• Large training usually combines DP with one or more of TP, PP, and FSDP.

Conclusion

Both training and inference “use several GPUs,” but they distribute different things and therefore have different communication patterns and bottlenecks.

Training in One Sentence

DDP is not fundamentally about averaging losses. It is a mechanism for agreeing on the averaged gradients used in the update.

Inference in One Sentence

Inference chooses among several methods to distribute memory or computation, avoid OOM, or increase speed.

• Layer placement or PP solves model placement by passing activations between GPUs.
• TP lets GPUs jointly compute large Linear and Attention operations, communicating through all-reduce or all-gather.
• Data-parallel inference replicates the model and distributes requests, offering high throughput with little communication.

Choosing a Method

• The model does not fit on one GPU: solve placement first with PP or TP.
• Single-request latency or token generation speed matters: TP is often advantageous.
• Throughput matters and the model fits on one GPU: replicated data-parallel inference is the simplest and strongest choice.