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

2020/06/30 by Andrei Ivanov, Nikoli Dryden, Ivanov, Andrei +8 · 24 citations
Computer Science · Engineering · Mathematics · #Advanced Neural Network Applications #Artificial intelligence #Artificial neural network #Bottleneck #Computer architecture #Computer engineering #Computer science #Deep learning #Embedded system #Engineering #FOS: Computer and information sciences #Implementation #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Machine learning #Parallel Computing and Optimization Techniques #Software engineering #Topic Modeling #Training set #Transformer #Voltage #cs.LG #stat.ML

paper · pdf · doi:10.48550/arxiv.2007.00072

published in arXiv (Cornell University) 3, 711-732 (Cornell University) · 22 pages, 8 figures; MLSys 2021 camera ready

openalex publication_date 2020/06/30 · arxiv created 2021/11/08 · arxiv updated 2021/11/09 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Transformers 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 a BERT encoder layer and 1.19x for the entire BERT. Our approach is applicable more broadly to optimizing deep neural networks, and offers insight into how to tackle emerging performance bottlenecks.

Citations

• cuDNN: Efficient Primitives for Deep Learning
• Sequence Transduction with Recurrent Neural Networks
• Long Short-Term Memory
• Distributed Representations of Words and Phrases and their Compositionality
• Learning Phrase Representations using RNN Encoder–Decoder for Statistical Machine Translation
• ImageNet classification with deep convolutional neural networks
• Training Deep Nets with Sublinear Memory Cost
• Geometric deep learning: going beyond Euclidean data
• Device Placement Optimization with Reinforcement Learning
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• DLVM: A modern compiler infrastructure for deep learning systems
• Intel nGraph: An Intermediate Representation, Compiler, and Executor for Deep Learning
• Tensor Comprehensions: Framework-Agnostic High-Performance Machine Learning Abstractions
• Glow: Graph Lowering Compiler Techniques for Neural Networks
• TF-Replicator: Distributed Machine Learning for Researchers
• Massively Distributed SGD: ImageNet/ResNet-50 Training in a Flash
• Generating Long Sequences with Sparse Transformers
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Demystifying Parallel and Distributed Deep Learning
• CTRL: A Conditional Transformer Language Model for Controllable Generation
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• LLM-Supported Manufacturing Mapping Generation
• Fast Transformer Decoding: One Write-Head is All You Need
• GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism
• Dissecting the Graphcore IPU Architecture via Microbenchmarking
• Reformer: The Efficient Transformer
• ALBERT: A Lite BERT for Self-supervised Learning of Language Representations
• Scaling Laws for Neural Language Models
• MLIR: A Compiler Infrastructure for the End of Moore's Law
• TIRAMISU: A Polyhedral Compiler for Dense and Sparse Deep Learning
• Language Models are Few-Shot Learners
• Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
• An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale

Cited by

• Energy Scaling Laws for Diffusion Models: Quantifying Compute and Carbon Emissions in Image Generation
• PICNIC: Silicon Photonic Interconnected Chiplets with Computational Network and In-memory Computing for LLM Inference Acceleration
• A Study on Inference Latency for Vision Transformers on Mobile Devices
• Tawa: Automatic Warp Specialization for Modern GPUs with Asynchronous References
• First Attentions Last: Better Exploiting First Attentions for Efficient Transformer Training
• AURA: Action-Gated Memory for Robot Policies at Constant VRAM
• Exact Causal Attention with 10% Fewer Operations
• Confidential LLM Inference: Performance and Cost Across CPU and GPU TEEs
• Zen-Attention: A Compiler Framework for Dynamic Attention Folding on AMD NPUs
• COMET: A Framework for Modeling Compound Operation Dataflows with Explicit Collectives
• EMLIO: Minimizing I/O Latency and Energy Consumption for Large-Scale AI Training
• Improving OCR for Historical Texts of Multiple Languages
• Retrieval-Aware Distillation for Transformer-SSM Hybrids
• La RoSA: Enhancing LLM Efficiency via Layerwise Rotated Sparse Activation
• Power Law Guided Dynamic Sifting for Efficient Attention
• Hardware-Efficient Attention for Fast Decoding
• FlatAttention: Dataflow and Fabric Collectives Co-Optimization for Efficient Multi-Head Attention on Tile-Based Many-PE Accelerators
• Parallel Scaling Law for Language Models
• Nova: An End-to-End MLIR Compiler for Deep Learning
• Morello: Compiling Fast Neural Networks with Dynamic Programming and Spatial Compression
• Network Design for Wafer-Scale Systems with Wafer-on-Wafer Hybrid Bonding
• Nonlinear Computation with Linear Optics via Source-Position Encoding
• ComFuse: Fusing Complex Memory-Intensive Subgraphs with Compute-Intensive Kernels For Modern GPU Architectures
• Unified Lookup-Table Inference with Signed-Digit K/V Caches for Ternary LLMs
• MCHA: A Memory-Centric Hierarchical Architecture for Parallel-Sequential Computing

Related

• Tempo: Accelerating Transformer-Based Model Training through Memory Footprint Reduction
• READ: Recurrent Adaptation of Large Transformers
• Mesa: A Memory-saving Training Framework for Transformers
• LightSeq2: Accelerated Training for Transformer-based Models on GPUs
• A Multi-Level Framework for Accelerating Training Transformer Models