# Efficiently Scaling Transformer Inference

**Authors:** Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, Jeff Dean  
**arXiv:** [2211.05102](https://arxiv.org/abs/2211.05102)  
**Date:** 2022-11-09  
**Summary page:** https://www.summarizepaper.com/en/paper/2211.05102v1/

## Key points

- The paper investigates the problem of efficient generative inference for large deep Transformer models with tight latency targets and long sequence lengths.
- Challenges associated with generative inference of language models (LLMs) include a large memory footprint, low parallelizability, and high inference cost due to attention mechanisms scaling quadratically with input sequence length.
- The authors propose a powerful and abstract partitioning framework to optimize LLMs for inference efficiency, which enables reaching the limits of model parallel scaling given the limited parallelizability of Transformer inference.
- They develop a simple analytical model for inference efficiency to select the best multi-dimensional partitioning techniques optimized for TPU v4 slices based on application requirements.
- A suite of low-level optimizations is combined with the partitioning framework 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.
- Appropriate partitioning reduces memory requirements in multiquery attention, enabling scaling up to 32x larger context lengths.
- The authors 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.
- The paper discusses several approaches to improve ML inference efficiency such as efficient attention layers, distillation, pruning or quantization techniques which could be coupled with other model compression methods.
- In conclusion, this paper investigates the scaling properties of Transformer inference workloads and proposes practical partitioning approaches to meet challenging application requirements such as tight latency targets.

## Summary

In this paper, the authors investigate the problem of efficient generative inference for large deep Transformer models with tight latency targets and long sequence lengths. They highlight the challenges associated with generative inference of language models (LLMs), including a large memory footprint, low parallelizability, and high inference cost due to attention mechanisms scaling quadratically with input sequence length. To optimize LLMs for inference efficiency, they propose a powerful and abstract partitioning framework that enables reaching the limits of model parallel scaling given the limited parallelizability of Transformer inference. Within this framework, they develop a simple analytical model for inference efficiency to select the best multi-dimensional partitioning techniques optimized for TPU v4 slices based on application requirements. They 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. The authors further show that appropriate partitioning reduces memory requirements in multiquery attention (i.e., multiple query heads share single key/value head), enabling scaling up to 32x larger context lengths. Finally, they 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. The paper also discusses several approaches to improve ML inference efficiency such as efficient attention layers, distillation, pruning or quantization techniques which could be coupled with other model compression methods. In conclusion, this paper investigates the scaling properties of Transformer inference workloads and proposes practical partitioning approaches to meet challenging application requirements such as tight latency targets. The authors push beyond traditional paradigms of single-server inference by scaling up to 64+ chips and show that longer context lengths incur higher memory costs but multiquery attention with appropriate partitioning reduces this cost and makes long-context inference practical. They observe that FLOP count and communication volume can fundamentally limit inference performance of dense Transformer models and suggest sparsity techniques like task-based mixture of expert architectures or adaptive computation techniques that allocate different amounts of compute per input and generation timestep to reduce FLOPs per token in order to enable further gains in both cost and latency.