Distributed Training Strategies for GPU-Bound Models

The right distributed training strategy depends on what's actually out of room: throughput, optimizer state, or the model itself. If a single GPU can't hold the parameters, gradients, and optimizer state, shard the optimizer state first with ZeRO or FSDP, and only split the model itself with tensor or pipeline parallelism once that's not enough.

Why "Just Add More GPUs" Doesn't Automatically Work

Adding GPUs only helps if the problem you have is one that more GPUs actually solve. There are two distinct problems that get lumped together under "distributed training," and picking the wrong strategy for the one you have wastes both time and money.

The first problem is throughput: the model fits on one GPU, but training on the full dataset takes too long. The second is capacity: the model itself, or its optimizer state, is larger than what a single GPU's memory can hold. Data parallelism solves the first problem. It does nothing for the second, because every GPU still needs to hold a full copy of the model.

Confusing these two is the single most common distributed-training mistake. Teams wrap a 30B-parameter model in DistributedDataParallel, watch it OOM on GPU 0 before training even starts, and conclude they need "more compute" when what they actually need is a different parallelism strategy, not more of the same one.

Data Parallelism: The Default, Until the Model Doesn't Fit

Data parallelism (DP, or its more efficient form DistributedDataParallel in PyTorch) puts an identical copy of the model on every GPU, splits each batch across them, and synchronizes gradients after the backward pass. It's the right first choice whenever the model comfortably fits on one device, because it's simple, well-tested, and scales close to linearly on decent interconnects.

Its limit is exactly its mechanism: every GPU carries a full copy of the model, gradients, and optimizer state. If that state doesn't fit on one GPU, adding more GPUs under plain DP doesn't help. You need to shard that state instead of replicating it, or split the model itself.

Do the Memory Math Before You Pick a Strategy

Before choosing a strategy, work out whether the problem is really "model too big" or "optimizer state too big," because the fix is different. For mixed-precision training with Adam, the standard accounting (from the ZeRO paper's memory breakdown) is roughly 16 bytes per parameter:

• 2 bytes per parameter for fp16 weights

• 2 bytes per parameter for fp16 gradients

• 4 bytes per parameter for the fp32 master copy of the weights

• 4 bytes per parameter for Adam's momentum term

• 4 bytes per parameter for Adam's variance term

A 7B-parameter model needs roughly 112GB just for this state, before you've accounted for activations, which scale with batch size and sequence length on top of it. An 80GB GPU can't hold that alone. That's not a "the model is too big" problem in the naive sense; the weights themselves are only 14GB in fp16. It's an optimizer-state problem, and sharding the optimizer state (ZeRO/FSDP) solves it without touching model parallelism at all.

ZeRO and FSDP: Shard the State Instead of Copying It

ZeRO (Zero Redundancy Optimizer, introduced by Microsoft's DeepSpeed team) and PyTorch's Fully Sharded Data Parallel (FSDP) attack the capacity problem directly. Instead of every GPU holding a full copy of optimizer state, gradients, and parameters, each GPU holds only a shard, and the missing pieces are gathered on demand via collective communication during the forward and backward pass.

ZeRO ships in three stages, each sharding more: stage 1 shards only optimizer state, stage 2 adds gradient sharding, and stage 3 also shards the parameters themselves. FSDP is PyTorch's native implementation of roughly the ZeRO-3 idea, wrapped so it looks like ordinary data parallel training from the training-loop's perspective.

A minimal FSDP setup looks like this:

import torch
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from functools import partial

model = build_model()  # e.g. a HF transformer, still on CPU/meta device

auto_wrap_policy = partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={TransformerBlock},
)

model = FSDP(
    model,
    auto_wrap_policy=auto_wrap_policy,
    device_id=torch.cuda.current_device(),
)

This gets you most of the way for models that don't fit under plain DDP but still fit, once sharded, across the GPUs you have in a node. It does not split individual layers across devices, so if a single layer is itself too large (huge embedding tables, very wide MLP blocks), sharding the optimizer state alone won't save you.

When You Have to Split the Model Itself: Tensor and Pipeline Parallelism

Once ZeRO/FSDP sharding isn't enough, because even a single layer's activations or weights don't fit, you need to split the model's computation graph across GPUs. There are two ways to do that, and they solve different bottlenecks.

Tensor parallelism (popularized by Megatron-LM) splits individual weight matrices across GPUs, so a single matrix multiply is computed as a set of smaller multiplies that get combined with an all-reduce. This requires GPUs to communicate on every layer, which means it only performs well over very fast interconnects like NVLink, i.e. within a single node.

Pipeline parallelism (as in GPipe) instead splits the model into sequential stages, with each stage's layers living on a different GPU or set of GPUs. Data flows through the stages like an assembly line. Because stages only need to pass activations forward and gradients backward at stage boundaries, pipeline parallelism tolerates much slower interconnects, which makes it the right tool for spanning multiple nodes over standard networking.

The trade-off with naive pipelining is the "bubble": early stages sit idle waiting for later stages to finish the first micro-batch. Micro-batching (splitting each batch into smaller chunks that flow through the pipeline back-to-back) shrinks that idle time but doesn't eliminate it.

How Do I Choose a Distributed Training Strategy?

Work through this in order; stop at the first strategy that solves your actual bottleneck.

1. Does the model fit on one GPU, with room for activations and a reasonable batch size? If yes, use data parallelism (DDP) and stop. You don't need anything more complex.

2. Does the model fit, but the optimizer state doesn't? Use ZeRO stage 1 or 2, or FSDP with parameter sharding disabled for the parts that don't need it. This is the common case for models in the 1B-20B range on 40-80GB GPUs.

3. Does even the model's parameters and gradients not fit per GPU, once sharded? Move to ZeRO stage 3 or full FSDP, which shards parameters too and gathers them just-in-time for each layer's forward/backward pass.

4. Does a single layer not fit, even fully sharded? You need tensor parallelism within a node for that layer's matrix operations.

5. Does the full model, even split across a node's GPUs, still not fit, or do you need to scale across nodes? Add pipeline parallelism across nodes, and combine it with tensor parallelism inside each node and data parallelism across pipeline replicas. This three-way combination is usually called 3D parallelism, and it's how models with tens or hundreds of billions of parameters get trained in practice.

The mistake to avoid is jumping straight to step 5 because it's what large-model papers describe. Most teams training models under ~13B parameters never need pipeline or tensor parallelism at all; FSDP alone gets them there with far less engineering overhead and far fewer ways for a training run to silently deadlock on a collective operation.

A Starting DeepSpeed Config for ZeRO Stage 2

If you're using DeepSpeed instead of native FSDP, the whole strategy for the common "optimizer state doesn't fit" case is a JSON config, not a code rewrite:

{
  "train_batch_size": 256,
  "gradient_accumulation_steps": 4,
  "fp16": { "enabled": true },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": { "device": "cpu" },
    "contiguous_gradients": true,
    "overlap_comm": true
  }
}

offload_optimizer pushes optimizer state to CPU RAM when even sharded GPU memory isn't enough, at the cost of PCIe transfer latency. It's worth trying before jumping to tensor/pipeline parallelism, since it adds zero model-code changes and can be the difference between fitting on the GPUs you already have and requesting more.

The Takeaway

Match the strategy to the specific resource that's actually exhausted: batch throughput, optimizer state, parameter count, or per-layer size are four different problems with four different fixes, and only the last two require touching the model's code. Work through the decision order above before reaching for tensor or pipeline parallelism; most "the model is too big" problems turn out to be "the optimizer state is too big" problems that FSDP or ZeRO solves with a config change.

Sources

• PyTorch Fully Sharded Data Parallel documentation

• PyTorch DistributedDataParallel documentation

• Rajbhandari et al., "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models"

• Shoeybi et al., "Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism"

• Huang et al., "GPipe: Efficient Training of Giant Neural Networks Using Pipeline Parallelism"

• DeepSpeed documentation

FAQ

What's the actual difference between data parallelism and model parallelism?

Data parallelism copies the whole model onto every GPU and splits the batch, so it only helps once the model itself already fits in memory. Model parallelism splits the model's layers or tensors across GPUs, which is the only option once a single copy no longer fits.

Should I use tensor parallelism or pipeline parallelism?

Tensor parallelism splits individual layers across GPUs and needs fast interconnects like NVLink, so keep it inside a node. Pipeline parallelism splits the model into sequential stages and tolerates slower links, so use it to span multiple nodes.

Do I need DeepSpeed, or does PyTorch's built-in FSDP cover this?

FSDP now covers most single-framework use cases: sharded parameters, gradients, and optimizer states with a config-free API. Reach for DeepSpeed when you need ZeRO-Infinity's CPU/NVMe offload, specific ZeRO stage tuning, or 3D parallelism recipes that FSDP doesn't expose directly.

Can I combine multiple parallelism strategies at once?

Yes, and at large scale you usually have to. Tensor parallelism inside a node, pipeline parallelism across nodes, and data parallelism across pipeline replicas ('3D parallelism') is the standard recipe for training models with tens of billions of parameters or more.

How do I estimate GPU memory before choosing a strategy?

For mixed-precision training with the Adam optimizer, budget roughly 16 bytes per parameter: 2 bytes each for fp16 weights and gradients, plus 4 bytes each for the fp32 master weights, momentum, and variance the optimizer keeps. A 7B-parameter model needs on the order of 112GB just for this state, before activations.