Scaling ML Models: Data, Tensor, and Pipeline Parallelism Explained

Introduction

A single NVIDIA A100 has 80GB of memory. A 70-billion-parameter model, stored just as FP32 weights, needs 70B × 4 bytes = 280GB - it does not fit on one GPU, and that is before you count the optimizer states and activations needed to actually train it. Training GPT-scale models is therefore not a "run it and wait" problem; it is a systems problem of splitting a single mathematical computation across dozens, hundreds, or thousands of GPUs while keeping every one of them busy.

This post walks through the four main strategies used to make that split - data parallelism, tensor (model) parallelism, pipeline parallelism, and ZeRO - plus the two techniques that make all of them cheaper to run: activation checkpointing and mixed precision training. The goal is to go from "what does data parallel mean" to "why does Megatron-LM split the weight matrix column-wise instead of row-wise" in one sitting - starting from the simplest idea and building up, step by step, to the exact configuration real clusters run today.

Background: Where Does All the Memory Go?

Before splitting anything, it helps to know exactly what is competing for GPU memory during training. The DeepSpeed ZeRO paper splits this into two buckets.

Model states - memory needed regardless of batch size: the parameters themselves, their gradients, and the optimizer states (for Adam: momentum and variance, one value each per parameter).

Residual states - everything else: activations saved for the backward pass, temporary buffers, and memory fragmentation.

Here is the concrete number that motivates this entire post. Training with Adam in mixed precision, a common setup keeps two copies of the parameters (one FP16 copy for the fast forward/backward math, one FP32 "master" copy for the accurate weight update) plus FP16 gradients plus FP32 Adam momentum and variance:

A 7.5-billion-parameter model - roughly GPT-2 XL's scale - therefore needs 7.5B × 16 bytes = 120GB just for model states, before a single activation is stored. That single number is why none of the techniques below are optional at this scale; they are the only way the model fits at all.

A quick note on terminology before we begin: replicating copies the entire tensor identically onto every GPU (high memory redundancy, zero communication during forward compute), whereas sharding splits a tensor into N distinct slices across N GPUs so each GPU stores only 1/N of the data.

Collective Communication Primitives & Terms

Before splitting models across multiple GPUs, it helps to establish the fundamental communication building blocks. Distributed training libraries like NVIDIA NCCL rely on a set of standardized collective communication primitives to exchange data between devices:

Key Terms to Know

• Ring / Tree All-Reduce: Efficient topologies implemented by libraries like NCCL so that every GPU exchanges data with only a few immediate neighbors rather than sending N² messages across the network.
• NVLink vs. InfiniBand: NVLink is an ultra-fast intra-node GPU interconnect (up to 900 GB/s on H100), whereas InfiniBand connects separate server chassis/nodes over networking switches (50–100 GB/s).
• Pipeline Bubble: The idle execution time where downstream GPUs wait for upstream GPUs to finish processing micro-batches during pipeline fill and drain phases.
• Sequence Parallelism (Megatron-SP): Sharding non-tensor-parallel activation regions (LayerNorm & Dropout) along the sequence dimension (S/N) to eliminate activation memory redundancy.
• Expert Parallelism (MoE): Partitioning sparse "Expert" feed-forward layers across GPUs in Mixture-of-Experts models, using All-to-All collectives to dispatch tokens to assigned experts.

Data Parallelism

Data parallelism (DP) is the simplest strategy and the one every distributed training run starts from. The idea: if the model fits on one GPU, just make copies of it.

1. Split the global batch into N smaller batches (N = number of GPUs).
2. Copy the identical model weights onto all N GPUs.
3. Each GPU runs a full forward and backward pass on its own batch slice, independently.
4. Because each GPU saw different data, each GPU now holds a different gradient.
5. All-reduce: average the gradients across every GPU (sum, then divide by N), so every GPU ends up with the identical, averaged gradient.
6. Every GPU applies the same weight update, so all copies stay in sync for the next iteration.

The step that makes this work is the all-reduce - a collective communication operation where every GPU sends its gradient to every other GPU and receives the sum back. Modern libraries such as NCCL implement all-reduce using efficient communication topologies like rings or trees so every GPU exchanges data with only a few neighbors rather than every other GPU. This is the only communication DP needs per step, which is why it scales so well up to a point.

Tensor (Model) Parallelism: Megatron-LM

Tensor parallelism, introduced in the Megatron-LM paper, takes the opposite approach from DP: instead of copying the whole model, split individual weight matrices across GPUs, so no single GPU ever holds the full matrix. This is called intra-layer parallelism, because the split happens inside one layer's math, not between layers.

Say a linear layer computes Y = GELU(X · A), where X is the input activations and A is a weight matrix. There are two ways to split A across 2 GPUs, and only one of them avoids an expensive synchronization in the middle of the layer.

Method 1: Row-wise split of A (needs a mid-layer sync)

Split A by rows into A1 (top half) and A2 (bottom half), one on each GPU. To multiply correctly, X must be split by columns to match: X = [X1, X2]. Each GPU computes a partial product, but neither GPU's partial product is a valid final answer on its own - X1·A1 and X2·A2 must be summed together before the GELU nonlinearity can be applied, since GELU is nonlinear and cannot be split across a sum:

GPU 1: partial_1 = X1 · A1
GPU 2: partial_2 = X2 · A2
[ALL-REDUCE: sum partial_1 + partial_2 across GPUs]
Y = GELU(partial_1 + partial_2)   ← needs the full sum on every GPU first

This works, but it forces a synchronization point inside the layer, before the activation function can even run.

Method 2: Column-wise split of A (no mid-layer sync)

Instead, split A by columns: A = [A1, A2]. Now X does not need to be split at all - every GPU keeps the full X. Each GPU computes:

GPU 1: Y1 = GELU(X · A1)
GPU 2: Y2 = GELU(X · A2)

Because GELU is applied elementwise, GELU(X·A1) and GELU(X·A2) are each already complete, correct results for their own columns of Y - no partial sum needed, no sync required mid-layer. Y = [Y1, Y2] is simply the column-wise concatenation of the two GPUs' outputs. This is why Megatron-LM's MLP block splits the first linear layer column-wise: it lets the nonlinearity run independently on each GPU.

The second linear layer in the MLP block (the one projecting back down) is then split row-wise, because its input (Y, already split by columns across GPUs) lines up perfectly with a row-wise split of the second weight matrix - and this second layer has no nonlinearity in the way, so the required all-reduce happens exactly once, at the very end of the MLP block, instead of in the middle of it. Combining a column-split first layer with a row-split second layer means the entire two-layer MLP needs only one all-reduce total per forward pass (and one per backward pass), not one per layer.

Seeing it with real dimensions

Labels like A1 and A2 can hide what's actually happening on each device. Picture a batch of activations X with shape [batch, 4096] feeding into a weight matrix A with shape [4096, 16384] - a realistic MLP up-projection ratio. Splitting A by columns across 2 GPUs looks like this:

Each GPU keeps the full, unsplit X, but holds only half of A's columns - 8192 out of 16384. Each GPU's output is a genuine, complete [batch, 8192] slice of the final [batch, 16384] answer, computed entirely independently. Concatenating Y1 and Y2 along the column dimension gives back exactly what a single GPU holding the full 16384-wide A would have produced - just computed by two GPUs working on non-overlapping columns at the same time.

Splitting self-attention

Multi-head attention splits even more naturally: each attention head is already an independent computation, so Megatron-LM simply assigns different heads to different GPUs - head 1 and head 2 to GPU 1, head 3 and head 4 to GPU 2, and so on. The Q, K, V projection matrices are split column-wise (same reasoning as Method 2 above), each GPU computes attention for its own heads completely independently, and the output projection matrix is split row-wise so the final all-reduce happens once, after the heads are concatenated back together.

The gap basic tensor parallelism leaves: sequence parallelism

Splitting the attention and MLP blocks column/row-wise covers most of a transformer layer's heavy compute, but not all of it. Operations like LayerNorm and Dropout sit between these blocks and are normally left replicated - every GPU redundantly runs and stores the exact same LayerNorm output, quietly giving back some of the activation memory tensor parallelism was supposed to be saving.

Sequence parallelism closes this gap by sharding those replicated regions along the sequence dimension instead of duplicating them: each GPU handles LayerNorm/Dropout for only its own slice of the sequence, using an all-gather to reassemble the full sequence right before a tensor-parallel region needs it, and a reduce-scatter to shard it again immediately after. This keeps the memory savings intact across the entire layer, not just inside the attention and MLP sub-blocks.

Why this communication pattern forces tensor parallelism to stay local

Two all-reduces per layer sounds cheap until real bandwidth numbers are attached to it - this communication fires on every forward and backward pass, for every layer, on every training step, so the link it travels over matters enormously:

Peak theoretical bandwidths; achieved bandwidth in practice is lower and depends on topology and workload.

This is the concrete reason tensor parallelism almost always stays confined to the GPUs inside a single node: firing an all-reduce roughly every layer needs NVLink's ~900GB/s, not the ~50-100GB/s available between nodes over InfiniBand. Pipeline parallelism, covered next, only needs to pass one activation tensor at each stage boundary - far less frequent traffic - which is exactly why it is the strategy used to cross node boundaries instead.

Pipeline Parallelism: GPipe

Where tensor parallelism splits inside a layer, pipeline parallelism (from the GPipe paper) splits between layers - an inter-layer approach. If a model has 32 transformer layers and you have 4 GPUs, GPU 1 gets layers 1-8, GPU 2 gets layers 9-16, GPU 3 gets layers 17-24, and GPU 4 gets layers 25-32.

The bubble problem

The naive version of this is wasteful. During the forward pass, GPU 1 must finish processing the batch through layers 1-8 before GPU 2 can even start on layers 9-16 - GPU 2, 3, and 4 sit completely idle while GPU 1 works. Then GPU 1 sits idle while GPU 2 works, and so on. The same problem occurs in reverse during the backward pass: GPU N must wait for gradients from GPU N+1 before it can compute its own gradients. These idle gaps are called bubbles, and in the naive scheme they can waste more than half of every GPU's time.

The fix: micro-batching

GPipe's solution is to split each mini-batch into smaller micro-batches and pipeline them through the GPUs like an assembly line. Instead of waiting for the entire batch to clear GPU 1 before GPU 2 starts, GPU 2 can start on micro-batch 1 the moment GPU 1 finishes it - while GPU 1 immediately moves on to micro-batch 2.

Time step:    1     2     3     4     5     6
GPU 1 (L1-8): mb1   mb2   mb3   mb4    -     -
GPU 2 (L9-16): -    mb1   mb2   mb3   mb4    -
GPU 3 (L17-24): -    -    mb1   mb2   mb3   mb4
GPU 4 (L25-32): -    -     -    mb1   mb2   mb3

There is still a bubble at the start (while the pipeline "fills up") and at the end (while it "drains"), but once the pipeline is full, all 4 GPUs are busy simultaneously on different micro-batches. GPipe's analysis shows the bubble fraction shrinks as the number of micro-batches M grows relative to the number of pipeline stages K - specifically, the paper recommends M ≥ 4K to keep the bubble overhead small enough that throughput scales close to linearly with the number of GPUs.

Illustrative bubble-fraction values for K=4 stages, following the (K−1)/(M+K−1) relationship from the GPipe paper - the exact fraction shrinks as M grows relative to K.

Activation Checkpointing

Micro-batching solves the bubble problem, but pipelining also means more activations must be kept in memory at once (multiple micro-batches are in flight simultaneously). Activation checkpointing is the technique that makes this affordable.

Normally, the backward pass needs the activations from every layer's forward pass to compute gradients, so all of them must be kept in memory until backprop reaches that layer. Activation checkpointing instead:

1. Only stores activations at a handful of checkpoint points - for example, at the end of every transformer block, rather than after every single operation inside it.
2. Discards everything in between the checkpoints during the forward pass.
3. During the backward pass, when gradients for a discarded region are needed, recomputes that region's forward pass on the fly, starting from the nearest earlier checkpoint.

This is a direct memory-for-compute trade: you pay for a second forward pass through the recomputed sections, but you no longer need to store every intermediate activation in memory simultaneously. For deep transformer stacks, where activation memory can dwarf parameter memory, this trade is usually worth it.

Mixed Precision Training

Historically, training used FP32 (32-bit floating point) throughout. The mixed precision training paper showed that running the forward pass, backward pass, and gradient computation in FP16 (16-bit) - while keeping a single FP32 "master copy" of the weights for the actual update - preserves accuracy while roughly halving memory use and substantially speeding up the matrix math.

FP32 vs FP16 vs BF16

BF16 keeps FP32's 8 exponent bits (the same dynamic range) but shrinks the mantissa to 7 bits, trading precision for range. FP16 does the opposite: it keeps more mantissa bits (10) but shrinks the exponent to 5 bits, trading range for precision - which is exactly why FP16 needs the loss scaling trick described below, while BF16 mostly does not.

The training loop with master weights

FP32 master weights
        │  (cast down)
        ▼
FP16 working copy  ──► forward pass (FP16) ──► loss
                                                  │
                        FP16 gradients  ◄─── backward pass (FP16)
                                                  │
                        (cast up to FP32)         ▼
                                          FP32 weight update (Adam)
                                                  │
                                                  ▼
                                     updated FP32 master weights

A natural question: if you are keeping both an FP32 master copy and an FP16 working copy, doesn't that add memory rather than save it? It does add the 4-byte master copy, but the FP16 copy only costs 2 bytes instead of 4 - and, more importantly, every activation saved during the forward pass for use in the backward pass is now stored in FP16 instead of FP32. Since activations, not the weights themselves, dominate memory for large batch sizes, the net effect is close to a 2x memory reduction despite carrying the extra master copy.

Loss scaling

FP16's 5-bit exponent means any gradient value smaller than roughly 2⁻²⁴ gets flushed to exactly zero (this region is called subnormal/denormal, and FP16 cannot represent it usefully). In deep networks, many gradients naturally fall in this tiny range and would simply vanish if computed directly in FP16.

The fix: multiply the loss by a large scaling factor (for example, 1024 or higher) right before the backward pass. Since gradients scale linearly with the loss, every gradient in the backward pass is scaled up by the same factor, pushing values that would have underflowed to zero back into FP16's representable range. Right before the optimizer step, the gradients are divided by the same factor to undo the scaling - so the final weight update is mathematically identical to running in full FP32, but no gradient information was lost to underflow along the way.

DeepSpeed ZeRO: Removing Redundancy in Data Parallelism

Recall the 120GB model-state example from earlier - and recall that plain data parallelism copies the full model states onto every single GPU. With N GPUs doing data parallelism, that 120GB is redundantly stored N times, even though every GPU only ever needs its own local shard of the optimizer computation. The ZeRO paper (Zero Redundancy Optimizer, from DeepSpeed) removes exactly this redundancy by partitioning model states across the data-parallel GPUs instead of replicating them.

ZeRO has three progressive stages, each partitioning one more category of model state across the N GPUs:

The exact memory formulas depend on the bookkeeping convention used (e.g., how FP32 master weights are accounted for), but the qualitative trend is consistent: each ZeRO stage shards more of the model state, reducing per-GPU memory, with ZeRO-3 approaching a 1/N reduction in model-state memory.

At Stage 3, memory per GPU shrinks in direct proportion to N - the more GPUs you add, the less memory each one needs for model states, since every GPU only ever materializes the full parameter set momentarily (via an all-gather) right when it needs it for computation, then releases it again.

Because full parameters are only reassembled transiently, ZeRO relies heavily on the same collective operations as DP and tensor parallelism - all-gather to reconstruct a full tensor from shards right before it's needed, and reduce-scatter to both average gradients and immediately shard the result, combining what would otherwise be an all-reduce followed by a separate partition step into one communication call.

Putting It Together: 3D Parallelism

In practice, large model training combines all three axes rather than picking one:

• Tensor parallelism within a node (fast NVLink), splitting individual layers across a handful of GPUs to fit a single transformer block's weights.
• Pipeline parallelism across nodes, splitting the model's depth (its stack of layers) so each node only needs to hold a fraction of the total layers.
• Data parallelism across replicas of the entire tensor+pipeline group, scaling throughput by processing more batches in parallel - often combined with ZeRO to avoid redundantly storing optimizer states within each data-parallel replica.

Each axis solves a different bottleneck: tensor parallelism solves "this one layer doesn't fit," pipeline parallelism solves "this whole model doesn't fit," and data parallelism (with ZeRO) solves "we want more throughput without redundantly storing what we don't have to." Combining all three is often called 3D parallelism, and it's the standard recipe behind training the largest publicly known language models.

A concrete cluster configuration

These three axes multiply together rather than simply add up, and that multiplication is what turns a cluster size into an explainable engineering decision instead of a headline number. A plausible configuration on a cluster built from 8-GPU nodes:

Reading it this way - TP degree × PP degree × DP degree = total GPUs - makes the choice legible: 2-way TP because that's how many GPUs can share fast NVLink within a small group, 4-way PP because that's how many stages the model's depth was cut into, and 32-way DP layered on top purely for raw throughput, ideally paired with ZeRO so those 32 replicas aren't redundantly storing full optimizer states.

Key Takeaways

1. Model states, not just weights, are the real memory bill. A 7.5B parameter model in mixed precision with Adam needs roughly 16 bytes per parameter across FP16 weights, FP16 gradients, and FP32 master weights + momentum + variance - 120GB before a single activation is stored.

2. Data parallelism scales throughput but not model size. It replicates the entire model on every GPU and averages gradients via all-reduce - simple and effective, but useless once the model itself doesn't fit on one GPU.

3. Tensor parallelism splits inside a layer. Column-wise splitting of the first MLP matrix avoids a mid-layer synchronization by keeping GELU's nonlinearity independent per GPU; pairing it with a row-wise split of the second matrix collapses the whole MLP block down to a single all-reduce - and sequence parallelism extends the same savings to the LayerNorm/Dropout regions in between.

4. Pipeline parallelism splits across layers, and lives or dies on micro-batching. Naive layer-splitting wastes most of your GPU time in bubbles; GPipe's micro-batch pipelining with M ≥ 4K keeps the bubble overhead small enough for near-linear scaling.

5. Activation checkpointing trades compute for memory by only storing activations at block boundaries and recomputing everything in between during the backward pass - essential once pipelining keeps multiple micro-batches in flight simultaneously.

6. Mixed precision training roughly halves memory and speeds up matmuls, provided you keep an FP32 master weight copy and use loss scaling to stop small FP16 gradients from underflowing to zero.

7. ZeRO removes data parallelism's biggest inefficiency - full model-state redundancy across GPUs - by partitioning optimizer states, gradients, and eventually parameters across the data-parallel group, shrinking per-GPU memory in direct proportion to the number of GPUs, with zero change to the training math.

8. Real large-scale training combines all of the above as 3D parallelism: tensor parallelism within a node (bounded by NVLink bandwidth), pipeline parallelism across nodes, and data parallelism (usually with ZeRO) across replicas of the whole setup - with the TP × PP × DP product explaining exactly why a cluster is sized the way it is.

References

• Video by Aleksa Gordic: Scaling ML Models (YouTube)
• Shoeybi et al.: Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Rajbhandari et al.: ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
• Micikevicius et al.: Mixed Precision Training
• Huang et al.: GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism
• Wikipedia: Collective Operations
• Wikipedia: IEEE Half-Precision Floating-Point Format
• Google Cloud Blog: BF16: The Secret to High Performance on Cloud TPUs

Have questions, corrections, or want to discuss distributed training? Feel free to reach out via email at vinayrjumani@gmail.com or connect with me on LinkedIn.