AI and Memory Wall

In plain words

Computer chips have become very good at arithmetic, but the wires and memory that feed them data have not kept up. Much of a large language model's run time can go to waiting for numbers to arrive from memory, not doing the multiplications. The authors call this the "memory wall".

The problem is worst when a model writes its answer one token at a time. Each new token needs a small amount of math but has to pull in the model's weights, so the chip can spend its time moving data. An encoder model that reads the whole input at once does more math per byte it loads, which keeps the chip busier.

The paper does not build a new model. It collects trend data on hardware and models, measures a few Transformer models on one CPU, and then surveys ways to cope: training methods that need less memory, compressing models so they fit on one chip, and chips that spend more of their budget on memory. Its main message is that these ideas, rather than faster arithmetic alone, will decide how large models can be trained and served.

TL;DR

The paper argues that peak hardware compute has grown much faster than memory and interconnect bandwidth, so memory bandwidth now limits AI workloads, most of all decoder-only Transformer inference §1 fig. 1 fig. 3. Its tool is arithmetic intensity, FLOPs divided by bytes moved from memory §2.1 eq. 1. It is an analysis with no new algorithm: it lists directions in training, quantization and pruning, and in accelerator design §3. All numbers are author-reported trend fits or citations to other work.

Problem

Training and serving large Transformer models needs more memory capacity, more memory bandwidth and more chip-to-chip communication, and these grow more slowly than peak FLOPS §1 fig. 1. The paper argues that for single-chip decoder inference, even a model that fits on one chip is limited by memory transfers §1 §2.2 fig. 3. Under these trends, FLOP counts are misleading on their own, because two models with similar FLOPs can differ greatly in run time §2.1 §2.2. The paper's scope is server-class hardware and the Transformer models of roughly 2018 to 2022 §1 fig. 2.

The paper's own results depend on several assumptions: profiling on one CPU with 8-bit values and batch size 1 §2.2, footnote, and trend fits over a fixed window §1, footnote.

Method

The paper's core technique is a measurement, not an algorithm.

1. Trend analysis. Fit growth rates to peak FLOPS, DRAM and interconnect bandwidth over 20 years, and to model size and training compute for Transformers, using linear regression on the plotted data §1 fig. 1 fig. 2 §1, footnote.
2. Arithmetic intensity. For each operation, divide its total FLOPs by the total bytes it loads or stores in memory. A multiply-add in a MAC is counted as two FLOPs. The bytes figure assumes local memory can hold both operand matrices, so the result is an upper bound on data reuse §2.1 eq. 1 §2.1, footnote.
3. Profiling. Run BERT-Base, BERT-Large and GPT-2 on one Intel CPU at batch size 1, with 8-bit parameters and activations, and record FLOPs, memory operations (MOPs), arithmetic intensity and latency for several sequence lengths §2.2 fig. 3.
4. Directions. Discuss existing research for training algorithms, memory reduction, low-precision training, quantization, pruning, small models and accelerator design. The paper does not reimplement these methods §3.

The per-operation check from step 2:

# Arithmetic intensity of one operation  [§2.1] [eq. 1]
# Input:  operation (matrix multiply or other op), batch size 1
# Output: FLOPs per byte moved from memory

def arithmetic_intensity(operation):
    # Count FLOPs; multiply and add in a MAC count separately  [§2.1, fn.]
    total_flops = count_flops(operation)
    # Count all bytes read or written to memory  [§2.1] [eq. 1]
    # Assumes local memory holds both matrices: an upper bound  [§2.1, fn.]
    total_bytes = count_memory_bytes(operation)
    # Intensity is FLOPs divided by bytes moved  [eq. 1]
    return total_flops / total_bytes

Key equation:

\text{Arithmetic Intensity} = \frac{\text{\# FLOPs}}{\text{\# MOPs}} eq. 1

Intuition from the paper: encoders run matrix-matrix products, while autoregressive decoding runs matrix-vector products, which have much lower arithmetic intensity §2.1. The paper notes that the decoder penalty shrinks for batched decoding and for processing a long prompt, where matrix-matrix products dominate §2.2, footnote.

Hyperparameters

The paper has no training hyperparameters. The settings below are the measurement and fit settings it reports.

Results

Author-reported. The paper has no benchmark table. Its results are growth rates from trend fits and profiling of three models on one CPU. Numbers from cited work are stated as the paper states them. The "Baseline" column names the comparison the paper uses; "n/a" means the paper gives no baseline.

Paper vs code

Official code found: the first author's repository, github.com/amirgholami/ai_and_memory_wall, at commit a1f1d676a6e54413b57a7f89374b0dddeab5f356, licensed MIT LICENSE. The paper names no repository itself. The owner is the first author, which makes it an official repository.

The repository is a data release, not the paper's code. Its README describes tables of parameters, feature sizes and FLOPs for NLP and CV models and a memory breakdown table for training README.md. The repository contains no CPU profiling scripts, no code for the arithmetic-intensity measurement §2.1 and no regression code for the growth rates §1, fn.. Those results cannot be reproduced from the repository alone.

Two differences to note when reading the tables:

• The repository's memory breakdown lists total training memory per model README.md. The paper's figure 2a instead plots GPU capacity divided by 6 as an upper bound fig. 2. These are different quantities.
• The repository's GPT-2 row lists an input length of "2014" and a batch size of 1 in a column labelled "Input Resolution (Sentence length)" [code: README.md]. The text does not state this sequence length for the profiling run §2.2.

No discrepancies found in the parts that can be compared.

Pitfalls

• No algorithm to implement. Section 3 discusses cited work and does not test any method here §3. Treat its numbers for quantization, pruning, rematerialization and optimizers as citations, not results from this paper §3.1 §3.2.
• Trend numbers are fits. Every "x per two years" rate comes from a linear regression on plotted points, and the points are not listed in the text §1, fn.. Recomputing a rate needs the figure data. Hardware and model trends cover different windows: 20 years for hardware, 2018 to 2022 for models §1.
• Reference point matters. Peak FLOPS is normalized to an R10000 system in figure 1 fig. 1. Absolute ratios in the text depend on that choice.
• Training compute excludes RL. The training-compute series does not include reinforcement-learning models fig. 2, fn. §1, fn.. The paper also says reported training cost is a lower bound §3, intro.
• The memory-bound claim is conditional. The authors say decoder inference is memory-bound at low batch sizes. They say it does not hold for processing a long prompt or for large-batch inference §2.2, fn.. Do not generalise to training or large-batch serving without measuring.
• Intensity is an upper bound. The bytes count assumes local memory holds both operands of each operation §2.1, fn.. Per the footnote, the value is therefore an upper bound on reuse; kernels that reload data would have lower intensity than this count gives.
• No threshold given. The paper cites the roofline model §2.1 but states no hardware peak-to-bandwidth ratio and no cut-off above or below which an operation counts as memory-bound. Compare against your own hardware numbers; the paper does not supply them.
• Single CPU profile. The latency and intensity comparisons use an Intel CPU, 8-bit values and batch size 1 §2.2, fn.. No GPU or accelerator profile is reported in the text. Results on other hardware are not stated.
• Figure values are only in plots. Panel-level values for the profiling figure are not given in the text. Cite the panel, not a number from it fig. 3.
• Figure numbers are derived from the LaTeX source. The numbering used here follows the order of figure environments in the source, and the compiled PDF was not checked against it. If in doubt, verify with the published version.
• Two versions of this work. The arXiv first page states this version is an extended version of a paper published in IEEE Micro (arXiv first-page notice). The card describes the arXiv v1 text only; the published article may differ.
• Historical illustration. The 5-cycle example in the introduction illustrates the memory-wall argument, which the paper attributes to 1995 work §1. It is not a measurement from this work.

Lineage

• Builds on: the roofline model for arithmetic intensity (Williams et al., 2009), cited in §2.1 (no card).
• Builds on: the memory-wall argument of Wulf and McKee (1995) and related earlier observations, cited in §1 (no card).
• Cited methods discussed in §3: Checkmate for rematerialization §3.1, ZeRO for sharding §3.1, several post-training quantization methods §3.2 and a pruning method §3.2. No cards exist for these.
• Compared against: none. The paper contains no benchmark comparison between methods.
• Extended version of: a paper published in IEEE Micro (per the arXiv first-page notice).
• Superseded by: not stated.