LLM Training Series: A Systems Map for Distributed Transformers

LLM training stops being “run the model on more GPUs” the moment one replica no longer fits, one link becomes the step time, or one schedule leaves half the cluster idle. From that point on, the training run is a placement problem.

The stable question is simple: which bytes are replicated, which bytes are sharded, and which link moves them on the critical path? Data parallelism, tensor parallelism, pipeline parallelism, ZeRO, sequence/context parallelism, and expert parallelism are different answers to that question. The public systems that matter - GPipe, Megatron-LM, DeepSpeed ZeRO, PyTorch DDP, GShard, Switch Transformer, and modern Megatron-Core - all expose the same constraints in different shapes (GPipe, Megatron-LM, ZeRO, PyTorch Distributed, GShard, Switch Transformer).

TL;DR

• Data parallelism replicates the model and shards the batch. The hard part is reducing gradients without turning one server or one slow rank into the step time.
• Tensor parallelism splits the algebra inside a Transformer block. It moves activation-sized tensors every layer, so it belongs on fast local links.
• Pipeline parallelism splits the layer stack. It buys capacity and another scale axis, then charges you in bubbles, micro-batches, and schedule complexity.
• ZeRO/FSDP-style sharding keeps data-parallel semantics while partitioning optimizer state, gradients, and parameters.
• Sequence/context parallelism appears when long context makes activation and attention memory larger than the model-state problem.
• MoE expert parallelism shards experts and routes tokens. Sparse compute reduces FLOPs per token, but all-to-all, load balance, and small GEMMs become first-class.
• The production answer is hybrid. Single-axis scaling is a warmup exercise, not a large-model recipe.

Reading Order

Read these in order if you want the stack to build cleanly. Each post still stands alone.

1. Pipeline Parallelism from First Principles: Why GPipe Split the Batch
2. Data Parallelism: From Parameter Server to Ring All-Reduce
3. ZeRO: Partitioning Optimizer State, Gradients, and Parameters
4. Tensor Parallelism in Megatron-LM: Splitting Layers, Not Stacks
5. Megatron Internals I: Building the DP / TP / PP Process Groups
6. Megatron Internals II: Column/Row Parallel Linear and Vocab Parallel Embedding
7. Megatron Internals III: Mixed Precision, Loss Scaling, and Grad Clipping
8. MoE Parallelism Principles: GShard, Expert Parallel, and All-to-All
9. MoE Internals: DeepSpeed-Megatron Expert Parallel Implementation
10. Sequence Parallelism I: Megatron SP
11. Sequence Parallelism II: DeepSpeed Ulysses
12. Sequence Parallelism III: Ring Attention
13. Sequence Parallelism IV: Megatron Context Parallel
14. Hiding Tensor-Parallel Collectives: AG/RS Overlap in Megatron
15. The ZeRO-3 Diagram Most People Remember Is Wrong

The Map

Start by naming the partition. Most confusion in distributed training comes from mixing these rows.

These are not competing features. They compose because they attack different tensors.

A dense 70B-class run may use tensor parallelism inside an NVSwitch domain, data parallelism across replicas, ZeRO or a distributed optimizer across the DP dimension, and pipeline parallelism only if depth still does not fit. A long-context run adds sequence or context parallelism. An MoE run adds expert parallelism, and suddenly all-to-all placement matters as much as GEMM throughput.

Canonical Axes

Data parallelism

The invariant is that every rank applies the same update. Parameter servers made that explicit by centralizing gradient aggregation, but they concentrate network traffic at the server (Scaling Distributed Machine Learning with the Parameter Server). Ring all-reduce removes that hot spot by decomposing all-reduce into reduce-scatter plus all-gather; Baidu popularized the pattern for deep learning clusters (arXiv:1702.05847), Horovod made it easy to use (arXiv:1802.05799), and PyTorch DDP wraps the same idea around autograd buckets (arXiv:2006.15704).

Pipeline parallelism

Pipeline parallelism treats the model as a chain. GPipe’s key mechanism is micro-batching: split one mini-batch into M pieces so K pipeline stages do useful work instead of waiting through a (K - 1) / K bubble (arXiv:1811.06965). PipeDream then showed what changes when the schedule is asynchronous and weight versions can be stale (arXiv:1806.03377). Megatron-LM’s 1F1B and interleaved schedules are the production descendants for dense Transformer training (arXiv:2104.04473).

ZeRO and sharded data parallel

ZeRO starts from one fact: Adam state, gradients, master weights, and parameters are replicated across data-parallel ranks even though each rank only needs some of them at a given moment. ZeRO-1 shards optimizer state, ZeRO-2 also shards gradients, and ZeRO-3 shards parameters too (arXiv:1910.02054). ZeRO-Offload and ZeRO-Infinity extend the storage hierarchy to CPU DRAM and NVMe, which helps only when prefetch and overlap keep the GPU fed (arXiv:2101.06840, arXiv:2104.07857).

Tensor parallelism

Megatron-LM’s tensor parallelism is not arbitrary matrix slicing. It uses a column-parallel first MLP projection, local GeLU, and a row-parallel second projection so communication lands after the nonlinearity, not before it (arXiv:1909.08053). Attention is split by heads. Vocab embeddings and logits are sharded by token range. The later Megatron paper shows how this TP axis composes with DP and PP at cluster scale (arXiv:2104.04473).

Sequence, context, and MoE

Long context shifts the bottleneck from model state to activations and attention. Sequence parallelism shards sequence-dimension activations; ring attention shards the context blocks and circulates K/V blocks instead of materializing the full attention problem on every rank. MoE shifts the problem again: GShard and Switch Transformer show that sparse experts can scale parameter count without dense FLOPs, but the system pays in routing, expert placement, all-to-all, and load balance (GShard, Switch Transformer).

Code

The papers are useful, but the contracts are clearest in code:

• NVIDIA Megatron-LM and Megatron-Core: tensor parallel layers, pipeline schedules, distributed optimizer, sequence/context parallelism, and MoE.
• DeepSpeed, especially deepspeed/runtime/zero/: ZeRO stages, offload, partitioning, and optimizer-state orchestration.
• PyTorch distributed and DistributedDataParallel: process groups, c10d collectives, autograd hooks, and gradient buckets.
• PyTorch pipelining: the modern PyTorch lineage for pipeline schedules.
• Horovod: a compact reference for ring-allreduce-centered data parallelism across frameworks.

The Recurring Checks

Before choosing a parallel recipe, answer these in this order:

1. What must be resident? Parameters, gradients, optimizer state, activations, KV/cache-like temporaries, communication buffers.
2. What can be sharded without changing semantics? Model states are easier than activations; activations are easier than arbitrary dynamic routing.
3. Which collective moves the hot bytes? All-reduce, reduce-scatter, all-gather, all-to-all, send/recv, or a custom ring.
4. Can communication overlap compute? Exposed communication is the cost that matters.
5. Which physical link carries it? NVLink/NVSwitch, PCIe, InfiniBand, Ethernet, CPU DRAM, and NVMe are not interchangeable.
6. What new failure mode appears? Bubbles, stale weights, small collectives, load imbalance, memory fragmentation, checkpoint complexity, or graph breaks.

Practical Order

The least painful plan usually looks like this:

1. Estimate model-state memory and activation memory.
2. Use plain data parallelism while one replica fits.
3. Add ZeRO/FSDP when replicated model states are the blocker.
4. Add tensor parallelism when individual layers are too wide or too slow.
5. Add sequence/context parallelism when long context dominates activations or attention.
6. Add pipeline parallelism when depth still does not fit or the cluster needs another scale axis.
7. Add expert parallelism only when the model architecture is sparse.

This is not a law. It is a bias toward the smallest system that fits the constraint in front of you.

Related Posts Already on This Blog

• From Scaling Laws to Cluster Size - tokens, FLOPs, and GPU-hours before the parallel recipe.
• Large MoE Performance: The Three Walls After Sparsity - production MegaScale / Megatron-Core MoE co-design.
• Learning PyTorch DDP Performance Tuning on a One-GPU Machine - DDP pathologies you can reproduce locally.
• Why Variable Sequence Length Breaks DDP Throughput - token imbalance under data parallelism.

References

• Yanping Huang et al., GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism, NeurIPS 2019.
• Aaron Harlap et al., PipeDream: Fast and Efficient Pipeline Parallel DNN Training, SOSP 2019.
• Mohammad Shoeybi et al., Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism, 2019.
• Deepak Narayanan et al., Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM, SC 2021.
• Samyam Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC 2020.
• Jie Ren et al., ZeRO-Offload: Democratizing Billion-Scale Model Training, USENIX ATC 2021.
• Samyam Rajbhandari et al., ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning, SC 2021.
• Mu Li et al., Scaling Distributed Machine Learning with the Parameter Server, OSDI 2014.
• Andrew Gibiansky et al., Bringing HPC Techniques to Deep Learning, 2017.
• Alexander Sergeev and Mike Del Balso, Horovod: Fast and Easy Distributed Deep Learning in TensorFlow, 2018.
• Shen Li et al., PyTorch Distributed: Experiences on Accelerating Data Parallel Training, VLDB 2020.
• Noam Shazeer et al., GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, 2020.
• William Fedus et al., Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, JMLR 2022.
• Code: NVIDIA Megatron-LM, DeepSpeed, PyTorch distributed, Horovod.