Distributed Training Frameworks
A one-page reference — how large models get split across many GPUs, the parallelism strategies that define the space, and the three frameworks that matter most.
Why it exists
A single GPU can't hold a large model. A 70-billion-parameter model needs roughly 280 GB just for its weights in fp32 — well past an H100's 80 GB, before you add gradients and optimizer state (which can triple the footprint). Distributed training spreads the model, its gradients, and its optimizer state across many GPUs. Every framework is a different answer to the same question: how do you split the work?
The parallelism toolbox
These five strategies are the real vocabulary — frameworks are mostly different packagings of them, and frontier runs combine several ("3D parallelism"). Each animation loops; watch what moves.
1. Data parallel
Every GPU holds a full copy of the model and processes a different slice of the batch; an all-reduce then averages the gradients so all copies stay identical. Simple and fast — but every GPU pays the full memory cost.
2. Tensor parallel
A single layer is split across GPUs. Each computes a partial result, then an all-gather stitches them back together. Powerful but communication-heavy — it wants fast NVLink between GPUs.
3. Pipeline parallel
Different layers live on different GPUs; micro-batches march through the stages like an assembly line. The empty corners are the pipeline bubble — idle time while the line fills and drains, which is why you feed many micro-batches.
4. Sharded data parallel (ZeRO / FSDP)
Parameters, gradients, and optimizer state are sharded across GPUs and all-gathered on demand — assembled into the full layer just before it runs, then freed. ZeRO stages 1/2/3 shard progressively more; offload spills to CPU/NVMe. Slashes memory at the cost of extra communication.
5. Sequence & expert parallel
Expert parallel routes each token to the expert that handles it (experts live on different GPUs, tokens shuffle via all-to-all) — the traffic pattern that dominates Mixture-of-Experts training. Sequence parallel is the sibling idea: split one long sequence's tokens across GPUs so no single device holds the whole context.
The three frameworks at a glance
| Framework | Best for | Strengths | Trade-off |
|---|---|---|---|
| PyTorch FSDP / FSDP2 | The default. PyTorch-native shops, up to ~one node of dense training. | Cleanest upstream integration; sharded data parallelism; composes with torch.compile. | No built-in tensor/pipeline parallel — you compose those yourself. |
| NVIDIA Megatron-Core | Frontier scale (70B+), where throughput & efficiency are everything. | Best-in-class tensor + expert + pipeline parallel; Transformer Engine FP8 kernels; peak NVIDIA throughput. | Complex, rigid config, NVIDIA-locked. |
| Microsoft DeepSpeed | Memory-constrained runs, heavy offload, or MoE. | ZeRO 1/2/3 + ZeRO-Infinity (CPU/NVMe offload) and MoE primitives in one package. | Config complexity; in 2026 wins neither integration nor FP8 outright. |
How to choose
- → FSDP by default — modern PyTorch, mid-scale, minimal abstraction.
- → Megatron-Core at frontier scale, when squeezing every FLOP matters most.
- → DeepSpeed when you're memory-bound, offloading, or training MoE.
Also worth knowing: Torchtitan and Nanotron are readable reference implementations of 3D parallelism (learn from them, don't deploy them), and MaxText is the JAX/XLA equivalent for TPU-native training — powerful, but it requires rewriting your model in JAX.
Where it's heading (2025–26)
- Low-precision training goes mainstream. FP8 (via NVIDIA Transformer Engine) is now routine at scale, and Blackwell-class hardware pushes toward FP4 / microscaling — cutting memory and boosting throughput, with the framework handling loss-scaling and numerics.
- Mixture-of-Experts is the frontier default. Sparse MoE models dominate, making expert parallelism first-class. Work like Megatron-Core's "MoE parallel folding" and hybrid expert-parallel schemes optimize the all-to-all communication that MoE routing demands.
- Long context forces context/sequence parallelism. As windows stretch into the hundreds of thousands of tokens, splitting the sequence dimension moves from niche to standard, alongside memory-saving attention.
- RL post-training became an infrastructure problem. RLHF/RL frameworks — verl, OpenRLHF, TRL — now co-locate training and inference (rollout) engines and lean on asynchronous pipelines, a distinct scaling challenge from pre-training.
- Convergence on native-PyTorch 3D parallelism. The field is settling on composable FSDP2 + tensor + pipeline patterns (à la Torchtitan), with communication–computation overlap and fast interconnect (NVLink, InfiniBand/RDMA) as the real performance frontier.
Sources: MegaCpp — Framework Survey: FSDP2 vs Megatron-Core vs DeepSpeed vs Torchtitan vs Nanotron vs MaxText · pdpspectra — Distributed Training in 2026: DeepSpeed vs Megatron vs FSDP · Spheron — Distributed LLM Training: FSDP, DeepSpeed ZeRO-3 & Megatron-Core Guide · NVIDIA — Optimizing Communication for MoE Training with Hybrid Expert Parallel · Megatron-Core — MoE Parallel Folding · verl — RL Post-Training Framework
Figures are order-of-magnitude; exact memory depends on precision, optimizer, and activation checkpointing.
Animations embedded from assets/ — keep the folder alongside this file so they render.
react & discuss
EOF · cd ~