LLM Training on Duo's Tech Blog

The ZeRO-3 Diagram Most People Remember Is Wrong

Mon, 16 Jun 2025 00:00:00 +0000

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.

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

Mon, 09 Jun 2025 00:00:00 +0000

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.

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

Mon, 02 Jun 2025 00:00:00 +0000

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).

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

Mon, 26 May 2025 00:00:00 +0000

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).

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

Mon, 19 May 2025 00:00:00 +0000

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.

Tensor Parallelism in Megatron-LM: Splitting Layers, Not Stacks

Sun, 18 May 2025 00:00:00 +0000

Tensor Parallelism in Megatron-LM: Splitting Layers, Not Stacks

Pipeline parallelism splits a model by depth. Tensor parallelism splits the math inside one layer.

Megatron-LM made tensor parallelism practical for Transformers by choosing split points that preserve local GEMMs and put collectives only where the algebra requires them. The original Megatron paper introduced the core intra-layer pattern (arXiv:1909.08053). The later Megatron-LM systems paper showed how that TP axis composes with data and pipeline parallelism at cluster scale (arXiv:2104.04473).

MoE Internals: DeepSpeed-Megatron Expert Parallel Implementation

Sat, 17 May 2025 00:00:00 +0000

MoE Internals: DeepSpeed-Megatron Expert Parallel Implementation

DeepSpeed-Megatron MoE is easy to misread if you only inspect one initialization function. Megatron builds the usual TP, PP, and DP topology. DeepSpeed then wraps the model, finds MoE modules, and gives those modules expert-parallel groups. The topology is split across the training framework and the MoE layer implementation.

This post connects the pieces: deepspeed.initialize(), EP groups, expert-DP groups, MoELayer, and the difference between a DeepSpeed MoE wrapper and a Megatron-integrated SwitchMLP-style module. For routing, capacity, and All-to-All fundamentals, start with MoE Parallelism Principles.

Sequence Parallelism I: Megatron SP Cuts Activation Memory Along the Sequence

Mon, 12 May 2025 00:00:00 +0000

Sequence Parallelism I: Megatron SP Cuts Activation Memory Along the Sequence

Tensor parallelism is usually introduced as a way to make matrix multiplications fit. That is true, but it hides a second problem. After the weights are split, many activations are still replicated on every tensor-parallel rank. For short contexts this is tolerable. For long contexts it becomes one of the reasons training throughput collapses into activation checkpointing.

Megatron sequence parallelism, usually shortened to Megatron SP, is a targeted fix from Reducing Activation Recomputation in Large Transformer Models (arXiv:2205.05198). It does not replace tensor parallelism. It keeps Megatron’s column-parallel and row-parallel linear layers, then shards the sequence-local regions that tensor parallelism had left replicated. The trick is small enough to miss and important enough to change the memory budget of a whole Transformer block.

MoE Parallelism Principles: GShard, Expert Parallel, and All-to-All

Sat, 10 May 2025 00:00:00 +0000

MoE Parallelism Principles: GShard, Expert Parallel, and All-to-All

Mixture-of-Experts is sparse compute wrapped in a communication problem. The model idea is simple: replace a dense FFN with many expert FFNs and route each token to a few of them. The distributed-training problem starts when those experts live on different GPUs.

Shazeer et al. introduced the sparsely-gated MoE layer in Outrageously Large Neural Networks. GShard made the pattern practical at scale with top-2 routing, capacity, load-balancing loss, and all-to-all dispatch in GShard. Switch simplified the router to top-1 in Switch Transformers. This post explains the shared systems model underneath them.

ZeRO: Partitioning Optimizer State, Gradients, and Parameters

Sun, 27 Apr 2025 00:00:00 +0000

ZeRO: Partitioning Optimizer State, Gradients, and Parameters

Plain DDP is clean but wasteful. Every data-parallel rank stores the same parameters, gradients, fp32 master weights, and optimizer moments. At small scale that redundancy is convenient. At LLM scale it is the memory wall.

ZeRO, the Zero Redundancy Optimizer, keeps data-parallel semantics and removes the redundant storage one category at a time. ZeRO-1 shards optimizer state. ZeRO-2 shards optimizer state and gradients. ZeRO-3 shards optimizer state, gradients, and parameters (arXiv:1910.02054).

Megatron Internals III: Mixed Precision, Loss Scaling, and Grad Clipping

Sat, 26 Apr 2025 00:00:00 +0000

Megatron Internals III: Mixed Precision, Loss Scaling, and Grad Clipping

Mixed precision training is not just “turn on fp16.” It is a state machine: fast low-precision tensors do forward and backward, stable fp32 tensors receive optimizer updates, and every rank agrees whether the step is valid before any shard changes.

The classic recipe comes from Micikevicius et al., Mixed Precision Training: use lower precision where hardware is fast, keep fp32 master weights for updates, and use loss scaling when fp16 gradients underflow. Megatron wraps that recipe around tensor, pipeline, and data parallelism.

Megatron Internals II: Column/Row Parallel Linear and Vocab Parallel Embedding

Sat, 19 Apr 2025 00:00:00 +0000

Megatron Internals II: Column/Row Parallel Linear and Vocab Parallel Embedding

Tensor parallelism is not “split every tensor somehow.” In Megatron, it is a small set of layer contracts: which dimension is local, which collective completes the dense math, and which gradient path communicates.

The original Megatron-LM paper is still the cleanest starting point: split transformer matrix multiplies so each GPU does useful dense GEMM, then communicate only where the algebra requires it. This post walks the implementation-level contracts behind ColumnParallelLinear, RowParallelLinear, VocabParallelEmbedding, and parallel cross entropy.

Megatron Internals I: Building the DP / TP / PP Process Groups

Sat, 12 Apr 2025 00:00:00 +0000

Megatron Internals I: Building the DP / TP / PP Process Groups

Megatron-LM’s first trick is not tensor-parallel matmul. It is rank bookkeeping. Before the model runs, every process must know the small set of peers it will communicate with for tensor parallelism, pipeline parallelism, data parallelism, embeddings, and optimizer state.

The original Megatron-LM paper introduced intra-layer tensor parallelism for transformer training. The later Megatron-LM scaling paper put tensor parallelism, pipeline parallelism, and data parallelism into one training system. This post is the control plane underneath that system: how a flat list of ranks becomes a DP / PP / TP mesh.

Data Parallelism: From Parameter Server to Ring All-Reduce

Sun, 06 Apr 2025 00:00:00 +0000

Data Parallelism: From Parameter Server to Ring All-Reduce

Data parallelism is the default scaling move because it preserves the model program. Every rank owns the same model, sees different examples, computes gradients, and applies the same update.

The algorithm is simple. The system is not. The whole post is about removing one bottleneck: do not push all gradient traffic through one server when every GPU could be moving bytes at the same time.

Pipeline Parallelism from First Principles: Why GPipe Split the Batch

Sun, 16 Mar 2025 00:00:00 +0000

Pipeline Parallelism from First Principles: Why GPipe Split the Batch

Pipeline parallelism starts with a useful accident: a Transformer is already a chain. If one GPU cannot hold the whole stack, put consecutive blocks on consecutive GPUs and send activations across stage boundaries.

That solves capacity. It does not solve throughput. A naive layer split turns an expensive cluster into a queue where most devices wait. GPipe’s contribution was to make the queue busy by slicing the mini-batch into micro-batches and recomputing activations instead of storing everything (arXiv:1811.06965).

LLM Training Series: A Systems Map for Distributed Transformers

Sun, 02 Mar 2025 00:00:00 +0000

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).