Data Movement Is All You Need: A Case Study on Optimizing Transformers

By Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, Torsten Hoefler. 2020

190 citations recorded by the source corpus.

Transformers have become widely used for language modeling and sequence learning tasks, and are one of the most important machine learning workloads today. Training one is a very compute-intensive task, often taking days or weeks, and significant attention has been given to optimizing transformers. Despite this, existing implementations do not efficiently utilize GPUs. We find that data movement is the key bottleneck when training. Due to Amdahl's Law and massive improvements in compute performance, training has now become memory-bound. Further, existing frameworks use suboptimal data layouts. Using these insights, we present a recipe for globally optimizing data movement in transformers. We reduce data movement by up to 22.91% and overall achieve a 1.30x performance improvement over state-of-the-art frameworks when training BERT. Our approach is applicable more broadly to optimizing deep neural networks, and offers insight into how to tackle emerging performance bottlenecks.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Hardware and Architecture, Model Architectures, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Christopher R'e
2. Sergey Levine
3. Stefano Ermon
4. Yoshua Bengio
5. Ruslan Salakhutdinov
6. Tri Dao
7. Tom Goldstein
8. Pieter Abbeel

In Computer Vision and Pattern Recognition

1. Luc Van Gool
2. Dacheng Tao
3. Jitendra Malik
4. Xiaogang Wang
5. Trevor Darrell
6. Ming-Hsuan Yang
7. Andrea Vedaldi
8. Chunhua Shen

In Hardware and Architecture

1. Tri Dao
2. Christopher R'e
3. Atri Rudra
4. Stefano Ermon
5. Daniel Y. Fu
6. Yuxiong He
7. Beidi Chen
8. Ion Stoica

In Model Architectures

1. Tri Dao
2. Christopher R'e
3. Daniel Y. Fu
4. Stefano Ermon
5. Atri Rudra
6. Beidi Chen
7. Yuandong Tian
8. Ion Stoica

In Efficiency & Compression

1. Tri Dao
2. Yuxiong He
3. Samyam Rajbhandari
4. Reza Yazdani Aminabadi
5. Minjia Zhang
6. Ion Stoica
7. Shaden Smith
8. Jeff Rasley

Citations in the corpus

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, background): “On modern GPUs, compute speed has out-paced memory speed [64–66], and most operations in Transformers are bottlenecked by memory accesses [45].”
• DeepSpeed- Inference: Enabling Efficient Inference of Transformer Models at Unprecedented Scale (2022, background): “Transformer Kernels: There is also a suite of work focused on accelerating the performance of transformer kernels [24], [25], [26].”
• A Cookbook of Self-Supervised Learning (2023, background): “Another approach, is to resort to IO-aware optimizations [Ivanov et al., 2021], the most known one perhaps being FlashAttention [Dao et al., 2022].”
• Wavelet: Efficient DNN Training with Tick-Tock Scheduling (2021, methodology): “Although recent literature advances model parallel training performance by incorporating pipeline parallelism (Huang et al., 2019; Narayanan et al., 2019), the issue of GPU under-utilization is still not eliminated (Ivanov et al., 2021).”
• Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time (2023, background): “We highlight some hardware characteristics of GPUs: • Small-batch generation is bottlenecked by GPU memory I/Os (NVIDIA, 2022; Ivanov et al., 2021; Dao et al., 2022).”
• EfficientViT: Memory Efficient Vision Transformer with Cascaded Group Attention (2023, background): “Recent studies reveal that memory-inefficient operations are mainly located in MHSA rather than FFN layers [31,33].”