PaperlayerView exact revision on arXiv

Source-linked AI summary

AI and Memory Wall

Amir Gholami, Zhewei Yao, Sehoon Kim, Coleman Hooper, Michael W. Mahoney, Kurt Keutzer

arXiv:2403.14123v1cs.LGcs.ARcs.DC

TL;DR

Rapid growth in AI model size and training compute has exposed widening gaps between computation and memory or communication capabilities. The paper analyzes Transformer inference using arithmetic intensity and profiling, finding that decoder models can be dominated by memory bandwidth, and argues for redesigning models and training strategies.

  • Problem

    AI model growth has outpaced memory capacity, bandwidth, and communication improvements, creating bottlenecks that compute-focused hardware scaling does not address.

  • Method

    The paper re-examines hardware trends and performs a Transformer case study comparing encoder and decoder inference through FLOPs, memory operations, arithmetic intensity, and latency.

  • Results

    Memory increasingly became a bottleneck rather than compute, and decoder inference showed substantially longer latency because of higher memory operations and lower arithmetic intensity.

  • Takeaways & Limitations

    Breaking the memory wall requires rethinking AI model architecture, data-efficient training, optimization, and deployment strategies.

Abstract

from arXiv · show

The availability of unprecedented unsupervised training data, along with neural scaling laws, has resulted in an unprecedented surge in model size and compute requirements for serving/training LLMs. However, the main performance bottleneck is increasingly shifting to memory bandwidth. Over the past 20 years, peak server hardware FLOPS has been scaling at 3.0x/2yrs, outpacing the growth of DRAM and interconnect bandwidth, which have only scaled at 1.6 and 1.4 times every 2 years, respectively. This disparity has made memory, rather than compute, the primary bottleneck in AI applications, particularly in serving. Here, we analyze encoder and decoder Transformer models and show how memory bandwidth can become the dominant bottleneck for decoder models. We argue for a redesign in model architecture, training, and deployment strategies to overcome this memory limitation.

I. INTRODUCTION

AI models and training costs have scaled rapidly while compute hardware has outpaced memory and interconnect bandwidth, making memory and communication increasingly dominant bottlenecks. The paper revisits this memory wall for modern AI workloads and motivates architectural, training, and deployment changes.

  • Contribution: The paper re-examines recent hardware and AI computations, concluding that the memory wall is increasingly the dominant bottleneck for a range of AI tasks.It frames the issue as an emerging challenge for both training and serving.
  • Memory wall: Memory-wall limits span capacity, bandwidth, and latency across on-chip memory, DRAM, and communication between processors.Data-transfer capabilities have lagged behind hardware compute capabilities at each of these levels.
  • Motivation: 750×/2yrs: Compute needed to train recent LLMs increased at this rate during 2018–2022, encouraging accelerators to prioritize peak compute.This focus often simplified other hardware components, including memory hierarchy.
  • Motivation: 410×/2yrs: LLM sizes grew at this rate, exceeding the memory available on a single chip.Distributed-memory parallelism avoids single-chip capacity limits but can introduce additional memory-wall challenges.
  • Memory wall: Communication between neural-network accelerators is slower and less efficient than on-chip movement, while network bandwidth remains difficult to scale.This creates a bottleneck when work is distributed across multiple processes or accelerators.
  • Hardware trends: 3.0×/2yrs versus 1.6 and 1.4 times every 2 years: peak server FLOPS outpaced DRAM and interconnect bandwidth, respectively.The disparity increasingly made memory a bottleneck even when models fit within one chip.

II. CASE STUDY

The paper studies Transformer inference through encoder and decoder architectures. Encoders process all tokens concurrently, whereas decoders generate one token at a time autoregressively.

  • Transformer architectures: Encoder Transformers such as BERT concurrently process all tokens during inference.Decoder Transformers such as GPT instead process and generate one token at each iteration.
  • Transformer architectures: Decoder Transformers such as GPT run autoregressively, processing and generating one token at each iteration.This inference pattern distinguishes them from encoder architectures.

A. Arithmetic Intensity

Total FLOPs alone can misidentify Transformer bottlenecks, so the paper evaluates arithmetic intensity alongside memory operations. A case study compares encoder matrix-matrix operations with decoder matrix-vector operations.

  • Arithmetic intensity: Arithmetic intensity is the number of FLOPs performed per byte loaded from memory.It is computed as total FLOPs divided by total bytes accessed, also called memory operations or MOPs.
  • Arithmetic intensity: Arithmetic intensity complements total FLOPs by measuring how much computation is performed for each byte transferred from memory.The paper uses this measure to assess performance bottlenecks in encoder-only and decoder-only Transformers.
  • Model comparison: BERT-Base and BERT-Large use matrix-matrix operations, whereas GPT-2 uses repeated matrix-vector multiplications during autoregressive inference.These models illustrate why operation type matters when assessing memory-related bottlenecks.

B. Profiling

Profiling on an Intel Gold 6242 CPU shows that GPT-2 has longer inference latency than similarly configured BERT models despite comparable end-to-end FLOPs. Higher memory operations and lower arithmetic intensity in autoregressive matrix-vector operations explain this result, especially at low batch sizes.

  • Profiling results: GPT-2 latency was significantly longer than BERT-Base or BERT-Large at every tested sequence length, despite similar configuration and end-to-end FLOPs to BERT-Base.The models were profiled on an Intel Gold 6242 CPU.
  • Profiling results: Higher memory operations and lower arithmetic intensity in GPT matrix-vector operations account for its longer latency relative to the BERT models.A model with higher arithmetic intensity can run faster with the same or more FLOPs than one with lower arithmetic intensity.
  • Profiling results: Decoder models can become primarily memory-bandwidth bound rather than compute bound at low batch sizes.The profiling evidence connects this bottleneck to autoregressive inference.

III. PROMISING SOLUTIONS FOR BREAKING THE WALL

The paper argues that continued model scaling will make training increasingly costly as compute growth diverges from bandwidth growth, motivating changes to model design, training, and deployment.

  • III. PROMISING SOLUTIONS FOR BREAKING THE WALL: 410×/2yrs scaling of LLM size, combined with widening compute–bandwidth disparity, is expected to make training larger models increasingly costly.The paper notes that exponential scaling cannot continue indefinitely.
  • III. PROMISING SOLUTIONS FOR BREAKING THE WALL: Breaking the memory wall requires rethinking AI model design, improving data efficiency, reducing hyperparameter-tuning overhead, and addressing training-state size.The paper presents these as fundamental challenges motivating promising solution directions.
  • III. PROMISING SOLUTIONS FOR BREAKING THE WALL: The paper reviews recent research targeting data-efficient training, lower optimization overhead, and reduced memory requirements.This discussion focuses on the last three challenges identified in the preceding motivation.

A. Efficient Training Algorithms

Efficient training strategies target tuning overhead and memory footprint through more robust optimization, rematerialization, and precision-aware methods, while trading additional computation or memory for lower training cost.

  • A. Efficient Training Algorithms: Brute-force hyperparameter tuning adds substantial overhead because SGD variants are difficult to tune for new models.The relevant parameters include learning rate, annealing schedule, and convergence iterations.
  • A. Efficient Training Algorithms: Second-order stochastic optimization can improve tuning robustness and achieve SOTA, but its memory footprint is 3–4× higher than current methods.The paper identifies removing this overhead as an important direction.
  • A. Efficient Training Algorithms: Memory-local optimization algorithms can trade additional computation for reduced memory footprint and greater data locality.Communication-avoiding algorithms are cited as an example from numerical linear algebra.
  • A. Efficient Training Algorithms: Rematerialization stores only selected forward-pass activations and recomputes the rest when needed, reducing memory footprint by up to 5× with 20% more compute.This can support training large models within single-chip memory.
  • A. Efficient Training Algorithms: FP16 arithmetic enabled more than a 10× increase in hardware compute capability, while reducing training precision below FP16 remains challenging without accuracy degradation.The paper identifies mixed FP8/FP16 and newer FP4 approaches as promising directions.

B. Efficient Deployment

Efficient deployment focuses on compressing models so they require less distributed memory, using quantization, pruning, and smaller language models.

  • B. Efficient Deployment: Distributed-memory inference makes deploying state-of-the-art models such as GPT-3 and large recommendation systems challenging.The paper proposes compression and smaller models as ways to address this deployment burden.
  • B. Efficient Deployment: INT4 inference quantization can reduce model footprint and latency by up to 8× with minimal accuracy impact.The paper distinguishes the greater difficulty of reducing training precision below FP16.
  • B. Efficient Deployment: Structured sparsity can prune up to 30% of neurons and unstructured sparsity up to 80%, with minimal accuracy impact.Pushing beyond these levels is described as challenging and often accuracy-damaging.
  • B. Efficient Deployment: Small language models are presented as a potential route toward broader AI adoption beyond continued scaling of model data and size.The paper describes recent results on their capabilities as promising.

C. Rethinking the Design of AI Accelerators

The paper proposes accelerator designs between CPU and GPU extremes, combining more efficient caching with higher-capacity, hierarchical DRAM to improve compute–bandwidth trade-offs.

  • C. Rethinking the Design of AI Accelerators: 43. Accelerator designs that maximize peak compute often sacrifice cache hierarchy, whereas CPUs perform better on bandwidth-bound problems.The paper notes that CPUs have roughly an order of magnitude less peak compute than GPUs or TPUs.
  • C. Rethinking the Design of AI Accelerators: An intermediate architecture with efficient caching and higher-capacity DRAM could mitigate distributed-memory communication bottlenecks.The proposed DRAM organization may use multiple levels with different bandwidths.

IV. CONCLUSION

Training compute and Transformer model size have scaled far faster than hardware FLOPS, while DRAM and interconnect bandwidth lag substantially. These trends indicate that memory, particularly intra/inter-chip transfer, is becoming the main limitation.

  • 750×/2yrs: Training compute for recent SOTA Transformer models in NLP has scaled at this rate.
  • 410×/2yrs: Model parameter size for recent SOTA Transformer models in NLP has scaled at this rate.
  • 3.0×/2yrs versus 1.6×/2yrs and 1.4×/2yrs: Hardware FLOPS growth has outpaced DRAM and interconnect bandwidth growth.Over 20 years, peak hardware FLOPS increased by 60,000×, compared with 100× for DRAM bandwidth and 30× for interconnect bandwidth.
  • Memory, particularly intra/inter-chip memory transfer, is becoming the main limitation as bandwidth falls behind model and compute scaling.
Loading 2403.14123v1…