Source-linked AI summary
Data Movement Is All You Need: A Case Study on Optimizing Transformers
Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, Torsten Hoefler
TL;DR
Transformer training is memory-bound and existing implementations underutilize GPUs because data movement and layouts are not globally optimized. The paper uses dataflow analysis, fusion, layout selection, and configuration optimization, achieving lower movement and faster BERT training while acknowledging a layout-selection scope limitation.
Problem
Training transformers is increasingly expensive, and existing implementations do not efficiently utilize GPUs because data movement is the key training bottleneck.
Method
The paper constructs data-centric dataflow graphs to analyze movement, maximize reuse through fusion, select performant layouts, and optimize end-to-end configurations.
Results
Up to 22.91% less data movement and at least 1.30× faster BERT training than general-purpose frameworks are achieved, with 1.08× over DeepSpeed.
Takeaways & Limitations
The recipe is applicable to other deep neural networks and emphasizes considering data movement across the training stack from applications to hardware.
Takeaways & Limitations
The configuration-selection algorithm omits dataflow connections between forward and backpropagation, so globally optimal layouts are not guaranteed.
Abstract
from arXiv · showhide
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.
1 INTRODUCTION
Transformer training is increasingly constrained by data movement rather than arithmetic, motivating a global optimization recipe that improves GPU utilization and training performance.
- Training transformers is highly compute-intensive, often requiring days on hundreds of GPUs or TPUs.
- Data movement is identified as the key bottleneck because accelerated tensor contractions have made transformer training memory-bound.
- 37% of BERT iteration runtime is spent in memory-bound operators, although tensor contractions account for over 99% of arithmetic operations.
- Up to 22.91% less data movement and at least 1.30× faster training than general-purpose frameworks are achieved for the evaluated BERT workloads.
- The approach combines dataflow analysis, data reuse, layout selection, and end-to-end configuration optimization.
2 BACKGROUND
The paper frames transformer optimization as a data-movement problem and uses data-centric programming to expose and transform computation separately from movement.
- Transformers combine multi-head attention with feed-forward layers, normalization, dropout, and residual connections in encoder layers.
- Their dataflow graph exposes parallelism and data reuse because edges represent exact data movement whose volume and access sets can be inspected.
- Modern optimization techniques largely tune data movement and maximize data reuse through layouts, tiling, specialized units, buffering, and vectorization.
- DaCe separates computation from data movement using Stateful DataFlow multiGraphs that represent computations and containers as nodes and movement as edges.
- Its white-box transformations can change schedules, layouts, mappings, and other movement properties without changing the underlying computation.
3 OPTIMIZING TRANSFORMERS
The optimization recipe analyzes transformer dataflow, classifies operators by computation and movement, then targets reuse, layouts, and end-to-end configurations.
- The recipe constructs a training dataflow graph, reduces movement using data reuse, evaluates layouts, and selects configurations for end-to-end performance.
- 3.1 Dataflow Analysis: SDFGs represent operators and their input/output data, allowing data-access volume and floating-point operation counts to be estimated.
- 3.1 Dataflow Analysis: Data movement often dominates operator runtime, making it the principal optimization target even when computation is substantial.
- 3.1 Dataflow Analysis: Transformer operators are classified as tensor contractions, statistical normalizations, or element-wise operations based on movement-to-operation ratios and computation structure.
- 3.2 Operators in Transformers: Tensor contractions are compute-intensive matrix multiplications where layouts and tiling strategies are critical, while normalizations depend on reductions, layout, and vectorization.
4 FUSION
Fusion reduces data movement by combining compatible operators and avoiding intermediate loads and stores, while results show both strong gains and configuration-sensitive trade-offs.
- 4 FUSION: Normalization and element-wise operators consume a significant share of runtime despite being less computationally intensive than tensor contractions.
- Fusion rules: Operators are fused when their iteration-space implementations are compatible, with partial fusion sharing outer dimensions while sequencing inner spaces.
- Fusion rules: Fusion is selected to reduce kernel launches or avoid intermediate memory loads and stores, and maximal fusion continues until reductions or iteration spaces change.
- Results: Every forward fused operator outperforms PyTorch, while two backward operators are slower because layout selection can be suboptimal for global performance.
- Results: High MUE often accompanies fused operators, but 100% MUE may be unattainable with multiple differently shaped tensors and irregular DRAM access.
- Additional Fusion Approaches: Fusing into tensor contractions is not profitable, whereas algebraic fusion of MHA input projections is profitable.
5 DATA LAYOUT
The paper selects data layouts empirically because layout choices affect vectorization, tiling, memory access, and operator runtime. Exhaustive configuration search reveals that intuitive layout rules are useful but insufficient for consistently identifying fast configurations.
- Layout selection: Data layouts enable efficient access patterns, vectorization, and tiling for tensor contractions and statistical normalization operators.The authors benchmark feasible layouts and other parameters because the best choice depends on GPU model and tensor sizes.
- Tensor contractions: Tensor contraction performance varies across layouts, tensor sizes, and Tensor Core usage, with some cases failing to saturate Tensor Cores.The analysis includes algebraic fusion variants and both Tensor Core and regular floating-point execution.
- Tensor contractions: 14.24%: cuBLAS’s heuristic algorithm was worse than the best algorithm for a half-precision dX1 QKT contraction.At single precision, the heuristic was up to 7.18% worse, demonstrating hardware- and workload-specific tuning requirements.
- Fused operators: Vectorized-memory-access layouts provide the most noticeable performance improvement, indicating that moved-data volume is the main performance limitation.Layouts joining reduction and vector dimensions are another notable category.
- Fused operators: Joining reduction and vector dimensions reduces partially reduced values from eight FP16 registers to one.This is identified as a notable layout category for fused operators.
- Fused operators: 65 µs versus 771 µs: configurations satisfying intuitive layout rules can still differ substantially in runtime for the AIB kernel.Exhaustive search is therefore needed to identify performant configurations reliably.
6 END-TO-END OPTIMIZATION
The paper globally assembles fused operators and data layouts using performance measurements and graph-based configuration selection. Evaluations show improvements over optimized frameworks for attention, BERT encoder layers, and end-to-end BERT training, while the recipe also applies beyond BERT.
- 6 END-TO-END OPTIMIZATION: Global optimization assembles fused components and selects layouts jointly, allowing downstream gains to outweigh transpose overheads.The implementation uses dataflow graphs and integrates with PyTorch through its C++ operator API.
- 6.1 Configuration Selection: SSSP selects the final configuration path through a DAG whose edges represent minimum runtimes for input-output layout pairs.For BERT, the algorithm takes seconds and saves the selected path for automatic layout definition at training start.
- 6.1 Configuration Selection: The configuration runtime is within 6% of an ideal layout configuration, although omitting forward-backpropagation dataflow prevents a global-optimality guarantee.The authors state that this assumption could be relaxed in future work.
- 6.3 End-to-End Performance: 1.30× faster: the optimized BERT encoder layer outperforms PyTorch for combined forward and backpropagation, including framework overheads.It is also 1.20× faster than TensorFlow+XLA and 1.08× faster than DeepSpeed.
- 6.3 End-to-End Performance: 22.91%: the optimized implementation reduces total data movement over the standard implementation.The comparison attributes the advantage to fused-kernel I/O accounting across the implementation.
- 6.3 End-to-End Performance: 1.19×: end-to-end BERT training achieves this overall speedup despite additional training overheads.The result is reported for BERT pretraining at scale.
- 6.3 End-to-End Performance: The recipe is directly applicable to similar transformer layers, including GPT-2/3 decoder layers, with few changes.The implementation can also be extended by stacking optimized layers into a full training pipeline.
7 RELATED WORK
Prior work accelerates transformer training through distributed-memory scaling, large batches, and general compiler optimizations, but does not provide the paper’s complete, systematic treatment of data movement.
- Distributed-memory methods such as ZeRO, Megatron, and Mesh-TensorFlow scale transformer training across many GPUs.
- DeepSpeed is the closest compared approach, but it performs optimizations and layout selections manually.
- The paper’s data-centric approach combines dataflow-based fusion and layout transformations across granularity levels, surpassing the cited tools’ optimization capabilities.
8 DISCUSSION
The discussion argues that the data-movement recipe extends beyond transformers and can guide automated optimization across diverse operators, architectures, and hardware designs.
- The recipe can be directly adopted for Megatron-LM, GPT-3, and other DNN architectures once their data-centric graphs are constructed.
- Operator classifications cover many DNN operators beyond transformers, including tensor contractions, convolutions, pooling, and normalization.
- Automated optimization selects schemes and configurations from data layouts and operator types, avoiding infeasible manual kernel optimization across the large configuration space.
- Different operators favor different data layouts, motivating hardware support for fast layout changes; tensor contractions can remain bounded by data transfer to Tensor Cores.
- The analysis is relevant to hardware trends involving reduced format conversion, more on-chip memory, low-latency interconnects, and coarse-grained spatial hardware.
9 CONCLUSIONS
The paper concludes that transformer training is memory-bound and GPU-underutilizing, while its data-movement recipe improves implementations and applies broadly across deep learning.
- Transformer training is memory-bound and underutilizes GPUs, despite the workload’s importance.
- The recipe identifies data-movement bottlenecks and optimizations that outperform highly tuned state-of-the-art implementations.
- The approach applies broadly because many DNNs fit within the proposed operator classification.
- The results emphasize considering data movement across the training stack, from application software through hardware.
A.1 Additional Background
The training setup uses data parallelism and separates backpropagation into input-gradient and parameter-gradient stages.
- Data parallelism partitions each mini-batch across many GPUs.
- Backpropagation distinguishes computing input gradients dX from computing parameter gradients dW.
- The parameter-gradient stage applies only to layers with learnable parameters.
A.1.1 Transformers
Transformers rely on parallel multi-head attention and are optimized here through dataflow-guided fusion, data layouts, and operator configurations. The approach reduces data movement and improves BERT training performance over optimized frameworks.
- Transformers: Transformers replace recurrent or convolutional sequence processing with attention, enabling parallel processing of every sequence element during training.
- Transformers: Multi-head attention uses parallel attention heads over learned projections of queries, keys, and values, followed by concatenation and linear projection.
- Transformers: The dataflow graph exposes operator parallelism, data reuse, access sets, and movement volume for guided bottleneck analysis and optimization.
- Transformers: The optimization recipe maximizes data reuse through fusion, selects performant layouts, and uses measured performance to choose end-to-end operator configurations.
- Transformers: Fully fusing the three self-attention input projections performed best, although cuBLAS kernel occupancy made task parallelism unprofitable.
- Transformers: 1.19× overall speedup was achieved for BERT pretraining, compared with 1.30× for a BERT encoder layer over PyTorch.The end-to-end speedup was 1.18× in phase 1 and 1.22× in phase 2; the lower overall gain reflects unoptimized operations and distributed-training and data-loading overheads.
- Transformers: The data-centric approach supports user-extensible fusion and data-layout transformations across multiple granularity levels in a parametric dataflow graph.