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

May 26, 2025 · 8 min · Duo An
Column-parallel linear layer splitting output features across tensor-parallel ranks

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

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

April 19, 2025 · 7 min · Duo An