Title: Blockwise Parallel Transformers for Large Context Models
Venue: NeurIPS 2023
Accepted: yes
Keywords: Language Model

Long Context Modeling

Reinforcement Learning
Abstract: Transformers have emerged as the cornerstone of state-of-the-art natural language processing models, showcasing exceptional performance across a wide range of AI applications. However, the memory demands posed by the self-attention mechanism and the large feedforward network in Transformers limit their ability to handle long sequences, thereby creating challenges for tasks involving multiple long sequences or long-term dependencies. We present a distinct approach, Blockwise Parallel Transformer (BPT), that leverages blockwise computation of self-attention and feedforward network fusion to minimize memory costs. By processing longer input sequences while maintaining memory efficiency, BPT enables training sequences 32 times longer than vanilla Transformers and up to 4 times longer than previous memory-efficient methods. Extensive experiments on language modeling and reinforcement learning tasks demonstrate the effectiveness of BPT in reducing memory requirements and improving performance.

Blockwise Parallel Transformers for Large Context Models

Abstract

Transformers have emerged as the cornerstone of state-of-the-art natural language processing models, showcasing exceptional performance across a wide range of AI applications. However, the memory demands posed by the self-attention mechanism and the large feedforward network in Transformers limit their ability to handle long sequences, thereby creating challenges for tasks involving multiple long sequences or long-term dependencies. We present a distinct approach, Blockwise Parallel transformers (BPT), that leverages blockwise computation of self-attention and feedforward network fusion to minimize memory costs. By processing longer input sequences while maintaining memory efficiency, BPT enables training sequences 32 times longer than vanilla Transformers and up to 4 times longer than previous memory-efficient methods. Extensive experiments on language modeling and reinforcement learning tasks demonstrate the effectiveness of BPT in reducing memory requirements and improving performance.

1 Introduction

Transformers [54] have become the backbone of many state-of-the-art natural language processing models [15, 45, 5, 37]. They have demonstrated impressive performance across a wide range of AI problems, including language modeling, machine translation, image captioning, and protein folding [41, 49, 32, 45, 5, 47, 9]. Transformers achieve this success through their architecture design that uses self-attention and position-wise feedforward mechanisms. These components facilitate the efficient capture of long-range dependencies between input tokens, enabling scalability in terms of context length and model size through highly parallel computations.

However, the memory requirements of Transformers limit their ability to handle long sequences, which is necessary for many AI problems, such as high-resolution images, podcasts, code, or books and especially those that involve multiple long sequences or long-term dependencies [10, 7, 41, 7, 36, 29, 49, 32, 1]. The quadratic self-attention and the large feed forward network of Transformers require a large amount of memory, which makes it challenging to scale to longer input sequences. This limitation has led to various techniques proposed to reduce the memory requirements of Transformers, including sparse-approximation, low-rank approximation, and low precision approximation [see e.g. 53, 24, 22, 11, 25, 38, 56].

One distinct line of research does not rely on approximation but instead focuses on computing exact self-attention with linear memory complexity. This approach leverages the observation that the softmax matrix in self-attention can be computed without materializing the full matrix [39]. This technique has led to the development of FlashAttention [14] and Memory Efficient Attention [44]. Both methods propose a blockwise computation of the self-attention softmax, demonstrating reduced memory requirements.

VanillaAttentionFlashAttention/MemoryEfficientAttentionBlockWiseParallelTransformer

150000

150000

40000

30000

100000

100000

20000

50000

50000

10000

1B

3B

3B

7B

13B

30B

13B

30B

70B

1A100

8 A100

64TPUv4

Despite the resulting reduced memory requirements of the self-attention block in transformers models, a significant challenge still arises from the feedforward network. This network contains a large number of parameters and produces high-dimensional intermediate vectors, resulting in substantial memory requirements. This issue is becomes the key memory challenge once employing memory-efficient attention mechanisms. Consequently, training Transformers on longer context lengths and scaling up transformers models become significantly hindered due to the overwhelming memory demands imposed by the feedforward network.

To address this challenge, we make an important observation: when self-attention is computed in a blockwise manner to reduce memory requirements, it becomes feasible to merge the computation of the feedforward network. This eliminates the need to wait for the self-attention computation to finish before performing the feedforward step on the entire sequence. By computing the feedforward network on a block-by-block basis, we effectively reduce the memory cost associated with the feedforward network. This process involves the utilization of two nested loops over the input sequence blocks. In the outer loop, we iterate over each block and compute the query. In the inner loop, we iterate over each block to calculate the key and value. These key-value pairs, along with the query, are then used to compute the blockwise attention specific to the corresponding input block. This blockwise attention is subsequently used to calculate the output of the feedforward network, followed by a residual connection. This approach enables us to process longer input sequences while maintaining lower memory budget. Since our approach performs blockwise parallel computation and fuses the feedforward and self-attention computations, we name our method the Blockwise Parallel transformers (BPT).

We evaluate the effectiveness of our approach on several benchmarks, including language modeling and reinforcement learning. Our experiments show that BPT can reduce the memory requirements of Transformers, enabling us to train 32 times longer sequence than vanilla attention [54] based GPT models and up to 4 times longer sequence than prior state-of-the-arts FlashAttention [14] and Memory Efficient Attention [44]. Furthermore, we demonstrate the application of BPT on the task of traning transformers based RL agent. By conditioning on multiple trajectories, BPT significantly improves the performance and achieves better results on challenging RL benchmarks. We believe that our approach has the potential to enable the training and evaluation of more complex models that require longer input sequences, which could lead to further breakthroughs in AI research.

Our contributions are twofold: (a) proposing a blockwise computation of self-attention and feedforward approach that enables 32 times longer and up to 4 times longer context lengths than vanilla transformers and previous memory-efficient Transformers, respectively, and (b) demonstrating the effectiveness of our approach through extensive experiments.

2 Memory Bottleneck of Transformer

Given input sequences Q,K,V ∈ R s × d where s is the sequence length and d is the head dimension. We compute the matrix of outputs as:

where softmax is applied row-wise. Standard attention implementations materialize the matrices QK T and softmax( QK T √ d ) to HBM, which takes O ( s 2 ) memory, so the overall space complexity is O ( s 2 ) . There has been a large body of work trying to reduce memory usage of self-attention by using online softmax [39, 44, 14] to reduce memory cost of self-attention by preventing it from full materialization. And these approaches reduce memory footprint from O ( s 2 ) to O ( s ) . However, the large feedforward layers have been overlooked.

In addition to attention sub-layers, each of the attention layers is accomplished with a fully connected feedforward network, which is applied to each position separately and identically. This consists of two linear transformations with a ReLU activation in between.

While the linear transformations are the same across different positions, they use different parameters from layer to layer. The large size of the feedforward network requires substantial memory resources, and this becomes even more pronounced when dealing with large context sizes. See Section 3.1 for analysis of memory cost associated with transformers.

3 Blockwise Parallel for Large Context Models

Self-attention can be computed in a blockwise manner without materializing the softmax attention matrix softmax( QK T ) [39, 14, 44]. This approach involves splitting the sequences Q ∈ R s × d into B q blocks and sequences K,V ∈ R s × d into B kv blocks. For each query block, the blockwise attention Attention( Q,K,V ) can be computed by iterating over all key-value blocks. Once the blockwise attention is computed, the global attention matrix can be obtained by scaling the blockwise attention using the difference between the blockwise and global softmax normalization constants [39]. This is achieved by keeping track of normalization statistics and combining them from all blocks to scale each block accordingly. For a specific query block Q i , 1 ≤ i ≤ B q , the corresponding attention output can be computed by scaling each blockwise attention as follows:

The scaling operation scales each blockwise attention based on the difference between the blockwise maximum and the global maximum:

This blockwise self-attention computation eliminates the need to materialize the full attention matrix of size O ( n 2 ) , resulting in significant memory savings.

We observe that the blockwise computation is not limited to self-attention but can also be applied to the feedforward network. For each query block, after iterating over the key and value blocks, the feedforward network can be computed along with a residual connection, completing the attention and feedforward network computation for that query block. This means that the model does not need to compute the feedforward network on the full sequence, but rather on intermediate blocks, resulting in memory savings. The computation for a query block is given by:

Therefore, the output for each block consists of the feedforward network, self-attention, and residual connection computed in a blockwise manner.

It is worth mentioning that for large models, the memory cost of the feedforward network on the full sequence can be much larger than the memory efficient attention. Therefore computing the feedforward network on the same block as attention can significantly reduce memory cost, and it also reduces data movements, contributing to overall computational efficiency. Moreover, we should remark that blockwise parallelism can be directly applied to the final cross entropy loss, which can further minimize memory cost. The full process of our framework, coined as BPT, is summarized in Algorithm 1.

Softmax

Linear

Feed

Forward

Q x

Nx

Masked

Attention

KV x

Projection

Input

3.1 Analysis of Memory Cost

We present an analysis of memory cost across different transformers architectures: the Vanilla Transformer, the memory-efficient / Flash Attention variant, and BPT.

Vanilla Transformers:

Attention: For Q , K , V , saving their input x needs 2 bsh bytes, where b is batch size, s is sequence length, and h is hidden dimension. For QK T matmul, saving activations Q and K needs 4 bsh bytes. For softmax ( QK T ) , saving input QK T needs 2 bs 2 a bytes, where a is the number of attention heads. For mask and dropout, saving mask needs bs 2 a bytes. For score × V , saving score needs 2 bs 2 a bytes, and saving V needs 2 bsh bytes. For output projection and dropout, saving the input needs 2 bsh bytes, and saving dropout mask needs bsh bytes. The maximum attention activation size of attention is O ( s 2 ) with checkpointing.

Algorithm 1 Reduce memory cost with BPT.

Required: Input sequence x . Number of query blocks B q . Number of key and value blocks B kv . Initialize

Project input sequence x into query, key and value.

Split query sequence into B q of query input blocks.

Split key and value sequences into B kv of key-value input blocks.

for outer = 1 to B q do

Choose the outer -th query.

for inner = 1 to B kv do

Choose the inner -th key and inner -th value block.

Compute attention using query, key and value, and record normalization statistics.

end for

Combine each blocks by scaling them to get attention output for the outer -th input block.

Compute feedforward on attention output and add residual connection.

end for

FFN: For the first linear layer, saving input needs 2 bsh bytes. For activation, saving input needs 8 bsh bytes. For the second linear layer, saving input needs 8 bsh bytes. For dropout, saving the mask needs bsh bytes. With checkpointing, the maximum activation size of FFN is 8 bsh bytes.

Consequently, for a large context length, the memory cost of activation in vanilla transformers is O ( s 2 ) .

BPT:

Attention: Since BPT does not materialize full attention and instead computes it blockwise, it needs to store intermediate blockwise activations in the key-value loop, which has a maximum activation size of 4 bch with checkpointing. Additionally, it needs to store q output activations for the query loop, which requires 2 bsh bytes. Since s ≫ c , the maximum activation size is 2 bsh .

FFN: When iterating the FFN over blocks, BPT needs to save the following activations: For the first linear layer, saving input needs 2 bch bytes. For activation, saving input needs 8 bch bytes. For the second linear layer, saving input needs 8 bch bytes. For dropout, saving the mask needs bch bytes. In total, 19 bch bytes are needed. Additionally, storing the output of the for loop requires 2 bsh bytes. Therefore, the maximum FFN activation size is 2 bsh .

Consequently, each BPT layer's memory cost of activation is 2 bsh .

Memory-Efficient / Flash Attention:

Attention: Similar to BPT attention, the maximum activation size is 2 bsh .

FFN: Similar to the vanilla FFN, the maximum activation size is 8 bsh .

Consequently, each Flash Attention layer's memory cost is 8 bsh .

Comparing the activation memory of Flash Attention/Memory-Efficient transformers with BPT, we see that BPT offers 8 bsh/ 2 bsh = 4 times memory saving. By taking into account other factors of memory cost such as model parameters and optimizer states, BPT allows training with context lengths 2-4 times larger than prior state-of-the-arts.

3.2 Why Blockwise Parallel

The utilization of blockwise parallelization may raise questions about the effectiveness of running parallel computers, as computation can become sequential between blocks. However, the benefits of blockwise parallelization depend on the model size and hardware configuration. In cases where the model is large or the context length is extremely long, a block may reach its maximum arithmetic density, making it impractical to execute the original full-length sequence in parallel. In such scenarios, blockwise parallelization treats the long sequence as short ones, allowing dealing with large models and effectively enabling large context size. Moreover, using blockwise parallelization allows us to avoid waiting for the completion of self-attention and allocating a significant amount of memory solely for feed-forward network computation.

Another notable advantage of blockwise parallelization is its ability to leverage hardware with significantly faster SRAM speed compared to HBM speed. For instance, in Nvidia GPUs, SRAM is an order of magnitude faster than HBM, while in Google TPUs, SRAM also offers higher speed than HBM. By utilizing blockwise parallelization, we can tap into the increased speed of SRAM, thereby reducing communication costs and increasing throughput. This advantage aligns with memory efficient self-attention approaches [14, 44].

3.3 Implementation

Algorithm 1 provides the pseudocode of the algorithm. Figure 3 in Appendix shows a Jax implementation optimized for simplicity. The full code of BPT is provided at GitHub 1 which supports large-scale distributed training of large context models using BPT.

The blockwise_ffn function begins by accepting a rematerialized feed forward module, inputs and chunk size. The remat_ffn compute feedforward on inputs with checkpointing, i . e .without saving intermediates. The scan_ffn function is then used to scan over input sequences and generate outputs.

The blockwise_attn function process query, key, and value to produce attention blockwise. The scan_attention function is defined, which computes the attention weights between the query vector and key-value pairs from another chunk. This is done by applying the scan_kv_block function to the key-value chunk, calculating the dot product between the query and key vectors, and then adding a bias term. The bias term introduces a positional bias between different chunks based on their indices without materializing the full matrix. The softmax function is then applied to the attention weights in a numerically stable manner, using the max-score trick to avoid large exponentiation results.

Finally, BPT combines the outputs from all chunks, normalizes them using their max-score-adjusted weights, and passes them through a feed-forward neural network ( blockwise_ffn ). The final output is the sum of the feed-forward output, the attention output, and the original input.

4 Setting

We evaluate the impact of using BPT in improving large transformers models by benchmarking memory requirement, maximum sequence length and throughout speed. We show apply BPT to reinforcement learning as an application.

Model Configuration. Our study is built upon the GPT architecture. Table 1 provides a overview of the model sizes considered in our experiments.

Baselines. We evaluate our method by comparing it with vanilla Transformer [54] which is denoted as 'Vanilla', and FlashAttention [14] and Memory Efficient Attention [44] which are state-of-the-art memory efficient attention, we denote them as 'MemoryEfficient' in our experiments. All methods use the same gradient checkpointing in the experiments.

Datasets. We consider two datasets for evaluation purposes. Including pretraining on OpenWebText dataset and large context reinforcement learning on ExoRL.

OpenWebText. The OpenWebText dataset [18] is a large and diverse collection of web pages that has been filtered and cleaned for use in natural language processing (NLP) tasks. The dataset consists of over 6 billion tokens from more than 40 million web pages, covering a wide range of topics and genres.

ExoRL. The ExoRL [58] dataset is based on unlabeled exploratory data collected by running unsupervised RL algorithms. For each environment, it comes with eight different unsupervised data collection algorithms, taken from from URLB [28]. The datasets are collected by unsupervised RL and then relabeled using task reward function. The resulting mixed dataset consists of 8 millions timesteps (8000 episodes), with each episode spanning a length of 1000 steps.

Training Configuration. Our main baselines are vanilla attention [54], which computes self-attention by materializing the attention matrix and computes the feedforward network normally. We also

consider two prior state-of-the-art memory-efficient methods, namely FlashAttention [14], which focuses on GPU efficiency, and Memory Efficient Attention [44], which focuses on TPU efficiency. Since they share a similar idea, for notation simplicity, we refer to them as FlashAttention in our experiments. We tune the block size for both the baselines and BPT, and report the best results achieved by each. The experiments are on NVIDIA 80GB A100 GPUs, we consider both single GPU for smaller model training and 8 GPUs settings for model parallel training. We also experiment with scaling up model on 64 TPUv4.

We note that no data parallelism is considered in our evaluations since our approach is independent of data parallelism. As a result, the batch sizes used in our analysis are much lower than the ones used for the end-to-end training. All of our results are obtained using full precision instead of mixed precision.

5 Results

In our experiments, our primary objective is to comprehensively evaluate the performance of BPT across multiple key metrics, including maximum sequence length, memory usage, and throughput. Moreover, we extend the applicability of BPT to reinforcement learning and evaluate its effectiveness in large context application.

5.1 Evaluation of Context Length

We present experimental results comparing the maximum training sequence lengths achieved using three different attention mechanisms: Vanilla, MemoryEfficient, and Blockwise Parallel. Table 2 summarizes the findings. On one A100 GPU, Vanilla transformers supports a maximum training sequence length of 16K for 1B parameters and 8K for 3B parameters. In contrast, MemoryEfficient enables longer sequences of 65K for 1B parameters and 16K for 3B parameters. Notably, our proposed method, Blockwise Parallel, surpasses both methods, achieving a maximum sequence

length of 131K for 1B parameters and 3B parameters. Moving on larger models, Blockwise Parallel again outperforms the other two methods, allowing training sequences of 65K for 30B large model on 8 GPUs and 8K for 70B large model on 64 TPUv4, which are two and four times longer than MemoryEfficient, respectively.

Table 3 shows the analysis of memory usage across different settings with three distinct approaches: Vanilla Transformer, MemoryEfficient, and our proposed method, BPT. It is evident that Vanilla transformers consumes the highest amount of memory, while MemoryEfficient and BPT offer notable improvements in memory optimization. Notably, our BPT technique consistently outperforms both Vanilla transformers and MemoryEfficient in all settings, showcasing memory efficiency.

5.2 Evaluation on Throughput and Speed

In Table 4, we present a comparison of the throughput achieved by different attention mechanisms on the GPT-XL (1B) model trained on the OpenWebText dataset using 8 GPUs. Throughput is measured as number of tokens processed per device per second. We evaluate the performance at various context lengths, including 2K, 8K, 16K, 33K, and 65K tokens. Our proposed method achieves competitive throughput as MemeoryEfficient mechanism, and surpasses the Vanilla transformer, achieving 1.17x speedup at context length 8k and 1.2x speedup at context length 16k. At context length 32K and 64K, our method maintains high throughput and training speed, while the alternatives cannot train due to running out of memory. This demonstrates the scalability and efficiency of our proposed method, allowing it to effectively handle large context lengths without compromising on throughput and training speed.

5.3 Evaluation on Reinforcement Learning

In this section, we present the results of applying BPT to improve the performance of transformers in reinforcement learning (RL). We report our results in Table 5, where we evaluate our proposed model on the ExoRL benchmark across six different tasks. On ExoRL, we report the cumulative return, as per ExoRL [58]. The numbers of BC, DT [6] and AT [34] are from the ExoRL and AT paper. AT + ME and AT + BPT numbers are run by ourselves. Since the ExoRL data is significantly more diverse than D4RL because it is collected using various unsupervised RL algorithms [28, 35, 33], it is found that TD learning performs best while behavior cloning struggles [58]. AT [34] shows that conditioning transformers on multiple trajectories with relabeled target return can significantly outperforms behavior cloning approaches BC-10% and DT, and achieves competitive results with TD learning. For more details, please refer to their papers. We are interested in applying BPT to improve the performance of AT by conditioning on a 32 trajectories rather than 4 trajectories in original work. It is worth noting that each trajectory has 1000 × 4 length where 1000 is sequence length while 4 is return-state-action-reward, making training 32 trajectories with 350M size model infeasible for both Vanilla and MemoryEfficient. Results in Table 5 show that, by scaling the sequence length, AT + BPT consistently outperforms the original transformers model in all six tasks, achieving a total average return of 111.13 compared to the original transformers model's total average return of 83.02

6 Related Work

Transformers have garnered significant attention in the field of natural language processing (NLP) and have become the basis for numerous state-of-the-art models. Several works have explored memory-

efficient techniques to address the memory limitations of Transformers and enable their application to longer input sequences. One line of research focuses on various approximation techniques or compressing along the sequence dimension [see e.g. 24, 12, 14, 4, 44, 56, 38, 25]. Other works explored replacing attention [19, 20, 43, 23, 3, 59, 42, 55]. Another line of work explores partitioning the large hidden dimension of the feedforward network into parts and retrieving only one part per token [30, 50, 17, 26, 60, 62]. Additionally, extending the context by attending over states from previous sequences has been explored [13, 46], as well as combining local and global contexts [21, 11]. For a comprehensive review of these techniques, we recommend referring to the surveys by Tay et al. [53], Narang et al. [40], Tay et al. [52]. Several studies explored sharding large model on distributed devices tensor, data, or sequence parallelism [51, 16, 57, 27, 61, 31, 48]. Ours shares similarities with the sequence parallelism [27] where sequences are distributed across devices, in contrast, ours implements blockwise computation on sequences for each device. This creates an orthogonal relationship between our method and sequence parallelism, allowing for straightforward combination. In addition, our methodology is compatible with both tensor and data parallelism. Another direction involves computing exact self-attention in a blockwise manner using the tiling technique [39]. This approach has led to the development of memory efficient attention mechanisms [14, 44]. In line with these advancements, our work falls into this category. We propose computing both the feedforward

network and self-attention in a blockwise manner, resulting in a significant reduction in memory requirements.

7 Conclusion

In conclusion, we propose a blockwise parallelization approach to reduce the memory requirements of Transformers, the backbone of state-of-the-art AI models. Our approach enables processing longer input sequences while maintaining or improving performance. Through extensive experiments, we demonstrate its effectiveness, achieving up to 4x memory reduction than memory-efficient Transformers. Our contributions include a practical method for large context sizes in large transformers models. With the increasing capability of hardware, larger models and longer context length are widely used in AI research. At the same time, as we are pushing up against physics and fabrication limits, it is more important to design scaling approaches as efficient as possible to scale up large models and large context size. Our approach holds promise for training and evaluating complex models with longer input sequences, potentially driving new breakthroughs in machine learning research.

Limitations and Future Work. Although our method achieves state-of-the-art low memory usage for transformers models, it does have some limitations that need to be addressed:

Optimal performance. While our implementation prioritizes simplicity with high-level Jax operations, optimizing low-level operations is crucial for achieving optimal performance. In future work, we suggest considering porting our method to CUDA and OpenAI Triton to achieve minimal memory cost and maximum speedup.

References

J.-B. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y . Hasson, K. Lenc, A. Mensch, K. Millican, M. Reynolds, et al. Flamingo: a visual language model for few-shot learning. Advances in Neural Information Processing Systems , 35:23716-23736, 2022.

K. Andrej. GitHub - karpathy/nanoGPT: The simplest, fastest repository for training/finetuning medium-sized GPTs. - github.com. https://github.com/karpathy/nanoGPT, 2023. [Accessed 16-May-2023].

I. Bello. Lambdanetworks: Modeling long-range interactions without attention. arXiv preprint arXiv:2102.08602 , 2021.

I. Beltagy, M. E. Peters, and A. Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 , 2020.

T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural information processing systems , 33:1877-1901, 2020.

L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems , 34:15084-15097, 2021.

M. Chen, J. Tworek, H. Jun, Q. Yuan, H. P. d. O. Pinto, J. Kaplan, H. Edwards, Y. Burda, N. Joseph, G. Brockman, et al. Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374 , 2021.

T. Chen, B. Xu, C. Zhang, and C. Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174 , 2016.

X. Chen, X. Wang, S. Changpinyo, A. Piergiovanni, P. Padlewski, D. Salz, S. Goodman, A. Grycner, B. Mustafa, L. Beyer, et al. Pali: A jointly-scaled multilingual language-image model. arXiv preprint arXiv:2209.06794 , 2022.

X. Chen, M. Lin, N. Schärli, and D. Zhou. Teaching large language models to self-debug. arXiv preprint arXiv:2304.05128 , 2023.

R. Child, S. Gray, A. Radford, and I. Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

K. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Davis, A. Mohiuddin, L. Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794 , 2020.

Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. V. Le, and R. Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860 , 2019.

T. Dao, D. Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems , 35: 16344-16359, 2022.

J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805 , 2018.

Facebook. Fully Sharded Data Parallel: faster AI training with fewer GPUs - engineering.fb.com. https://engineering.fb.com/2021/07/15/open-source/fsdp/ , 2023. [Accessed 16-May-2023].

W. Fedus, B. Zoph, and N. Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. The Journal of Machine Learning Research , 23(1):5232-5270, 2022.

A. Gokaslan and V. Cohen. Openwebtext corpus, 2019. URL http://Skylion007.github. io/OpenWebTextCorpus .

A. Gu, T. Dao, S. Ermon, A. Rudra, and C. Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems , 33:1474-1487, 2020.

A. Gu, K. Goel, and C. Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396 , 2021.

J. Ho, N. Kalchbrenner, D. Weissenborn, and T. Salimans. Axial attention in multidimensional transformers. arXiv preprint arXiv:1912.12180 , 2019.

E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021.

W. Hua, Z. Dai, H. Liu, and Q. Le. Transformer quality in linear time. In International Conference on Machine Learning , pages 9099-9117. PMLR, 2022.

A. Jaegle, F. Gimeno, A. Brock, O. Vinyals, A. Zisserman, and J. Carreira. Perceiver: General perception with iterative attention. In International conference on machine learning , pages 4651-4664. PMLR, 2021.

N. Kitaev, Ł. Kaiser, and A. Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451 , 2020.

A. Komatsuzaki, J. Puigcerver, J. Lee-Thorp, C. R. Ruiz, B. Mustafa, J. Ainslie, Y. Tay, M. Dehghani, and N. Houlsby. Sparse upcycling: Training mixture-of-experts from dense checkpoints. arXiv preprint arXiv:2212.05055 , 2022.

V. Korthikanti, J. Casper, S. Lym, L. McAfee, M. Andersch, M. Shoeybi, and B. Catanzaro. Reducing activation recomputation in large transformer models. arXiv preprint arXiv:2205.05198 , 2022.

M. Laskin, D. Yarats, H. Liu, K. Lee, A. Zhan, K. Lu, C. Cang, L. Pinto, and P. Abbeel. Urlb: Unsupervised reinforcement learning benchmark. arXiv preprint arXiv:2110.15191 , 2021.

M. Laskin, L. Wang, J. Oh, E. Parisotto, S. Spencer, R. Steigerwald, D. Strouse, S. Hansen, A. Filos, E. Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:2210.14215 , 2022.

D. Lepikhin, H. Lee, Y. Xu, D. Chen, O. Firat, Y. Huang, M. Krikun, N. Shazeer, and Z. Chen. Gshard: Scaling giant models with conditional computation and automatic sharding. arXiv preprint arXiv:2006.16668 , 2020.

S. Li, J. Fang, Z. Bian, H. Liu, Y . Liu, H. Huang, B. Wang, and Y. You. Colossal-ai: A unified deep learning system for large-scale parallel training. arXiv preprint arXiv:2110.14883 , 2021.

Y. Li, D. Choi, J. Chung, N. Kushman, J. Schrittwieser, R. Leblond, T. Eccles, J. Keeling, F. Gimeno, A. Dal Lago, et al. Competition-level code generation with alphacode. Science , 378 (6624):1092-1097, 2022.

H. Liu and P. Abbeel. Aps: Active pre-training with successor features. In International Conference on Machine Learning , 2021.

H. Liu and P. Abbeel. Emergent agentic transformer from chain of hindsight experience. International Conference on Machine Learning , 2023.

H. Liu, T. Zahavy, V. Mnih, and S. Singh. Palm up: Playing in the latent manifold for unsupervised pretraining. Advances in Neural Information Processing Systems , 35:3588035893, 2022.

H. Liu, C. Sferrazza, and P. Abbeel. Chain of hindsight aligns language models with feedback. arXiv preprint arXiv:2302.02676 , 2023.

Y. Liu, M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and V. Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv:1907.11692 , 2019.

X. Ma, C. Zhou, X. Kong, J. He, L. Gui, G. Neubig, J. May, and L. Zettlemoyer. Mega: moving average equipped gated attention. arXiv preprint arXiv:2209.10655 , 2022.

M. Milakov and N. Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv:1805.02867 , 2018.

S. Narang, H. W. Chung, Y. Tay, W. Fedus, T. Fevry, M. Matena, K. Malkan, N. Fiedel, N. Shazeer, Z. Lan, et al. Do transformer modifications transfer across implementations and applications? arXiv preprint arXiv:2102.11972 , 2021.

OpenAI. Gpt-4 technical report, 2023.

A. Orvieto, S. L. Smith, A. Gu, A. Fernando, C. Gulcehre, R. Pascanu, and S. De. Resurrecting recurrent neural networks for long sequences. arXiv preprint arXiv:2303.06349 , 2023.

M. Poli, S. Massaroli, E. Nguyen, D. Y. Fu, T. Dao, S. Baccus, Y. Bengio, S. Ermon, and C. Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint arXiv:2302.10866 , 2023.

M. N. Rabe and C. Staats. Self-attention does not need o(n2) memory. arXiv preprint arXiv:2112.05682 , 2021.

A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

J. W. Rae, A. Potapenko, S. M. Jayakumar, and T. P. Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint arXiv:1911.05507 , 2019.

R. M. Rao, J. Liu, R. Verkuil, J. Meier, J. Canny, P. Abbeel, T. Sercu, and A. Rives. Msa transformer. In International Conference on Machine Learning , pages 8844-8856. PMLR, 2021.

J. Rasley, S. Rajbhandari, O. Ruwase, and Y. He. Deepspeed: System optimizations enable training deep learning models with over 100 billion parameters. In Proceedings of the 26th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining , pages 3505-3506, 2020.

K. M. Ruff and R. V. Pappu. Alphafold and implications for intrinsically disordered proteins. Journal of Molecular Biology , 433(20):167208, 2021.

N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean. Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. arXiv preprint arXiv:1701.06538 , 2017.

M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-lm: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053 , 2019.

Y. Tay, M. Dehghani, S. Abnar, H. W. Chung, W. Fedus, J. Rao, S. Narang, V. Q. Tran, D. Yogatama, and D. Metzler. Scaling laws vs model architectures: How does inductive bias influence scaling? arXiv preprint arXiv:2207.10551 , 2022.

Y. Tay, M. Dehghani, D. Bahri, and D. Metzler. Efficient transformers: A survey. ACM Computing Surveys , 55(6):1-28, 2022.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

J. Wang, J. N. Yan, A. Gu, and A. M. Rush. Pretraining without attention. arXiv preprint arXiv:2212.10544 , 2022.

S. Wang, B. Z. Li, M. Khabsa, H. Fang, and H. Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 , 2020.

Y. Xu, H. Lee, D. Chen, B. Hechtman, Y. Huang, R. Joshi, M. Krikun, D. Lepikhin, A. Ly, M. Maggioni, et al. Gspmd: general and scalable parallelization for ml computation graphs. arXiv preprint arXiv:2105.04663 , 2021.

D. Yarats, D. Brandfonbrener, H. Liu, M. Laskin, P. Abbeel, A. Lazaric, and L. Pinto. Don't change the algorithm, change the data: Exploratory data for offline reinforcement learning. arXiv preprint arXiv:2201.13425 , 2022.

S. Zhai, W. Talbott, N. Srivastava, C. Huang, H. Goh, R. Zhang, and J. Susskind. An attention free transformer. arXiv preprint arXiv:2105.14103 , 2021.

Z. Zhang, Y. Lin, Z. Liu, P. Li, M. Sun, and J. Zhou. Moefication: Transformer feed-forward layers are mixtures of experts. arXiv preprint arXiv:2110.01786 , 2021.

L. Zheng, Z. Li, H. Zhang, Y. Zhuang, Z. Chen, Y. Huang, Y. Wang, Y. Xu, D. Zhuo, E. P. Xing, et al. Alpa: Automating inter-and { Intra-Operator } parallelism for distributed deep learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22) , pages 559-578, 2022.

S. Zuo, Q. Zhang, C. Liang, P. He, T. Zhao, and W. Chen. Moebert: from bert to mixture-ofexperts via importance-guided adaptation. arXiv preprint arXiv:2204.07675 , 2022.

A Experiment Details

A.1 Evaluation of Memory

In the experimental results presented in Section 5.1, we used model parallelism to partition the model across 8 GPUs or 64 TPUv4 units. Our evaluation focused on determining the maximum achievable sequence length, using a sequence number of one. For TPUs, we utilized its default training configuration, which involved performing matmul operations in bfloat16 format with weight accumulation in float32 . On the other hand, for GPUs, we adopted the default setup, where all operations were performed in float32 .

To profile memory usage, we utilized jax.profile and repeated the evaluation 100 times, reporting the average results. We conducted a grid search for the optimal query block size and key-value block size, considering values from the set [16 , 64 , 128 , 512 , 1024 , 2048 , 4096] . For each method, we reported the lowest memory achieved.

A.2 Evaluation of Throughput

In the evaluation presented in Section 5.2, we split OpenWebText following the methodology of [2]. Throughput is measured as tokens per device per second. To ensure a fair comparison, we performed a grid search for the optimal query block size and key-value block size, considering values from the set [16 , 64 , 128 , 512 , 1024 , 2048 , 4096] . For gradient checkpointing [8], we additionally grid search among three commonly used checkpointing policies including noth-ing_saveable , dots_saveable , and dots_with_no_batch_dims_saveable for attention and use nothing_saveable for feedforward network (FFN). For more details, please refer to Jax documentation. We selected the best performing configuration for both baselines and our method.

The training was conducted using FSDP [16] and gradient accumulation. We used weight decay of 0.1 and utilized cosine learning rate decay with a maximum learning rate of 2 . 0 × e -4 . For sequence lengths of 2048 , 4096 , 8192 , 16384 , the batch sizes in trajectories were set as 8 , 4 , 2 , 1 , 1 respectively. We use gradient accumulation to accumulate batch size in tokens to 1 million per batch.

A.3 Evaluation on RL

In the experiment presented in Section 5.3, we followed the prior work's setting for learning rate, batch size, and other hyperparameters, while modifying the number of trajectories. The specific hyperparameters are provided in Table 6. The original agentic transformers used 4 trajectories during training, we increase the number to 32.

During testing, increasing the number of trajectories has been shown to improve performance. However, performing autoregressive sampling over a large number of trajectories (e.g., 64 × 1000 × 4

total number of tokens) can be computationally slow. To reduce the sampling time, we limited the rollout to 16 trajectories.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper extends the work in Flash Attention to apply block wise attention to not just the self attention block, but also to the feedback forward part of the transformer block. While doing so, the output of the self attention block is not stored for backprop, and recalculated. This leads to very efficient gain in memory requirements for training on large sequences.\n\nAs already shown by FlashAttention work, performing these operations at block level can take advantage of limited SRAM by scheduling operations in the right order to reduce the amount of communication thereby not just leading to memory efficiency, but also improved throughput and latency.\n\nExperiments on multiple datasets show that this improves upon the foundation laid by FlashAttention by being able to support longer sequences along with throughput improvements.\n\nSTRENGTHS:\n- Efficiency of Transformer Architectures is a pretty important field, and any improvement in this area has a wide impact on various applications of Transformers.\n- Memory Efficient architectures are key to the future of both research and deployment.\n- The paper does a reasonable number of experiments to prove that even though the operations are done blockwise, the throughput improves because of better SRAM utilization with reduced communication.\n- Table 4: By increasing the sequence length, they are also able to get better validation loss.\n\nWEAKNESSES:\n- Line 23: \"such as\" is repeated.\n- The contribution on top of Flash Attention is pretty limited, but still the impact is pretty significant.\n\nQUESTIONS:\n- In the RL task, the AT baseline is for NTrajs=4, and compared with N Trajs =64 which OOMs for AT. Is it possible to evaluate on N Trajs = something in between 4 and 64 that allows AT to run without OOM but perform better than N Trajs =4 for a better comparison?\n- Why is stop_gradient needed in the implementation for `max_score`?\n\nLIMITATIONS:\nLimitions are well noted in the paper.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: SUMMARY:
This paper proposed to merge the computation of the FFN layers in Transformer into the block-wise computation of self-attention. This eliminates the need to wait for the self-attention computation to finish before performing the feedforward step on the entire sequence, because FFN is entirely position-wise and hence block-parallelizable.

Experiments on large language pretraining and reinforcement learning demonstrate that BPT reduces the peak memory costs and allows training large-scale Transformer models on longer sequences.

STRENGTHS:
The proposed method is straight forward and well-motivated. The paper is writing well, easy to follow. Experiments show the benefits of BPT on long sequence training, especially the cost of memory.

WEAKNESSES:
The novelty of the proposed method is not that significant, and there are still some parts unclear (see questions).

QUESTIONS:
If I understand correctly, even the FFNs are computed in a block-wise manner, we still need to keep all the output vectors of each FFN layer (and self-attention layer) to do back-propagation. So the memory cost still increases linearly w.r.t the length of inputs. By using the activation checkpointing, we can reduce the memory cost of the hidden vectors inside FFNs (4 x model dim), but not the output of each layer. Then, why BPT is able to train 2-4 times longer sequences than FlashAttention? Just because of the memory reduction of the FFN hidden vectors?

LIMITATIONS:
NA.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-11

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents a novel variant of the Transformer model, the Blockwise Feedforward Transformer (BFT), which is designed to address the memory inefficiency issue in standard Transformer models. The authors propose a blockwise computation method that significantly reduces the memory footprint without compromising the model's performance.\n\nThe key innovation of the BFT is the introduction of a blockwise computation method. Instead of computing the feedforward network on the full sequence, the BFT computes it on intermediate blocks, resulting in substantial memory savings. The computation for a query block is given by: Outputi = FFN(Attention(Qi, K, V ) + Qi) + Attention(Qi, K, V ) + Qi.\n\nSTRENGTHS:\n1. The Blockwise Feedforward Transformer (BFT) presents a novel approach to address the memory inefficiency issue in standard Transformer models. This is a significant contribution as memory efficiency is a critical factor for training long context language models\n2. The authors conducted extensive experiments on both reinforcement learning and language modeling tasks, providing a comprehensive evaluation of the model's performance across different domains.\n\nWEAKNESSES:\n1. The paper lacks a comprehensive analysis of the model's performance, particularly in terms of perplexity or its impact on downstream tasks. This omission makes it challenging to assess the effectiveness of the BFT in maintaining performance when compared to other models.\n\nQUESTIONS:\n1. Implementation Details: More information on the implementation details of the BFT would be appreciated. Are there any specific hardware/software requirements/dependencies for implementing the BFT?\n\nLIMITATIONS:\nThe authors argue that \"The application of long context models to NLP training remains uncertain due to the scarcity of large datasets encompassing extensive context information.\" However, resources such as the Books Corpus, Github, and Arxiv offer a wealth of long documents. It would be beneficial if the authors could experiment with training their proposed models using these resources and conduct a more thorough evaluation of the long context capabilities of Language Models.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n\"Blockwise Parallel Transformer for Large Models\" presents a novel approach to handle the memory demands of Transformer models, mainly when dealing with long sequences or tasks involving multiple sequences or long-term dependencies. The authors propose a method called Blockwise Parallel Transformer (BPT) that leverages blockwise computation of self-attention and feedforward network fusion to minimize memory costs. \n\nThe critical contributions of the paper are:\n\n1. The authors propose a blockwise computation of self-attention and feedforward approach that enables 16 to 64 times longer and 2 to 4 times longer context lengths than vanilla Transformer and previous memory-efficient Transformers, respectively.\n2. They demonstrate the effectiveness of their approach through empirical experiments. BPT can reduce the memory requirements of Transformers, enabling training at least 8 to 64 times longer sequence than vanilla attention-based GPT models and at least 2 to 4 times longer sequence than prior state-of-the-art FlashAttention and Memory Efficient Attention.\n3. The authors also show that BPT significantly improves performance and achieves better results on challenging reinforcement learning benchmarks by conditioning on multiple trajectories.\n\nThe authors argue that their approach has the potential to enable the training and evaluation of more complex models that require longer input sequences, which could lead to further breakthroughs in AI research.\n\nThe authors also explain the memory bottleneck of Transformer models and how their proposed method addresses this issue. They explain how self-attention can be computed blockwise without materializing the softmax attention matrix, resulting in significant memory savings. They also show how this blockwise computation can be applied to the feedforward network, resulting in further memory savings. The paper includes a detailed algorithm of the BPT method and key parts of its implementation in Jax which is nice to see.\n\nSTRENGTHS:\n**Originality:** BPT enables processing longer input sequences while maintaining or improving performance. This is achieved by computing both the feedforward network and self-attention blockwise, significantly reducing memory requirements (Page 9). This approach is original as it combines the techniques of blockwise computation of self-attention and feedforward networks, which is a distinct approach from previous works.\n\n**Quality:** The paper demonstrates the effectiveness of the BPT approach from the perspective of speed/throughput (but not quality/capability of the model) through empirical experiments. It shows that the BPT approach can train on 8 to 64 times longer sequences than vanilla attention-based GPT models and 2 to 4 times longer sequences than prior state-of-the-art models like FlashAttention and Memory Efficient Attention.\n\n**Clarity:** The paper is well-structured and clear in its presentation. It provides a comprehensive background on Transformers and their challenges in handling long sequences. The paper also clearly explains the proposed BPT approach and its advantages, such as its ability to leverage hardware with significantly faster SRAM speed compared to HBM speed. The evaluation of the method is also clearly presented with comparisons to other attention mechanisms.\n\n**Significance:** The paper's significance lies in its potential to enable the training and evaluation of more complex models that require longer input sequences.\n\nWEAKNESSES:\nThe use of RL as a benchmark for testing the quality of BPT, is a weird choice to say the least, especially given how noisy we know RL is. Language modeling is the primary use case for transformers and by extension BPT. Therefore, evaluating BPT's performance (perplexity) on language modeling tasks would be more appropriate.\n\nI initially misunderstood that this is not an approximation but an exact implementation of attn+ffn (similar to flash attention). Therefore it's not imperative that there's LLM benchmarks but regardless would be nice to see that this implementation gives exact same output as standard attention in all precisions int8-fp32.\n\nQUESTIONS:\n- How does it work with parallelism schemes that chunk on sequence or FFN?\n- Does this type of parallelism work with lower precision since it seems like there's more steps for accumulation and hence likely more noisy.\n\nLIMITATIONS:\n**Dependence on Hardware:** The paper mentions that the BPT can leverage hardware with significantly faster SRAM speed compared to HBM speed (Page 5). This implies that the performance of BPT might be dependent on the specific hardware configuration, and may not perform as well on hardware with slower SRAM speed.\n\n**Sequence/Tensor Level Parallelism:** How does it work with parallelism schemes that chunk on sequence or FFN.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}