SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Hao Liu
hao.liu@cs.berkeley.edu
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.
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.
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 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 4 times longer than FlashAttention [14] and Memory Efficient Attention [44]. Section 3.1 provides a memory cost breakdown.
- Vanilla Attention - FlashAttention / Memory Efficient Attention - BlockWise Parallel Transformer
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 training 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 \in \mathbb{R}^{s \times d}$  where  $s$  is the sequence length and  $d$  is the head dimension. We compute the matrix of outputs as:
$$
\operatorname {A t t e n t i o n} (Q, K, V) = \operatorname {s o f t m a x} \left(\frac {Q K ^ {T}}{\sqrt {d}}\right) V, \tag {1}
$$
where softmax is applied row-wise. Standard attention implementations materialize the matrices  $QK^T$  and  $\mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)$  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.
$$
\operatorname {F F N} (x) = \max  \left(0, x W _ {1} + b _ {1}\right) W _ {2} + b _ {2} \tag {2}
$$
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  $\mathrm{softmax}(QK^T)$  [39, 14, 44]. This approach involves splitting the sequences  $Q\in \mathbb{R}^{s\times d}$  into  $B_{q}$  blocks and sequences  $K,V\in \mathbb{R}^{s\times d}$  into  $B_{kv}$  blocks. For each query block, the blockwise attention  $\mathrm{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\leq i\leq B_{q}$ , the corresponding attention output can be computed by scaling each blockwise attention as follows:
$$
\operatorname {A t t e n t i o n} \left(Q _ {i}, K, V\right) = \operatorname {S c a l i n g} \left(\left\{\exp \left(Q _ {i} K _ {j} ^ {T}\right) V _ {j} \right\} _ {j = 1} ^ {B _ {k v}}\right). \tag {3}
$$
The scaling operation scales each blockwise attention based on the difference between the blockwise maximum and the global maximum:
$$
\operatorname {A t t e n t i o n} \left(Q _ {i}, K _ {j}, V _ {j}\right) = \exp \left(Q _ {i} K _ {j} ^ {T} - \max  \left(Q _ {i} K _ {j} ^ {T}\right)\right) / \sum \exp \left(Q _ {i} K _ {j} ^ {T} - \max  \left(Q _ {i} K _ {j} ^ {T}\right)\right)
$$
$$
\operatorname * {m a x} _ {i} = \max  \left(\max  (Q _ {i} K _ {1} ^ {T}), \dots , \max  (Q _ {i} K _ {B} ^ {T})\right)
$$
$$
\operatorname {A t t e n t i o n} (Q _ {i}, K, V) = \left[ \exp (Q _ {i} K _ {j} ^ {T} - \max _ {i}) \operatorname {A t t e n t i o n} (Q _ {i}, K _ {j}, V _ {j}) \right] _ {j = 1} ^ {B _ {k v}}.
$$
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:
$$
\operatorname {O u t p u t} _ {i} = \operatorname {F F N} \left(\operatorname {A t t e n t i o n} \left(Q _ {i}, K, V\right) + Q _ {i}\right) + \operatorname {A t t e n t i o n} \left(Q _ {i}, K, V\right) + Q _ {i}.
$$
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.
Figure 2: We use the same model architecture as the original transformers but with a different way of organizing the compute. In the diagram, we explain this by showing that for the bottom first incoming input block, we project it into query; then we 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 a residual connection. In our proposed approach, this process is then repeated for the other incoming input blocks.
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  $2bsh$  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  $4bsh$  bytes. For softmax( $QK^T$ ), saving input  $QK^T$  needs  $2bs^2a$  bytes, where  $a$  is the number of attention heads. For mask and dropout, saving mask needs  $bs^2a$  bytes. For score  $\times V$ , saving score needs  $2bs^2a$  bytes, and saving  $V$  needs  $2bsh$  bytes. For output projection and dropout, saving the input needs  $2bsh$  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  $2bsh$  bytes. For activation, saving input needs  $8bsh$  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 is  $8bsh$  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  $4bch$  with checkpointing. Additionally, it needs to store  $q$  output activations for the query loop, which requires  $2bsh$  bytes. Since  $s \gg c$ , the maximum activation size is  $2bsh$ .
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  $8bch$  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-Efficient / Flash Attention:
Attention: Similar to BPT attention, the maximum activation size is  $2bsh$ .
FFN: Similar to the vanilla FFN, the maximum activation size is  $8bsh$ .
Consequently, each Flash Attention layer's memory cost is  $8bsh$ .
Comparing the activation memory of Flash Attention/Memory-Efficient transformers 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 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].
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<sup>1</sup> 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 attentive 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 OpenText 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 URLB [28]. The datasets are collected by unsupervised RL and then relabeled using task reward function. The resulting mixed dataset consists of 8 millionstimesteps (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
<https://github.com/lhao499/11m_large_context
Table 1: Sizes and architectures of the models which we evaluated in experiments.
Model Name	nparams	nlayers	dmodel	nhheads	dhead
GPT 1B	1.3B	24	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-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.
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.
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.
1 A100	PartitionSpec	Vanilla Attention	MemoryEfficient	Blockwise Parallel
350M	(1,1,1)	16K (16384)	65K (65536)	131K (131072)
1B	(1,1,1)	16K (16384)	65K (65536)	131K (131072)
3B	(1,1,1)	8K (8192)	16K (16384)	65K (65536)
8 A100	PartitionSpec	Vanilla Attention	MemoryEfficient	Blockwise Parallel
3B	(1,1,8)	16K (16384)	65K (65536)	131K (131072)
7B	(1,1,8)	16K (16384)	65K (65536)	131K (131072)
13B	(1,1,8)	8K (8192)	33K (32768)	65K (65536)
30B	(1,1,8)	8K (8192)	16K (16384)	65K (65536)
64 TPUv4	PartitionSpec	Vanilla Attention	MemoryEfficient	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)	1k (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, 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.