GQA: Training Generalized Multi-Query Transformer Models from
Multi-Head Checkpoints
Joshua Ainslie, James Lee-Thorp, Michiel de Jongyy
Yury Zemlyanskiy ,Federico Lebrón ,Sumit Sanghai
Google Research
Abstract
Multi-query attention (MQA), which only
uses a single key-value head, drastically
speeds up decoder inference. However, MQA
can lead to quality degradation, and moreover
it may not be desirable to train a separate
model just for faster inference. We (1) propose
a recipe for uptraining existing multi-head language model checkpoints into models with
MQA using 5% of original pre-training compute, and (2) introduce grouped-query attention (GQA), a generalization of multi-query attention which uses an intermediate (more than
one, less than number of query heads) number
of key-value heads. We show that uptrained
GQA achieves quality close to multi-head attention with comparable speed to MQA.
1 Introduction
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). The memory bandwidth from loading
Multi-query attention (MQA), which only uses a single key-value head,
drastically speeds up decoder inference. However, MQA can lead to quality
degradation, and moreover it may not be desirable to train a separate model
just for faster inference. We (1) propose a recipe for uptraining existing
multi-head language model checkpoints into models with MQA using 5% of original
pre-training compute, and (2) introduce grouped-query attention (GQA), a
generalization of multi-query attention which uses an intermediate (more than
one, less than number of query heads) number of key-value heads. We show that
uptrained GQA achieves quality close to multi-head attention with comparable
speed to MQA.
Joshua Ainslie James Lee-Thorp Michiel de Jong Yury Zemlyanskiy Federico Lebrón Sumit Sanghai