Blockwise Parallel Transformers for Large Context Models
Authors: Hao Liu
Pieter Abbeel
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Language Model
Long Context Modeling
Reinforcement Learning
TL;DR: We present an efficient method of computing the standard transformer architecture, enabling effective processing of long contextual information.

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
Hao Liu UC Berkeley hao. liu@cs. berkeley. edu
Pieter Abbeel
UC Berkeley berkeley edu
pabbeel@cs wtr.oleeeeebeeeeeee........eeeee
Abstract
Transformers have emerged as the cornerstone of state-of-the-art natural language processing models, showcasing exceptional performance across wide range of Al applications. However, the memory demands posed by the self-attention mecha- nism and the large feedforward network in Transformers limit their ability to handle long sequences, thereby creating challenges for tasks involving multiple long se- quences 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 se- quences 32 times longer than vanilla Transformers and up times longer than previous memory-efficiett 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-arr natural language processing They have demonstrated impressive performance across wide range of Al problems, including language modeling, machine translation, image captioning, and protein folding
models 15 45 37
Transformers achieve this success through their architecture design that uses self self-attention and position-wis feedforward mechanisms. These components facilitate the efficient capture long-range dependencies between input tokens, enabling scalability terms of context length and model size through highly parallel computations.
47
However, the memory requirements of Transformers limit their ability to handle long sequences, which is necessary for many Al problems, such as high-resolution images, podcasts, code, or books and especially those that involve multiple long sequences long-term dependencies [IO 41 36
229 49 [ The quadratic self-attention and the large feed forward network of Transformers require a large amount of memory, which makes it challenging to scale longer input sequences. This limitation has led to various techniques proposed to reduce the memory requirements of Transformers, including sparse-approximmtion low-rank approximation, and low precision approximation [see e.g.
53
24
25 38 56
One distinct line of research does not rely approximation but instead focuses computing exact self-attention with linear memory complexity. This approach leverages the observation that the softmax matrix 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 blockwise computation of the self-attention softmax, demonstrating reduced memory requirements.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Vanilla Attention
FlashAtteniion Memorry Efficient Attention
BlockWise Parallel Transformer
5000
3003
300M
7B 138 50B
13R 300 70R
Figure 1: Maximum context length during training time with the GPT model using different methods. Model sizes range from 1B to 70B. Figures (A), (B), and (C) show evaluation using one, eight A A100, and 64 TPUv4, respectively, with a single sequence. Our method enables training sequences 32 times longer than vanilla attention-based transformers [54), and 2 to times longer than FlashAttention and Memory Efficient Attention 44
14
Section
provides memory cost breakdown.
Despite the resulting reduced memory requirements of the self-attention block in transformers models, significant challenge still arises from the feedforward network. This network contains 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, make important observation: when self-attention is computed in blockwise manner to reduce memory requirements, it becomes feasible to merge the computation of the feedforward network. This eliminates the need wait for the self-attention computation to finish before performing the feedforward step on the entire sequence. By computing the feedforward network on 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 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 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 saae--f-the-aaee 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 Al research.
Our contributions are twofold: (a) proposing blockwise computation of self-attention and feedfor- ward approach that enables 32 times longer and up to 4 times longer context lengths than vanilla transformers and previous memory- -fffeieetf Transformers, respectively, and (b) demonstrating the effectiveness of our approach through extensive experiments.
2 Memory Bottleneck of Transformer
Given input sequences Q, K, V E Bxd where s is the sequence length and dl is the head dimension. We compute the matrix of outputs as:
Attention(Q. K, = softmax QKT
Vd VV,

---

where softmax is applied row-wise. Standard attention implementations materialize the matrices QKT and softmax( QK" is (s))). There has been large body of work trying to reduce memory usage of self-attention by using online softmax (39 14 to reduce memory cost of self-attention by preventing it from full materialization. And these approaches reduce memory footprint from Os²2) to O(s). However, the large feedforward layers have been overlooked.
va to HBM, which takes ( memory, so the overall space complexity
In addition to attention sub-layers, each of the attention layers is accomplished with fully connected feedforward network, which is applied to each position separately and identically, This consists of two linear transformations with ReL.U activation in between.
(FN = m0 T:1 1 + b1 W + l2
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?) )
This approach involves splitting the sequences Q E Rxx into By blocks and sequences K, Vi R into By 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 specific query block Q, 1 < i S Bq. the corresponding attention output can be computed by scaling each blockwise attention as follows:
(39 14
Attention (Q;, K, V) =
VV B.v) (r))
(3)
The scaling operation scales each blockwise attention based on the difference between the blockwise maximum and the global maximum:
Autention(Q;; Kj, V2) = exp(Q:K7 - max (Q:K3)) Σ exp(Q:K] - max max((:
max; = max max(max(Q.KT)).. wwx.maaaaammmm max(Q:
Attention(Q,,0 Q,, V) = [exp(Q.K]] max;) Attention(Q::. K2, V2) BEs
This blockwise self-attention computation eliminates the need to materialize the full attention matrix of size O(n²), resulting in significant memory savings.
We observe that the blockwise computation is not limited to self-attention but can also be applied 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 query block is given by:
Output; FFN (Attention( Qu, K, V) + Q.) + Attention(Q,, K, K) Q++
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, summarized in Algorithm I
3

---

Softmax
Linear
Feed
Forward
Q
Nx
Masked Attention
KV
Projection
Input
We use the same model architecture the original transformers but with different way of organizing the compute. In the diagram, we explain this by showing that for the bottom first incoming input block, project it into query; then iterate over the same input sequence positioned above the bottom row, and project it to key and value. These query, key and value are used to compute self-attention (yellow box), whose output is pass to feedforward network (cyan box), followed by residual connection. In our proposed approach, this process is then repeated for the other incoming input blocks.
Figure 2:
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 Transtormers:
Attention: For Q, K, V, saving their input r needs 2bsh bytes, where b is batch size, s is sequence length, and h is hidden dimension. For QKT matmul, saving activations Q and K needs 4bsh bytes. For softmax roftmax(QK"))) (QKT), saving input QK needs 26s²a a bytes, where a is the number of attention heads. For mask and dropout, saving mask needs as bytes. For score V, saving score needs 2bs bytes, and saving V needs 2bsh bytes. For output projection and dropout, saving the input needs 2hsh bytes, and saving dropout mask needs bsh bytes. The maximum attention activation size of attention is O(s²) with checkpointing.

---

Algorithm Reduce memory cost with BPT.
Required: Input sequence Number of query blocks Bg. Number of key and value blocks Bkv-
Initialize
Project imput sequence into query, key and value.
Split query sequence into Bg of query input blocks.
Split key and value sequences into Bk of key-value input blocks.
for ouler = to B, do
Choose the outer-th query. for inner = to Bk 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 2bsh bytes. For activation, saving input needs 8hsh bytes. For the second linear layer, saving input needs 8bsh bytes. For dropout, saving the mask needs bsh bytes. With checkpointing, the maximum activation size of FFN 8bsh bytes.
Consequently, for large context length, the memory cost of activation in vanilla transformers is Oss²).
BPT:
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 maximum activation size of 40ch with checkpointing. Additionally, it needs to store output activations for the query loop, which requires 2bsh bytes. Since s >> c, the maximum activation size 2bsh.
Attention:
FFN: When iterating the FFN over blocks, BPT needs to save the following activations: For the first linear layer, saving input needs 2bch bytes. For activation, saving input needs 8bch bytes. For the second linear layer, saving input needs 8èch bytes. For dropout, saving the mask needs bch bytes. In total, 19bch bytes are needed. Additionally, storing the output of the for loop requires 2bsh bytes. Therefore, the maximum FFN activation size is 2bsh.
Consequently, each BPT layer's memory cost of activation is 2bsh.
Memory-Efficies Flash Attention:
Attention: Similar to BPT attention, the maximum activation size is 2bsh.
FFN: Similar to the vanilla FFN, the maximum activation size is 8hsh.
Consequently, each Flash Attention layer's memory cost is 8bsh.
Comparing the activation memory of Flash MMention/Memory-Efiiie that BPT offers 8bsh /bbhh 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.
transformers with BPT, we times memory saving. By taking into account other factors
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 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 avoid waiting for the completion of self-attention and allocating a significant amount of memory solely for -erwwrd network computation.
5

---

Another notable advantage of blockwise parallelization is its ability to leverage hardware with significantly faster SRAM speed compared HBM speed. For instance, 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 provides the pseudocode of the algorithm. Figure in Appendix shows Jax imple- mentation optimized for simplicity. The full code of BPT is provided GitHub which supports large-scale distributed training of large context models using BPT.
The blockwise_fr fn function begins by accepting rematerialized feed forward module, inputs and chunk size. The remat_ffi compute feedforward on inputs with checkpointing, without saving intermediates. The sfn function is then used scan over input sequences and generate outputs.
The blockwise_atee saatttntttttteeeeeeeeeeeeeeeeeet function process query, key, and value to produce attention blockwise. The ocan_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_ _alnck___coc bblock function to the key-value chunk, calculating the dot product between the query and key vectors, and then adding bias term. The bias term introduces 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 numerically stable manner, using the max-score trick avoid large exponentiation results.
Finally, BPT combines the outputs from all chunks, normalizes them using their max-score-adjusse adjussedd weights, and passes them through a feed-forward neural network (blockwise_frr) iffn.. The final output is the sum of the
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 I 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 [ and Memory Efficient Attention 44 which are -tate-of-the-art
memory efficient attention, we denote them as "MemoryEfficiiet" in our experiments, All methods use the same gradient checkpointing the experiments.
Datasets. We consider two datasets for evaluation purposes. Including pretraining on WeenTeeText dataset and large context reinforcement learning on ExoRL.
Webn The OeenWetText dataset [8 is 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 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 unsuper- vised 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 by materializing the attention matrix and computes the feedforward network normally. We also
which computes self-attention
hatt:://ggthuu..cc//laagg//////
context

---

Table 1: Sizes and architectures of the models which we evaluated in experiments.
Model Name
T2 parauss
Taayrrs
dmudee
Theads
dhead
GPT 1B
1.3B
2048
16
128
GPT 3B
2.7B
32
2560
32
80
GPT 7B
6.7B
32
4096
32
128
GPT 13B
13.0B
40
5140
40
128
GPT 30B
30.0B
48
7168
56
128
GPT 70B
70.0B
80
8192
64
128
consider two prior state-of-the-arr memory-efficient methods, namely FlashAttention 14], which focuses GPU efficiency, and Memory Efficient Attention <4] which focuses 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 A 100 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 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.
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 BPT to reinforcement learning and evaluate its effectiveness in large context application.
Table 2: Maximum context length during training with different methods. BPT enables training 2-4 times longer sequence length than FlashAttention Memory Efficient Attention, and up to 32 times longer sequence length than vanilla attention.
A100
PartitionSpec
Vanilla Attention
MemoryEfficient
Blockwise Parallel
350M
(1,1,))
16K (16384)
65K (65536)
13IK (131072)
IB
(1,1,)
16K (16384)
65K (65536)
131K (131072)
(1,1,)
8K (8192)
16K (16384)
65K (65536)
A100
PartitionSpec
Vanilla Attention
MemoryEfficient
Blockwise Parallel
3B
(1.1,))
16K (16384)
65K (65536)
131K (131072)
7B
(1,1,))
16K (16384)
65K (65536)
131K (131072)
13B
(1,1,))
8K (8 (8192)
33K (32768)
65K (65536)
30B
(1.1,)
8K (8192)
16K (16384)
65K (65536)
64 TPUv4 PartitionSpec
Vanilla Attention
Memory
Blockwise Parallel
13B
(1,1,64)
4K (4096)
16K (16384)
33K (32768)
30B
(1,1,64)
2K (2048)
4K (4096)
16K (16384)
70B
(1,1,64)
Ik (1024)
2K (2048)
8K (8192)
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, MemoryEfffcient enables longer sequences of 65K for IB parameters and 16K for 3B parameters. Notably, our proposed method, Blockwise Parallel, surpasses both methods, achieving maximum sequence
7

---

length of 131K for IB 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 MemoryFfficient, respectively.
Table Shows the analysis of memory usage across different settings with three distinct approaches: Vanilla Transformer, MemoryEfficient, and our proposed method, BPT. 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.
Table 3: Memory usage comparison for different settings. "oom" denotes out of memory.
Setting
3B on A100
13B on 8 A100
Context Length
Vanilla
EemoeyEfficient
BPT
Vanilla
MemoryEfficient
BPT
8192
64GB
44GB
43GB
59GB
44GB
42GB
16384
oom
47GB
45GB
oom
46GB
45GB
32768
oom
55GB
52GB
com
55GB
52GB
65536
00m
75GB
70GB
oom
75GB
68GB
131072
00m
oom
79GB
00m
com
78GB
5.2 Evaluation on Throughput and Speed
In Table present a comparison of the throughput achieved by different attention mechanisms on the GPT-XL (1B) model trained on the OpenWebText WebText dataset using 8 GPUs. Throughput 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 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 to effectively handle large context lengths without compromising on throughput and training speed.
5.3 Evaluation on Reinforcement Learning
In this section, present the results of applying BPT to improve the performance of transformers in reinforcement learning (RL). We report our results in Table 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 (44 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 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 32 trajectories rather than 4 trajectories in original work. It is worth noting that each trajectory has 1000 X 4 length where 1000 is sequence length while 4 return-state- petate-attion-rewaaae
making training 32 trajectories with 350M size model infeasible for both Vanilla and MemoryEfffiient. 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--rrttrttt models. Several works have explored memory-

---

Table 4: Throughput comparison on GPT-XL (IB) using Open WebTet dataset. Throughput is measured tokens processed per second. 'oom' denotes running out of memory, 'na' denotes results not available because we early terminated these runs to reduce compute cost.
Model
Context Len
Val Loss
Throughput
Speed up
Vanila transformers
2048
2.46
3827
1x
MemoryEfficient
2048
2.46
4371
1.14x
Blockwise Parallel
2048
2.46
3985
1.04x
Vanila transformers
4096
2.44
2340
Ix
MemoryEfficient
4096
2.44
2567
1.1x
Blockwise Parallel
4096
2.44
2687
1.15x
Vanila transformers
8192
2.43
2455
1x
MemoryEfficient
8192
2.43
2781
1.13x
Blockwise Parallel
8192
2.43
2875
1.17x
Vanila transformers
16384
2,41
1701
1x
MemoryEfficient
16384
2.41
1889
1.11x
Blockwise Parallel
16384
2.41
2045
1.2x
Vanila transformers
32768
oom
oom
oom
MemoryEfficient
32768
810
1x
Blockwise Parallel
32768
na
857
1.1x
Vanila transformers
65536
oom
oom
oom
MemoryEfficient
65536
oom
oom
oom
Blockwise Parallel
65536
na
600
1x
Table 5: Application BPT on improving transformers in RL. All the baselines use vanilla attention. AT + ME denotes using "MemoryEfficien"" AT BPT denotes using Blockwise Parallel.
ExoRL benchmark
BC-10%
DT
AT
AT
AT + ME AT + BPT
Task
N Trajs 4
N Trajs = 32
Trajs = 32
Trajs = 32
Walker Stand
52.91
34.54
68.55
oom
00m
95.45
Walker Run
34.81
49.82
88.56
00m
105.88
Walker Walk
13.53
34.94
64.56
00m
00m
78.56
Cheetah Run
34.66
67.53
125.68
00m
60m
178.75
Jaco Reach
23.95
18.64
52.98
87.56
Cartpole Swingup
56.82
67.56
97.81
00m
0om
120.56
Total Average
36.11
45.51
83.02
00m
0om
111.13
efficient techniques to address the memory limitations Transformers and enable their application to longer input sequences. One line of research focuses various approximation techniques
Other works Another line of work explores partitioning the large hidden dimension of the feedforward network into parts and retrieving only one part per 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
compressing along the sequence dimension [see e.g. 24
12 14
56
38 25
explored replacing attention [19
20
43
23
3
59
42
55
token [30 50
26
60
62
For 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, sequence parallelism [5]] 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 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 ad advancements, our work falls into this category. We propose computing both the feedforward

---

network and self-attention in a blockwise manner, resulting in significant reduction in memory requirements,
7 Conclusion
In conclusion, propose blockwise parallelization approach to reduce the memory requirements of Transformers, the backbone of state-of-the- tett-offthe-rrttt Al models. Our approach enables processing longer input sequences while maintaining improving performance. Through extensive experiments, demonstrate its effectiveness, achieving up to 4x memory reduction than memory-efficient Transform- Our contributions include 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 Al research. At the same time, 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.
are pushing up against physics and fabrication limits, it is
Limitations and Future Work. Although method achieves state-of-the-art low memory usage for transformers models, it does have some limitations that need to 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 OpenAl Triton to achieve minimal memory cost and maximum speedup.
Acknowledgements
This project is supported in part by ONR under N00014-21-1-2769 We would like to express our sincere appreciation to the members the RLL Lab and Berkeley Al Lab, as well as Anselm Levskaya, Markus Rabe, Federico Lebron, Sharad Vikram, and Tri Dao for their valuable insights and contributions to this paper. We thank Google TPU Research Cloud for granting us access to TPUs.
10

---

References
[1] J.-B. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y. Hasson, K. Lenc, A. Mensch, K. Millican, M. Reynolds, al. Flamingo: visual language model for few-shot learning. Advances in Neural Information Processing Systems, 55:23716-23753 2022.
[2] K. Andrej. GitHub karpathy/nanooGP: The simplest, fastest repository for training/finetuning medium-sized GPTs. - github.com.
2023. [Accessed
16-May-2023).
[3] Bello. Lambdanetworks: Modeling long-range interactions without attention. arXiv preprint arXiv:2102.08602, 2021.
[4] I. Beltagy, M. E. Peters, and A. Cohan. Longformer: The long-document transformer. arXiv preprint 055::2000.0055000 2020.
[5] 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.
[6] L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-11099 2021.
[7] 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.03377 2021.
[8] Chen, B. Xu, Zhang, and C. Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv: 1604.06//44 2016.
[9] X. Chen, X. Wang, S. Changpinyo, Piergiovanni, P. Padlewski, D. Salz, Goodman, A. Gryener, B. Mustafa, L. Beyer, et al. Pali: A jointly-scaled multilingual language-image model. arXiv preprint arXiv:2209.06799 2022.
[10] X. Chen, M. Lin, N. Schärli, and D. Zhou. Teaching large language models to self-debug. arXiv preprint arXiv:2304.05 2023.
[11] R. Child, S. Gray, A. Radford, and Sutskever. Generating long sequences with sparse transformers, arXiv preprint arXiv:1904.1080 2019.
[12] K. Choromanski, V. Likhosherstov, D. Dohan, Song, A. Gane, T. Sarlos, P. Hawkins, J. Davis, A. Mohiuddin, I.. Kaiser, et al. Rethinking attention with performers. arXiv preprint
arXiv:2009.14794 2020.
[13] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. V. Le. and R. Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint arXiv: 000000060, 2019.
[14] 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.
[15] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04880) 2018
[16] Facebook. Fully Sharded Data Parallel: faster Al training with fewer GPUs engineer- ing.fb.com. [Accessed 16-May-2023).
fb. http:///001111111/////////aaaee0
2023.
[17] 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()):523--5270, 2022.

---

[18] A. Gokaslan and V. Cohen. Openwebtext corpus, 2019. URL http://SSky/ioo0007.. io/ So/DpenWeeBeextorpuus
[19] 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- 1114774-487, 2020.
[20] A. Gu, K. Goel, and C. Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396 2021.
[21] Ho, N. Kalchbrenner, D. Weissenborn, and T. Salimans. Axial attention in multidimensional transformers, arXiv preprint arXiv: 2019.
[22] 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: 106.0985, 2021.
[23] 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.
[24] A. Jaegle, E. Gimeno, A. Brock, O. Vinyals, A. Zisserman, and J. Carreira. Perceiver: General perception with iterative attention. In International conference machine learning, pages 4651-4664. PMLR, 2021.
[25] N. Kitaev, E. Kaiser, and A. Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451, 2020.
[26] A. Komatsuzaki, J. Puigcerver, J. Lee-Thorp, C. R. Ruiz, B. Mustafa, Ainslie, Y. Tay, M. Dehghani, and N. Houlsby. Sparse upcycling: Training mixture-of-experts from dense checkpoints. arXiv preprint arXiv:2212.05083 2022.
[27] V. Korthikanti, J. Casper, S. Lym, L. McAfee, M. Andersch, M. Shoeybi, and B. Catanzaro. Re- ducing activation recomputation in large transformer models. arXiv preprint 285.5.2205.05/98 2022.
[28] 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:21 10. 15/91, 2021.
[29] 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 2022.
preprint arXiv:2210.1422
[30] 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.
[31] 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:2 10. 14883, 2021.
[32] Y. Li, D. Choi, J. Chung, N. Kushman, J. Schrittwieser, R. Leblond, T. Eccles, J. Keeling, Gimeno, A. Dal Lago, al. Competition-level code generation with alphacode. Science, 378 (6624)::092-1097, 2022.
[33] H. Liu and P. Abbeel. Aps: Active pre-training with successor features. In International Conference on Machine Learning, 2021.
[34] H. Liu and P. Abbeel. Emergent agentic transformer from chain of hindsight experience. International Conference Machine Learning, 2023.
[35] 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:35880- 35893, 2022.
[36] H. Liu, C. Sferazza, and Abbeel. Chain of hindsight aligns language models with feedback. arXiv preprint dtXt::2302.0266,, 2023.
12

---

[37] Y. Liu. M. Ott, N. Goyal, J. Du, M. Joshi, D. Chen, O. Levy, M. Lewis, L. Zettlemoyer, and Stoyanov. Roberta: robustly optimized bert pretraining approach. arXiv preprint 1eXiv:99072 /692, 2019.
[38] 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.
[39] M. Milakov and N. Gimelshein. Online normalizer calculation for softmax. arXiv preprint arXiv: 02867, 2018.
[40] S. Narang, H. W. Chung, Y. Tay. W. Fedus, T. Fevry, M. Matena, K. Malkan, N. Fiedel, N. Shazeer, Lan, et Do transformer modifications transfer across implementations and applications? arXiv preprint (tti::2002.//97// 2021.
[41] OpenAl. Gpt-4 technical report, 2023.
[42] 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.06:: 230..03333 2023.
[43] 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.
[44] M. N. Rabe and Staats. Self-attention does not need o(n2) memory. arXiv preprint ar arXiv:2112.05682 2021.
[45] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[46] W. Rae, A. Potapenko, S. M. Jayakumar, and T. P. Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprint arXiv: 0e/////11.05507 2019.
[47] 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.
[48] 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.
[49] K. M. Ruff and R. V. Pappu. Alphafold and implications for intrinsically disordered proteins. Journal of Molecular Biology, 433(20):167208, 2021.
[50] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean. Outra- geously large neural networks: The sparsely- gated mixture-of-experts layer. arXiv preprint arXiv: 17110033,, 2017.
[51] M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-Im: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv: 08053, 2019.
[52] 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 2022.
influence scaling? arXiv preprint 2277
[53] Y. Tay, M. Dehghani, D. Bahri, and D. Metzler. Efficient transformers: A survey, ACM Computing Surveys, 55(6):1-28, 2022.
[54] A. Vaswani, 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.
[55] Wang, J. N. Yan, A. Gu, and A. M. Rush. Pretraining without attention. arXiv preprint 212.52221,.1544, 2022.
13

---

[56] S. Wang, B. Z. Li, M. Khabsa, H. Fang. and H. Ma. Linformer: Self-attention with linear complexity, arXiv preprint orXiv:2006.04700 2020.
[57] Y. Xu, Lee, D. Chen, B. Hechtman, 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.
[58] D. Yarats, D. Brandffonbrener, 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 preprini arXiv:2201. 13425, 2022.
[59] S.
Zhai, W. Talbott, N. Srivastava, C. Huang, H. Goh, R. Zhang, and J. Susskind. An attention free transformer. arXiv preprint WxXi:2205.141033 2021.
[60] 7. Zhang, Lin, Z. Liu, Li, M. Sun, and J. Zhou. Moefication: Transformer feed-forward layers are mixtures of experts. arXiv preprint 1rX:::2110.01766, 2021.
[61] 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 16th USENIX Symposiun Operating Systems Design and Implementation (OSDI 22), pages 559-578, 2022.
parallelism for distributed deep learning. In
[62] S. Zuo, Q. Zhang, C. Liang, P. He, T. Zhao, and W. Chen. Moebert: from bert to mixture-of- experts via importance-guiin guided adaptation. arXiv preprint arXiv:2204.07675, 2022.
14

---

A Experiment Details
A.1 Evaluation of Memory
In the experimental results presented in Section 5.T we used model parallelism to partition the model across 8 GPUs 64 TPUv4 units. Our evaluation focused determining the maximum achievable sequence length, using sequence number of onc. 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 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 compari- son, performed grid search for the optimal query block size and key-value block size, con- sidering values from the set [16, 64, 128, 512, 1024, 2048, 4096). For gradient checkpointing [8], additionally grid search among three commonly used checkpointing policies including noth- iigg saveable,, dots sateaaaa,,bbe,,eeeeeeeeeeeeeeeee and dots with _ooo__n_____ batthhnooohhhhhhaeeaaeaae siims_savesl saveable for feedforward network (FFN). For more details, please refer Jax documentation. We selected the best performing configuration for both baselines and our method.
for attention and
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 x e e4, For sequence lengths of 2048, 4096, 8192, 16384, the batch sizes in trajectories were set 8, 4, 2, 1, respectively. We use gradient accumulation to accumulate batch size tokens to million per batch.
A.3 Evaluation RL
Table 6: Hyperparameters used in RL evaluation.
Hyperparameter
Value
Number of layers
Number attention heads
Embedding dimension
128
Activation function
ReL.U
Batch size
64
Dropout
Learning rate
10-4
Leamning rate decay
Lánear warmup for 10" steps
Grad norm clip
0.25
Weight decay
10
Initial desired target return at test time
850 Walker Stand
400 Walker Run
900 Walker Walk
350 Cheetah Run
300 Jaco Reach
800 Cartpole Swingup
Number of trajectories during training
32
Number of trajectories at test time
+ 16
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 The original agentic transformers used 4 trajectories during training, we increase the number 32.
During testing, increasing the number of trajectories has been shown to improve performance. However, performing autoregressive sampling over a large number of trajectories (c.g., 64 x 1000 x 4
15

---

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

---

def Blockvise_fffrremme (remat_fff., imputa, chunk_ size, deterministic).
# remat_ffn: rematerialized ffn
inputs rearrangeeinnuut 'b (c n) -> b cn d', c=chunkssizo
def scan_ifu(remat_fff carry, hidden_states):
outputs remat_ remat_ wtw.ffniddddnssttteee. deferiihisticcdddemmhitttttt return carry, outputs
scan_axis = inputs.ndin ndim
scanc
scan_ffn,
10
crrtable_broadcastt"ppaaaa
"olit__rng==(parrmms.. False, "dropout": True),
in_exes=scan_axis.
13
out_axes==san_axx1s
14
(remat_ffn, None, inputs)
15
rea rearrange(res) (res, 'b bend>> (c n) d!)
return res
18
def blockvise_attn(quer) key, value, query_chunk_size,
19
key_chunk_siee size, dtype, policy, precision, prevent_cse): shape [-1]) astype (dtype)
query query / jup sqrt sqretyerryyyy (huer..shape
query rearrange(query, 'b (c n) h d -> b c h sqquery_chunk_stt key, value map(lamdda
22
rearrange(t, siey_chhnk.ssze), (key, value))
23
24
num_q. batch, num num_heada, dim_per_hent head query... shape num_kv key.shape[o]
def scan_attention (args):
27
query_chunk, query_chunk_iee
args
28
patactooss.patiinn (jaz. checkpoint,
policy-policy)
def scan_kv_bii bean__kv_blockcarrr args):
key. chunk, value_ key. coy_chhuk_iix idx arge
(numerator, denominator, prev prev_max_ scyore) carry
32
attn_weights jup.einsume einsum ('bqhd, bbchk'> query_chunk, key_chunk, precision-precisionl
35
idx, key_chunk_idy, bias bias_chunk jup moveaxis (bias_ _acnass__k,,,,___h 1, 2)
chunk_ _.cunk__bbba (querrynnuuuuuuuuurue
36
attn_weights attn_weights bias_chunk
37
38
max_score - jnp max (etttntwiighhts axis--1, keepdims-
max_acore = jnp maximmm(pree sareee__aa__oreeeeeeeeeeeeeeeeee max_score)
acore jax.
stop gradient (max_score)
exp_weights jup exp.exp(ate exptattn_ weights max_score)
42
exp_values jup einsum(
'bqhv, bwhfbbvhhhthroo exp_weights, value_chunk, precision-precision
correction jup.exp
max max_score)
46
numerator numerator correction + exp_values
47
denominator = denominator correction + exp_ sumaaxis-. keepdims-Trus T return Carry (mumerator, denominator, max_score), None
init_carry Carry
50
jnp. zeros (bbath, query_chunk_size sizun num num_heads, hem_per_headd, caree-uury.dyyppee jnp.zeros zeros (batch, Querrycchnk_siie num_heads, sin_per_heady dtype (ueeeeuuurpee))) dtype), (- (-jnp.inf) inf)
51
52
jnp. ones (batch, query_chunk_siee num nea_heads, 1), dtype stype-query dtype),
(numerator, denominator, max_score),
lax.scam.
55
scan_kv_block, init_carry, xs= (kee, value, arange(O, num_kv))
56
outputs = (numerator denominator) astype (dtype)
return outputs
59
res lax.
60
lambda -, x: (O, scan_attention (x))
(), xs= xs= (query, jnp (r, num_q))
res
rearrange (res, 'nbchd
64
return res
Figure 3: Key parts of the implementation of Blockwise Parallel in Jax. The full code is available battps://githii com homm/hhao99// 1a1ao449911mm.arggg context


Reviews:
[{"review_id": "sAUrexIG2r", "reviewer": "Reviewer_wLgp", "summary": "\"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.", "questions": "- 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.", "limitations": "**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.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "**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.", "weaknesses": "The 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and for a positive assessment of the work. We found the reviewer’s questions and suggestions insightful, and we list our plans to incorporate them below. Please let us know if our answers address your questions. \n\n\n> Q: The 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\n\nA: The primary reason for using this benchmark is training transformers across trajectories provides an easy way to set up a meaningful problem setting that requires challenging long context length.  While RL tasks in general exhibit significant noise, it's important to note that the ExoRL benchmark employed within this paper distinguishes itself by possessing a relatively lower level of noise. This attribute is attributed to the benchmark's utilization of *millions of diverse trajectories* for training, alongside the paper’s focus on training an agentic transformer using MLE loss. \n\nFor reviewer requested language model tasks, we have attached results in the table. The results are validation loss of GPT-XL (1B) on OpenWebText dataset. For shorter context length, BPT shows the validation loss as vanilla transformers and flashattention baselines, and BPT surpasses these baselines in terms of achieving a lower validation loss with extended context length. \n\n| Model              | Context Len | Val Loss | Throughput | SpeedUp |\n| ------------------ | ----------- | -------- | ---------- | ------- |\n| Vanila Attention   | 2048        | 2.46     | 3827       | 1x      |\n| FlashAttention     | 2048        | 2.46     | 4371       | 1.14x   |\n| Blockwise Parallel | 2048        | 2.46     | 3985       | 1.04x   |\n| Vanila Attention   | 4096        | 2.44     | 2340       | 1x      |\n| FlashAttention     | 4096        | 2.44     | 2567       | 1.10x   |\n| Blockwise Parallel | 4096        | 2.44     | 2687       | 1.15x   |\n| Vanila Attention   | 8192        | oom      | 2455       | 1x      |\n| FlashAttention     | 8192        | 2.43     | 2781       | 1.13x   |\n| Blockwise Parallel | 8192        | 2.43     | 2875       | 1.17x   |\n| Vanila Attention   | 16384       | oom      | 1701       | 1x      |\n| FlashAttention     | 16384       | oom      | 1889       | 1.11x   |\n| Blockwise Parallel | 16384       | 2.41     | 2045       | 1.20x   |\n\n> Q: Does BPT work with lower precision since it seems like there's more steps for accumulation and hence likely more noisy? \n\nA: BPT works with lower precision, this is because in the attention part, BPT employs a numerically stable softmax, similar to the traditional standard transformer. It achieves this by maintaining a record of the largest value within each block during the attention phase and subsequently re-normalizing it using the overall maximum value; and in the FFN part, the blockwise computation for MLP does not incur potential numerical errors because it’s position wise. Our experimentation involved a comparison between BPT and a conventional standard transformer using various precisions, ranging from int8 to fp32. This evaluation was done with the provided code snippet. \n\n```\nfor dtype in [jnp.float16, jnp.float32, jnp.int8, jnp.bfloat16]:\n        for seed in range(10):\n            batch, seq_len, n_heads, head_dim = 4, 1024, 48, 64\n            key = jax.random.PRNGKey(seed)\n            if dtype == jnp.int8:\n                inputs = jnp.random.randint(key, (batch, seq_len, n_heads * head_dim), -256, 256, dtype=dtype)\n            else:\n                inputs = jnp.random.normal(key, (batch, seq_len, n_heads * head_dim), dtype=dtype)\n            variables = bpt_block.init(key, inputs)\n            o_bpt = bpt_block.apply(variables, inputs).block_until_ready()\n            variables = vanilla_block.init(key, inputs)\n            o_vanilla = vanilla_block.apply(variables, inputs).block_until_ready()\n            np.testing.assert_allclose(o_bpt, o_vanilla, rtol=1e-5, atol=1e-5)\n```\n\n> Q: Does BPT work with sequence / tensor parallelism schemes? \n\nA: BPT is compatible with sequence parallelism and tensor parallelism. The sequence parallelism distributes sequence dimension across devices, in contrast, BPT implements blockwise computation on sequences for each device. This creates an orthogonal relationship between our method and sequence parallelism, allowing for straightforward combination. BPT is compatible with both tensor and data parallelism since BPT does not change the computation along tensor and batch dimensions.  \n\n\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\nA: On modern accelerators, the speed of SRAM vastly surpasses that of HBM. Consequently, BPT offers the dual benefits of accelerated speed and efficient memory utilization. However, it's important to note that in cases where the hardware features slower SRAM speed, BPT still remains advantageous in terms of memory optimization, even if the degree of speed enhancement is not as pronounced."}, {"review_id": "yDMwdj81BL", "reviewer": "Reviewer_DQeN", "summary": "The 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\n", "questions": "1. 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", "limitations": "The 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.\n\n\n\n\n", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. 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", "weaknesses": "1. 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and for a positive assessment of the work. We found the reviewer’s questions and suggestions insightful, and we list our plans to incorporate them below. Please let us know if our answers address your questions. \n\n\n> Q: 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\nA: For reviewer requested language model tasks, we have attached results in the table. The results are validation loss of GPT-XL (1B) on OpenWebText dataset. For shorter context length, BPT shows the validation loss as vanilla transformers and flashattention baselines, and BPT surpasses these baselines in terms of achieving a lower validation loss with extended context length. \n\n| Model              | Context Len | Val Loss | Throughput | SpeedUp |\n| ------------------ | ----------- | -------- | ---------- | ------- |\n| Vanila Attention   | 2048        | 2.46     | 3827       | 1x      |\n| FlashAttention     | 2048        | 2.46     | 4371       | 1.14x   |\n| Blockwise Parallel | 2048        | 2.46     | 3985       | 1.04x   |\n| Vanila Attention   | 4096        | 2.44     | 2340       | 1x      |\n| FlashAttention     | 4096        | 2.44     | 2567       | 1.10x   |\n| Blockwise Parallel | 4096        | 2.44     | 2687       | 1.15x   |\n| Vanila Attention   | 8192        | oom      | 2455       | 1x      |\n| FlashAttention     | 8192        | 2.43     | 2781       | 1.13x   |\n| Blockwise Parallel | 8192        | 2.43     | 2875       | 1.17x   |\n| Vanila Attention   | 16384       | oom      | 1701       | 1x      |\n| FlashAttention     | 16384       | oom      | 1889       | 1.11x   |\n| Blockwise Parallel | 16384       | 2.41     | 2045       | 1.20x   |\n\n\n> Q: 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\nA: We plan to make all BPT code available upon acceptance. As for the hardware prerequisites, our experiments were conducted on various setups, including single-device and multi-host configurations, for both GPUs (A100) and TPUs (v3 and v4), therefore BPT should work on most GPUs and TPUs. \n\n\n> Q: The 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.\n\n\nA: Undoubtedly, the potential to leverage extensive long context data is substantial. However, the current landscape presents specific challenges that warrant careful consideration. While resources like the Books Corpus, Github, and Arxiv offer a plethora of lengthy documents, effectively harnessing these resources demands addressing some key points. \n\nFor instance, take the case of code repositories like Jax and Pytorch, which encompass a remarkable volume of approximately **120 million and 150 million tokens** respectively, excluding dependencies. These repositories harbor a wealth of contextual information. However, public Github datasets often contain only individual files rather than complete repositories, leading to an average document length significantly shorter than what models like BPT are capable of handling. Consider the 'starcoderdata' dataset, where the average document length is approximately 1500 tokens. \n\nSince the authors' primary focus is on addressing architecture bottlenecks, dealing with dataset challenges remains an ancillary pursuit.\n"}, {"review_id": "KINbuQKnDU", "reviewer": "Reviewer_t1yD", "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.\n\nExperiments 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.", "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.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the review. We found the reviewer’s questions and suggestions insightful, and we list our plans to incorporate them below. You raised a good question about how exactly does BPT help save memory. Please let us know if our answers address your questions. \n\n\n> Q: 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?\n\n\nA: We provide an analysis of the memory costs associated with vanilla Transformer, the memory-efficient/flash attention variant, and BPT. This breakdown will provide you with a clear understanding of the memory usage of each architecture.\n\n**Vanilla Transformer**\n\n*Attention:*\nFor Q, K, V, saving their input x needs 2bsh bytes, where b is batch size, s is sequence length, and h is hidden dimension \nFor QK^T matmul, saving activation Q and K needs 4bsh bytes \nFor softmax(QK^T), saving input QK^T needs 2bs ** 2a, a is number of attention heads\nFor mask and dropout, saving mask needs bs ** 2a\nFor score * V, saving score needs 2bs ** 2a and saving V needs 2bsh \nFor output projection and dropout, saving the input needs 2bsh, saving dropout mask needs bsh\nIn total attention activation size 11bsh + 5bs ** 2a \nThus the peak attention activation size of attention is O(s ** 2) with checkpointing.\n\n*FFN:*\nFor first linear layer, saving input needs 2bsh \nFor activation, saving input needs 8bsh\nFor second linear layer, saving input needs 8bsh\nFor dropout, saving mask needs bsh\nIn total MLP activation size is 19bsh\nWith checkpointing, the peak activation size of FFN is 8bsh. \n\nConsequently, for large context length, the memory cost of activation in vanilla Transformer is O(s**2).\n\n**BPT**\n\n*Attention:*\nSince BPT doesn’t materialize full attention and instead computes it blockwise, it needs to store intermediate blockwise activations in the key-value for loop which has 4bch peak activation size with checkpointing, and it also needs to store q output activations for query for loop which requires 2bsh. Since s >> c, the peak activation size is 2bsh. \n\n*FFN:*\nIterating ffn over blocks, BPT needs to save the following activations:  \nFor first linear layer, saving input needs 2bch \nFor activation, saving input needs 8bch\nFor second linear layer, saving input needs 8bch\nFor dropout, saving mask needs bch\nIn total 19bch, in addition, storing the output of the for loop needs 2bsh. Therefore, the peak FFN activation size is 2bsh.\n \nConsequently, each BPT layer’s memory cost of activation is 2bsh. \n\n**Memory efficient  / flash attention transformer**\n\n*Attention:*\nSame as BPT attention, the memory cost is 2bsh + 4bch.\n\n*FFN:*\nSame as vanilla FFN, the peak activation size is 8bsh. \n\nConsequently, each flashattention layer’s memory cost is 8bsh\n\n---\nComparing activation memory of FlashAttention/MemEff with BPT, we see that BPT offers 8bsh/2bsh = 4 times memory saving. By taking into account other factors of memory cost such as model and optimizer state, hardware memory, and distributed partitioning, BPT allows training 2-4x larger context length. We plan to include the analysis into the revision. \n"}, {"review_id": "wBsfF5tOpO", "reviewer": "Reviewer_KeWq", "summary": "The 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.", "questions": "- 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", "limitations": "Limitions are well noted in the paper.", "rating": 7, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- 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.", "weaknesses": "- Line 23: \"such as\" is repeated.\n- The contribution on top of Flash Attention is pretty limited, but still the impact is pretty significant.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and for a positive assessment of the work. We found the reviewer’s questions and suggestions insightful, and we list our plans to incorporate them below. Please let us know if our answers address your questions. \n\n\n> Q: 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\nA: AT with vanilla transformer architecture runs out of memory for more than 16 trajectories. We have conducted experiments for AT with 16 trajectories. On ExoRL, increasing to 16 trajectories improves results significantly compared to 4 trajectories. AT + BPT trained with 64 trajectories performs the best. \n\n| ExoRL benchmark | BC-10% | DT | AT | AT | AT | AT + FA | AT + BPT |\n|-----------------|--------|----|----|----|----|---------|---------|\n| Task            |        |    | N Trajs = 4 | N Trajs = 16 | N Trajs = 64 | N Trajs = 64 | N Trajs = 64 |\n| Walker Stand    | 52.91  | 34.54  | 148.55   | 161.45   | oom   | oom   | 195.45     |\n| Walker Run      | 34.81  | 49.82  | 98.56    | 125.44   | oom   | oom   | 145.98     |\n| Walker Walk     | 13.53  | 34.94  | 84.56    | 89.66    | oom   | oom   | 103.88     |\n| Cheetah Run     | 34.66  | 67.53  | 145.68   | 164.31   | oom   | oom   | 178.75     |\n| Jaco Reach      | 23.95  | 18.64  | 58.98    | 60.95    | oom   | oom   | 87.56      |\n| Cartpole Swingup| 56.82  | 67.56  | 187.56   | 191.34   | oom   | oom   | 220.56     |\n| Total Average   | 36.11  | 45.51  | 120.65   | 132.19   | oom   | oom   | 155.36     |\n\n\n> Q: Why is stop_gradient needed in the implementation for max_score?\n\nA: Jax compiler and XLA use heuristics to optimize computation graphs created by users code, but such optimization is oftentimes inefficient due to combinatorial possibilities. Without using the stop_gradient, we found that the compilation process becomes time-consuming. Given that the gradient is unnecessary for max_score, utilizing stop_gradient serves as a clear directive to the compiler to disregard its gradient during optimization.\n"}]

Author rebuttal:

Meta-review: The paper introduces an efficient blockwise computation scheme for transformers. Instead of fully materializing the NxN attention matrix, the computation for queries is fused with subsequent feedforward layers, which significantly improves memory efficiency. Importantly, the method outputs the exact same result as standard transformers. All the reviewers agreed that this is a straightforward but well motivated contribution, and it seems very practically useful, so this is a clear accept.
Decision: Accept (spotlight)