Distributed Deep Learning & Model Sharding

The constraint is rarely FLOPs

A 70B parameter model in BF16 is roughly 140 GB of weights alone. Train it with Adam in mixed precision and you are looking at on the order of 16 bytes per parameter before activations: FP16/BF16 weights (~2), gradients (~2), and Adam’s FP32 master weights plus first and second moments (~12). That is about 1.1 TB of persistent state for 70B—before a single token of activation lands. One H100-80GB does not hold that. Eight of them barely do, and only if you stop replicating the full training state on every rank.

What most teams get wrong is treating “distributed training” as one technique. It is four different cuts through the same graph, each solving a different scarcity: data volume, parameter memory, activation memory, or interconnect bandwidth. Pick the wrong cut, and you burn cluster time on AllReduce traffic that never needed to leave the NVLink island. Pick the right one and a model that OOMs on a single node, trains cleanly with a config change.

By the end of this, you should be able to decide, for a given model and cluster, whether to stay on data-parallel sharding (FSDP / ZeRO), add tensor parallelism inside the node, add pipeline parallelism across nodes, or stop and quantize instead.

Data parallelism only works until the replica does not fit

Classic data parallelism (PyTorch DDP and its ancestors) copies the full model onto every GPU, splits the batch, and AllReduce gradients each step. Communication volume scales with parameter count, not batch size. That is fine for a 7B model on eight GPUs. It is fatal once a single replica exceeds device memory.

Two systems removed the replica without forcing you to rewrite the model as column- and row-parallel layers: Microsoft’s ZeRO (Zero Redundancy Optimizer) and PyTorch’s Fully Sharded Data Parallel (FSDP), which is explicitly inspired by ZeRO Stage 3.

ZeRO stages are cumulative:

DeepSpeed’s own tutorial is blunt about the optimizer bill: for a 1.5B GPT-2, Adam states alone were ~18 GB on a 32 GB V100; ZeRO-1 across eight ranks dropped that partition to ~2.25 GB per device and made the run fit. Stage 2 on 32 V100s is what they used to demonstrate a 10B GPT-2 under pure data parallelism. Stage 3 partitions the 16-bit parameters as well and, with ZeRO-Infinity, can offload parameters and optimizer state to CPU or NVMe when even aggregate GPU memory is insufficient.

FSDP’s default FULL_SHARD matches that Stage-3 pattern: parameters live sharded outside compute; an all-gather materializes the full layer for forward; they reshard; backward all-gathers again; gradients are reduce-scattered back to shards; the optimizer steps on shards only. PyTorch also exposes SHARD_GRAD_OP (shard gradients and optimizer, keep params unsharded longer) and HYBRID_SHARD (shard inside the node, replicate across nodes)—the last one is the practical answer when inter-node bandwidth cannot sustain a full cross-cluster all-gather every layer.

I treat FSDP/ZeRO-3 as the default for fine-tunes and for any dense model that still fits in aggregate GPU memory with room for activations. You keep a normal nn.Module, you do not hand-split every Linear, and you scale memory almost linearly with DP degree. The cost is collective latency on the critical path. If your interconnect is weak and your model is huge, that cost dominates—and that is when tensor and pipeline parallelism stop being optional.

Tensor parallelism buys layer width; it taxes the link every layer

Tensor parallelism (TP), in the Megatron-LM form that became the industry default, splits individual weight matrices across GPUs so each rank holds a shard of every layer. For a transformer block the standard pattern is:

• Column-parallel QKV (and the MLP up-projection): split along the output dimension; each rank multiplies the full input by its column shard; no communicate on the way in.
• Local attention or MLP nonlinearity on the partial results (attention heads partition naturally when you split QKV by head).
• Row-parallel output projection (and MLP down-projection): each rank holds a row shard, produces a partial sum, then AllReduce (or reduce-scatter under sequence parallel) to form the full residual stream.

Hugging Face’s Megatron-LM notes put the tax clearly: in a simple transformer layer you pay two AllReduces in forward and two in backward. Every layer. That is why NVIDIA’s TensorRT-LLM sharding guide keys the decision off interconnect, not parameter count:

If your GPUs have fast connections between them like NVLink then tensor parallel is likely a good choice. However if the communication will go over slow connections (across nodes for example) pipeline parallel is likely better.

Their decision tree is the one I use:

• Fits one GPU — do not shard. Communication overhead of zero beats every clever schedule.
• Fits one node — start with TP across the NVLink domain. Sanity-check pure pipeline only if the node has no high-bandwidth GPU fabric (some L40S boxes).
• Spans nodes — TP within the node, pipeline between nodes. Exception: NVL36/NVL72-class multi-node NVLink domains, where TP can stay wide without falling onto the slow fabric.

vLLM’s serving docs give the same layout for inference: tensor_parallel_size = GPUs per node, pipeline_parallel_size = nodes, with tp * pp = world_size.

One more TP detail that saves real memory: sequence parallelism (Megatron / the Korthikanti line of work). After the TP AllReduce, LayerNorm and dropout activations would otherwise be fully replicated on every TP rank. Sequence parallel replaces the all-reduce with reduce-scatter, shards those activations along the sequence axis, and all-gathers only when the next column-parallel matmul needs the full hidden state. Net communication equals a normal all-reduce (reduce-scatter + all-gather), but activation memory for those ops drops by the TP degree. Enable it whenever TP > 1; there is no good reason not to on modern stacks.

TP degree is not free even on NVLink. Larger TP means smaller GEMMs per rank, and past a point you under-occupy the tensor cores. In practice TP=2/4/8 inside one node is the useful range; TP=16 across a weak link is how people discover their MFU collapsed.

Pipeline parallelism is the cross-node tool, and the bubble is the bill

Pipeline parallelism (PP) assigns contiguous blocks of layers to ranks (or to TP groups). Stage $i$ sends activations to stage $i+1$; gradients flow the other way. Communication is point-to-point and proportional to activation size, not to the full parameter set—exactly why it survives inter-node Ethernet or InfiniBand when TP’s per-layer AllReduce does not.

Naive PP leaves most GPUs idle while the single batch walks the stages. GPipe’s fix is micro-batching: split the global batch into $M$ micro-batches and fill the pipe. The idle fraction—the pipeline bubble—scales as roughly $(P-1)/(M+P-1)$ for $P$ stages (forms vary slightly by schedule; the shape is what matters). More micro-batches shrink the bubble; they also raise the number of in-flight activations and couple you to a larger effective batch before the optimizer steps.

Schedules matter more than the Wikipedia diagram:

• GPipe-style fills all forwards then all backwards; simple, larger activation footprint.
• 1F1B (PipeDream-Flush / Megatron) interleaves one forward and one backward after the pipeline fills; steady-state bubble is small and activation memory is lower because completed micro-batches release sooner.
• Interleaved 1F1B assigns each rank multiple non-contiguous stage “chunks,” which further reduces bubble at the cost of more P2P traffic and nastier load balance.

Hugging Face’s Megatron notes: a 24-layer model with PP=4 puts 6 layers per stage; micro-batches in the PP setting are what other codebases call gradient accumulation. If you are already using Megatron-style PP, set gradient accumulation with that equivalence in mind rather than stacking another accumulation loop on top.

PP’s failure mode is load imbalance, not bandwidth. An uneven layer split (embeddings + early layers on stage 0, a heavy final LN and LM head on the last stage) makes the slowest stage the clock. Profile tokens/s per stage before you chase kernel fusion. The second failure mode is too-deep PP with too few micro-batches: bubble dominates, utilization looks like a staircase, and someone “fixes” it by buying more GPUs.

Hybrid parallelism is not a slogan; it is a placement rule

Production runs almost never use one axis. The composition that actually shows up in Megatron-LM, DeepSpeed, and NeMo looks like this:

with optional sequence/context parallel folded into the TP or as its own group, and expert parallel for MoEs.

Placement rule I will defend:

1. TP = width of the NVLink/NVSwitch island (often 8 on a DGX/HGX node, sometimes 4).
2. PP = number of such islands you need so that parameters + optimizer shards + working activations fit.
3. DP (or FSDP/ZeRO over the DP group) = whatever remains, used to scale throughput and to shard optimizer state with a distributed Adam.

Example: 16 GPUs across 2 nodes × 8 GPUs. TensorRT-LLM and vLLM both land on TP=8, PP=2 as the starting point—not TP=16. Optimizer state then shards across the DP dimension if you still have one; if the whole world is one model replica (DP=1), you are pure model-parallel and throughput only moves with micro-batch packing and kernel quality.

For Mixture-of-Experts, add expert parallelism (EP): different experts live on different ranks; tokens are dispatched to the ranks that own the chosen experts and results return. Capacity is no longer “all weights on every GPU,” which is why MoE trains at parameter counts that would be absurd if dense. EP traffic is all-to-all and sensitive to expert imbalance; it composes with TP (often TP within attention, EP across the MoE MLP) rather than replacing it. If you are not running MoE, ignore EP entirely—do not cargo-cult it into a dense LLaMA fine-tune.

Context / long-sequence parallelism is the other axis people reach for when the OOM is activations from sequence length, not weights. Megatron-style sequence parallel (above) is the cheap companion to TP. Full context parallel (ring attention and relatives) partitions the sequence through attention itself so activation memory for long context drops across the CP group. Use it when you are pushing 128k–1M contexts and TP+checkpointing is not enough; it is not a substitute for weight sharding.

Memory levers that are not parallelism

Sharding decides where state lives. These decide how much state exists:

Activation (gradient) checkpointing discards most forward activations and recomputes them in backward. Selective checkpointing on transformer blocks is the usual compromise: Hugging Face documents ~20% slowdown for a large cut in activation memory; Megatron’s full recompute figures for GPT-3-class models have been reported around 70% activation memory reduction for a few percent FLOPs overhead depending on what you recompute. Turn this on before you add another PP stage.

Gradient accumulation does not reduce the peak memory of a single micro-batch. It only simulates a larger global batch by summing grads over steps. If micro-batch=1 still OOMs, accumulation will not save you; checkpointing, sharding, or a smaller micro-batch will.

Mixed precision (BF16/FP16 compute with FP32 master weights) is already assumed in the ~16 bytes/parameter training figure. BF16 is the less fragile default on Ampere and newer; FP16 still wants loss scaling.

Optimizer choice moves the 12-byte Adam term. bitsandbytes 8-bit Adam and similar cut optimizer state hard; for full pretrains most serious runs still want standard Adam/AdamW numerics and pay the memory, then shard it with ZeRO-1/2 or a Megatron distributed optimizer.

Offload (FSDP CPU offload, ZeRO-Offload, ZeRO-Infinity NVMe) extends capacity when you are willing to trade PCIe/NVMe bandwidth for “it fits.” Fine for fitting a larger model on a single node overnight; a poor substitute for correct TP/PP placement if you care about sustained MFU.

Rough decision order when you OOM:

1. BF16 + activation checkpointing + smaller micro-batch.
2. ZeRO-2 or FSDP (or ZeRO-1 if only optimizer state is tight).
3. TP within the node if a layer’s weights or activations still blow memory or if you need lower latency per step for large batches.
4. PP across nodes when the model no longer fits the node even sharded.
5. Offload last for training throughput; acceptable earlier for offline jobs.

Serving is a different objective function

Training hides latency behind large batches and pipeline fill. Serving optimizes tokens/s under latency SLOs, KV-cache capacity, and concurrent users.

For a single replica, vLLM and TensorRT-LLM converge on the same guidance as training placement: prefer TP inside the high-bandwidth domain; add PP only when the model will not fit one node (or when GPUs lack NVLink and PP’s cheaper communication wins). vLLM logs GPU KV cache size and Maximum concurrency after load—if concurrency is below target, add GPUs or nodes and reshard; do not assume a bigger TP always helps, because TP increases the collective work per token decode step.

Inference does not carry Adam state, so weight memory is ~2 bytes/parameter in BF16 (plus KV cache, which often dominates at long context and high batch). That is why a model that needed a 16-GPU train job may serve on 2–4 GPUs with TP—or on one GPU after INT8/INT4 quantization. I would rather quantize a 70B to fit one or two GPUs for a latency-sensitive service than PP it across a weak multi-node link and gift every request a pipeline hop. Multi-replica data parallelism (independent copies behind a router) is how you scale QPS after one replica meets the latency budget; that is not model sharding, and it should not be planned as if it were.

Prefill versus decode can prefer different shardings (prefill is compute-heavy and more TP-friendly; decode is memory-bandwidth heavy). Research systems re-shard between phases; production stacks you actually run still mostly fix TP/PP at launch. Design for the decode path you will live in.

What I would configure on Monday

For a dense transformer fine-tune that fits in the aggregate memory of one node: FSDP FULL_SHARD or ZeRO-3, BF16, activation checkpointing, no TP. Measure. Only introduce TP=2/4 if profiler or memory traces show you still bound on unsharded layer peaks or if you need the Megatron fused stack for throughput.

For pretrain or full fine-tune of something in the 30B–100B class on multi-node HGX: TP = GPUs per node (or half if GEMMs get too thin), PP = nodes required to fit, DP = remainder with distributed optimizer / ZeRO-1. Sequence parallel on. 1F1B schedule. Micro-batches sized so bubble is a few percent, not so many that you distort the optimization recipe.

For serving the same weights: requantize if quality allows; else TP within node; PP only cross-node; read the KV-cache concurrency line before you buy another chassis.

The expensive mistake is not picking ZeRO instead of Megatron. The expensive mistake is TP across nodes over Ethernet because a blog said “tensor parallel is faster,” or PP=8 with two micro-batches, because the model card said the model is large. Communication topology is part of the model config. Set TP × PP × DP to match the fabric you actually have, then spend the memory you freed on batch size—not on a wider shard that your links cannot feed.