LLM Training

Current DeepSpeed ZeRO-3 partitions each parameter as flattened intra-layer slices across ranks

The ZeRO-3 Diagram Most People Remember Is Wrong

The ZeRO-3 Diagram Most People Remember Is Wrong Many engineers first learned ZeRO-3 from an animation that looked like pipeline parallelism. One GPU owned early layers, another GPU owned later layers, and the active layer block appeared to be broadcast to all other GPUs when needed. That picture is memorable. It is also not the right mental model for current DeepSpeed ZeRO-3 training. The current steady-state model is intra-layer partitioning. Each parameter is flattened, padded if necessary, and split across data-parallel ranks. Before forward or backward compute needs that parameter, ranks AllGather the full parameter. After gradients are produced, ranks ReduceScatter gradients back to the owning shards. That is the operational reading of ZeRO-3 from the ZeRO paper (arXiv:1910.02054) and the current DeepSpeed runtime. ...

Naive AllGather waits before GEMM while overlapped tensor-parallel communication starts GEMM as chunks arrive

Hiding Tensor-Parallel Collectives: AG/RS Overlap in Megatron

Hiding Tensor-Parallel Collectives: AG/RS Overlap in Megatron Once Megatron SP is enabled, tensor-parallel communication often appears as AllGather and ReduceScatter rather than a single AllReduce. That is a memory win because activations can stay sequence-sharded between tensor-parallel regions. It also creates a scheduling question. Can the collectives be hidden under GEMM work? That question sits directly on top of Megatron’s tensor/sequence-parallel training path (arXiv:2104.04473, arXiv:2205.05198). This post is about that scheduling question. It complements Megatron tensor parallelism and Megatron SP. The goal is not to memorize every flag in Megatron or Transformer Engine. The goal is to recognize where the dependency graph allows overlap and where it does not. ...

Megatron Context Parallel balances causal ring attention by pairing early and late sequence chunks

Sequence Parallelism IV: Megatron Context Parallel and Load-Balanced Rings

Sequence Parallelism IV: Megatron Context Parallel and Load-Balanced Rings Megatron Context Parallel, usually shortened to CP, brings long-context attention into Megatron’s hybrid-parallel world. It keeps the familiar tensor, pipeline, and data-parallel axes, then adds a context axis that shards the sequence. Inside that context group, attention behaves like a ring, borrowing the same blockwise K/V circulation idea as Ring Attention (arXiv:2310.01889) while staying compatible with Megatron-style tensor and pipeline parallelism (arXiv:2104.04473). ...

Ring Attention keeps query blocks fixed on each GPU while key/value blocks move around a communication ring

Sequence Parallelism III: Ring Attention for Context That Does Not Fit

Sequence Parallelism III: Ring Attention for Context That Does Not Fit Megatron SP reduces replicated activation memory around tensor-parallel blocks. DeepSpeed Ulysses uses All-to-All to turn sequence shards into head shards for attention. Ring Attention changes the unit of work again. It asks each rank to keep a block of queries fixed, then circulate key/value blocks around a ring until every query block has seen every key/value block it needs. That is the core idea in Ring Attention with Blockwise Transformers for Near-Infinite Context (arXiv:2310.01889). ...

Ulysses sequence-shards activations and uses All-to-All to make each rank own all tokens for one attention head

Sequence Parallelism II: DeepSpeed Ulysses and All-to-All Attention

Sequence Parallelism II: DeepSpeed Ulysses and All-to-All Attention Megatron SP is a careful memory optimization around an existing tensor-parallel block. DeepSpeed Ulysses starts from a different question. What if each device owns a sequence slice most of the time, but attention temporarily wants each device to own a head slice instead? The answer in DeepSpeed Ulysses is an All-to-All transpose (arXiv:2309.14509). Before attention, every rank has all heads for a subset of tokens. After All-to-All, every rank has all tokens for a subset of heads. That one layout change lets local attention run per head while the rest of the layer can remain sequence-sharded. ...