DSP: Dynamic Sequence Parallelism for Multi-Dimensional Transformers

## Abstract

Abstract Scaling multi-dimensional transformers to long sequences is indispensable across various domains. However, the challenges of large memory requirements and slow speeds of such sequences necessitate sequence parallelism. All existing approaches fall under the category of embedded sequence parallelism, which are limited to shard along a single sequence dimension, thereby introducing significant communication overhead. However, the nature of multi-dimensional transformers involves independent calculations across multiple sequence dimensions. To this end, we propose Dynamic Sequence Parallelism (DSP) as a novel abstraction of sequence parallelism. DSP dynamically switches the parallel dimension among all sequences according to the computation stage with efficient resharding strategy. DSP offers significant reductions in communication costs, adaptability across modules, and ease of implementation with minimal constraints. Experimental evaluations demonstrate DSP’s superiority over state-of-the-art embedded sequence parallelism methods by remarkable throughput improvements ranging from 32.2% to 10x, with at least 50% communication volume reduction.

## 1 Introduction

Efficiently scaling multi-dimensional transformers to accommodate long sequences is necessary across diverse domains, including video generation (Singer et al., ; Blattmann et al., ; Ma et al., ), image generation (Ramesh et al., ; Rombach et al., ; Liu et al., ), protein structure prediction (Jumper et al., ), spatial-temporal information processing (Cong et al., ), and beyond. The long length of sequences in these applications entails substantial activation memory costs and a notable slowdown in processing speeds, underscoring the need for employing sequence parallelism.

Current methods for sequence parallelism, such as Megatron-SP (Korthikanti et al., ), DeepSpeed-Ulysses (Jacobs et al., ), and Ring-Attention (Li et al., ; Liu et al., ) are all embedded sequence parallelism methods. As shown in Figure 1, these embedded methods shard along a single sequence dimension, which are tailored to the task-specific pattern of modules and introduce intricate communication and complex code modification. However, multi-dimensional neural networks often operate independent computations across multiple sequence dimensions. For instance, for video generation models like OpenSora (Zangwei Zheng, ) and Latte (Ma et al., ), Spatial-Temporal Attention (Yan et al., ) is favored over full Self-Attention (Vaswani et al., ), facilitating separate attention computations for temporal and spatial dimensions.

Figure: Figure 1: Comparison of Embedded and Dynamic Sequence Parallelism. Reshard means the communication to change sequence parallel layout. The blue arrow represents communication. The number and width of the arrows indicate the volume and frequency of communication, respectively.
Refer to caption: x1.png

Therefore, there exists an explorable space for a general sequence parallelism abstraction. To adapt to the flexible patterns of multi-dimensional transformers, we introduce Dynamic Sequence Parallelism (DSP) as a novel abstraction of sequence parallelism, characterized by its elegance, high effectiveness, and excellent compatibility with popular transformers. Unlike embedded sequence parallelism, DSP dynamically switches the parallel dimension of sequences during the computation stage with an efficient resharding strategy, a process completely decoupled from the computation module. DSP offers several advantages over embedded sequence parallelism: 1) Efficient Communication: DSP incurs significantly lower communication costs due to its simplified communication patterns and reduced frequency of exchanges. 2) Adaptability: DSP seamlessly adapts to most modules without necessitating specific modifications and imposes few limitations on its usage. 3) Ease of Use: DSP is remarkably easy to implement, and also provides a simple API shown in Appendix A.4 for users to enable it effortlessly.

Our experiments yield promising results, showcasing DSP’s superiority over state-of-the-art embedded sequence parallelism methods. It achieves an end-to-end throughput improvement ranging from 32.2% to 10x and reduces communication volume by at least 75%.

We summarize our contributions as follows:

- •
We introduce DSP as a novel abstraction of sequence parallelism aimed at effectively scaling multi-dimensional transformers. DSP dynamically switches the parallel dimension of sequences during the computation stage, offering high effectiveness, elegant formalism, and excellent compatibility.
- •
By significantly reducing communication volume and frequency, DSP improves end-to-end throughput by 32.2% to 10x and reduces communication volume by at least 50% compared to state-of-the-art methods.
- •
DSP seamlessly integrates with various modifications without requiring specific modifications and imposes few limitations. Its ease of use is highlighted by the minimal code changes needed to incorporate it into existing frameworks with our high-level API.

## 2 Background and Related Work

**Table 1: Meanings of the symbols that are used in this paper.**
| $B$ | The number of batch sizes | $N$ | The number of GPUs |
| --- | --- | --- | --- |
| $C$ | The size of hidden states | $n$ | The $n$-th GPU |
| $S_{i}$ | The $i$-th sequence dimension | $\mathbf{X}$ | The a multi-dimensional sequence |
| $s_{i}$ | The status that sequence is sharded from |  | tensor |
|  | dimension $S_{i}$ | $\mathbf{X}_{p,n}$ | The partition of X assigned to GPU $n$ |
| $\hat{s}$ | The status that sequence is not sharded |  | for sequence parallel |
| $M$ | The volume of a sequence tensor |  |  |

### 2.1 Background

Transformer Architecture.
Transformer (Vaswani et al., ) is a type of neural network architecture that has become highly influential in natural language processing (Devlin et al., ; Brown et al., ; Reid et al., ) and other domains (Dosovitskiy et al., ; Jumper et al., ; Peebles & Xie, ). The Transformer is composed of a stack of layers, each consisting of a multi-head attention (MHA) and a position-wise feed-forward network (FFN). Specifically, the MHA comprises H independently parameterized attention heads, formulated as:

$$ $\mathrm{MHA}(x)=\mathrm{Concat}(\mathrm{head}_{1},\dots,\mathrm{head}_{H}) \mathbf{W}^{O},\quad\mathrm{head}_{i}=\mathrm{Att}(\mathbf{Q}_{i},\mathbf{K}_{ i},\mathbf{V}_{i}),$ (1) $$

$$ $\mathrm{Att}(\mathbf{Q},\mathbf{K},\mathbf{V})=\mathrm{softmax}\left(\frac{ \mathbf{Q}\mathbf{K}^{\top}}{\sqrt{d_{k}}}\right)\mathbf{V},\quad x_{\mathrm{ MHA}}=\mathrm{LayerNorm}(x+\mathrm{MHA}(x)),$ (2) $$

where $\mathrm{Att}(\cdot)$ denotes the scaled dot-product attention, $\mathbf{Q},\mathbf{K},\mathbf{V}$ are query, key, value projections, and LayerNorm is the layer normalization. The output $x_{\mathrm{MHA}}$ is fed into the FFN, which consists of two linear transformations with a ReLU activation in between, computed as:

$$ $\mathrm{FFN}(x)=\max(0,x\mathbf{W}_{1}+\mathbf{b}_{1})\mathbf{W}_{2}+\mathbf{b }_{2},\quad x_{\mathrm{out}}=\mathrm{LayerNorm}(x_{\mathrm{MHA}}+\mathrm{FFN}( x_{\mathrm{MHA}})),$ (3) $$

where $\mathbf{W}_{1},\mathbf{W}_{2},\mathbf{b}_{1},\mathbf{b}_{2}$ are the parameters of the FFN.

Figure: Figure 2: Illustration of Multi-Dimensional (2D) Transformer.
Refer to caption: x2.png

Multi-Dimensional Transformer. Multi-dimensional transformers (Ho et al., ; Yang et al., ) extend the self-attention mechanism of standard transformers to operate over multiple dimensions beyond just one sequence. An example of 2D-Transformer is shown in Figure 2. Let the input multi-dimensional sequence be denoted as $\mathbf{X}\in\mathbb{R}^{[B,S_{1},S_{2},\dots,S_{K},C]}$, where $B$ is the batch size, $S_{1},S_{2},\dots,S_{K}$ are the sequence lengths along $K$ different sequence dimensions, and $C$ is the hidden size. Multi-dimensional transformer can be formatted as:

$$ $\mathbf{X}_{\mathrm{reshape}}=\mathrm{Reshape}(\mathbf{X},[B\times\prod_{j\neq i }S_{j},S_{i},C]).$ (4) $$

The transformer block operation is then applied along the $i$-th sequence dimension of $\mathbf{X}_{\mathrm{reshape}}$.

$$ $\mathbf{X}_{\mathrm{out}}=\mathrm{transformer_block}(\mathbf{X}_{\mathrm{ reshape}}).$ (5) $$

After applying the transformer block operation along all $N$ dimensions, the final output tensor $\mathbf{X}_{\mathrm{out}}$ has the same shape as the input tensor $\mathbf{X}$. Multi-dimensional Transformer is widely used for applications with multi-dimensional inputs including video data (Xu et al., ; He et al., ; Geng et al., ; Ma et al., ), 3D data (Zheng et al., ; Chen et al., ), protein structure prediction (Jumper et al., ; Mirdita et al., ), time-series data (Pan et al., ; Huang et al., ; Deihim et al., ) and beyond.

### 2.2 Related Work

In this section, we discuss the four main parallelism techniques employed in deep learning: data parallelism, tensor parallelism, pipeline parallelism, and sequence parallelism.

Data parallelism (Hillis & Steele Jr, ; Li et al., ) is one of the most widely adopted parallelism techniques. The input data is partitioned across devices, each processing a subset. Model parameters are replicated, and gradients are summed. ZeRO (Rajbhandari et al., ; ) optimizes memory by partitioning parameters, states, and gradients across devices, enabling the training of larger models. Tensor parallelism (Shazeer et al., ; Shoeybi et al., ), or model parallelism, partitions model parameters across devices. Different model parts are assigned to different devices. Pipeline parallelism (Huang et al., ; Narayanan et al., ; Li & Hoefler, ; Liu et al., ) partitions the model into stages executed in parallel across devices. Activations are passed between devices in a pipeline style.

Unlike parameter parallelism discussed earlier, sequence parallelism is a technique specifically designed for distributing long sequences and activation across multiple devices. Here are three main methods of sequence parallelism:

Ring Attention (Li et al., ) employs an innovative approach to partitioning the sequence dimension using a ring-style peer-to-peer (P2P) communication pattern to transfer keys and values across GPUs. (Liu et al., ) enhance this method with an online softmax mechanism, allowing for the computation of attention scores without retaining the full sequence length. However, Ring Attention’s reliance on P2P communication can be less efficient in high-latency environments.
Megatron-SP (Korthikanti et al., ) further optimizes activation usage in the attention based on tensor parallelism. To transit between tensor parallelism and sequence parallelism in the Transformer block, additional all-gather and reduce-scatter operations are introduced. And it’s constrained by the number of attention heads, as self-attention relies on the parallelism of the head dimension of the sequence.
DeepSpeed-Ulysses (Jacobs et al., ) introduces an innovative approach for training long sequences by utilizing all-to-all collective communication. This method partitions the query, key, and value matrices across attention heads while preserving the original attention computation structure. The process is facilitated by two sets of all-to-all communications that alternate between sequence splitting and attention head splitting. Nevertheless, it is constrained by the number of attention heads as well.

Moreover, these sequence parallelism methods are designed for parallelism within a single sequence dimension. For multi-dimensional transformers, this strategy becomes inefficient due to unnecessary communication. While specialized parallelism for multi-dimensional sequences has been explored in specific domains (Cheng et al., ), their applicability remains limited.

## 3 Dynamic Sequence Parallelism

### 3.1 Overview

Figure: Figure 3: System overview of Dynamic Sequence Parallelism.
Refer to caption: x3.png

In the realm of multi-dimensional transformers, computation occurs independently for each sequence dimension. To harness this inherent feature, we introduce Dynamic Sequence Parallelism (DSP), an efficient, adaptive and ease-of-use sequence parallelism abstraction for multi-dimensional transformers.

To ensure correct computation logic with sequence parallelism, embedded methods typically require complex and time-consuming communications within computation modules to change the parallel dimension.
However, as illustrated in Figure 3, the key feature of DSP is its dynamic switch of parallel dimension at the intersections of computation stages. By resharding only between computation stages dynamically, rather than within them, this approach allows DSP to remain independent of the computation logic within the module. Therefore, DSP eliminates numerous unnecessary communications within modules, and is able to utilize efficient all-to-all operations to switch parallelism dimensions for the intermediate sequence.

For operations involving all sequence dimensions, including the beginning and end of the model, DSP handles them by split and gather operation. Furthermore, we also propose a high-level, user-friendly implementation of DSP compatible with all distributed frameworks based on PyTorch, details in Appendix A.4.

### 3.2 Problem Definition

In sequence parallelism, the objective is to distribute activation computations across multiple GPUs to reduce the memory overhead caused by long sequences. This approach, however, incurs additional communication costs between GPUs. Our goal is to optimize this trade-off in the context of multi-dimensional transformers.

Given a multi-dimensional sequence $\mathbf{X}\in\mathbb{R}^{[B,S_{1},S_{2},\dots,S_{K},C]}$ and a set of $N$ GPUs, where $S_{1},\dots,S_{K}$ are the sequence along $K$ different sequence dimensions, we aim to partition the computation such that the memory usage per GPU is under capacity while maintaining acceptable communication costs. Let $\mathbf{X}_{p,n}$ denote the partition of $\mathbf{X}$ assigned to GPU $n$, where $p$ represents the partition strategy. The optimization problem is formulated as:

$$ $\min_{p}\sum_{n=1}^{N}\mathrm{CommCost}(\mathbf{X}_{p,n})\ \ s.t.\ \ \mathrm{ Memory}(\mathbf{X}_{p,n})<Capacity.$ (6) $$

Here, $\mathrm{Memory}(\mathbf{X}{p,n})$ denotes the memory usage of partition $\mathbf{X}_{p,n}$ on GPU $n$, $\mathrm{CommCost}(\mathbf{X}_{p,n})$ represents the communication cost. We aim to achieve a balance that minimizes the overall computational overhead while optimizing GPU resource utilization.

### 3.3 Dynamic Primitives

**Table 2: Definition of Dynamic Primitives for DSP. $s_{i}$ denotes the sequence sharded from dimension $i$, while $\hat{s}$ indicates the sequence is not sharded. $M$ represents the sequence size, and $N$ signifies the sequence parallel size.**
| Source | Target | Primitives | Comm | Comm | Freq |
| --- | --- | --- | --- | --- | --- |
| Shard | Shard | Operation | Volume |  |  |
| $s_{i}$ | $s_{i}$ | / | / | / | / |
| $s_{i}$ | $s_{j}$ | Switch | all-to-all | $M/N$ | High |
| $\hat{s}$ | $s_{i}$ | Split | / | 0 0 | Low |
| $s_{i}$ | $\hat{s}$ | Gather | all-gather | $M$ | Low |

Figure: Figure 4: Illustration of dynamic switch.
Refer to caption: x4.png

In this section, we introduce the key dynamic primitives of DSP, as outlined in Table 2. These three primitives form the cornerstone for implementing DSP across a spectrum of multi-dimensional transformers.

The first condition is that when there is no need to alter sequence parallelism between computation stages, we maintain the shard status of the sequence. This approach significantly reduces unnecessary communication overhead. However, when it becomes necessary to transit parallelism between dimensions, we employ dynamic switch to efficiently transform parallelism. Specifically, as depicted in Figure 4, dynamic switching adjusts the parallelism to a dimension unrelated to the ongoing computation, utilizing highly efficient all-to-all operations.

Assume $\mathbf{X}\in\mathbb{R}^{[B,S_{1},\dots,S_{i}/N,\dots,S_{j},\dots,S_{K},C]}$ represents the input. The current parallel dimension is $i$ so its sequence length is $S_{i}/N$ on each device, where $N$ is sequence parallel size. If we want to switch the shard dimension from $i$ to $j$, the operation can be formulated as follows:

$$ $\mathbf{Y}=\text{DynamicSwitch}(\mathbf{X},i,j),$ (7) $$

where the resulting tensor $\mathbf{Y}$ has the shape $\mathbb{R}^{[B,S_{1},\dots,S_{i},\dots,S_{j}/N,\dots,S_{K},C]}$. Furthermore, Split and Gather operations facilitate smooth transitions between sharded and non-sharded states. Although these operations may involve increased communication compared to Switch operations, they are primarily utilized at the onset and conclusion of most networks, and also for some global operations in very rare conditions, rendering their costs negligible.

### 3.4 Adaptability and Flexibility

Given its decoupling from the computation of modules, DSP exhibits remarkable adaptability, making it compatible with a wide array of transformer variants such as Cross Attention (Hertz et al., ; Ma et al., ); specialized kernels like FlashAttention (Dao et al., ); special attention mechanisms including multi-query attention (Shazeer, ) and grouped-query attention (Ainslie et al., ); and even beyond like Mamba (Gu & Dao, ) and RWKV (Peng et al., ). This inherent flexibility enables DSP to seamlessly integrate into diverse transformers without specific modification. Furthermore, while DeepSpeed-Ulysses and Megatron-LM necessitate attention head splitting, DSP’s scalability is significantly better because it shards on sequence length, which much larger especially when scaling sequences.

Moreover, DSP’s adaptability extends beyond module compatibility to encompass various parallelism methodologies. From conventional data parallelism to more sophisticated approaches like ZeRO and pipeline parallelism, DSP effortlessly integrates with diverse parallel computing paradigms, thereby enhancing scalability and performance across distributed computing environments.

As shown in Appendix A.4, we demonstrate the API usage of DSP. By calling just four functions without knowing the detailed implementation, DSP can be enabled on PyTorch and is compatible with various distributed frameworks, including FSDP (Zhao et al., ), Accelerate (Gugger et al., ), DeepSpeed (Rasley et al., ), and Megatron-LM (Shoeybi et al., ).

## 4 Theoretical Analysis

We choose 2D-Transformer as described in Equation 5 as our base model, which is widely employed in real-world applications. To be specific, we use the OpenSora (Zangwei Zheng, ) variant of 2D-Transformer, an open-source video generation model, where there are two transformer blocks for two sequence dimensions separately. More details can be found in Appendix A.1. We choose DeepSpeed-Ulysses (Jacobs et al., ), Megatron-SP (Korthikanti et al., ) and RingAttention (Liu et al., ) as baselines, which represent the state-of-the-art sequence parallelism methods employed for processing long sequences with transformers.

### 4.1 Communication Analysis

The primary advantage of DSP lies in its ability to minimize communication costs and enable scalable communication operations. DSP exploits the inherent characteristics of multi-dimensional transformers to eliminate unnecessary communication, compared with embedded approaches such as Megatron-SP (Korthikanti et al., ), RingAttention (Liu et al., ) and DeepSpeed-Ulysses (Jacobs et al., ). Consider an activation size of $M$ and a sequence parallel size of $N$. In a 2D-Transformer, there is one transformer block for each sequence dimension per layer, resulting in two transformer blocks per layer. More details of the implementation of each method are demonstrated in Appendix A.2.

Both Megatron-LM and DeepSpeed-Ulysses require transforming sequence parallelism from a sequence-shard to a head-shard layout. Megatron-SP employs 2 all-gather operations to aggregate the entire sequence and 2 reduce-scatter operations to distribute results in the attention and MLP layers for one transformer block. This results in a total of 8 collective communication operations, leading to a total communication volume of $8M$. Conversely, DeepSpeed-Ulysses incurs 4 communication operations in temporal block for the query, key, value, and output layout transformations. Consequently, the communication volume transmitted per device for an all-to-all communication of size $M$ across $N$ GPUs is $4M/N$. Ring-Attention needs to communicate the entire key and value in the temporal as well, resulting in a total communication volume of $2M$.

In contrast, DSP mitigates communication cost by employing only two all-to-all operations in total two blocks per layer. As shown in Table 3, it reduces the communication volume to $2M/N$, significantly outperforming other sequence parallelism techniques. Notably, with merely two all-to-all operations, DSP exhibits efficient scalability even in super-large clusters for training and inference on extremely long sequences because the communication volume decreases as the number of nodes increases, rendering DSP an exceptional choice for large-scale distributed training and inference tasks involving extreme long sequences.

**Table 3: Comparison of DSP with other sequence parallelism methods for 2D-Transformer architectures. $M$ denotes the activation size, and $N$ represents the number of devices. Communication volume refers to the per-layer volume.**
| Method | Communication | Activation | Parameter | Ease |
| --- | --- | --- | --- | --- |
| Volume | Memory | Memory | of Use |  |
| Ring Attention | $2M$          \faThumbsDown | \faThumbsOUp | \faThumbsOUp | \faThumbsDown\faThumbsDown |
| Megatron-SP | $8M$        \faThumbsDown\faThumbsDown | \faThumbsDown | \faThumbsOUp | \faThumbsDown |
| DeepSpeed-Ulysses | $4M/N$   \faThumbsOUp | \faThumbsOUp | \faThumbsOUp | \faThumbsOUp |
| DSP (ours) | $2M/N$ \faThumbsOUp\faThumbsOUp | \faThumbsOUp\faThumbsOUp | \faThumbsOUp | \faThumbsOUp\faThumbsOUp |

### 4.2 Memory Analysis

Regarding activation memory, since we shard every tensor in the transformer, we are theoretically able to achieve the minimum activation cost, similar to DeepSpeed-Ulysses and Ring Attention. In practice, however, our approach requires less shape transformation and communication overhead, allowing us to further reduce intermediate activation memory compared to other methods. Megatron-SP, on the other hand, needs to hold the entire activation after the all-gather operation, resulting in higher memory requirements for hosting the activation.

As for parameter memory, as discussed in Section 3.4, DSP is compatible with most parameter parallelism strategies. In this analysis, we utilize the ZeRO technique (Rajbhandari et al., ) to evenly shard all parameters across different devices. Consequently, our parameter memory footprint can be kept low.

## 5 Experiments

Experiments are conducted on 128 NVIDIA H100 GPUs, interconnected via NVLink within nodes and InfiniBand across nodes. Our methods and implementations are not dependent on specific hardware architectures and can generalize to other devices, particularly those with less efficient interconnects. We follow the same baseline and settings as discussed in Section 4, utilizing 720M and 3B size Transformer-2D models in our experiments. Despite the existence of various 2D-Transformer variants, their architectures are fundamentally similar. Consequently, we select one base model similar to OpenSora (Zangwei Zheng, ), an open-source video generation model, for our study.
The code is implemented using PyTorch (Paszke et al., ).

In the following evaluations, we focus on addressing the following
questions:
1) How is DSP’s end-to-end performance compared with other SOTA sequence parallelism?
2) How is DSP’s scaling ability when scale to many GPUs?
3) What is DSP’s memory consumption like in practice?

### 5.1 End-to-End Performance

In this section, we compare the end-to-end performance of different sequence parallelism methods on 128 NVIDIA H100 GPUs. We use a combination of sequence parallelism and data parallelism, with the sequence parallelism set to the minimum size for each method. We evaluate across different sequence lengths ranging from 0.5 million to 4 million tokens, which are common usages for video generation. Details can be found in Appendix A.3.2. As shown in Figure 5, DSP is able to outperform DeepSpeed-Ulysses by 32% to 75%, and other methods by up to 10x due to its communication efficiency. As the sequence length becomes longer and the sequence parallel size increases, as DSP’s communication volume decreases as the device number increases, our performance’s advantage over the baselines becomes even more pronounced. When scaling from 0.5M to 4M tokens, our FLOPS drops by at most 23%, while other methods experience at least a 40% drop.

Figure: Figure 5: End-to-end performance comparison of different sequence parallel methods combined with data parallel on 128 NVIDIA H100 GPUs. The sequence parallel size is set to minimum for different method.
Refer to caption: x5.png

### 5.2 Scaling Ability

This section evaluates the scaling ability of DSP from two perspectives: weak scaling and strong scaling. Weak scaling refers to scenarios where the computational workload per device remains constant while incrementally increasing the number of devices. This setup is analogous to the training stage, where the goal is to scale longer sequences over more GPUs. Strong scaling, on the other hand, is more challenging as it requires keeping the total computational workload constant while incrementally increasing the number of devices. In this case, the computation becomes more sparse on each device. Strong scaling is often employed when the objective is to infer an input sequence rapidly across many GPUs for low-latency applications. The experiments are divided into intra-node and inter-node evaluations due to the different interconnection conditions. Intra-node experiments leverage NVLink interconnect for communication, while inter-node experiments utilize a combination of NVLink and InfiniBand interconnect. More details can be found in Appendix A.3.3.

Weak Scaling. In the weak scaling experiments, to maintain a consistent computational workload for each GPU, the batch size is linearly increased proportional to the number of GPUs, while the sequence length is fixed. As shown in Figures 6, DSP significantly outperforms other methods by more than 80.7%. Moreover, DSP can scale up to 64 GPUs without being limited by the number of attention heads, unlike DeepSpeed-Ulysses and Megatron-SP. Despite scaling to 64 GPUs, DSP maintains an almost linear throughput increase, with only a 15% performance loss from 8 GPUs to 64 GPUs. Additionally, DSP can achieve super-linear scaling for intra-node setups due to reduced optimizer and gradient computation overhead, and its highly efficient performance within a single node.

Figure: Figure 6: Weak scaling ability evaluation of different methods with sequence parallelism only. “×” denotes out of memory or head. Black boxes represent linear scaling.
Refer to caption: x6.png

Strong Scaling. In the strong scaling experiments, both the batch size and sequence length are fixed. As shown in Figure 7, DSP can maintain linear scalability when scaling up to 8 GPUs for the 720 M model and 4 GPUs for the 3B model, which covers most practical scenarios. To evaluate the extreme performance capabilities of DSP, we further scale up to 64 GPUs with very little workload per device. Although there is an inevitable performance drop, DSP’s throughput remains significantly better than the baselines, outperforming Ulysses by up to 62%. As shown in Figure 9, we also compare the inference latency of all baselines with the same workload. DSP can speed up 29% to 63% compared with baselines.

Figure: Figure 7: Strong scaling ability evaluation of different methods with sequence parallelism only. “×” denotes out of memory or head. Black boxes represent linear scaling.
Refer to caption: x7.png

### 5.3 Memory Consumption

Figure: Figure 8: Inference latency comparison of different sequence parallelism methods.
Refer to caption: x8.png

Figure 9 demonstrates the memory consumption comparison of different baselines in the weak scaling setting. The semi-transparent bar represents the cache memory, while the solid bar represents the allocated memory. The total memory usage is the sum of them. Our approach exhibits the lowest memory usage, scaling efficiently for longer sequences. Furthermore, DSP’s memory usage is compact without excessive cache memory bloat, unlike Ring-Attention and Megatron-SP.

## 6 Conclusion and Discussion

In this work, we introduced Dynamic Sequence Parallelism (DSP), a novel sequence parallel abstraction for effectively scaling multi-dimensional transformers to long sequences. Unlike current embedded sequence parallelism methods that only shard on single sequence dimension and are tailored to specific tasks and patterns, DSP offers a general and elegant solution by dynamically switching the parallel dimension during computation, using an efficient resharding strategy decoupled from the computation module.
The key advantages of DSP are: 1) substantially reduced communication costs, 2) adaptability across modules without specialized modifications, and 3) remarkable ease of implementation enabled by a simple high-level API. Our experiments demonstrated DSP’s superiority, achieving from 32.2% to 10x higher end-to-end throughput and at least 75% lower communication volume compared to state-of-the-art methods.
Its elegance and ease of use make it a promising solution for efficient sequence parallelism across a wide range of applications.

One limitation of this work is that DSP is specifically designed for multi-dimensional transformers and may not adapt well to single-dimensional ones like language models. Additionally, while there are global operations that involve all sequence dimensions, which are rare in transformer, DSP may not be of optimal efficiency.

In the future, DSP could expand its scope beyond transformer architectures to architectures including convolution, recurrent, and graph neural networks to utilize its potential across various tasks. Furthermore, automated optimization techniques could enable DSP to dynamically and autonomously determine the most effective switching strategy based on network analysis, thereby optimizing overall system efficiency and efficacy.

## Acknowledgements

Yang You’s research group is being sponsored by NUS startup grant (Presidential Young Professorship), Singapore MOE Tier-1 grant, ByteDance grant, ARCTIC grant, SMI grant (WBS number: A-8001104-00-00), Alibaba grant, and Google grant for TPU usage.

## Appendix A Appendix

### A.1 Model Details

In the theoretical analyses and evaluation section, we use a Transformer-2D model as our base model, similar to OpenSora (Zangwei Zheng, ). However, it is not exactly OpenSora; we have removed its specific cross-attention module to ensure that the performance can be generalized to other models. Therefore, in each layer, there are only two transformer blocks that process two sequence dimensions separately, as shown in Figure 10. Specifically, the two dimensions are temporal $t$ and spatial $s$ for a sequence. Each dimension is processed by a corresponding transformer block, which is a common strategy in many applications.

### A.2 Parallelism Implementation

Figure: Figure 10: Overview of different sequence parallelism methods for 2D-Transformer.
Refer to caption: x10.png

In Figure 10, we demonstrate the detailed implementation of different sequence parallel methods on 2D-Transformer. The implementation of DeepSpeed-Ulysses (Rasley et al., ) is directly adopted from the official OpenSora (Zangwei Zheng, ) implementation, while Megatron-SP (Korthikanti et al., ) is adopted based on its official implementation. For Ring-Attention (Liu et al., ), we adopt an [unofficial implementation](https://github.com/zhuzilin/ring-flash-attention) for the 2D-Transformer.

Megatron-SP employs four resource-intensive collective communication operations per transformer block. Specifically, it initiates an AllGather operation to aggregate the entire input $x$, succeeded by ReduceScatter operations at the output for both attention and MLP modules, culminating in a total communication volume of $8M$.

DeepSpeed-Ulysses adopts the more efficient AlltoAll approach. It leverages all-to-all for query, key, value to transform their shard dimension before attention, and a all-to-all for output after attention. And it only need to communicate in temporal transformer block. Consequently, the communication volume transmitted per device for an AlltoAll communication of size $M$ across $N$ GPUs is $4M/N$.

Ring-Attention is not shown in the figure because it does not require resharding. We implement sequence communication in the temporal transformer as the time axis is split. In the attention module, it needs to pass the key and value to all other devices, resulting in a total communication volume of $2M$.

DSP applies dynamic switching between stages to switch the parallel dimension, which involves two AlltoAll operations, totaling $2M/N$ communication.

### A.3 Experiment Settings

#### A.3.1 Model Size

In the experiments, we use 720M and 3B size for 2D-Transformer. There specific model settings are shown in Table 4.

**Table 4: Model settings of 720M and 3B 2D-Transformer.**
| Model Name | Layers | Hidden States | Attention Heads | Patch Size |
| --- | --- | --- | --- | --- |
| 720M | 28 | 1152 | 16 | (1, 2, 2) |
| 3B | 36 | 2038 | 32 | (1, 2, 2) |

#### A.3.2 End-to-end Performance

Here is the polished text in a more formal format without using bullet points:
In end-to-end performance experiments, 128 GPUs were utilized for all methods. For each method, the minimum sequence parallel size that would not result in out-of-memory errors was employed to reduce communication overhead, with data parallelism employed for the remaining size. ZeRO-2 was used for all methods except Megatron-SP. The specific parallel size is detailed in Table 5.

The accumulated sequence length ranged from 0.5M to 4M, which appears significantly larger than typical text lengths. However, such lengths are common for multi-dimensional tasks. In this case, we followed the workload of video generation. The spatial sequence, representing video resolution, was fixed at 1024x1024. After applying the Variational Autoencoder (VAE) and Patch Embedding, the final length for the spatial sequence was 4096. The temporal sequence, representing video length, scales linearly in the test.

**Table 5: End-to-end performance parallel settings. Tuple for methods denotes (sequence_parallel_size, data_parallel_size).**
| Model | Sequence | Temporal | Spatial | DeepSpeed | Megatron | Ring | DSP |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Size | Length | Sequence | Sequence | Ulysses | SP | Attention |  |
| 720M | 0.5M | 128 | 4096 | (2, 64) | (2,64) | (2, 46) | (2, 64) |
| 1M | 256 | 4096 | (4, 32) | (4,32) | (4, 32) | (4, 32) |  |
| 2M | 512 | 4096 | (8, 16) | (16,8) | (8, 16) | (8, 16) |  |
| 4M | 1024 | 4096 | (16, 8) | / | (16, 8) | (16, 8) |  |
| 3B | 0.5M | 128 | 4096 | (4, 32) | (4, 32) | (4, 32) | (4, 32) |
| 1M | 256 | 4096 | (8, 16) | (16,8) | (8, 16) | (8, 16) |  |
| 2M | 512 | 4096 | (16, 8) | / | (16, 8) | (16, 8) |  |
| 4M | 1024 | 4096 | (32, 4) | / | (32, 4) | (32, 4) |  |

#### A.3.3 Scaling Ability

**Table 6: Strong scaling experiment settings.**
| Model Size | Type | Batch Size | Temporal | Spatial |
| --- | --- | --- | --- | --- |
| 720M | Intra-Node | 1 | 64 | 4096 |
| Inter-Node | 1 | 256 | 4096 |  |
| 3B | Intra-Node | 1 | 16 | 4096 |
| Inter-Node | 1 | 128 | 4096 |  |

**Table 7: Weak scaling experiment settings.**
| Model Size | Type | GPU Number | Batch Size | Temporal | Spatial |
| --- | --- | --- | --- | --- | --- |
| 720M | Intra-Node | 1 | 1 | 64 | 4096 |
| 2 | 2 | 64 | 4096 |  |  |
| 4 | 4 | 64 | 4096 |  |  |
| 8 | 8 | 64 | 4096 |  |  |
| Inter-Node | 8 | 1 | 256 | 4096 |  |
| 16 | 2 | 256 | 4096 |  |  |
| 32 | 4 | 256 | 4096 |  |  |
| 3B | Intra-Node | 1 | 1 | 16 | 4096 |
| 2 | 2 | 16 | 4096 |  |  |
| 4 | 4 | 16 | 4096 |  |  |
| 8 | 8 | 16 | 4096 |  |  |
| Inter-Node | 8 | 1 | 128 | 4096 |  |
| 16 | 2 | 128 | 4096 |  |  |
| 32 | 4 | 128 | 4096 |  |  |
| 64 | 8 | 128 | 4096 |  |  |

In weak scaling experiments, as shown in Figure 7 we fix the sequence length and linearly increase the batch size, ensuring that the workload on each device remains constant as the number of devices scales. In strong scaling experiments, as shown in Figure 6, we fix both the sequence length and batch size, keeping the total computation constant. For each experiment, we set the sequence length to the maximum for the least GPU case to fully utilize the computational resources. Specifically, we use the same spatial sequence length and adjust the temporal sequence length to its maximum for each test and sequence parallel size is set to GPU number.

### A.4 API

Figure: Figure 11: An example to demonstrate DSP’s API for Pytorch. Developers can use only four functions to enable DSP. $shard$ refers to the sequence parallel dimension.
Refer to caption: x11.png

In the API, as illustrated in Figure 11, we provide four functions for users: $dsp_dataloader$, $split$, $gather$, and $dynamic_switch$. The $dynamic_switch$ function will reshard sequences between the computation stages. Users only need to input the source shard dimension and target shard dimension. The $split$ and $gather$ functions assist users in distributing and retrieving the entire sequence when needed. The $dsp_dataloader$ function helps users process their dataloader for sequence parallelism, which requires the same data within a single sequence parallel group.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: - How does DSP impact training phase? Will it take longer to converge with DSP?
- How does DSP perform on the larger models with more parameters?
- If the experimental environment changes, for example, with less or more GPUs, how will DSP perform over the other works?
- Will DSP perform better or worse for relatively shorter sequence length?
soundness: 3
strengths: + This paper targets an important topic -- the low efficiency of transformers on long sequences, and proposes effective solution that yields significant end-to-end throughput improvement;
+ Well written paper with clear logic, the illustration and presentation are helpful for understanding the paper's idea.
summary: This paper proposes to improve the efficiency of multi-dimensional transformers for long sequences via dynamic sequence parallelism (DSP). Unlike the conventional data parallelism, model parallelism, pipeline parallelism, etc. which are designed only within a single dimension and have limited flexibility, DSP can adaptively switch between dimensions and therefore, can minimize the communication overhead. The evaluation results show that DSP can improve the E2E throughput by 10x and at least 50\% communication. The authors also make DSP a user-friendly API that can be easily integrated into the existing transformer training/inference frameworks.
weaknesses: - The evaluation section can be further elaborated. I assume that authors are evaluating the inference tasks only. If yes, this should be explicitly pointed out. Also, the discussions for DSP on training phase are missing. For example, will DSP affect accuracy or the training convergence cycles?
- Extending the experiments on more configurations could be beneficial. See questions for details.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Please see the weakness section.
soundness: 3
strengths: 1. The paper is well organized and easy to follow.
2. The illustration with the tensor shape helps the understanding of the proposed solution.
3. Applicability to the emerging ND transformer.
4. Easy to use API design.
summary: The paper presents Dynamic Sequence Parallelism (DSP), a novel method to scale multi dimensional transformers with long sequence length efficiently. The proposed solution dynamically switches the parallel computation dimension and minimizes communication costs and enhances adaptability. The experimental results show that communication reduction, and throughput improvement compared to state-of-the-art embedded sequence parallelism methods.
weaknesses: 1. Can you please show the breakdown for communication overhead/actual computation etc?
2. Can you please show the performance with a larger parameter/multi axis model?

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: Why do the authors compare their work only with the techniques that exploit intra-sequence parallelism? Are there no other techniques for parallelising multi-dimensional transformers? Are there no other techniques for exploiting inter-sequence parallelism? A more comprehensive literature review and comparisons would be very useful. If such techniques do not exist, it would be valuable for the authors to explicitly state this and discuss why their approach is novel in this context.
soundness: 1
strengths: The experiments demonstrate that exploiting inter-sequence parallelism instead of intra-sequence parallelism leads to significant performance advantages when parallelising multi-dimensional transformers.
summary: This paper presents a method for parallelising the inference computation of multi-dimensional transformers, i.e., transformers that operate on multi-dimensional data such as video streams. In a Multi-GPU setup, the authors argue that it is sufficient to exploit the inter-sequence (e.g., inter-column and inter-row) parallelism available in multi-dimensional data instead of embarking on more advanced methods that can parallelize the processing of a single sequence across several GPUs  (referred to as embedded sequence parallelism in the paper).  The approach proposed by the authors requires all-to-all communication to reshuffle the data across GPUs after processing each dimension. The main argument is that because there is substantial inter-sequence parallelism in multi-dimensional data, there is no need to exploit intra-sequence parallelism, which leads to fewer communication rounds than the approaches that support intra-sequence parallelism.
weaknesses: 1) The technical contribution is quite limited. The main contribution of the solution presented appears to be insertion of all-to-all communication rounds to re-arrange the data before switching to a different dimension. This mechanism is similar to 2D FFT processing, where transformations are applied row-wise and column-wise with transpositions in between. 
2) The theoretical and the practical comparisons with the methods that focus on intra-sequence parallelism (Megatron-SP, RingAttention, and DeepSpeed-Ulysses) are not fair because:
- These comparison points are not designed to support multi-dimensional data. They are designed to support large single-dimensional sequences. This submission does not necessarily offer better solutions for supporting large single-dimensional sequences, which should be better clarified in the paper. 
- The comparison points focus on both training and inference and involve optimisations for reducing the activation memory usage during training. This submission appears to be covering only the inference part without covering any activation memory minimisation techniques, which should be acknowledged in the paper, perhaps in a limitations or discussions section.
- RingAttention can be efficiently implemented on simple ring-like topologies without requiring all-to-all interconnects. It is not fair to assume an all-to-all interconnect and then compare the communication complexity per link (or per device) with a solution that assumes a ring topology, which uses fewer links. It would be great if the authors could clarify their assumptions about the network topology and justify their choice of comparison methods.
- The authors build on some assumptions made in DeepSpeed-Ulysses, which this reviewer believes to be problematic: there is no significant difference between the complexity of all-reduce and all-to-all operations when the underlying topology is all-to-all. Because this submission makes the same assumptions made by DeepSpeed-Ulysses, the communication complexity comparisons given in Table III of this submission are also questionable. In particular, the authors appear to be reporting the total communication volume for Megatron-SP and RingAttention while reporting per-device communication volume for DeepSpeed-Ulysses and their work (DSP).

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: How does the resharding process impact performance in practice, especially as the number of GPUs or sequence dimensions increases? Quantifying the computational or memory overhead of these operations would clarify DSP’s efficiency as scales grow and reveal any potential trade-offs.

Appendix A.4 briefly introduces the API, but more detailed examples, including guidance for integrating DSP with popular frameworks like TensorFlow or JAX, would be helpful. Could you add further explanation on API usage and potential challenges in integrating DSP across different environments?
soundness: 3
strengths: Dynamic Sequence Parallelism (DSP) is a novel approach that overcomes limitations of existing sequence parallelism by dynamically switching dimensions based on computation stages, enhancing efficiency in multi-dimensional transformers. The technical foundations are solid, with thorough mathematical definitions and a comprehensive communication and memory analysis that demonstrate DSP’s advantages over state-of-the-art methods. The clarity of the paper is bolstered by well-structured figures, tables, and logical organization, making complex concepts accessible. DSP’s adaptability and scalability in high-dimensional transformer applications make it a significant contribution, with broad relevance for scaling long sequences across various transformer architectures and distributed deep learning systems.
summary: This paper introduces Dynamic Sequence Parallelism (DSP), a novel approach to sequence parallelism in multi-dimensional transformers. Unlike existing methods that shard along a single sequence dimension, DSP dynamically switches the parallel dimension based on the computation stage, utilizing an efficient resharding strategy that minimizes communication overhead. DSP’s adaptable and easy-to-implement design allows it to operate across various modules without specialized modifications. Experimental results demonstrate DSP's performance advantages over state-of-the-art methods, achieving throughput improvements of 32.2% to 10x and reducing communication volume by at least 50%. The paper’s contributions include a formal definition of DSP, a comprehensive analysis of its communication and memory efficiency, and experimental validations that highlight its scalability and ease of integration with existing distributed frameworks.
weaknesses: The paper could benefit from broader evaluation across additional transformer architectures, including single-dimensional applications, to strengthen its generalizability claims. A deeper analysis of DSP’s computational overhead, particularly with frequent resharding on large-scale setups, would clarify if its efficiency holds consistently as dimensions and devices scale. Practical limitations, such as handling global operations involving all sequence dimensions, are briefly mentioned but would benefit from expanded discussion or mitigation strategies. More detailed examples of the API and implementation details, especially for integration with popular frameworks, would improve usability for practitioners. Additionally, weak scaling experiments indicate challenges in inter-node communication as GPU counts increase; exploring optimized strategies for inter-node setups could further enhance DSP’s scalability.

---

confidence: 1
contribution: 3
initial_rating: 6
presentation: 2
questions: DSP requires data resharding between stages, which means that the memory layout of data changes dynamically. This process can complicate memory management and may lead to increased memory fragmentation or inefficient memory use. How to solve this case in the proposed appraoch?

Could the authors provide more details on the computational and memory overhead associated with the resharding operation in DSP? How this overhead scales with an increasing number of GPUs or longer sequences?

The paper mentions altering sequence parallelism between computation stages, but it is unclear under what specific circumstances or computations this alteration becomes necessary. Could the authors clarify when and why it is necessary to change the sequence parallelism strategy between different stages of computation?
soundness: 3
strengths: ++ The paper targets important problem for multi-dimensional transformers scalability.

++ The proposed approach effectively reduced communication overhead.
summary: The paper presents Dynamic Sequence Parallelism, a approach for scaling multi-dimensional transformers efficiently. DSP dynamically switches the parallel dimension between computation stages using a resharding strategy, which reduces communication overhead and simplifies implementation. DSP provides substantial performance improvements over existing sequence parallelism methods, such as Megatron-SP and DeepSpeed-Ulysses, achieving higher throughput (32.2% to 10x increase) and reducing communication volume by up to 75%.
weaknesses: -- The overhead associated with resharding in DSP is not clear.

Meta-review: The paper extends the idea of embedded sequence parallelism to multiple sequence dimensions, as seen in video generation models. In particular, the proposed method is motivated by the challenges presented by long sequences.

Reviewers agreed that the paper tackles an emerging problem of interest, and that the paper was clear in its technical presentation. However, discussions with Reviewer EErB revealed that the performance gains were - on video generation models - limited to shorter spatial sequence lengths, and that these shorter sequence lengths were the result of reducing the resolution of the input data; no evidence was provided to show that the accuracy loss was small enough to be acceptable. Moreover, the only multi-dimensional transformer models studied were the aforementioned video generation models. The paper's claims on generalizability to all multi-dimensional transformers (e.g. AlphaFold2, which was discussed in the rebuttal but no results were provided) and long sequence lengths (which was claimed in the abstract) are not fully substantiated by the empirical evidence.

Decision: Reject