# USP: A Unified Sequence Parallelism Approach for Long Context Generative AI

Jiarui Fang
Tencent
<EMAIL_ADDRESS>
&Shangchun Zhao
Tencent
<EMAIL_ADDRESS>

###### Abstract

Sequence parallelism (SP), which divides the sequence dimension of input
tensors across multiple computational devices, is becoming key to unlocking
the long-context capabilities of generative AI models. This paper investigates
the state-of-the-art SP approaches, i.e. DeepSpeed-Ulysses and Ring-Attention,
and proposes a unified SP approach, which is more robust to transformer model
architectures and network hardware topology. This paper compares the
communication and memory cost of SP and existing parallelism, including
data/tensor/zero/pipeline parallelism, and discusses the best practices for
designing hybrid 4D parallelism involving SP. We achieved 47% MFU on two
8xA800 nodes using SP for the LLAMA3-8B model training using sequence length
208K. Our code is publicly available at https://github.com/feifeibear/long-
context-attention.

##  1 Introduction

The field of artificial intelligence is witnessing a trend as the context
length in generative AI models grows ever longer. Claude has pioneered this
trend in large language models (LLMs) by extending the sequence length to 100K
tokens. Following closely in its footsteps, OpenAI’s GPT-4 has expanded the
context length to 128K tokens. The advent of multi-modality models is
propelling this trend forward, with Gemini 1.5 Pro boasting a context length
of a staggering 10 million tokens, and OpenAI’s Sora, a Diffusion Model,
accommodating at least 1 million visual tokens. These breakthroughs underscore
the imperative for generative AI techniques to adeptly handle a larger context
length.

Sequence Parallelism (SP), a technique that partitions input sequences, has
emerged as a promising approach for the training or inference of longer
sequences. Following an initial exploration period of two years, by the latter
of 2023, two landmark works, DeepSpeed-Ulysses [1] and Ring-Attention [2],
marked the maturation of the SP technique. DeepSpeed-Ulysses maintains
constant communication volume when sequence length and compute devices are
increased proportionally, while Ring-Attention hides P2P communication costs
introduced by SP through overlapping computation and communication. However,
challenges remain, such as the SP parallel degree of DeepSpeed-Ulysses is
limited to less than the number of attention heads, and the computational
efficiency of Ring-Attention degrading due to the subdivision of matrix
multiplications. These limitations currently hinder the broader adoption of
Sequence Parallelism in distributed Transformer computation.

In this paper, we delve deeper into the realm of SP. We begin by highlighting
that Ulysses and Ring are not mutually exclusive approaches; they can be
combined through a hybrid parallel strategy to mitigate their drawbacks. Then,
we discussed the relationship between SP and data/tensor/zero/expert/pipeline
parallelism. The most complex among these is the relationship between SP and
tensor parallelism. Since tensor parallelism also has its specific sequence
parallel optimizations to reduce activation memory cost [3]. For each
parallelism approach, whether SP should replace it or is there some issue to
using SP with it together, remains an open question. After addressing these
questions, we have provided a list of best practices for building a 4D hybrid
parallelism system.

The primary contributions of this paper include:

  * •

We propose a unified sequence parallel method that integrates DeepSpeed-
Ulysses and Ring-Attention, overcoming the shortcomings of both and
demonstrating greater robustness to model architecture and network hardware.

  * •

We systematically analyze the application of SP in conjunction with Tensor
Parallelism, ZeRO, and Pipeline Parallelism as 4D parallelism, and provide a
list of best practices to apply SP.

##  2 Sequence Parallelism Approaches

Before delving into the concept of SP, let’s review the computational process
of the standard Transformer Block. The Notion used in this paper is shown in
Table 1, where $d=hc\times hs$.

$L$ | Sequence Length | $d$ | Hidden Dimension | $hc$ | Head Count
---|---|---|---|---|---
$hs$ | Head Size | $bs$ | Batch Size | $N$ | Device Number
Table 1: Notation for Transformer Parameters

Given input sequences $Q,K,V\in\mathbb{R}^{L\times d}$, where $L$ is the
sequence length and $d$ is the head dimension, we compute the matrix of
outputs as follows:

$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{QK^{T}}{\sqrt{d}}\right)V,$
(1)

Each self-attention sub-layer is accompanied by a feedforward network (FFN),
which is applied to each position separately and identically.

$\text{FFN}(x)=\max(0,xW_{1}+b_{1})W_{2}+b_{2}.$ (2)

Compared to Tensor Parallelism (TP) [4] and ZeRO [5], the research on SP for
Transformer models has been relatively underdeveloped for a long time. The
challenge lies in the characteristic of attention computation, where the
sequence dimension serves as a common dimension in the matrix multiplication
after softmax, making it difficult to partition the tensors and distribute the
computation across multiple nodes after slicing the sequence dimension.

Early attempts [6, 7, 8] at SP were not successful, often leading to redundant
memory consumption [6] and inefficient communication pattern [8]. For long
input sequences, the best practice is to adopt the Sequence Parallelism of
Megatron-LM. This method optimizes the AllReduce operation of TP, reducing the
memory cost of activations while maintaining the same communication overhead.
As shown in Figure 1, the principle of Megatron-LM Sequence Parallelism is the
similar to ZeRO-2 [5]. It replaces the AllReduce operation on replicated
tensors (the left figure) into equivalent allgather and reduce-scatter
operations on partitioned data (the right figure). Since an AllReduce
operation is exactly a combination of Allgather and ReduceScatter, the
communication cost remains the same. The size of the input and output tensors
is reduced by a factor of 1/$N$ across $N$ computational devices. Because of
the sequence dimension in input/output tensors is partitioned, it is named as
Sequence Parallelism. However, this form of Sequence Parallelism cannot be
used independently without tensor parallelism and communication volume remains
constant regardless of the degree of parallelism.

Figure 1: The Principle of Megatron-LM Sequence Parallelism.

The maturity of the standalone SP technology is marked by the publication of
two milestone papers in late 2023. The DeepSpeed-Ulysses [1], named as SP-
Ulysses, and Ring-Attention [2], named as SP-Ring, solving the longstanding
memory and communication issues inherent to SP from two distinct perspectives.
For both methodologies, each computational device is allocated a distinct
segment of the $Q$ (query), $K$ (key), $V$ (value), and $O$ (output) tensors,
which are segregated along the dimensions of the sequence. There is no
redundancy in their storage between devices, which is a primary distinction
from the early SP design [6].

SP-Ring can be viewed as a distributed version of FlashAttention [9]. As shown
in the right part of Figure 2, SP-Ring employs a nested two-level loop that
orchestrates the communication and computation in a blockwise fashion. When
computing for the blocks of the tensor $O$ segment, if the required tensor $K$
and $V$ blocks are not locally available, Peer-to-Peer (P2P) communication is
utilized to fetch them from other devices. Communication can be organized in a
Ring fashion, where each device simultaneously sends and receives $K,V$
blocks, allowing communication to overlap computation.

Figure 2: SP-Ulysses and SP-Ring.

SP-Ulysses leverages All2All communication for segments of the $Q$, $K$, $V$,
and $O$ tensors, as shown in the left part of Figure 2. After the All2All
operation, the partitioning of these four tensors changes from the sequence
dimension $L$ to the dimension of the attention number heads $hc$. Therefore,
the computation of $softmax(QK^{T})V$ for each attention head is maintained in
its entirety, and can be implemented using the under the hood Attention
operator library, like FlashAttention.

##  3 Unified Ulysses-Ring Sequence Parallelism

Currently, both SP-Ulysses and SP-Ring are facing certain issues that hinder
their effectiveness in practical applications.

SP-Ulysses is sensitive to the number of attention heads. The parallelism
degree 111The parallelism degree is the number of devices that participate in
parallel computation. In other words, it is the number of processes in a
parallel process group. of DS-Ulysses cannot exceed the number of attention
heads $hc$. Consequently, it is not suitable for the GQA (Grouped Query
Attention) [10] and MQA (Multi-Query Attention) [11] scenarios. For instance,
Llama3-8B employs GQA with a KV head number of 8, which means that when using
DS-Ulysses SP, the maximum SP degree is 8. However, if MQA is used and the KV
head number is 1, DS-Ulysses will not function. In addition, since Tensor
Parallelism also requires division across the $hc$ dimension, SP-Ulysses, and
TP are in conflict.

SP-Ring is inefficient in computation and communication. Ring-Attention
segments the $Q,K,V,O$ tensors into smaller blocks, which can lead to a
decrease in computation efficiency of the fused operator $Softmax(QK^{T})V$.
Even if communication and computation fully overlap, the total execution time
lags behind that of DS-Ulysses. When using a causal mask has not been
addressed, the DS-Ring has a load-unbalancing issue. DS-Ring does not impose
any restrictions on $hc$.

SP-Ulysses and SP-Ring are currently considered alternative strategies for SP,
with the choice of one precluding the other. This point of view was reinforced
by the SP-Ring authors in their ICLR open review
rebuttal222https://openreview.net/forum?id=WsRHpHH4s0&noteId=HIY0tae4Gz.
Currently, Megatron-DeepSpeed utilizes SP-Ulysses, while Megatron-LM opts for
SP-Ring for its SP implementation. However, we claim that, rather than viewing
them as rivals, they can work together as a unified SP approach.

As shown in Algorithm 1, SP-Ring and SP-Ulysses are organized in a hybrid
parallel manner named as USP-Attention to work together in partitioning the
sequence dimension. The SP process group is segmented into two orthogonal
process group sets: a set of SP-Ring process groups and a set of SP-Ulysses
process groups. For a more intuitive understanding, an SP process group can be
viewed as a 2D mesh, SP-Ring operates across each column of the mesh, while
SP-Ulysses runs across each row. For example, a process group containing 8
processes could viewed as $2\times 4$, where a SP-Ulysses process group
$ulysses\\_pg$ of size 2 and a SP-Ring process group $ulysses\\_pg$ of size 4.
This is the same as how data parallelism and tensor parallelism process groups
are partitioned.

The inputs to the USP-Attention include the segment of the $Q,K,andV$ tensors
after being partitioned along the sequence dimension and the output is a
tensor $O$ shard. The size of a tensor segment is ($bs$, $L/N$, $hs$, $hd$).
Note that when using MAQ, the shape of the $hc$ for the $K$ and $V$ tensor
segments differs from that of the Q tensor. During forward propagation,
$scatter\\_idx$ is set to 1, and $gather\\_idx$ is set to 2. AllToAll4D merges
the dimension $L$ and partitions the dimension $hc$ of the tensors $Q$, $K$,
and $V$, transforming them into ($hc/N$, $bs$, $L$, $d$). They also partition
the $O$ tensor along the $L$ dimension and merge the $hc$ dimension. During
backward propagation, $scatter\\_idx$ is set to 2, and $gather\\_idx$ is set
to 1.

Algorithm 1 Unified Sequence Parallelism Attention Implementation

1:function USP-Attn($ulysses\\_pg$, $ring\\_pg$, $Q$, $K$, $V$,
$scatter\\_idx$, $gather\\_idx$)

2:
$Q\leftarrow\textbf{AllToAll4D}{(Q,scatter\\_idx,gather\\_idx,group=ulysses\\_pg)}$

3:
$K\leftarrow\textbf{AllToAll4D}{(K,scatter\\_idx,gather\\_idx,group=ulysses\\_pg)}$

4:
$V\leftarrow\textbf{AllToAll4D}{(V,scatter\\_idx,gather\\_idx,group=ulysses\\_pg)}$

5: $O\leftarrow\textbf{LoadBalance-RingAttention}(Q,K,V,group=ring\\_pg)$

6:
$O\leftarrow\textbf{AllToAll4D}{(O,gather\\_idx,scatter\\_idx,group=ulysses\\_pg)}$

7: return $O$

8:end function

The vanilla SP-Ring introduces load-unbalancing issues when applying causal
attention, as only the lower triangular matrix of $QK^{T}$ needs to be
computed. As shown in Figure 3, if the sequence dimension is divided evenly,
the computational tasks are not evenly distributed among the devices. As shown
in the left side of the figure, on 4 GPUs the computation load of GPU3 is
nearly 7 times that of GPU0.

Figure 3: Loading Balancing for SP-Ring.

The solution to the load-unbalancing issue is to reorder the input sequence
tokens along the sequence dimension, as depicted on the right side of Figure
3. In the figure, the input sequence consists of 16 tokens. Under even
partition, GPU0 processes tokens 0-3, while GPU3 handles tokens 12-15. After
the reorder operation for load balance partitioning, GPU0 now processes tokens
0, 1, 14, 15, and GPU3 processes tokens 5,6,11,12. The workloads handled by
each GPU are perfectly balanced, which is a superior solution to the striped
attention [12].

The aforementioned load balancing method is quite straightforward and a
similar implementation has already been applied in Megatron-LM. We present its
application for Unified SP method. The Algorithm 2 extracts and reorders the
input sequence from the global input sequence in hybrid Ulysses and Ring
parallelism. For the positional encoding only involves element-wise
operations, i.e. RoPE [13], the same reordering operation also needs to be
applied to the positional encoding parameters of the model. Since the
operation is applied to the model’s input token sequence, which is an integer
vector of length $bs*L$, the additional overhead for load balancing is
negligible.

Algorithm 2 Prepare Load Balance Sequence Segment from The Global Input
Sequence

1:function LocalBalanceLocalSeq($seq$, $ring\\_process\\_group$,
$ulysses\\_process\\_group$)

2: $ring\\_degree\leftarrow ring\\_process\\_group.get\\_world\\_size()$

3: $ring\\_rank\leftarrow ring\\_process\\_group.get\\_rank()$

4: $ulysses\\_rank\leftarrow ulysses\\_process\\_group.get\\_rank()$

5: $seq\\_chunks\leftarrow seq.chunk(2\times ring\\_degree)$

6: $reorder\\_seq\leftarrow
concat([seq\\_chunks[r\\_rank],seq\\_chunks[2\times rd-r\\_rank-1]])$

7: $local\\_seq\leftarrow reorder\\_seq.chunk(ud)[u\\_rank]$

8: return $local\\_seq$

9:end function

Unified SP is highly flexible and robust, allowing for the various
combinations of the Ulysses degree and the Ring degree, as long as the product
of the Ulysses degree and the Ring degree equals the SP degree. When the
Ulysses degree equals $N$, it becomes SP-Ulysses, and when the ring degree
equals $N$, it becomes SP-Ring. Therefore, it covers the ability of both SP-
Ulysses and SP-Ring.

Firstly, Unified SP can remove the head number limitation of SP-Ulysses. For
example, to run llama3-8B 333https://github.com/meta-llama/llama3 with $hc$=8
in a 16-degree SP, one can set the Ulysses degree to 8 and the ring degree to
2. Secondly, Unified SP allows lower bandwidth and topology requirements for
network infrastructure by offering a more robust communication pattern. By
setting the Ulysses degree to a value between 1 and N, the Attention
communication pattern will be a mix of P2P and All2All as shown in Figure 4.
Such a communication pattern is particularly well-suited for heterogeneous
communication networks, allowing All2All operations to operate in high-
bandwidth interconnections while asynchronous P2P communications operate in
lower-bandwidth sections. This is applicable to scenarios such as a node in
which GPUs are connected via PCIe Switch or a cluster of GPU nodes between
which the network is Ethernet connected.

Tip 1: We suggest using Unified-SP in place of SP-Ring and SP-Ulysses, as it
encompasses the capabilities of both while offering additional benefits.

Figure 4: The Unified-SP is more robust to network hardware topology.

##  4 SP in 4D Parallelism

SP, as a newly emerged parallelism method, how to integrate it into the
existing hybrid parallelism framework of Data Parallelism (DP), Tensor
Parallelism (TP), and Pipeline Parallelism (PP) has not been thoroughly
studied. This section will analyze the relationships between SP and DP, TP,
and PP, and discuss the best practices for designing 4D parallelism involving
sequence dimension parallelism.

As shown in Table 2, we analyze the communication and memory cost of a
standard transformer block for different parallelism. The impact of GQA is not
reflected in the table, but we will analyze it later. Entries in the
Communications Params columns indicate the collective communication operations
are conducted on the parameters and gradients of the transformer block. It
includes the parameters/gradients of the weight and bias tensors of 4 Linear
Layers in the Self-Attention layer, as well as 2 Linear Layers of the FFN
layer, amounting to a volume of $12\times O(d^{2})$ elements in the GPT-2
model. Note that llama3 and llama2 has $9.37\times O(d^{2})$ elements since
the intermediate size is lower. Entries in Communications Act(ivatiom) columns
indicate the collective communications are conducted on a single hidden states
tensor belonging to activations, containing $bs\times L\times d$ elements. The
memory cost is broken down into model parameters/gradients (P/G), Optimizer
States (OS), and intermediate activation tensors (Act).

The Cost of the Communication represents the bandwidth requirements. It is
calculated by the product of the communicated element number by an algorithm
bandwidth (algobw) factor related to the algorithm of collective communication
444https://github.com/NVIDIA/nccl-tests/blob/master/doc/PERFORMANCE.md. For
the collective communication algorithm of AllReduce, AllGather, ReduceScatter,
and AllToAll, the respective algobw factors are $2\frac{n-1}{n}$,
$\frac{n-1}{n}$, $\frac{n-1}{n}$, and 1. In the table, we approximate the term
$O(\frac{n-1}{n})$ to $O(1)$ for simplicity.

The table is built for the mix-precision training using the fp16(bf16) format.
The memory requirement for the model parameters and gradients are $P$ and $G$
Bytes. The Optimizer States (OS), which includes the parameter in fp32,
momentum, and variance in the Adam optimizer, is 6 times that of the fp16
parameters. The memory requirement for the peak activation is $A$ bytes. The
parallel degree is $N$.

Table 2: Comparison of Communications and Memory Cost of SP, DP, TP, and ZeRO for a standard transformer block | Communication (FWD+BWD) | Split | Memory
---|---|---|---
| Param | Cost | Act | Cost | Dim | P/G | OS | Act
SP-Ulysses | allreduce | $12O(d^{2})$ | 8*all2all | $\frac{8}{N}O(bs*L*d)$ |  $hc$/$L$ | $P+G$ | $6P$ | $A/N$
SP-Ring | allreduce | $12O(d^{2})$ | P2P | $4O(bs*L*d)$ |  $L$/$L$ | $P+G$ | $6P$ | $A/N$
DP | allreduce | $12O(d^{2})$ | 0 | 0 |  $bs$/$bs$ | $P+G$ | $6P$ | $A/N$
ZeRO1 | allgather+ | $12O(d^{2})$ | 0 | 0 | $hc$/$L$ | $P+G$ | $6P/N$ | $A/N$
reducescatter
SP-Unified+ | allgather+ | $12O(d^{2})$ | P2P+8*all2all | $\leq 4O(bs*L*d)$ | $hc$/$L$ | $P+G$ | $6P/N$ | $A/N$
ZeRO1 | reducescatter
SP-Unified+ | allgather+ | $12O(d^{2})$ | P2P+8*all2all | $\leq 4O(bs*L*d)$ | $hc$/$L$ | $P+\frac{G}{N}$ | $6P/N$ | $A/N$
ZeRO2 | reducescatter
SP-Unified+ | 2*allgather+ | $18O(d^{2})$ | P2P+8*all2all | $\leq 4O(bs*L*d)$ | $hc$/$L$ | $\frac{P+G}{N}$ | $6P/N$ | $A/N$
ZeRO3 | reducescatter
TP | 0 | 0 | 4*allreduce | $8O(bs*L*d)$ |  $hc$/$d$ | $\frac{P+G}{N}$ | $6P/N$ | $\alpha A$
TP-sp | 0 | 0 | 6*allgather+ | $10O(bs*L*d)$ | $hc$/$d$ | $\frac{P+G}{N}$ | $6P/N$ | $A/N$
4*reducescatter

#### Data Parallelism (DP):

In terms of communication cost, SP is inferior to DP. Both SP and DP require
the allreduce operation on gradients during backward propagation. The
difference in their communication performance lies in the attention module,
where SP introduces additional communications overhead for activations. When
the Ulysses degree is set to be greater than 1, the communication overhead of
SP will be larger than that of DP due to the all2all operations. When using
the Ring method, although the additional P2P communication for attention is
overlapped, it introduces extra performance issues. The ideal performance is
only to reach the performance of DP without communication for attention. In
terms of memory performance, both SP and DP are equivalent, as they all can
reduce the activation footprint to 1/$N$.

Tip 2: We suggest prioritizing the use of DP over SP if possible. Only when
the batch size (bs) is insufficient for partitioning should one consider
whether to employ SP.

#### ZeRO:

ZeRO [14] is a distributed parameter management method that reduces the
storage space requirements of each computing device by sharding the Optimizer
States (ZeRO-1), Gradients (ZeRO-2), and Parameters (ZeRO-3) across multiple
devices. The memory cost for Optimizer States, Gradients, and Parameters is
reduced to 1/$N$ of the original. ZeRO can also operate within an SP process
group because partitioning along the batch dimension ($bs$) or the sequence
dimension ($L$) is equivalent to ZeRO’s approach. ZeRO is working on the
unified process group of size $N_{sp}\times N_{dp}$, which combines the SP and
DP process groups.

Tip 3: We suggest that when utilizing SP, it should always be used in
conjunction wit ZeRO-1/2. One can also consider employing ZeRO-3, and Offload
techniques [5, 15] to trade off communication cost for memory savings.

#### Tensor Parallelism (TP):

The Tensor Parallelism (TP) approach, pioneered by Megatron-LM [4], shards the
parameters of models across computing devices. In the TP part of activation
tensors, not all are partitioned and distributed across multiple computational
devices. Consequently, the memory cost for activations, as reflected in the
Table 2, is denoted by the $\alpha A$, where 0<$\alpha$<1\. Please refer to
the Equation(2) of paper [3] for a more precise $\alpha$. TP has been further
refined by Megatron-LM Sequence Parallelism [3], which replaces the an
allreduce in TP with an allgather and a reducescatter, and therefore reduces
the activation memory cost to $A/P$ at the cost of redo two allgathers for
attention and FFN. To distinguish Megatron-LM Sequence Parallelism from Ring
and Ulysses Sequence Parallelism, we named it TP-sp here.

As shown in Table 2, in terms of communication cost, TP-sp is higher than SP-
Ulysses and SP-Ring. Firstly, the communication volume of TP-sp is greater
than that of SP-Ring, and the latter can be overlapped with computation.
Additionally, the communication volume of TP-sp does not decrease with an
increase in parallelism, whereas SP-Ulysses can achieve this. Therefore, TP-sp
is inferior to any form of sequence parallelism in terms of communication. The
SP has a lower communication cost for activations, but it requires
synchronizing gradients and parameters. However, the parameter communication
volume is small compared to the activation communication volume, and it can be
overlapped by computation. GQA/MQA can reduce SP communication costs, while
the communication cost of TP-sp remains unchanged. Assuming the GQA group
number is $G$, the Ulysses and Ring communication cost for $K$ and $V$ is
reduced to 1/$G$, and the activation communication cost is reduced to
$\frac{4}{N}O(bs*L*d)+\frac{4}{N}O(bs*L*d/G)$ and
$\frac{4}{(}bs*L*\frac{d}{G})$.

Tip 4: We suggest that SP has an advantage over TP-sp in terms of
communication cost on a large scale. GQA can further reduce the communication
cost of SP.

In terms of memory cost, TP-sp holds an advantage over SP. Even when SP
employs ZeRO-1/2 to align the memory cost of OS and activations with that of
TP-sp, the memory cost of the parameter remains more substantial. However,
when SP employs ZeRO3, it achieves a similar memory cost with TP-sp. The SP-
Ulysses+ZeRO3 is the exact strategy to scale sequence length employed by the
authors of DS-Ulysses [1].

Based on the above analysis, under the same model and input configurations 1,
simply switching parallelism from TP-sp to SP-Unified does not increase the
trainable sequence length; on the contrary, it is very likely to result in a
reduction of the sequence length.

However, SP can still extend the sequence length compared to TP-sp in high
parallel degree. When the sequence is very long, as previously analyzed, the
proportion of parameter communication volume in the total communication is
relatively small. Therefore, the additional communication overhead of an
allgather operation introduced by ZeRO has a limited impact.

Tip 5: We suggest switching TP-sp to SP cannot increase the sequence length in
training. SP+ZeRO3 can train a similar sequence length as TP-sp.

Due to the inherent limitation of TP-sp’s parallelism, which is capped at hc,
it is not possible to further reduce the memory cost of activations by
increasing the TP-sp paralleled degree. In contrast, SP can continue to expand
its parallel degree by leveraging the SP-Ring technique, thereby enabling the
training of larger models on a larger scale.

Tip 6: We suggest a higher degree of SP parallelism, which may need to set a
large ring degree when the head number is limited, to train a long sequence
across a greater number of computational devices. This is an advantage that
cannot be achieved with TP-sp approaches.

#### Pipeline Parallelism (PP):

PP partitions transformer blocks across layers, so it is complementary with
the SP, which partitions tensors inside a transformer block. Therefore, SP and
PP are fully compatible. Since SP can form a unified parallel group with ZeRO
for collective communication, we believe TP should still be placed at the
lowest dimension of the 4D parallel group partitioning.

Tip 7: We suggest that in 4D hybrid parallelism, the order of process group
dimensionality from low to high is TP, SP-Ulysses, SP-Ring, ZeRO-DP, PP.

##  5 Experimental Results

###  5.1 Performance of Unified SP

We evaluated the attention module’s performance using SP-Unified on an L20
PCIe GPU cluster, benchmarking the throughput in TFLOPS. We replicated the
parameter settings of the llama3-8B model. Table 3 illustrates that on an
8xL20 setup, optimal throughput for 32K and 128K sequences was achieved with a
ulysses-degree of 4 and ring-degree set to 2. This finding supports Tip 1: SP-
Unified is well-suited for heterogeneous networks. The lb-ring, an enhanced
version of the standard Ring-Attention with load-balancing, outperformed the
original. It consistently outperforms the basic ring attention implementation,
with the advantage becoming more pronounced as the sequence length increases.

Table 3: Throughput (iters/sec) of SP-Unified on 8xL20 PCIe fwd-only (Ring-Degree $\times$ Ulysses-Degree=8) group_num | bs | seqlen | head_num | head_size | ulysses_degree | basic-ring | lb-ring
---|---|---|---|---|---|---|---
4 | 1 | 8K | 32 | 128 | 8 | 57.346 | 57.098
4 | 1 | 8K | 32 | 128 | 4 | 153.134 | 152.189
4 | 1 | 8K | 32 | 128 | 2 | 415.5 | 454.93
4 | 1 | 8K | 32 | 128 | 1 | 358.595 | 361.969
4 | 1 | 32K | 32 | 128 | 8 | 15.229 | 14.262
4 | 1 | 32K | 32 | 128 | 4 | 28.584 | 32.818
4 | 1 | 32K | 32 | 128 | 2 | 44.348 | 62.754
4 | 1 | 32K | 32 | 128 | 1 | 40.478 | 58.377
4 | 1 | 128K | 32 | 128 | 8 | 2.563 | 2.586
4 | 1 | 128K | 32 | 128 | 4 | 3.217 | 4.235
4 | 1 | 128K | 32 | 128 | 2 | 3.399 | 5.476
4 | 1 | 128K | 32 | 128 | 1 | 3.131 | 5.186

As shown in Table 4, we repeated the attention benchmarking on an 8xA100-SXM4
NVLink node, the highest throughput for both 32K and 128K sequence lengths was
achieved when the ulysses-degree is set to 8, that is the same as SP-Ulysses.
SP-Ulysses demonstrated a significant advantage over SP-Ring. This verified
the argument in Sec. 3 of SP-Ring, that although communication overhead can be
hidden through overlapping with computation, it results in a reduction in
computational efficiency.

Table 4: Throughput (iters/sec) of SP-Unified on 8xA100-SXM4 NVLink fwd-only (Ring-Degree $\times$ Ulysses-Degree=8) group_num | bs | seqlen | head_num | head_size | ulysses_degree | basic-ring | lb-ring
---|---|---|---|---|---|---|---
4 | 1 | 32K | 32 | 128 | 8 | 135.569 | 136.375
4 | 1 | 32K | 32 | 128 | 4 | 103.525 | 132.979
4 | 1 | 32K | 32 | 128 | 2 | 91.365 | 132.979
4 | 1 | 32K | 32 | 128 | 1 | 81.985 | 113.79
4 | 1 | 128K | 32 | 128 | 8 | 2.782 | 2.785
4 | 1 | 128K | 32 | 128 | 4 | 2.024 | 2.771
4 | 1 | 128K | 32 | 128 | 2 | 1.73 | 2.89
4 | 1 | 128K | 32 | 128 | 1 | 1.628 | 2.91

###  5.2 End-to-end SP Performance in Megatron-LM

We have incorporated the SP-Unified method into Megatron-LM. Currently,
Megatron-LM includes a preliminary version of SP-Ring, but lacks an
implementation of SP-Ulysses. Our software is based on Megatron-LM commit
2196398, dated April 12, 2024. The SP-Ring is implemented using Megatron-LM’s
native Context Parallel, while SP-Ulysses is developed from our repository
code. We conducted experiments using the docker image
nvcr.io/nvidia/nemo:24.03. By default, we utilize ZeRO-1 for both Data
Parallelism (DP) and Sequence Parallelism (SP), and consistently apply
Sequence Parallelism Optimization for Tensor Parallelism (TP-sp). We do not
employ gradient accumulation or activation recomputation. Our experimental
setup includes two GPU nodes, each equipped with 8xA800 GPUs connected via
400GB/s NVLink, and node-to-node communication via 1.6 Tbps RDMA. We have
adjusted the MFU computation in Megatron-LM to account only for effective
computations under causal masking. As a result, the MFU is notably lower when
training with long sequences compared to the MFU figures printed by Megatron-
LM.

###  5.3 SP vs. DP

Firstly, we compare the performance of SP and DP under the same LLAMA2-7B
workload on a single node of A800 GPUs. The global batch size is 8. We use the
SP-Unified, and pick the best ulysses-degree and ring-degree settings. In a
single node, ulysses-degree = 8 usually achieves the best performance. As
shown in Figure 6, DP outperforms SP-Unified across various input sequence
lengths, which confirms our conclusion in Tip 2.

Figure 5: Comparison DP and SP on LLAMA2-7B Task with global bs=8.

###  5.4 Hybrid SP and TP

Table 5 presents the performance of the llama2-7B model on a single node of
8xA800 GPU 80GB. The longest sequence length that can be achieved is 64K, and
the best performance is obtained when tp-degree=4 and ulysses-degree=2. It
outperforms TP-sp-only by 10%. When seqlen=64K, SP-only will bring OOM issue,
which echos our conclusion in Tip 5 that SP is less memory-efficient than TP-
sp.

When global-bs=16 and seqlen=30K, SP-Ulysses delivers the optimal performance,
significantly surpassing other SP and TP hybrid strategies. It is also 26%
better than TP-sp-only in throughput. This indicates that, despite the similar
communication cost of SP-Ulysses and TP-sp, SP-Ulysses is more communication
efficient than TP-sp in practice in an NVLINK-connected setting, and
communication patterns of hybrid SP-Ulysses and SP-Ring sometimes bring better
performance.

Table 5: LLAMA2-7B FLOPS per GPU of hybrid parallelism using TP and SP-Unified on a single 8xA800 NVLink node. seqlen | global-bs | tp-degree | ulysses-degree | ring-degree | FLOPS/GPU | MFU
---|---|---|---|---|---|---
64K | 1 | 4 | 2 | 1 | 154.49 | 0.50
64K | 1 | 4 | 1 | 2 | 151.40 | 0.49
64K | 1 | 8 | 1 | 1 | 141.85 | 0.45
30K | 16 | 2 | 4 | 1 | 155.98 | 0.50
30K | 16 | 2 | 1 | 4 | 147.77 | 0.47
30K | 16 | 4 | 1 | 1 | 150.05 | 0.48
30K | 16 | 1 | 8 | 1 | 163.42 | 0.52
30K | 16 | 1 | 1 | 8 | 142.16 | 0.46
30K | 16 | 8 | 1 | 1 | 129.12 | 0.41

We employed a hybrid parallel strategy combining TP and SP-Unified to
benchmark the training throughput of LLAMA3-8B across two nodes. The results
are shown in Table 6. In most of cases, the global batch size is fixed to be
1, aiming to maximize the sequence length. Since llama3-8B has only 8 heads,
the maximum product of ulysses-degree and tp-degree is 8. Table 6 presents
throughput in FLOPS on 64K, 80K and 120K.

Table 6: LLAMA3-8B FLOPS per GPU of hybrid parallelism using TP and SP-Unified on two RDMA-connected 8xA800 NVLink nodes. seqlen | global-bs | tp-degree | ulysses-degree | ring-degree | FLOPS/GPU | MFU
---|---|---|---|---|---|---
64K | 1 | 1 | 8 | 2 | 136.31 | 0.44
64K | 1 | 1 | 4 | 4 | 137.48 | 0.44
64K | 1 | 1 | 2 | 8 | 129.44 | 0.41
64K | 1 | 1 | 1 | 16 | 121.83 | 0.39
64K | 1 | 8 | 1 | 2 | 129.75 | 0.42
64K | 1 | 4 | 2 | 2 | 122.45 | 0.39
64K | 1 | 2 | 4 | 2 | 87.67 | 0.28
64K | 1 | 2 | 2 | 4 | 89.35 | 0.29
64K | 1 | 4 | 1 | 4 | 122.57 | 0.39
64K | 1 | 2 | 1 | 8 | 101.35 | 0.32
64K | 2 | 8 | 1 | 1 | 141.20 | 0.45
80K | 1 | 1 | 8 | 2 | 147.46 | 0.47
80K | 1 | 1 | 4 | 4 | 148.90 | 0.48
80K | 1 | 1 | 2 | 8 | 140.13 | 0.45
80K | 1 | 1 | 1 | 16 | 132.86 | 0.43
80K | 1 | 8 | 1 | 2 | 136.16 | 0.44
80K | 1 | 4 | 2 | 2 | 137.49 | 0.44
80K | 1 | 2 | 4 | 2 | 111.05 | 0.36
80K | 1 | 2 | 2 | 4 | 110.81 | 0.36
80K | 1 | 4 | 1 | 4 | 130.27 | 0.42
80K | 1 | 2 | 1 | 8 | 121.14 | 0.39
80K | 2 | 8 | 1 | 1 | 144.40 | 0.46
120K | 1 | 4 | 2 | 2 | 152.51 | 0.49
120K | 1 | 2 | 4 | 2 | 136.63 | 0.44
120K | 1 | 8 | 1 | 2 | 145.92 | 0.47
120K | 1 | 4 | 1 | 4 | 150.96 | 0.48

The optimal performance for sequence lengths of 64K and 80K is achieved with
SP-only without TP-sp, the optimal setting of both is ulysses-degree=4 and
ring-degree=4. For sequence length is 64K and 80K, the unified SP outperforms
the SP-Ring by 13% and 12%, respectively. These results echo our conclusion in
Tip 1.

For sequence lengths are 64K and 80K, we can increase the global batch size
(global-bs) for the TP+SP hybrid setting when tp-degree=8. However, SP-only
always has an OOM issue with global-bs=2. Increasing the global-bs to 2, TP+SP
has a 2.7% improvement in throughput to SP-only at the sequence length 64K.
But at sequence length 80K, TP+SP with global-bs as 2 is still worse than SP-
only with global-bs as 1.

When the sequence length reaches 120K, the optimal performance is achieved
with tp-degree=4 and ulysses-degree=2, reaching 152.51 TFLOPS, with an MFU of
0.49. At this sequence length, SP-only meets an OOM issue. This confirms our
Tip 5 again, which implies that the memory efficiency of SP is inferior to
that of TP-sp. It is noteworthy that two other TP+SP hybrid settings yield
similar performance, 0.47 and 0.48 in MFU, which is quite close to the
optimal.

Table 7: Exploring Upper Bound of Sequence Length for LLAMA3-8B using TP and SP-Unified on 2 8xA800 NVLink nodes. seqlen | global-bs | tp-degree | ulysses-degree | ring-degree | FLOPS/GPU | MFU
---|---|---|---|---|---|---
160K | 1 | 4 | 2 | 2 | 158.64 | 0.51
160K | 1 | 8 | 1 | 2 | 156.63 | 0.50
208K | 1 | 8 | 1 | 2 | 147.26 | 0.47
160K | 1 | 4 | 1 | 4 | 159.37 | 0.51
190K | 1 | 4 | 1 | 4 | 157.08 | 0.50

We explored the upper bound of sequence length on 2 nodes, and the results are
presented in Table 7. The largest sequence length was achieved with tp-
degree=8 and ring-degree=2. At this point, SP-only could not run due to OOM.
Compared to SP, TP-sp is more memory-efficient, hence for training the longest
sequences, parallel degrees are limited by 8 attention head numbers should all
be assigned to TP-sp. It is noteworthy that we can achieve longer sequence
lengths through activation optimization, such as activation recomputation
[16], and offloading [5, 15], but these methods will harm MFU.

###  5.5 Convergence

We compared the convergence differences between USP and DP. Using the same
dataset, we tested the loss curves over 10K iterations on 4 GPUs. We found
that the curves for USP and DP completely overlapped, which validates that our
modifications to RoPE and SP-Ring for load balancing are correct.

Figure 6: Comparison the loss of DP and USP on LLAMA2-7B with global bs as 4.

##  6 Future Work

#### SP on Large Scale Cluster

We believe that SP is highly beneficial for extremely large scale LLM training
tasks. Currently, for publicly disclosed large-scale model training tasks [17,
18] over 10K GPUs, SP has not been utilized. This is because these training
tasks were started before November 2023, when the SP methods were not yet
mature.

Firstly, SP can introduce the dimension that can be partitioned from the
sequence length, alleviating the constraints hindered by batch size
limitations. MegaScale project [18] uses a large global batch size to increase
DP degree which impacts convergence. However, we can increase the SP degree
instead of the DP degree to reduce the global batch size, thereby avoiding the
convergence problem.

Secondly, increasing the SP degree can decrease the activation cost, allowing
for longer model context length during training. Theoretically, the SP-Ring
degree can be increased arbitrarily, whereas the TP degree is limited by the
number of heads. Our experimental results on two nodes have not yet
demonstrated this advantage of SP.

#### SP+ZeRO-3

Megatron-LM does not officially support ZeRO-3, possibly because the TP-sp
already reduces memory cost for parameters and gradients, which is also the
target of ZeRO-3. We have claimed that ZeRO-3 is highly compatible with SP in
Tip 2. Given that Megatron-LM already has an implementation of SP-Ring, ZeRO-3
becomes a necessary feature.

#### SP+MoE

As more and more models transition to a Mixture-of-Experts (MoE) architecture,
the research on combining SP with expert parallelism (EP) becomes increasingly
significant. The sequence parallelism employed on the attention module is
decoupled with the FFN module. Therefore, hybrid sequence parallelism can also
be compatible with Mixture of Experts (MoE), as long as the All2All
communication between the Attention and FFN modules is carefully designed.

##  7 Conclusion

In this paper, we propose an approach that unifies DeepSpeed-Ulysses and Ring-
Attention for sequence parallelism. This method encompasses the capabilities
of both techniques, broadening the applicability and delivering superior
performance in some cases. We systematically analyzed the interplay between
sequence parallelism and other established parallelism methods, deriving a set
of best practice suggestions. These suggestions have been validated through
experimental results obtained from two GPU nodes.

##  8 Acknowledgements

We express our gratitude to Zilin Zhu from Tencent. The code utilized in our
research was built from his GitHub repository, and he also owns the authorship
of Figure 3 presented in this paper.

## References

  * [1] Sam Ade Jacobs, Masahiro Tanaka, Chengming Zhang, Minjia Zhang, Leon Song, Samyam Rajbhandari, and Yuxiong He.  Deepspeed ulysses: System optimizations for enabling training of extreme long sequence transformer models.  arXiv preprint arXiv:2309.14509, 2023.
  * [2] Hao Liu, Matei Zaharia, and Pieter Abbeel.  Ring attention with blockwise transformers for near-infinite context.  arXiv preprint arXiv:2310.01889, 2023.
  * [3] Vijay Anand Korthikanti, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, and Bryan Catanzaro.  Reducing activation recomputation in large transformer models.  Proceedings of Machine Learning and Systems, 5, 2023.
  * [4] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro.  Megatron-lm: Training multi-billion parameter language models using model parallelism.  arXiv preprint arXiv:1909.08053, 2019.
  * [5] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He.  Zero-offload: Democratizing billion-scale model training.  In 2021 USENIX Annual Technical Conference (USENIX ATC 21), pages 551–564, 2021.
  * [6] Shenggui Li, Fuzhao Xue, Yongbin Li, and Yang You.  Sequence parallelism: Making 4d parallelism possible.  arXiv preprint arXiv:2105.13120, 2021.
  * [7] Shenggui Li, Hongxin Liu, Zhengda Bian, Jiarui Fang, Haichen Huang, Yuliang Liu, Boxiang Wang, and Yang You.  Colossal-ai: A unified deep learning system for large-scale parallel training.  In Proceedings of the 52nd International Conference on Parallel Processing, pages 766–775, 2023.
  * [8] Dacheng Li, Rulin Shao, Anze Xie, Eric P Xing, Joseph E Gonzalez, Ion Stoica, Xuezhe Ma, and Hao Zhang.  Lightseq: Sequence level parallelism for distributed training of long context transformers.  arXiv preprint arXiv:2310.03294, 2023.
  * [9] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré.  Flashattention: Fast and memory-efficient exact attention with io-awareness.  Advances in Neural Information Processing Systems, 35:16344–16359, 2022.
  * [10] Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai.  Gqa: Training generalized multi-query transformer models from multi-head checkpoints.  arXiv preprint arXiv:2305.13245, 2023.
  * [11] Noam Shazeer.  Fast transformer decoding: One write-head is all you need.  arXiv preprint arXiv:1911.02150, 2019.
  * [12] William Brandon, Aniruddha Nrusimha, Kevin Qian, Zachary Ankner, Tian Jin, Zhiye Song, and Jonathan Ragan-Kelley.  Striped attention: Faster ring attention for causal transformers.  arXiv preprint arXiv:2311.09431, 2023.
  * [13] Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu.  Roformer: Enhanced transformer with rotary position embedding.  Neurocomputing, 568:127063, 2024.
  * [14] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He.  Zero: Memory optimizations toward training trillion parameter models.  In SC20: International Conference for High Performance Computing, Networking, Storage and Analysis, pages 1–16. IEEE, 2020.
  * [15] Jiarui Fang, Zilin Zhu, Shenggui Li, Hui Su, Yang Yu, Jie Zhou, and Yang You.  Parallel training of pre-trained models via chunk-based dynamic memory management.  IEEE Transactions on Parallel and Distributed Systems, 34(1):304–315, 2022.
  * [16] Tianqi Chen, Bing Xu, Chiyuan Zhang, and Carlos Guestrin.  Training deep nets with sublinear memory cost.  arXiv preprint arXiv:1604.06174, 2016.
  * [17] Ziheng Jiang, Haibin Lin, Yinmin Zhong, Qi Huang, Yangrui Chen, Zhi Zhang, Yanghua Peng, Xiang Li, Cong Xie, Shibiao Nong, et al.  Megascale: Scaling large language model training to more than 10,000 gpus.  arXiv preprint arXiv:2402.15627, 2024.
  * [18] Meta AI Team.  Introducing meta llama 3: The most capable openly available llm to date.  https://ai.meta.com/blog/meta-llama-3/, 2024.  Accessed: May 2024.