WallFacer: Harnessing Multi-dimensional Ring Parallelism for Efficient Long Sequence Model Training

Transformer-based models have become the mainstream approach in both Natural Language Processing (NLP) and Computer Vision (CV). GPT-style models [36] lead the field in NLP tasks, generating tokens in an auto-regressive manner. Meanwhile, the Diffusion Transformer (DiT) [34], which employs pure Transformer blocks within traditional diffusion processes, has gained widespread use over the past year for image and video generation. The key mechanism behind these Transformer-based models is attention[41], which captures the text feature by calculating the attention score between every two single tokens. A standard attention function is given as:

where $Q$, $K$, and $V$ are the query, value, and key of input x with shape ($B$, $N$, $H$) in the case of self-attention, and $d_{k}$ represents the head dimension in commonly used multi-head attention. For convenience, we use the symbolic representations in Table 1. We can observe that in the original attention function, we need to store the intermediate value $QK^{T}$ of size $N^{2}$, which means that the memory capacity requirement and computation amount both grow in quadratic proportion with the sequence length. To solve the memory IO and capacity bottleneck, methods like flash-attention[8] have been proposed. Flash-attention is designed upon the idea of softmax decomposition [18, 29, 35], enabling online softmax computation by storing some extra statistics. Flash-attention is thus capable of segmenting the $Q$, $K$, and $V$ matrices into blocks, allowing the entire computation to be performed within a single kernel for each block sequentially. This strategy effectively circumvents redundant read and write operations to the High Bandwidth Memory (HBM).

However, the sequence length can reach hundreds of thousands, when dealing with multi-round chatting, or high-resolution long video generation. It then becomes necessary to distribute the sequence across multiple GPUs. This distribution helps to reduce both the memory and computation demands on any single device. This strategy is also known as sequence parallelism. We will explore these methods in further detail in the following section.

Sequence parallelism was first conceptualized in a study published in 2021 [23], which introduced the concept as a strategy to distribute computational sequences across multiple processing steps. Presently, Sequence parallelism can be divided into two main categories: attention-head-sharding-based and peer-to-peer-communication-based. The former involves distributing the attention heads of multi-head attention across multiple GPUs, whereas the latter resembles a distributed version of flash-attention, relying on peer-to-peer communication to transfer keys, values, and intermediate statistics.

Through observation, we identify two main drawbacks of Ring Attention. First, the communication overhead is exceedingly high because every GPU in the system must send and receive keys and values for nearly the entire sequence length before completing the attention computation. Second, variations in bandwidth between and within computing nodes can cause communication bottlenecks. As illustrated in Figure 2, the bandwidths between GPUs 3 and 4 and between GPUs 11 and 12 are lower than those between other GPUs in the ring. Despite this, the system is forced to operate in a complete circle, which can result in unnecessary idle times for other GPUs.

Inspired by techniques used in n-body simulations [13], we propose an approach that addresses both challenges simultaneously. In n-body simulations, such as those calculating direct forces between particles (e.g., gravitational forces), researchers discovered that communication overhead can be minimized by introducing additional dimensions of communication [9]. Similarly, we propose grouping GPUs into teams within the ring-style sequence parallelism framework and distributing communication tasks among team members. This approach reduces the overall communication volume and mitigates the bandwidth bottlenecks previously mentioned. Building upon this concept, we developed the WallFacer training system, which we will detail in the following section.

The WallFacer training system is comprised of five main components. At the core of the system is WallFacer Attention, which leverages multiple ring-style P2P (peer-to-peer) communication strategies to enhance the efficiency of distributed attention computation. The Communication Configuration Generator plays a critical role in initially assigning keys and values to their corresponding ranks. Meanwhile, the Communication Topology Scheduler outlines the placement of parallelism across various computing nodes. The Dataloader is designed to organize tokens within each sub-sequence according to their mask (causal or full) and distribute them across different GPUs. Finally, the WallFacer Runtime incorporates additional supporting techniques for the training process, such as gradient checkpointing. We will now explore the details of these components in depth.

As discussed in the previous section, a major limitation of Ring Attention is the extensive amount of peer-to-peer (P2P) communication required, which becomes problematic in environments with weak connections between computing nodes. We enhance the ring sequence parallelism by introducing an additional dimension. This is achieved by duplicating the queries, keys, and values, thus dividing the communication tasks and distributing them within each team. This segmentation markedly improves computational efficiency and scalability.

We will now go through the specifics of the WallFacer Parallelism and provide a theoretical analysis demonstrating its advantages over Ring Attention.

The Communication Configuration Generator, as introduced in the previous section, plays a crucial role in the initial setup of queries, keys, and values on GPUs and determining each GPU’s position within the sub-rings.

Initially, for the setup stage, it is essential to establish the sub-rings by rearranging the activation positions. Specifically, during forward propagation, the queries do not require rearrangement; however, the keys and values must be transmitted to their corresponding positions in the ring prior to commencing the loop. As illustrated in Figure 5, this initialization ensures that each team member holds a different shard of keys and values. Moreover, it guarantees that no two teams within the same ring possess identical keys and values.

Furthermore, Team 0 and Team 1 form what is termed a team group, where ring P2P communication occurs exclusively within this group, similarly for Team 2 and Team 3. A team group is defined as a collection of teams that participate in the same sub-rings. The number of teams within a team group can be calculated as $Team\ size=\frac{P}{C^{2}}$. We give the details of the determination of initial sending targets in the case of placing teams within the same computing nodes in Algorithm 2, and the receiving source can also be calculated similarly.

After the initialization of activations, we can set up the rings by providing the GPUs their last and next GPU within their rings, as is described in Algorithm 3

The communication topology scheduler employs a grid-search algorithm to discover the optimal parallelism configuration, tailored to the specifics of the current cluster, model, and input data.

When using Ring Attention, all GPUs form a single ring, which considerably limits the flexibility of adjusting the placement scheme. In contrast, WallFacer expands the tuning space with its unique multi-ring structure in two significant ways. First, the ring configuration in WallFacer is determined by the parallelism dimension $C\in[1,\sqrt{P}]$. Smaller $C$ results in less collective communication, a higher total communication volume, and more inter-node communication, while higher C leads to more communication volume during all-gather and reduce-scatter, which may exceed the overhead of the QKV matrix multiplication that we use for overlapping, but can highly reduce the overall communication volume.

Second, the additional dimension provided by WallFacer allows for two strategies in parallelism placement. One strategy is to keep the all-gather and reduce-scatter operations intra-node, which confines these processes within the same computing node. The other strategy is to keep the ring P2P communication intra-node, reducing latency and potentially enhancing performance. The choice between these strategies depends on the extent of overlapping achievable within the ring and the communication volume during collective communications.

Based on the above tuning space, we provide an auto-scheduler, grid-searching through all possible configurations for the one that provides the highest throughput, which can be described as:

The scheduler requires only a few iterations to profile the performance of automatically generated configurations. The time spent on this profiling is negligible compared to the entire training process, as it only needs to be performed once for each new computing cluster used for training.

To accommodate both full mask and causal mask configurations for self-attention, we have implemented two distinct data loading schemes for load-balancing. For full masks, sequences are straightforwardly divided into sub-sequences of equal lengths and distributed among the devices. However, causal masks present a challenge due to the unbalanced computational load across GPUs; sub-sequences at the beginning of a sequence require significantly more computation than those at the end. To address this imbalance and achieve load equilibrium among GPUs, we modify the ZigZag scheme introduced by [45], illustrated in Figure 6. The figure illustrates the simplest case of zigzag load-balancing. Notably, the effectiveness of this strategy improves as the number of GPUs increases. This improvement correlates with the expanding difference in computation volume between the first and the last token, which escalates as the sequence length extends. This approach ensures that the total workload on each GPU is balanced, eliminating the need for additional communication mechanisms like those employed in DistFlashAttention[22].

WallFacer is written in PyTorch[32] and uses the PyTorch torch.autograd.function and NCCL[31] backend for forward and backward implementation. WallFacer also employs multiple techniques during runtime to improve its overall training efficiency.

Ingetrate Flash Attention. The WallFacer attention mechanism involves multiple iterations that loop over Keys and Values, with each iteration still using traditional self-attention with corresponding Query, Key, and Value (QKV). This approach enables WallFacer to incorporate flash attention effectively, extending its capability by preserving intermediate states across iterations. Additionally, WallFacer enhances the efficiency of the forward process with the help of torch JIT to fuse kernels aside from flash attention.

Overlap communication with computing. In WallFacer attention, P2P communication and self-attention computing are interleaved across iterations, each incurring considerable time. To mitigate this, WallFacer employs a double buffering technique to asynchronously execute communication and computing kernels, effectively overlapping these processes and enhancing GPU utilization.

Save recomputation with checkpoints. WallFacer adopts the checkpointing strategy introduced by DistFlashAttn[22], placing checkpoints at the end of the self-attention phase rather than the FFN of each transformer layer. This checkpoint placement effectively obviates the need to recompute the self-attention forward process during the backward pass, avoiding redundant attention computation.

Our first experiment aims to assess the performance of WallFacer and Ring Attention across different clusters with varying environments, testing the adaptability of both methods. There are several factors influencing the efficiency of ring-style attention computation:

Theoretical Computation-Communication Volume Ratio: Primarily determined by the sequence length used during training. Attention computation exhibits a computational complexity of $O(N^{2}\cdot H)$, whereas P2P communication complexity is $O(N\cdot H)$. Thus, the model configuration does not impact this ratio; only the sequence length does. A larger $N$ increases the computation-communication ratio, facilitating easier overlap of communication with computation. We evaluate varying sequence lengths to explore performance under different ratios.

Compute Capability and Connectivity of GPUs: The computing overhead, given a specific volume, affects the computation-communication overhead ratio. Higher compute capabilities make overlapping more challenging. We utilize two sets of GPUs in this evaluation: Nvidia A100 40GB and Nvidia H100 80GB, with the latter offering significantly higher theoretical tflops on bf16 computations. Connectivity is considered in two parts: intra-node and inter-node. Our clusters are equipped with NVLink, ensuring robust intra-node communication. For inter-node communication, while InfiniBand offers high bandwidth at a higher commercial cost, some clusters utilize Ethernet for connectivity. Specifically, our H100 nodes leverage InfiniBand with eight adapters per node for superior inter-node bandwidth, whereas the Google Cloud servers use Ethernet. The diversity in node configurations (8-GPU and 16-GPU nodes) allows us to assess adaptability across different topologies.

This evaluation not only highlights the inherent differences between the schemes but also tests their flexibility in various hardware settings.

The results of our evaluation are illustrated in Figure 7. We measure throughput in thousands of tokens per second. To better demonstrate how to select the optimal configuration of WallFacer under each condition, we included two configurations, Wall-2 and Wall-4, in the figure. We omit configurations with lower performance for clarity. As indicated in the figure, in all six settings, at least one configuration of WallFacer achieves higher throughput than Ring Attention, with performance improvements of 114.33%, 41.4%, 62.87%, 42.45%, 35.98%, 19.9%, 77.12%, and 34.62%. This advantage is primarily due to the additional parallel dimension that WallFacer introduces. Unlike Ring Attention, which requires inter-node P2P communication in each iteration, WallFacer’s P2P communication is mostly confined intra-node, except for initial data transfers. This experiment clearly demonstrates WallFacer’s superior performance across various environments.

Another observation is that the optimal configuration for WallFacer may vary depending on the environment, reflecting differences in the computation-communication ratio and the trade-offs between collective and P2P communication. For example, the best $C$ value for A100_16 is 2, while for A100_8, it is 4. This variation can be attributed to the topological differences: the 16-GPU-node cluster benefits less from reduced P2P communication volume, making a smaller $C$ preferable to decrease collective communication volume while maintaining P2P communication at an acceptable level. This finding underscores the importance of the Communication Topology Scheduler we provide, which helps users avoid the complex process of manually analyzing these factors.

In summary, WallFacer demonstrates superior efficiency and adaptability in all six cases, thanks to its flexible parallelism scheme.

Before evaluating the memory consumption of WallFacer, we first compared the supported sequence lengths of WallFacer with those reported in the RingAttention paper [26]. As shown in Table 4, despite the slightly higher memory requirements of WallFacer, it still supports sequence lengths commonly used in training tasks.

Memory consumption in our experiments stems from both the model weights and activations, with additional costs for WallFacer arising from the duplication of QKV matrices prior to attention computation. To quantitatively assess the additional memory required by WallFacer, we monitored the maximum memory allocated by PyTorch [32] during our experiments. It is important to note that memory fragmentation in PyTorch can impact memory allocation efficiency. To mitigate this, we limited the sequence lengths in our training to prevent PyTorch from triggering cuda_free when the allocated memory approaches the GPU’s limit, which would otherwise introduce significant overhead.

The results, displayed in Figure 8, reveal that for the configurations yielding the highest throughput, Ring Attention consumes between 7.9% and 30.79% less GPU memory than WallFacer, except the DiT 512k Wall-4 setting. However, considering the substantial throughput gains provided by WallFacer, the additional memory usage is deemed acceptable. We observed that the maximum supported sequence lengths were not significantly impacted by this additional memory usage.

Moreover, in scenarios involving larger models, the relative increase in memory consumption due to QKV duplication diminishes. This reduction is explained by equation 7, which shows that as the model size increases—owing to more parameters and additional Transformer layers—the proportion of memory consumed by model weights and activations grows, whereas the absolute extra memory incurred by QKV multiplication remains constant. This phenomenon is further evidenced by the fact that the extra memory ratio for experiments with the 7B model is significantly smaller than that for the 3B and 1B model.

In the scaling tests we carry out experiments for both strong and weak scaling. Strong Scaling maintains the scale of the problem that we are trying to solve while increasing the computing resource that we use to speed it up. So in this experiment, we fix the sequence length to 128K while increasing the number of GPUs from 8 to 64 for the GPT model and from 8 to 32 for the DiT model. As id depicted in Figure LABEL:Fig.strong, WallFacer shows gradually more advantage over Ring Attention as we increase the number of GPUs. When we use 64 GPUs, WallFacer shows 60.71% higher throughput than Ring Attention for the GPT model and 114.33% for the DiT model on 32 A100 GPUs. This can also explained by the computation-communication ratio. When scaled to more GPUs, the local sequence length on each GPU becomes smaller, and as explained the previous sections, makes it harder to overlap the P2P communication with attention computation. This is exactly the main advantage of WallFacer, reduction of total P2P communication volume. So WallFacer can give a promising performance when we are trying to speedup the training of a fixed amount of data with more GPUs.

In our weak scaling experiments using DiT models, we scale the sequence length from 128k to 512k while proportionally increasing the number of GPUs from 8 to 32. Assuming a scaling factor of $k$, both the sequence length and the number of GPUs increase by $k$, resulting in a total computation volume that increases by approximately $k^{2}$. Since the computational resources also increase by $k$, the theoretical overall throughput (tokens per second) should remain constant during this weak scaling setting. As shown in Figure LABEL:Fig.weak, both WallFacer and Ring Attention exhibit this expected behavior, with throughput remaining relatively constant. However, WallFacer consistently achieves higher throughput across different numbers of GPUs. This is because, although the local computation-to-communication ratio stays the same in weak scaling, Ring Attention requires a higher proportion of inter-node communication as the number of GPUs increases, whereas WallFacer keeps most communication local.

In summary, WallFacer shows better scalability in both strong and weak scaling experiments, making it a better choice for large-scale Transformer model training.

In addition to the original FlashAttention [8] used in our experiments, FlashAttention3 [38] has been introduced, specifically designed for Hopper and newer Nvidia GPUs. For FP16 precision, which is utilized in this paper, FlashAttention3 achieves a 1.5-2.0x speedup on Hopper GPUs. As discussed in Section 2, reducing attention computation overhead results in more P2P communication not being overlapped, further emphasizing the need to reduce communication volume. With the increasing adoption of Hopper GPUs, the significance of the WallFacer system will also grow.

Model Parallelism As is well known, tensor parallelism shards activations by attention heads during attention computation, making it easily combinable with WallFacer with minimal effort. However, when combined with tensor parallelism, the need for attention heads can limit the scalability of head-based sequence parallel methods like Ulysses. Pipeline parallelism, on the other hand, divides the model across layers without altering the computation patterns within Transformer blocks, making WallFacer orthogonal to it.

Other Sequence Parallelism WallFacer is orthogonal with other attention-head-sharding-based sequence parallelism approaches, such as DeepSpeed Ulysses [15]. While Ulysses distributes attention heads across different devices, WallFacer can independently partition activations along the sequence length dimension. In future work, we can explore combining WallFacer with DeepSpeed Ulysses to expand the communication scheduling space, harnessing the scalability of WallFacer alongside the efficiency of Ulysses.

In summary, WallFacer can be seamlessly integrated with other parallel training techniques, enabling the creation of a hybrid distributed training system.