How a giant model is split across GPUs

There are three conceptually different axes for splitting training work:

\[ \boxed{\text{data} \quad \text{layers} \quad \text{matrix dimensions}} \]

These map roughly to data, pipeline, and tensor parallelism.

Data parallelism¶

Every GPU (or replica group) holds a model copy, but sees different examples.

Each computes local gradients:

\[ g_1,g_2,\ldots,g_n \]

Then the gradients are reduced/averaged:

\[ g=\frac1n\sum_i g_i \]

so every replica applies the same update.

This scales throughput well if the model fits, but naïvely replicates model state.

Tensor parallelism¶

Suppose one large projection is:

\[ Y=XW \]

and \(W\) is too large or expensive for one GPU.

Split:

\[ W=[W_1\ W_2\ W_3\ W_4] \]

Each GPU computes:

\[ Y_i=XW_i \]

and the partial outputs are combined.

This lets several GPUs cooperate on a single Transformer layer, but it requires frequent communication inside layers.

Pipeline parallelism¶

Split layers:

GPU1: layers 1–2
GPU2: layers 3–4
GPU3: layers 5–6
GPU4: layers 7–8

A full batch would leave most GPUs idle while it moves sequentially through the pipeline. So the batch is divided into microbatches:

time →  1  2  3  4  5  6  7

GPU1    A  B  C  D
GPU2       A  B  C  D
GPU3          A  B  C  D
GPU4             A  B  C  D

The idle regions at the start/end are the pipeline bubble.

Combine them¶

With 64 GPUs, one layout might be:

\[ 4\text{-way tensor}\times4\text{-way pipeline}\times4\text{-way data}=64 \]

One model replica spans \(4\times4=16\) GPUs; four replicas process different data.

Optimizer sharding¶

Ordinary data parallelism may replicate:

W = weights
G = gradients
O = optimizer state

on every worker.

Sharding moves toward each GPU owning only a fraction:

GPU1: ¼W ¼G ¼O
GPU2: ¼W ¼G ¼O
...

and communicating when missing pieces are needed.

The recurring trade-off is:

\[ \boxed{\text{save memory} \leftrightarrow \text{increase communication}} \]

At large scale, compute + memory + network jointly determine throughput.

Why different parallelisms compose¶

Each form of parallelism attacks a different scaling problem.

Data parallelism handles more examples at once.

Tensor parallelism makes a single enormous layer computable across several devices.

Pipeline parallelism makes a deep stack of layers fit across devices.

A frontier-scale training run may need all three because no single axis is sufficient.

Communication can dominate¶

Suppose tensor parallelism splits a matrix multiplication across 8 GPUs. Each GPU may perform a large local GEMM, then participate in a collective operation such as an all-reduce or all-gather.

If computation takes 2 ms but synchronization takes 1.5 ms, nearly half the layer time is no longer arithmetic.

This is why accelerator interconnect matters: - bandwidth, - latency, - topology, - collective implementation.

Two clusters with identical GPU FLOPS can have very different effective model-training throughput.

Pipeline bubbles and microbatching¶

Pipeline parallelism introduces idle periods. Increasing the number of microbatches reduces the fraction of time spent filling and draining the pipeline, but can interact with memory and optimization choices.

Thus distributed training is a scheduling problem as well as a numerical one.

The mental model to retain is:

The model graph has been stretched across a network. Every edge crossing a device boundary is now a distributed-systems operation.

Collective operations are the hidden instructions of distributed training¶

When a model spans GPUs, operations such as all-reduce, all-gather, and reduce-scatter become as fundamental as matrix multiplication.

For data parallelism, all-reduce can combine gradients.

For sharded training, reduce-scatter can both sum contributions and leave each worker owning only its shard.

For tensor parallelism, all-gather may reconstruct a distributed activation needed by the next operation.

The training graph therefore contains two classes of operations:

local numerical work
        +
distributed collective communication

High-performance implementations try to overlap them.

For example, while backward computes gradients for lower layers, communication for already-computed upper-layer gradients can be occurring concurrently.

A perfectly balanced system hides much of the network cost behind arithmetic. A poorly balanced one leaves expensive GPUs idle while they wait for bytes.

A small communication example¶

If four data-parallel workers each produce a 1 GB gradient shard and an all-reduce effectively requires moving comparable amounts of data across the fabric, communication time can become a meaningful fraction of each step. Faster GPUs can actually make this worse proportionally: arithmetic finishes sooner, so network synchronization becomes the exposed bottleneck.

This is why training systems care about bytes per parameter per step, not just FLOPs. Sharding strategies are judged by the communication pattern they induce as much as by the memory they save.