Efficiently Scaling Transformer Inference

By Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, Joseph Dean. 2022

688 citations recorded by the source corpus.

We study the problem of efficient generative inference for Transformer models, in one of its most challenging settings: large deep models, with tight latency targets and long sequence lengths. Better understanding of the engineering tradeoffs for inference for large Transformer-based models is important as use cases of these models are growing rapidly throughout application areas. We develop a simple analytical model for inference efficiency to select the best multi-dimensional partitioning techniques optimized for TPU v4 slices based on the application requirements. We combine these with a suite of low-level optimizations to achieve a new Pareto frontier on the latency and model FLOPS utilization (MFU) tradeoffs on 500B+ parameter models that outperforms the FasterTransformer suite of benchmarks. We further show that with appropriate partitioning, the lower memory requirements of multiquery attention (i.e. multiple query heads share single key/value head) enables scaling up to 32x larger context lengths. Finally, we achieve a low-batch-size latency of 29ms per token during generation (using int8 weight quantization) and a 76% MFU during large-batch-size processing of input tokens, while supporting a long 2048-token context length on the PaLM 540B parameter model.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, 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

In Computer Vision and Pattern Recognition

In Model Architectures

In Efficiency & Compression

Citations in the corpus

Llama 2: Open Foundation and Fine-Tuned Chat Models (2023, methodology): “Either you duplicate the KV values in all GPUs (making the KV cache size equal to GQA), or an alternative is to shard across the batch dimension instead (Pope et al., 2022).”
GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (2023, background): “Autoregressive decoder inference is a severe bottleneck for Transformer models due to the memory bandwidth overhead from loading decoder weights and all attention keys and values at every decoding step (Shazeer, 2019; Pope et al., 2022; de Jong et al., 2022).”
Efficient Memory Management for Large Language Model Serving with PagedAttention (2023, methodology): “For Transformers, these states consist of the key and value tensors associated with the attentionmechanism, commonly referred to asKV cache [41], which represent the context from earlier tokens to generate new output tokens in sequence.”
QLoRA: Efficient Finetuning of Quantized LLMs (2023, background): “Lossy quantization approaches study the trade-offs for regular rounding [13, 72, 48] or how to optimize rounding decisions to improve quantization precision [18].”
The case for 4-bit precision: k-bit Inference Scaling Laws (2022, background): “Work by Pope et al. (2022) looks at scaling inference in a production setting where large batch sizes are common.”
Accelerating Large Language Model Decoding with Speculative Sampling (2023, methodology): “Using a combination of the above techniques, in addition to a number of low-level optimisations to TPUs, Pope et al. (2022) have greatly improved the serving latency and eﬃciency of PaLM 540B.”
Splitwise: Efficient Generative LLM Inference Using Phase Splitting (2023, methodology): “LLM inference serving is a rapidly developing field, with several recent works optimizing batching [23], [25], [51], [53], [81], scheduling [22], [42], [51], [66], [73], [79], and memory usage [32], [35], [51], [74].”
Big Little Transformer Decoder (2023, methodology): “This requires the models to load weight matrices, as well as the cached keys and values of previously generated tokens [43], for each token generation, thus preventing parallelization of the loaded values across multiple tokens.”
FiDO: Fusion-in-Decoder optimized for stronger performance and faster inference (2022, background): “The first constraint is not binding for FiDO: due to FiDO’s memory efficiency we are able to fit larger batches even for the XL-XXL model, and if necessary model size can be further extended with quantization (Zeng et al., 2022) and parallelism (Pope et al., 2022).”
ChipNeMo: Domain-Adapted LLMs for Chip Design (2023, background): “Pope et al. demonstrate that inference costs on an 8B model are 8-12x lower than on a 62B model for equal latency targets [35].”
H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (2023, neutral): “In addition to the widely-studied bottlenecks of model size and the quadratic cost of attention layers, the problem of the size of the KV cache , which stores the intermediate attention key and values during generation to avoid re-computation, is becoming increasingly prominent (Pope et al., 2022).”
Learning to Compress Prompts with Gist Tokens (2023, methodology): “Note that at these (relatively) small scales, the wall time improvements are smaller than the FLOPs reductions because much of the inference latency is caused by moving tensors from high-bandwidth memory (HBM) to the chip compute cores, i.e. what Pope et al. [26] call the “memory time”.”
Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time (2023, methodology): “…of LLM, there is an emergence of systems that are specialized for LLM inference, such as Faster Transformer (NVIDIA), Orca (Yu et al., 2022), LightSeq (Wang et al., 2021), PaLM inference (Pope et al., 2022), TurboTransformers (Fang et al., 2021), and Deepspeed-Inference (Aminabadi et al., 2022).”
Baldur: Whole-Proof Generation and Repair with Large Language Models (2023, background): “[66]), and (2) when providing additional context to the model, the context now does not have to be reprocessed for each proof step, but only once per proof.”
Efficient Streaming Language Models with Attention Sinks (2023, methodology): “…expand this window size (Chen et al., 2023; kaiokendev, 2023; Peng et al., 2023) and improve training (Dao et al., 2022; Dao, 2023) and inference (Pope et al., 2022; Xiao et al., 2023; Anagnostidis et al., 2023; Wang et al., 2021; Zhang et al., 2023b) efficiency for lengthy inputs, the…”
The physics of optical computing (2023, background): “, in neural-network inference [137], where in many applications—such as language translation—we may require the latency to be < 1 second).”