Training

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