Long input sequences are a memory bottleneck of Transformer models: the activations of every layer grow linearly with the sequence length, and the attention scores quadratically. In 2021, the paper Sequence Parallelism: Long Sequence Training from System Perspective proposed to split the sequence dimension across GPUs, and to compute attention by passing the keys and the values from GPU to GPU in a ring. In 2023, DeepSpeed Ulysses and Ring Attention scaled sequence parallelism to sequences of over a million tokens.

This post discusses and implements context and sequence parallelism, techniques that parallelize computation across the sequence-length dimension. Sequence parallelism is another dimension of model parallelism, in the universe of data, tensor, pipeline, context, sequence and expert parallelism.

Many model inputs and activations include an extra dimension that represents an (un)ordered sequence of tokens. Examples include text and other temporal data, with shape (B, S, E) for batch size B, sequence length (or token count) S and embedding size E. Videos have shape (B, C, T, H, W) for C channels, T frames, and HxW pixels per frame. Attention scores and probabilities have shape (B, H, S, S) for H heads and sequences of length S.

In these scenarios, parallelism can be exploited along the sequence dimension S (or T for videos). Using the notation for text, sequence parallelism over P processes requires local storage of shape (B, S/P, E) per process. With that in mind, this post implements two existing techniques for sequence parallelism: Ulysses and Ring Attention. The use case is a model composed of a sequence of GPTlite blocks, where each block is a multi-head attention (MHA) module followed by a feed-forward network (FFN), with some normalization and skip connections.

Before continuing, it is important to emphasize the following:

  • This post calls its algorithms sequence parallelism, even though several sources (such as Megatron-Core) also use the term context parallelism. Strictly speaking, sequence parallelism splits the sequence (token) dimension across GPUs so that each device processes a different chunk of tokens, mainly to scale activation memory/compute with longer sequences. Context parallelism splits the attention context (keys/values) across GPUs, so that each device holds a slice of the KV context and computes partial attention that is combined via communication, mainly to scale attention/KV-cache memory for very large context windows. In many scenarios (including this post), both are effectively used: the input tokens are split, and the KV elements needed by attention are also distributed.
  • The focus is on training and inference prefill. This is not to be confused with sequence/context parallelism for the inference decode step, which is a different setting: decode processes one token at a time and typically shards the KV cache rather than the input tokens.

Data loader setup for sequence parallelism

To implement sequence parallelism, the data loader first has to be adapted to load data at the token level, in a way that matches the hybrid data- and sequence-parallel setup. All that is needed is to configure a DataLoader and a DistributedSampler that allocate chunks of sequences (instead of full sequences) to the data loader workers.

As an example, take 4 processes and a dataset of 8 samples. With data parallelism only, the sampler gives whole samples to the processes in turn: process 1 gets samples 1 and 5, process 2 gets samples 2 and 6, and so on. With 2 data-parallel groups of 2 sequence-parallel processes, a custom DistributedSampler gives the odd samples to the first group and the even samples to the second, and each process of a group gets one half of the tokens of each of its samples. With sequence parallelism only, all 4 processes read the same samples, and each process gets one quarter of the tokens of each sample.

Once this is in place, most layers (like the feed-forward network) operate on the last (embedding) dimension of the input, and can therefore run sequence-parallel training/prefill without changes. The attention layer is an exception, and is covered next.

Distributed attention for context parallelism

During training/inference, the model runtime can scale well with the sequence parallelism degree, except for the attention layer. This is because, to reduce the activation memory footprint, each GPU stores only a subset of the KV elements.

In practice, to compute attention, each process uses its own subset/chunk of queries, but needs the keys and values for the full sequence. The rationale is the following: take the attention head computation \(softmax \left(QK^T \right)V\) (without the scaling of the scores, for simplicity), where all tensors are shaped (B, S, E). If each process holds a subset of rows of \(Q\) as \(Q_p\) with shape (B, S/P, E), it needs to access all elements of \(K^T\) and \(V\) to perform the dot-product \(Q_p K^T\), the row-wise \(softmax\), and the dot-product with \(V\). The result is an attention output of shape (B, S/P, E) per process, as illustrated in the following picture, where the blue elements are the ones that one process uses or computes:

Therefore, sequence parallelism requires some communication to make the relevant KV elements accessible across all GPUs. With that in mind, the following sections look into two alternative algorithms for sequence/context-parallel attention.

Ulysses sequence parallelism

Using the previous notation, assume a mini-batch split across the sequence dimension, with local storage of (B, S/P, E) per process. If a tensor of this shape is passed to a feed-forward network, parallelism of order P is achieved, as the linear layers treat the S dimension as a batch dimension. However, multi-head attention requires communication between processes, because the time dependencies across tokens require access to the full attention context.

(DeepSpeed) Ulysses parallelism solves this by swapping the distributed view from sequence-split to head-split before the MHA module, and back after it, as in the following illustration:

Overview of Ulysses sequence parallelism. Left: the initial view of the input tensor, distributed across 4 color-coded GPUs and split along the sequence (S) dimension. Center: the first all-to-all changes the distributed view from sequence-split to head-split. Each process now holds full sequences for a subset of heads and computes attention locally. Right: the second all-to-all changes the view back from head-split to sequence-split. Inspired by: DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models

In practice, implementing Ulysses only requires extra steps that swap the distributed view of the input tensor from, e.g., (H, B, S/P, E) to (H/P, B, S, E) and vice versa. This is done by a function dist_view_swap() that, given a tensor distributed across the process group group, swaps the distributed view by changing the split dimension from old_split_dim to new_split_dim, with an all-to-all collective:

def dist_view_swap(tensor: torch.Tensor, old_split_dim: int, new_split_dim: int, group: dist.ProcessGroup):
    '''Swap the split dimension of a tensor of full shape (H, B, S, E) distributed across P processes, between the
    heads (dim 0) and the sequence (dim 2): a local shape (H, B, S/P, E) becomes (H/P, B, S, E), and vice versa.'''
    full_shape, P = list(tensor.shape), group.size()
    full_shape[old_split_dim] *= P  # full distributed shape
    H, B, S, E = full_shape

    send = torch.cat([tensor.chunk(P, dim=new_split_dim)[r].contiguous() for r in range(P)])
    recv = torch.zeros_like(send)
    dist.all_to_all_single(output=recv, input=send, group=group)

    recv = torch.cat(
        [recv.chunk(P)[r].view(H // P, B, S // P, E) for r in range(P)],
        dim=old_split_dim,
    )
    return recv

From here, the implementation is straightforward. In the forward pass, the first (leftmost) all-to-all in the Ulysses diagram above converts the distributed view from sequence-split to head-split, while the backward pass performs the inverse view swap:

class first_alltoall(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x, group=None):
        ctx.group = group  # save for backward pass
        return dist_view_swap(x, old_split_dim=2, new_split_dim=0, group=ctx.group)

    @staticmethod
    def backward(ctx, dout):
        dout = dist_view_swap(dout, old_split_dim=0, new_split_dim=2, group=ctx.group)
        return dout, None

The second all-to-all is analogous, except that it performs the opposite view changes in the forward and backward passes:

class second_alltoall(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x, group=None):
        ctx.group = group
        return dist_view_swap(x, old_split_dim=0, new_split_dim=2, group=ctx.group)

    @staticmethod
    def backward(ctx, dout):
        dout = dist_view_swap(dout, old_split_dim=2, new_split_dim=0, group=ctx.group)
        return dout, None

The MultiHeadAttention module can now be implemented by calling both all-to-all functions:

from flash_attn.flash_attn_interface import flash_attn_func

class MultiHeadAttention(nn.Module):

    def __init__(self, n_embd=256, d_head=128, n_heads=8, dropout_p=0.1, group=None):
        '''A Ulysses multi-head attention. Variable names follow GPTlite's post.'''

        super().__init__()
        self.d_head = d_head
        self.n_heads = n_heads
        self.keys = nn.ModuleList([nn.Linear(n_embd, d_head, bias=False) for _ in range(n_heads)])
        self.queries = nn.ModuleList([nn.Linear(n_embd, d_head, bias=False) for _ in range(n_heads)])
        self.values = nn.ModuleList([nn.Linear(n_embd, d_head, bias=False) for _ in range(n_heads)])
        self.proj = nn.Linear(n_heads * d_head, n_embd)
        self.dropout = nn.Dropout(dropout_p)

        self.group = group  # Ulysses group
        if self.group is None:
            self.group = dist.new_group(range(dist.get_world_size()))

    def forward(self, x):
        # x is local chunk: (B, S/P, E)
        P = self.group.size()
        B, S_local = x.shape[0], x.shape[1]
        S = S_local * P

        # Q, K and V embeddings: (B, S/P, E) -> (H, B, S/P, d_head)
        q = torch.stack([q(x) for q in self.queries], dim=0)
        k = torch.stack([k(x) for k in self.keys], dim=0)
        v = torch.stack([v(x) for v in self.values], dim=0)

        if P > 1:  # (H, B, S/P, d_head) -> (H/P, B, S, d_head)
            q = first_alltoall.apply(q, self.group)
            k = first_alltoall.apply(k, self.group)
            v = first_alltoall.apply(v, self.group)

        # flash_attn expects (B, S, H_local, d_head)
        q_ = q.permute(1, 2, 0, 3)
        k_ = k.permute(1, 2, 0, 3)
        v_ = v.permute(1, 2, 0, 3)

        dropout_p, softmax_scale = 0.0, q_.shape[-1] ** (-0.5)
        out = flash_attn_func(q_, k_, v_, dropout_p, softmax_scale)  # (B, S, H/P, d_head)

        out = out.permute(2, 0, 1, 3)  # (H/P, B, S, d_head)

        if P > 1:  # (H/P, B, S, d_head) -> (H, B, S/P, d_head)
            out = second_alltoall.apply(out, self.group)

        out = out.permute(1, 2, 0, 3)  # (H, B, S/P, d_head) -> (B, S/P, H, d_head)
        out = out.reshape(B, S // P, -1)  # (B, S/P, H, d_head) -> (B, S/P, H*d_head)
        out = self.proj(out)  # (B, S/P, H*d_head) -> (B, S/P, E)
        out = self.dropout(out)
        return out

And that is it. It is conceptually simple.

As an important note, there are other similar approaches to this problem, such as ColAI-SP and Megatron-SP. The big advantage of DeepSpeed Ulysses parallelism is that it requires less communication than the alternatives: per link, its all-to-alls send P times less data than the all-gathers and reduce-scatters of Megatron-SP, as shown in Section 3.2 of the paper.

Its main downside is that the maximum parallelism is dictated by the number of attention heads (or of key-value heads with Grouped Query Attention (GQA), typically 8). Also, the all-to-all is a blocking collective communication that may incur heavy overhead on slow interconnects. That is where Ring Attention comes into play.

Ring Attention

Ring attention was presented in the paper Ring Attention with Blockwise Transformers for Near-Infinite Context, building on Blockwise Parallel Transformer for Large Context Models. It performs a per-block computation of attention that allows one to compute \(softmax \left(QK^T \right)V\) without materializing the full attention matrix or requiring all keys/values to be present at once. It can be viewed as a distributed variant of Flash Attention. The core rationale is also related to the ideas in Self-attention Does Not Need \(O(n^2)\) Memory.

Given scores \(s_i = dot(q, k_i)\) for a single query row, the attention output can be written as:

\[\text{Attention}(q, K, V) = \sum_i \text{softmax}(s)_i \, v_i = \sum_i \left( \frac{\exp(s_i)}{\sum_j \exp(s_j)} \right) v_i = \frac{\sum_i \exp(s_i)\, v_i}{\sum_j \exp(s_j)}.\]

This formulation shows that the output is a fraction whose numerator and denominator are sums of terms. Therefore, the full attention matrix \(QK^T\) does not need to be stored. Instead:

  1. In the forward pass, the algorithm iterates over blocks of keys/values and accumulates \(v^{\star} \leftarrow v^{\star} + \exp(s_i)\, v_i\) (numerator) and \(s^{\star} \leftarrow s^{\star} + \exp(s_i)\) (denominator);
  2. After all keys and values are processed, \(\frac{v^{\star}}{s^{\star}}\) is computed to get the final attention output;
  3. In the backward pass, gradients are computed blockwise, so the whole attention matrix from the forward pass does not need to be stored.

Improving numerical stability

The previous formulation is not numerically stable in floating-point arithmetic, because the softmax exponentials can become very large numbers. A standard solution is to compute the softmax by subtracting the maximum score from all scores. This does not change the result of the softmax but avoids overflow/underflow.

This is called the safe softmax trick and is well explained in this post. The safe softmax formulation is:

\[safe\_softmax(s_i) = \frac{\exp (s_i - s_{max})}{ \sum_j \exp( s_j - s_{max})}.\]

The issue is how to compute the maximum when the computation is done in blocks. This is done by keeping the maximum over all blocks seen so far in the loop over blocks, a method known as the online softmax, from the paper Online normalizer calculation for softmax. Then, if \(m_j\) is the maximum over all blocks up to and including block \(j\), and \(m_{j-1}\) is the previous maximum (up to block \(j-1\)), the accumulators are updated as:

\[\begin{align*} v^{\star} \leftarrow & \, v^{\star} \cdot \exp ( m_{j-1} - m_{j}) + \sum_{i \in \text{block } j} \exp \left( s_i- m_{j} \right) \, v_i \\ s^{\star} \leftarrow & \, s^{\star} \cdot \exp( m_{j-1} - m_{j} ) + \sum_{i \in \text{block } j} \exp \left( s_i- m_{j} \right). \end{align*}\]

Finally, in practice, \(v^{\star}\) and \(s^{\star}\) are often not computed directly, but a small variation is applied that accumulates log-values (LogSumExp, or LSE) for better stability (described here), where the LSE is defined as:

\[LSE (x_{1},\dots ,x_{n})=\log \left(\exp(x_{1})+\cdots +\exp(x_{n})\right).\]

Distributed algorithm

Now that the attention output can be computed per block, attention can be parallelized by assigning a sub-block to each process. The algorithm starts with sequence parallelism of the tensors \(q\), \(k\) and \(v\) across \(P\) processes, i.e., each process holds a non-overlapping sequence block of the \(q\), \(k\) and \(v\) tensors. Similarly to Ulysses, this allows for direct P-order parallelism on the feed-forward network, but not on multi-head attention.

During ring attention, each process keeps its local chunk of queries and rotates the key/value blocks among all \(P\) processes in a ring. At each communication step (out of \(P\) steps), each process sends its block of keys/values to the next process and receives the block of keys/values from the previous process in the ring. After \(P\) steps, every process has seen all key/value blocks and can compute attention for its local queries.

This pattern can be illustrated as:

Overview of the Ring Attention algorithm. Before Ring Attention: the initial view of the input tensor, distributed across 4 color-coded GPUs and split along the sequence (S) dimension. 1st Ring Attention step: each process holds its part of the Query, Key and Value tensors (Q, K and V) and computes block attention on them. Meanwhile, each process asynchronously sends its Key and Value tensors to the next process in the ring (clockwise). 2nd, 3rd, and 4th steps: each process receives the Key and Value blocks of the previous process and computes the attention output for its original Query tensor and the received Key/Value tensors. After Ring Attention: the MHA output is already sequence-split across processes, as in the initial data layout. Inspired by: Ring Attention with Blockwise Transformers for Near-Infinite Context

From a process standpoint, after all ring steps, each process has kept its own sequence block of \(q\) and has iterated over the full \(k\) and \(v\) tensors (in chunks). For example, the third (red) process above sees the following data at each step:

An important point is that, while Ulysses communicates synchronously (all-to-all), forcing processes to wait for the communication to complete, ring attention can use asynchronous point-to-point communication. If the computation time of a block exceeds its transmission time, the communication overhead can be hidden behind computation. For the sake of comparison, ring attention sends about \(P/2\) times more data than Ulysses in the forward pass. Each process sends its key and value blocks at every step, while each of the four all-to-alls of Ulysses sends less than one block per process.

Implementation

The implementation starts with two helper functions. First, isend_k_and_v(k, v, recv_k, recv_v, group) asynchronously sends the tensors k and v to the next neighbour, asynchronously receives the previous neighbour’s k and v into recv_k and recv_v, and returns the send/receive communication futures, which must be waited on with req.wait(). Then, accumulate_out_and_lse(out, lse, block_out, block_lse) merges the attention output and LogSumExp of a block into the accumulated ones, rescaling both so that the softmax is normalized over all the blocks seen so far:

import flash_attn.flash_attn_interface as fa  # flash attention


def isend_k_and_v(k, v, recv_k, recv_v, group):
    P, rank = group.size(), group.rank()
    dst = (rank + 1) % P  # the rank of the next process
    src = (rank - 1) % P  # the rank of the previous process
    req_k_send = dist.P2POp(dist.isend, k, dst, group)
    req_v_send = dist.P2POp(dist.isend, v, dst, group)
    req_k_recv = dist.P2POp(dist.irecv, recv_k, src, group)
    req_v_recv = dist.P2POp(dist.irecv, recv_v, src, group)
    return dist.batch_isend_irecv([req_k_send, req_v_send, req_k_recv, req_v_recv])


def accumulate_out_and_lse(out, lse, block_out, block_lse):
    # source: https://github.com/zhuzilin/ring-flash-attention/pull/34#issuecomment-2076126795
    block_lse = block_lse.transpose(-2, -1).unsqueeze(dim=-1)  # (B, H, S/P) -> (B, S/P, H, 1)
    if out is None:  # first block
        return block_out, block_lse

    # out = (1-w) * out + w * block_out, with w = exp(block_lse) / (exp(lse) + exp(block_lse)) = sigmoid(block_lse - lse)
    out = out - F.sigmoid(block_lse - lse) * (out - block_out)
    # lse = log(exp(lse) + exp(block_lse))
    lse = lse - F.logsigmoid(lse - block_lse)
    return out, lse

The forward pass performs \(P\) ring communication steps. On each step, it computes the attention output and the per-row LogSumExp of the current block (variables block_out and block_lse), and accumulates them. The keys and values of the next step are received into new buffers, as the current ones are still being sent and used:

class RingAttention(torch.autograd.Function):

    @staticmethod
    def forward(ctx, q, k, v, group):

        P = group.size()
        q, k, v = q.contiguous(), k.contiguous(), v.contiguous()  # (B, S/P, H, E)

        out = lse = None  # accumulators

        for step in range(P):  # do P ring steps

            # send K and V for next step, asynchronously, into new buffers: the current ones are still in use
            recv_k, recv_v = torch.empty_like(k), torch.empty_like(v)
            reqs_k_v = isend_k_and_v(k, v, recv_k, recv_v, group)

            # compute attention output and softmax lse for current block
            dropout_p, softmax_scale = 0.0, q.shape[-1] ** (-0.5)
            kwargs = dict(causal=False, window_size=(-1, -1), softcap=0.0, alibi_slopes=None, return_softmax=False)
            block_out, _, _, _, _, block_lse, _, _ = fa._flash_attn_forward(q, k, v, dropout_p, softmax_scale, **kwargs)

            # update out and lse
            out, lse = accumulate_out_and_lse(out, lse, block_out, block_lse)

            # wait for new K and V before starting the next iteration
            [req.wait() for req in reqs_k_v]
            k, v = recv_k, recv_v

        ctx.group = group  # save for backward
        out = out.to(dtype=q.dtype)
        ctx.save_for_backward(q, k, v, out, lse)
        return out

The backward pass takes as input the gradient of the loss with respect to the output (\(\nabla_{out} Loss\) or dout in the code below), and returns the gradients with respect to the inputs (gradients \(\nabla_{q} Loss\), \(\nabla_{k} Loss\), \(\nabla_{v} Loss\), or dq, dk, dv):

    @staticmethod
    def backward(ctx, dout, *args):

        P = ctx.group.size()
        q, k, v, out, softmax_lse = ctx.saved_tensors
        softmax_lse = softmax_lse.squeeze(dim=-1).transpose(1, 2)

        block_dq = torch.empty_like(q)
        block_dk = torch.empty_like(k)
        block_dv = torch.empty_like(v)

        dq = torch.zeros_like(q)
        dk = torch.zeros_like(k)
        dv = torch.zeros_like(v)

        for step in range(P):

            # send K and V for next step, asynchronously, into new buffers (as in the forward pass)
            recv_k, recv_v = torch.empty_like(k), torch.empty_like(v)
            reqs_k_v = isend_k_and_v(k, v, recv_k, recv_v, ctx.group)

            # compute gradients for the current block K/V and local Q
            dropout_p, softmax_scale = 0.0, q.shape[-1] ** (-0.5)
            kwargs = dict(causal=False, window_size=(-1, -1), softcap=0.0, alibi_slopes=None,
                          deterministic=False, rng_state=None)
            fa._flash_attn_backward(
                dout, q, k, v, out, softmax_lse,
                block_dq, block_dk, block_dv,
                dropout_p, softmax_scale, **kwargs
            )

            if step > 0:
                # wait for dK and dV from the previous steps (rotated accumulators)
                [req.wait() for req in reqs_dk_dv]
                dk, dv = recv_dk, recv_dv

            dq += block_dq
            dk += block_dk
            dv += block_dv

            # rotate dK and dV alongside K and V, into new buffers
            recv_dk, recv_dv = torch.empty_like(dk), torch.empty_like(dv)
            reqs_dk_dv = isend_k_and_v(dk, dv, recv_dk, recv_dv, ctx.group)

            # wait for new K and V before the next iteration
            [req.wait() for req in reqs_k_v]
            k, v = recv_k, recv_v

        # before returning, wait for the last dK and dV for this process' original block
        [req.wait() for req in reqs_dk_dv]
        dk, dv = recv_dk, recv_dv
        return dq, dk, dv, None

A nuance in the backward pass is that the gradients of a given block refer to the current \(K\) and \(V\) blocks, which are being rotated across processes. Therefore, the gradients dk and dv are accumulated while being rotated alongside their k and v blocks.

Finally, the forward pass of multi-head attention is straightforward and calls ring attention instead of regular attention:

class MultiHeadAttention(nn.Module):
     def forward(self, x):
        P = self.group.size()
        B, S_local = x.shape[0], x.shape[1]
        S = S_local * P

        # take Q, K and V, and collect all head embeddings: (B, S/P, E) -> (H, B, S/P, E)
        q = torch.stack([q(x) for q in self.queries], dim=0)
        k = torch.stack([k(x) for k in self.keys], dim=0)
        v = torch.stack([v(x) for v in self.values], dim=0)

        # convert to flash-attn layout (B, S/P, H, E)
        q_ = q.permute(1, 2, 0, 3)
        k_ = k.permute(1, 2, 0, 3)
        v_ = v.permute(1, 2, 0, 3)

        if P == 1:
            out_ = fa.flash_attn_func(q_, k_, v_)
        else:
            out_ = RingAttention.apply(q_, k_, v_, self.group)

        out = out_.reshape(B, S // P, -1)  # (B, S/P, H, E) -> (B, S/P, H*E)
        out = self.proj(out)  # (B, S/P, H*E) -> (B, S/P, E)
        out = self.dropout(out)
        return out

And that is all. As can be seen, a key disadvantage of ring attention is that the number of communication steps equals the number of processes. This may be a limiting factor on large networks, where very fine-grained sequence splits lead to small per-process workloads, making it harder to fully hide communication behind computation. This is because the compute time of a step shrinks quadratically with the chunk length, and its communication time only linearly.

Note that deterministic behaviour of sequence parallelism across networks with different node counts is difficult to achieve, because random number generators produce different values during model initialization and dropout.

When to use Ulysses and Ring attention

Ulysses should be used when the sequence-parallel group sits on a fast interconnect, like NVLink within a node, and the parallelism degree is at most the number of attention heads. That means Ulysses cannot be used with Multi-Query Attention (MQA), and its use with GQA is limited. A benefit of Ulysses is that any attention kernel works unchanged, and that it works well for any mask. However, the all-to-all is hard to overlap with compute and slows down across nodes.

Ring attention should be used when a parallelism degree beyond the head count, or parallelism across nodes, is needed. It has no head limit and hides communication, but only if each local chunk is long enough for compute to cover the transfer.

However, Ring Attention is not ideal for custom or arbitrary attention masks: while the blockwise formulation can support any masking pattern, its implementation depends on overlapping communication with computation (so the mask should yield computation on each step). In practice, most standard Ring Attention libraries natively support only causal masks (the triangular masks used in autoregressive models) or full masks (unmasked bidirectional attention). To overcome this, variants like Striped Attention and the zigzag chunks of Llama 3’s context parallelism explicitly permute and optimize data chunks around the specific mathematical structure of a causal mask to balance the hardware workload. Note that Llama 3’s context parallelism all-gathers the keys and values instead of passing them in a ring.

Both methods have drawbacks. The rule of thumb is to use a hybrid algorithm once the run spans multiple nodes: Ulysses within each node and Ring across nodes, which usually beats either one alone at scale. Hybrid algorithms for sequence parallelism are also discussed in USP: A Unified Sequence Parallelism Approach for Long Context Generative AI and LoongTrain: Efficient Training of Long-Sequence LLMs with Head-Context Parallelism. Megatron’s a2a+p2p for Hierarchical Context Parallelism (HCP) is one such hybrid, where p2p is the ring exchange and a2a is Ulysses-style parallelism.

Training with sequence- and multi-dimensional parallelism

PyTorch does not have a notion of partial sequences in the same way that data parallelism has a notion of partial batches. Thus, samples processed in parallel are usually assumed to be full-length samples in a data-parallel execution. To overcome this, when sequence parallelism of order P is run, P gradient accumulation steps must be performed with the corresponding gradient scaling, so that the intended batch size is effectively processed and gradients are properly averaged. Also, the gradients are averaged over all processes, including those of a sequence-parallel group, as each process computes them from a different chunk of the samples.

Moreover, when multi-dimensional parallelism (e.g., data + sequence) is performed, process groups must be defined for the data-parallel processes (the ones that load different samples) and for the sequence-parallel processes (the ones that load different chunks of the same sample). This can be done with PyTorch’s DeviceMesh or by creating the process groups manually. For the sake of illustration, if \(2 \times 2\) data and Ulysses sequence parallelism were implemented on 4 GPUs, this would be the memory layout before and during multi-head attention:

Allocation of activations on 4 GPUs with 2-GPU data parallelism and 2-GPU Ulysses sequence parallelism. Left: the blue and green processes form one sequence-parallel group and share half of the batch; the red and yellow processes form the other group and share the other half. Right: the first all-to-all in Ulysses parallelism converts token-level distributed storage into head-level distributed storage. All four processes then compute attention over full sequences.

To find more about the topic, the DeepSpeed tutorial Getting Started with DeepSpeed-Ulysses shows how to replace the attention module of a model with the DistributedAttention of Ulysses. The ring-flash-attention repository implements ring attention and its zigzag and striped variants with FlashAttention, together with the context parallelism of Llama 3, and PyTorch provides ring attention in its Context Parallel API.