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

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.

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.

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

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