Sequence Parallelism in Large Language Models: How Megatron-SP, DeepSpeed Ulysses, and RingAttention Distribute Long Contexts

Training and serving frontier large language models on context windows spanning hundreds of thousands to millions of tokens introduces a fundamental memory barrier. While model parameters can be distributed across GPUs using Tensor Parallelism (TP) or Fully Sharded Data Parallelism (FSDP / ZeRO), activation memory scales directly with sequence length $S$. For sequence lengths exceeding 32,768 tokens, activation tensors during the forward and backward passes rapidly outstrip available GPU High Bandwidth Memory (HBM).

Standard Data Parallelism cannot partition a single continuous input sequence across multiple GPUs. Tensor Parallelism splits model weights along hidden dimensions and attention heads, but its scaling factor is strictly bounded by intra-node NVLink domains and communication latency. To push context lengths into millions of tokens without out-of-memory crashes or prohibitive recomputation costs, distributed systems employ Sequence Parallelism (SP).

Sequence Parallelism shards the sequence dimension across a cluster of GPUs. Modern distributed architectures rely on three primary paradigms: Megatron-LM Sequence Parallelism, DeepSpeed Ulysses, and RingAttention (Context Parallelism), alongside hybrid 2D formulations.

The Activation Memory Wall in Long-Context Scaling

During transformer training, intermediate activation tensors generated during the forward pass must be retained in memory or recomputed during the backward pass to calculate weight gradients. In a standard transformer layer, activation memory comprises:

1. Non-attention operations: Linear projections, LayerNorm / RMSNorm, activation functions (SwiGLU, GELU), and dropout layers. Activation memory in these layers scales linearly with sequence length: $O(b \cdot s \cdot h)$, where $b$ is micro-batch size, $s$ is sequence length, and $h$ is hidden dimension.
2. Attention computation: The storage of Query, Key, and Value states, attention score matrices, and softmax outputs. With exact attention, score matrices scale quadratically: $O(b \cdot a \cdot s^2)$, where $a$ is the number of attention heads. With fused kernels like FlashAttention, intermediate attention matrices are computed in SRAM, reducing HBM footprint to $O(b \cdot s \cdot h)$, but the overall activation volume still dominates total GPU memory at large $s$.

When sequence length $s$ grows from 4,096 to 128,000 tokens, activation memory increases by more than 30x. On an 80 GB NVIDIA H100 GPU, a 70-billion parameter model cannot execute a single forward-backward pass at 128k context without activation checkpointing and distributed sequence sharding.

Megatron-LM Sequence Parallelism (Megatron-SP)

Introduced by NVIDIA researchers Korthikanti et al. (2022), Megatron-SP addresses the activation redundancy inherent to standard Tensor Parallelism.

In classical Megatron-LM Tensor Parallelism:

• Multi-Head Attention and Feed-Forward Networks (MLPs) use column-parallel and row-parallel linear layers.
• In column-parallel layers, weight matrices are split along columns; in row-parallel layers, weight matrices are split along rows.
• Between these blocks, operations such as LayerNorm, RMSNorm, and residual dropout are replicated identically across all GPUs in the TP group.
• At the end of each row-parallel block, an All-Reduce collective operation synchronizes activations across all $TP$ ranks.

Consequently, while linear layer activations are sharded across $TP$ ranks, LayerNorm and Dropout activations remain fully duplicated on every GPU.

Standard Tensor Parallelism:
[Column Parallel GEMM] -> [Row Parallel GEMM] -> [All-Reduce] -> [Replicated LayerNorm / Dropout]

Megatron-LM Sequence Parallelism:
[Column Parallel GEMM] -> [Row Parallel GEMM] -> [Reduce-Scatter] -> [Sharded LayerNorm / Dropout] -> [All-Gather]

Megatron-SP splits the activations in LayerNorm, RMSNorm, and Dropout across the sequence dimension among the $TP$ GPUs.

Collective Communication Transformation

Megatron-SP avoids adding any additional communication volume by decomposing the standard All-Reduce operation:

• An All-Reduce operation consists of a Reduce-Scatter phase followed by an All-Gather phase.
• Megatron-SP executes Reduce-Scatter directly after the row-parallel GEMM. Each GPU receives only a $1/TP$ slice of the sequence activations.
• LayerNorm and Dropout execute independently on each GPU's local sequence slice of size $s/TP$.
• Before entering the subsequent column-parallel GEMM, an All-Gather collective reconstructs the full sequence across all TP ranks.

Because the total communication volume of Reduce-Scatter plus All-Gather exactly equals that of All-Reduce ($2 \cdot \frac{TP-1}{TP} \cdot b \cdot s \cdot h$), Megatron-SP reduces LayerNorm and Dropout activation memory by a factor of $TP$ with zero communication overhead. However, Megatron-SP remains coupled to the Tensor Parallelism group and does not scale beyond the local TP degree (typically 8 GPUs).

DeepSpeed Ulysses: Head-to-Sequence All-to-All Transposition

To decouple sequence parallelism from tensor parallelism and scale context across larger node counts, Microsoft introduced DeepSpeed Ulysses (Jacobs et al., 2023).

DeepSpeed Ulysses partitions the input sequence across $P$ sequence-parallel GPUs and relies on fast All-to-All collective communication primitives to swap between sequence partitioning and attention-head partitioning.

Input Sequence: [Batch, S/P, H, D]
        │
        ▼ (Local Q, K, V Projections)
  Q, K, V: [Batch, S/P, H, D]
        │
        ▼ (All-to-All Collective: Sequence Gather, Head Scatter)
  Q, K, V: [Batch, S, H/P, D]
        │
        ▼ (Local Attention Kernel / FlashAttention on Full Sequence)
  Attn Out: [Batch, S, H/P, D]
        │
        ▼ (All-to-All Collective: Head Gather, Sequence Scatter)
  Attn Out: [Batch, S/P, H, D]
        │
        ▼ (Local Output Projection and MLP)

Execution Flow

1. Local Projections: Each GPU receives a partition of the input sequence of length $s/P$ containing all $H$ attention heads. The Query, Key, and Value linear projections are computed locally on this $s/P$ slice.
2. First All-to-All Collective: The All-to-All collective redistributes the tensors. Each GPU sends slices of its sequence to other ranks while receiving all sequence chunks for a subset of attention heads $H/P$. The tensor shape transitions from [Batch, s/P, H, D] to [Batch, s, H/P, D].
3. Local Attention Computation: Each GPU now holds the complete sequence $s$ for its assigned $H/P$ heads. Standard, highly optimized local attention kernels (such as FlashAttention-2 or FlashAttention-3) run directly on the full sequence without requiring custom distributed attention logic.
4. Second All-to-All Collective: The attention outputs are transposed back from [Batch, s, H/P, D] to [Batch, s/P, H, D].
5. Feed-Forward Network: The projection and MLP layers execute locally on each GPU's $s/P$ sequence chunk.

Communication Cost and Constraints

The communication volume per GPU for each All-to-All operation is:

$$V_{\text{all-to-all}} = \frac{P - 1}{P} \cdot (b \cdot s \cdot h)$$

Total communication volume per layer for both forward All-to-All operations is $2 \cdot \frac{P-1}{P} \cdot b \cdot s \cdot h$, which remains constant per link as sequence length and GPU count scale proportionally.

Head Count Ceiling: DeepSpeed Ulysses requires that the number of attention heads $H$ be divisible by the sequence parallel degree $P$ ($P \le H$). In models employing Grouped-Query Attention (GQA) or Multi-Query Attention (MQA), the constraint becomes $P \le H_{KV}$, where $H_{KV}$ is the number of Key-Value heads. For a model with 8 KV heads, Ulysses cannot scale beyond 8 GPUs without splitting within individual heads.

RingAttention and Blockwise Context Parallelism

Proposed by Liu, Zaharia, and Abbeel (2023), RingAttention (often referred to as Context Parallelism in Megatron-Core) eliminates the head-count constraint by distributing attention computation across a 1D ring topology using asynchronous peer-to-peer (P2P) transfers.

Instead of gathering the full sequence onto each GPU via All-to-All, RingAttention keeps Query blocks stationary on each device and circulates Key and Value blocks in a ring.

Ring Topology (4 GPUs):
Step 0: GPU 0 computes Attn(Q0, K0, V0) | GPU 1 computes Attn(Q1, K1, V1) | GPU 2 computes Attn(Q2, K2, V2) | GPU 3 computes Attn(Q3, K3, V3)
        Async Send KV: 0 -> 1, 1 -> 2, 2 -> 3, 3 -> 0
Step 1: GPU 0 computes Attn(Q0, K3, V3) | GPU 1 computes Attn(Q1, K0, V0) | GPU 2 computes Attn(Q2, K1, V1) | GPU 3 computes Attn(Q3, K2, V2)
        Async Send KV: 3 -> 0, 0 -> 1, 1 -> 2, 2 -> 3
...
Step P-1: Output accumulator normalized via online softmax.

Algorithmic Mechanics

1. Partitioning: The sequence of length $s$ is divided into $P$ blocks of size $B = s/P$. Device $i \in \{0, \dots, P-1\}$ holds local Query block $Q_i$, Key block $K_i$, and Value block $V_i$.
2. Ring Iteration:

• In step $k=0$, device $i$ computes attention between local Query $Q_i$ and local Key-Value pair $(K_i, V_i)$.
• While computing step $k$, device $i$ asynchronously transmits its current $(K, V)$ block to neighbor $(i+1) \pmod P$ and receives the incoming block from $(i-1) \pmod P$ using non-blocking CUDA streams.
• In step $k$, device $i$ computes attention between $Q_i$ and the received $(K_{(i-k) \pmod P}, V_{(i-k) \pmod P})$.

1. Online Softmax Rescaling: Because attention is computed piecewise across blocks, each device tracks running statistics (local row-wise maximum $m_i$ and sum of exponentials $l_i$) derived from the FlashAttention online softmax formulation:

$$m_{\text{new}} = \max(m_{\text{prev}}, m_{\text{block}})$$

$$l_{\text{new}} = l_{\text{prev}} e^{m_{\text{prev}} - m_{\text{new}}} + l_{\text{block}} e^{m_{\text{block}} - m_{\text{new}}}$$

$$O_{\text{new}} = O_{\text{prev}} \left(\frac{l_{\text{prev}} e^{m_{\text{prev}} - m_{\text{new}}}}{l_{\text{new}}}\right) + O_{\text{block}} \left(\frac{e^{m_{\text{block}} - m_{\text{new}}}}{l_{\text{new}}}\right)$$

After $P$ ring shifts, every Query token has attended to every Key-Value token across the global sequence.

Overlapping Compute and Communication

The communication per ring step involves sending one $(K, V)$ block of size $2 \cdot b \cdot \frac{s}{P} \cdot h_{KV}$. The computation per step scales with the matrix multiplication of $Q_i \in \mathbb{R}^{\frac{s}{P} \times d}$ and $K_j^T \in \mathbb{R}^{d \times \frac{s}{P}}$, requiring $O\left(b \cdot \left(\frac{s}{P}\right)^2 \cdot h\right)$ FLOPs.

As the local block size $s/P$ increases, arithmetic intensity grows linearly. Once computation time exceeds the P2P transfer latency ($T_{\text{comp}} \ge T_{\text{comm}}$), RingAttention achieves 100% communication hiding.

Causal Masking and Load Balancing

In causal language modeling, tokens at position $i$ only attend to positions $j \le i$. A naive ring distribution results in severe load imbalance: device 0 (holding the earliest tokens) has zero causal attention work in subsequent steps, while device $P-1$ performs full computations.

To maintain perfect load balancing, RingAttention implementations use:

• Striped Partitioning: Tokens are assigned to ranks round-robin (token $t$ goes to rank $t \pmod P$) rather than contiguous chunks, ensuring equal causal density across ranks.
• Zigzag Ring Scheduling: Sequences are paired $(i, s - i - 1)$ across ranks, balancing the causal attention triangle across pairs of devices throughout ring iterations.

Unified Sequence Parallelism (USP / 2D SP)

While DeepSpeed Ulysses achieves minimal kernel overhead on intra-node NVLink fabrics, it is constrained by head counts ($P \le H_{KV}$). Conversely, RingAttention scales across arbitrary GPU counts but incurs latency overhead when local block sizes $s/P$ become too small to hide P2P communication.

Unified Sequence Parallelism (USP) (Fang et al., 2024) combines both approaches into a hierarchical 2D grid:

• Intra-Node (Ulysses Dimension): Within a server node (e.g., 8 GPUs connected via 900 GB/s NVLink), Ulysses uses All-to-All collectives to parallelize across heads.
• Inter-Node (Ring Dimension): Across server nodes connected via 400 Gbps InfiniBand or RoCE, RingAttention circulates $(K, V)$ blocks, hiding inter-node network latency behind intra-node computation.

2D Sequence Parallel Grid (e.g., 32 GPUs = 4 Nodes x 8 GPUs):
- Intra-Node (Size 8): DeepSpeed Ulysses All-to-All over NVLink
- Inter-Node (Size 4): RingAttention P2P Ring over InfiniBand

This 2D hierarchy allows scaling sequence parallelism to thousands of GPUs, unlocking multi-million token context training while matching hardware interconnect topologies.

Architectural Trade-Offs

Distributed training frameworks select sequence parallelism strategies based on cluster topology, model head configurations, and sequence lengths:

• Megatron-SP:
• Communication collective: Reduce-Scatter and All-Gather
• Maximum parallel degree: $TP$ (typically $\le 8$)
• Head count sensitivity: None
• Interconnect requirement: NVLink intra-node
• Primary benefit: Zero extra communication over standard Tensor Parallelism
• DeepSpeed Ulysses:
• Communication collective: All-to-All
• Maximum parallel degree: $H_{KV}$ (bounded by KV heads)
• Head count sensitivity: High (requires $P \le H_{KV}$)
• Interconnect requirement: High-bandwidth fabric (NVLink preferred)
• Primary benefit: Compatible with unmodified local FlashAttention kernels
• RingAttention / Context Parallelism:
• Communication collective: Asynchronous P2P Send/Recv
• Maximum parallel degree: Thousands of GPUs ($P \le s$)
• Head count sensitivity: None
• Interconnect requirement: Standard inter-node network (InfiniBand / RoCE)
• Primary benefit: Near-infinite sequence scaling with fully overlapped communication
• Unified Sequence Parallelism (2D SP):
• Communication collective: Hybrid All-to-All + Asynchronous P2P Ring
• Maximum parallel degree: $H_{KV} \times \text{Nodes}$
• Head count sensitivity: Moderate
• Interconnect requirement: Heterogeneous (NVLink + InfiniBand)
• Primary benefit: Optimal hardware utilization across large multi-node clusters

Through these sequence-sharding paradigms, distributed training clusters scale context windows beyond individual GPU memory limits, supporting long-document reasoning, multi-turn agent histories, and multi-modal sequence processing.

Sources

• Reducing Activation Recomputation in Large Transformer Models (Korthikanti et al., 2022)
• DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models (Jacobs et al., 2023)
• Ring Attention with Blockwise Transformers for Near-Infinite Context (Liu, Zaharia, & Abbeel, 2023)
• USP: A Unified Sequence Parallelism Approach for Long Context Generative AI (Fang et al., 2024)