Parallelism on Duo's Tech Blog

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.

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.

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