SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Table 3 shows the analysis of memory usage across different settings with three distinct approaches: Vanilla Transformer, MemoryEfficient, and our proposed method, BPT. It is evident that Vanilla transformers consumes the highest amount of memory, while MemoryEfficient and BPT offer notable improvements in memory optimization. Notably, our BPT technique consistently outperforms both Vanilla transformers and MemoryEfficient in all settings, showcasing memory efficiency.
Table 3: Memory usage comparison for different settings. "oom" denotes out of memory.
Setting	3B on A100	13B on 8 A100
Context Length	Vanilla	MemoryEfficient	BPT	Vanilla	MemoryEfficient	BPT
8192	64GB	44GB	43GB	59GB	44GB	42GB
16384	OOM	47GB	45GB	OOM	46GB	45GB
32768	OOM	55GB	52GB	OOM	55GB	52GB
65536	OOM	75GB	70GB	OOM	75GB	68GB
131072	OOM	OOM	79GB	OOM	OOM	78GB
5.2 Evaluation on Throughput and Speed
In Table 4, we present a comparison of the throughput achieved by different attention mechanisms on the GPT-XL (1B) model trained on the OpenWebText dataset using 8 GPUs. Throughput is measured as number of tokens processed per device per second. We evaluate the performance at various context lengths, including 2K, 8K, 16K, 33K, and 65K tokens. Our proposed method achieves competitive throughput as MemeoryEfficient mechanism, and surpasses the Vanilla transformer, achieving 1.17x speedup at context length 8k and 1.2x speedup at context length 16k. At context length 32K and 64K, our method maintains high throughput and training speed, while the alternatives cannot train due to running out of memory. This demonstrates the scalability and efficiency of our proposed method, allowing it to effectively handle large context lengths without compromising on throughput and training speed.
5.3 Evaluation on Reinforcement Learning
In this section, we present the results of applying BPT to improve the performance of transformers in reinforcement learning (RL). We report our results in Table 5, where we evaluate our proposed model on the ExoRL benchmark across six different tasks. On ExoRL, we report the cumulative return, as per ExoRL [58]. The numbers of BC, DT [6] and AT [34] are from the ExoRL and AT paper.  $\mathrm{AT} + \mathrm{ME}$  and  $\mathrm{AT} + \mathrm{BPT}$  numbers are run by ourselves. Since the ExoRL data is significantly more diverse than D4RL because it is collected using various unsupervised RL algorithms [28, 35, 33], it is found that TD learning performs best while behavior cloning struggles [58]. AT [34] shows that conditioning transformers on multiple trajectories with relabeled target return can significantly outperforms behavior cloning approaches BC-10% and DT, and achieves competitive results with TD learning. For more details, please refer to their papers. We are interested in applying BPT to improve the performance of AT by conditioning on a 32 trajectories rather than 4 trajectories in original work. It is worth noting that each trajectory has  $1000 \times 4$  length where 1000 is sequence length while 4 is return-state-action-reward, making training 32 trajectories with 350M size model infeasible for both Vanilla and MemoryEfficient. Results in Table 5 show that, by scaling the sequence length, AT + BPT consistently outperforms the original transformers model in all six tasks, achieving a total average return of 111.13 compared to the original transformers model's total average return of 83.02
Transformers have garnered significant attention in the field of natural language processing (NLP) and have become the basis for numerous state-of-the-art models. Several works have explored memory-
Table 4: Throughput comparison on GPT-XL (1B) using OpenWebText dataset. Throughput is measured as 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	1x
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	na	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 of BPT on improving transformers in RL. All the baselines use vanilla attention.  $\mathrm{{AT}} + \mathrm{{ME}}$  denotes using "MemoryEfficient".  $\mathrm{{AT}} + \mathrm{{BPT}}$  denotes using Blockwise Parallel.
ExoRL benchmark	BC-10%	DT	AT	AT	AT + ME	AT + BPT
Task	N Trajs = 4	N Trajs = 32	N Trajs = 32	N Trajs = 32
Walker Stand	52.91	34.54	68.55	OOM	OOM	95.45
Walker Run	34.81	49.82	88.56	OOM	OOM	105.88
Walker Walk	13.53	34.94	64.56	OOM	OOM	78.56
Cheetah Run	34.66	67.53	125.68	OOM	OOM	178.75
Jaco Reach	23.95	18.64	52.98	OOM	OOM	87.56
Cartpole Swingup	56.82	67.56	97.81	OOM	OOM	120.56
Total Average	36.11	45.51	83.02	OOM	OOM	111.13
efficient techniques to address the memory limitations of Transformers and enable their application to longer input sequences. One line of research focuses on various approximation techniques or compressing along the sequence dimension [see e.g. 24, 12, 14, 4, 44, 56, 38, 25]. Other works explored replacing attention [19, 20, 43, 23, 3, 59, 42, 55]. Another line of work explores partitioning the large hidden dimension of the feedforward network into parts and retrieving only one part per token [30, 50, 17, 26, 60, 62]. Additionally, extending the context by attending over states from previous sequences has been explored [13, 46], as well as combining local and global contexts [21, 11]. For a comprehensive review of these techniques, we recommend referring to the surveys by Tay et al. [53], Narang et al. [40], Tay et al. [52]. Several studies explored sharding large model on distributed devices tensor, data, or sequence parallelism [51, 16, 57, 27, 61, 31, 48]. Ours shares similarities with the sequence parallelism [27] where sequences are distributed across devices, in contrast, ours implements blockwise computation on sequences for each device. This creates an orthogonal relationship between our method and sequence parallelism, allowing for straightforward combination. In addition, our methodology is compatible with both tensor and data parallelism. Another direction involves computing exact self-attention in a blockwise manner using the tiling technique [39]. This approach has led to the development of memory efficient attention mechanisms [14, 44]. In line with these advancements, our work falls into this category. We propose computing both the feedforward
network and self-attention in a blockwise manner, resulting in a significant reduction in memory requirements.
In conclusion, we propose a blockwise parallelization approach to reduce the memory requirements of Transformers, the backbone of state-of-the-art AI models. Our approach enables processing longer input sequences while maintaining or improving performance. Through extensive experiments, we demonstrate its effectiveness, achieving up to 4x memory reduction than memory-efficient Transformers. Our contributions include a practical method for large context sizes in large transformers models. With the increasing capability of hardware, larger models and longer context length are widely used in AI research. At the same time, as we are pushing up against physics and fabrication limits, it is more important to design scaling approaches as efficient as possible to scale up large models and large context size. Our approach holds promise for training and evaluating complex models with longer input sequences, potentially driving new breakthroughs in machine learning research.
Limitations and Future Work. Although our method achieves state-of-the-art low memory usage for transformers models, it does have some limitations that need to be addressed:
- Optimal performance. While our implementation prioritizes simplicity with high-level Jax operations, optimizing low-level operations is crucial for achieving optimal performance. In future work, we suggest considering porting our method to CUDA and OpenAI Triton to achieve minimal memory cost and maximum speedup.
This project is supported in part by ONR under N00014-21-1-2769. We would like to express our sincere appreciation to the members of the RLL Lab and Berkeley AI 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.
[1] J.-B. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y. Hasson, K. Lenc, A. Mensch, K. Millican, M. Reynolds, et al. Flamingo: a visual language model for few-shot learning. Advances in Neural Information Processing Systems, 35:23716-23736, 2022.
[2] K. Andrej. GitHub - karpathy/nanoGPT: The simplest, fastest repository for training/finetuning medium-sized GPTs. — github.com. https://github.com/karpathy/nanoGPT, 2023. [Accessed 16-May-2023].
[3] I. Bello. Lambda networks: 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 arXiv:2004.05150, 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, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-15097, 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.03374, 2021.
[8] T. Chen, B. Xu, C. Zhang, and C. Guestrin. Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174, 2016.
[9] X. Chen, X. Wang, S. Changpinyo, A. Piergiovanni, P. Padlewski, D. Salz, S. Goodman, A. Grycner, B. Mustafa, L. Beyer, et al. Pali: A jointly-scaled multilingual language-image model. arXiv preprint arXiv:2209.06794, 2022.
[10] X. Chen, M. Lin, N. Scharli, and D. Zhou. Teaching large language models to self-debug. arXiv preprint arXiv:2304.05128, 2023.
[11] R. Child, S. Gray, A. Radford, and I. Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509, 2019.
[12] K. Choromanski, V. Likhosherstov, D. Dohan, X. Song, A. Gane, T. Sarlos, P. Hawkins, J. Davis, A. Mohiuddin, L. Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794, 2020.
[13] Z. Dai, Z. Yang, Y. Yang, J. Carbonell, Q. V. Le, and R. Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860, 2019.
[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.04805, 2018.
[16] Facebook. Fully Sharded Data Parallel: faster AI training with fewer GPUs — engineering.fb.com. https://engineering.fb.com/2021/07/15/open-source/fsdp/, 2023. [Accessed 16-May-2023].
[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(1):5232-5270, 2022.
[56] S. Wang, B. Z. Li, M. Khabsa, H. Fang, and H. Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768, 2020.
[57] Y. Xu, H. Lee, D. Chen, B. Hechtman, Y. Huang, R. Joshi, M. Krikun, D. Lepikhin, A. Ly, M. Maggioni, et al. Gspmd: general and scalable parallelization for ml computation graphs. arXiv preprint arXiv:2105.04663, 2021.
[58] D. Yarats, D. Brandfonbrener, H. Liu, M. Laskin, P. Abbeel, A. Lazaric, and L. Pinto. Don't change the algorithm, change the data: Exploratory data for offline reinforcement learning. arXiv preprint arXiv:2201.13425, 2022.
[59] S. Zhai, W. Talbott, N. Srivastava, C. Huang, H. Goh, R. Zhang, and J. Susskind. An attention-free transformer. arXiv preprint arXiv:2105.14103, 2021.
[60] Z. Zhang, Y. Lin, Z. Liu, P. Li, M. Sun, and J. Zhou. Moefication: Transformer feed-forward layers are mixtures of experts. arXiv preprint arXiv:2110.01786, 2021.
[61] L. Zheng, Z. Li, H. Zhang, Y. Zhuang, Z. Chen, Y. Huang, Y. Wang, Y. Xu, D. Zhuo, E. P. Xing, et al. Alpha: Automating inter-and {Intra-Operator} parallelism for distributed deep learning. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 559–578, 2022.
[62] S. Zuo, Q. Zhang, C. Liang, P. He, T. Zhao, and W. Chen. Moebert: from bert to mixture-of-experts via importance-guided adaptation. arXiv preprint arXiv:2204.07675, 2022.
A Experiment Details
A.1 Evaluation of Memory
In the experimental results presented in Section 5.1, we used model parallelism to partition the model across 8 GPUs or 64 TPUv4 units. Our evaluation focused on determining the maximum achievable sequence length, using a sequence number of one. For TPUs, we utilized its default training configuration, which involved performing matmul operations in bfloat16 format with weight accumulation in float32. On the other hand, for GPUs, we adopted the default setup, where all operations were performed in float32.
To profile memory usage, we utilized jax.profile and repeated the evaluation 100 times, reporting the average results. We conducted a grid search for the optimal query block size and key-value block size, considering values from the set [16, 64, 128, 512, 1024, 2048, 4096]. For each method, we reported the lowest memory achieved.
A.2 Evaluation of Throughput
In the evaluation presented in Section 5.2, we split OpenWebText following the methodology of [2]. Throughput is measured as tokens per device per second. To ensure a fair comparison, we performed a grid search for the optimal query block size and key-value block size, considering values from the set [16, 64, 128, 512, 1024, 2048, 4096]. For gradient checkpointing [8], we additionally grid search among three commonly used checkpointing policies including nothing_saveable, dots_saveable, and dots_with_no_batch_dims_saveable for attention and use nothing_saveable for feedforward network (FFN). For more details, please refer to Jax documentation. We selected the best performing configuration for both baselines and our method.
The training was conducted using FSDP [16] and gradient accumulation. We used weight decay of 0.1 and utilized cosine learning rate decay with a maximum learning rate of  $2.0 \times e^{-4}$ . For sequence lengths of 2048, 4096, 8192, 16384, the batch sizes in trajectories were set as 8, 4, 2, 1, 1 respectively. We use gradient accumulation to accumulate batch size in tokens to 1 million per batch.
A.3 Evaluation on RL
Table 6: Hyperparameters used in RL evaluation.
Hyperparameter	Value
Number of layers	3
Number of attention heads	1
Embedding dimension	128
Activation function	ReLU
Batch size	64
Dropout	0.1
Learning rate	10-4
Learning rate decay	Linear warmup for 105 steps
Grad norm clip	0.25
Weight decay	10-4
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	4 → 32
Number of trajectories at test time	4 → 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 6. The original agentic transformers used 4 trajectories during training, we increase the number to 32.
During testing, increasing the number of trajectories has been shown to improve performance. However, performing autoregressive sampling over a large number of trajectories (e.g.,  $64 \times 1000 \times 4$
total number of tokens) can be computationally slow. To reduce the sampling time, we limited the rollout to 16 trajectories.
def blockwise_ffn(remat_ffn, inputs, chunk_size, deterministic):
    # remat_ffn: a rematerialized ffn
    inputs = rearrange(inputs, 'b (c n) d -> b c n d', c=chunk_size)
    def scan_ffn(remat_ffn, carry, hidden_states):
        outputs = remat_ffn(hidden_states, deterministic=deterministic)
        return carry, outputs
    scan_axis = inputs.ndim - 2
    _, res = nn.scan(
        scan_ffn,
        variable_broadcast="params",
        split_rings={'params": False,"dropout": True},
        in Axes=scan_axis,
        out Axes=scan_axis,
    ) (remat_ffn, None, inputs)
    res = rearrange(res,'b c n d -> b (c n) d')
    return res
def blockwise_attn(query, key, value, query_chunk_size,
            key_chunk_size, dtype, policy, precision, prevent_cse):
    query = query / jnp.sqrt(query.shape[-1]).astypedtype)
    query = rearrange(query,'b (c n) h d -> nb c h d', c=query_chunk_size)
    key, value = map(lambda t: rearrange(t,'b (c n) h d -> nb c h d',
                    c=key_chunk_size), (key, value))
    num_q, batch_, num_heads, dim_per_head = query.shape
    num_kv = key.shape[0]
    def scanattention(args):
        query_chunk, query_chunk_idx = args
        @functools.partial(jax.checkpoint, prevent_cse=prevent_cse, policy=policy)
        def scan_kv_block(carry, args):
            key_chunk, value_chunk, key_chunk_idx = args
            (numerator, denominator, prev_max_score) = carry
            attnweights = jnp.einsum('bqhdt,bkhdt->bqhkt', query_chunk,
            key_chunk, precision=precision)
            bias_chunk = _chunk.bias_fn(query_chunk_idx, key_chunk_idx)
            bias_chunk = jnp.moveaxis(bias_chunk, 1, 2)
            attnweights = attnweights + bias_chunk
        max_score = jnp.max(attnweights, axis=-1, keepdims=True)
        max_score = jnpmaximum(prev_max_score, max_score)
        max_score = jax.lax.stop-gradient(max_score)
        expweights = jnp.exp(attnweights - max_score)
        exp_values = jnp.einsum(
            'bqhdt,bhdt->bqhdt', expweights, value_chunk, precision=precision
        )
        correction = jnp.exp Prev_max_score - max_score)
        numerator = numerator * correction + exp_values
        denominator = denominator * correction + expweights-sum(axis=-1, keepdims=True)
        return Carry(numerator, denominator, max_score), None
        init carrvy = Carry(
            jnp.zeros((batch, query_chunk_size, num_heads, dim_per_head), dtype=query_dtype),
            jnp.zeros((batch, query_chunk_size, num_heads, dim_per_head), dtype=query_dtype),
            (-jnp.inf) * jnp.ones((batch, query_chunk_size, num_heads, 1), dtype=query_dtype),
        )
        (numerator, denominator, max_score), _ = lax.scan(
            scan_kv_block, init carrvy, xs=(key, value, jnp.arange(0, num_kv))
        )
        outputs = (numerator / denominator).astypedtype)
        return outputs
    _, res = lax.scan(
            lambda_, x: ((), scan attentive(x)), 
            ), xs=(query,jnp.arange(0, num_q))
    )
res = rearrange(res,'n b c h d -> b (n c) h d')
return res
Figure 3: Key parts of the implementation of Blockwise Parallel in Jax. The full code is available at https://github.com/lhao499/llm_large_context
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64