Sequence Parallelism I: Megatron SP Cuts Activation Memory Along the Sequence
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. ...