prerequisite: 最强英文开源模型LLaMA架构探秘,从原理到源码
GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
引入GQA的主要目的是提升推理速度,这种注意力机制由transformer的Multi-head Attention简化而来,再辅以KV cache的checkpoint机制进一步提速。
如上图:
- 左边是transformer原始的Multi-head Attention,它有H个query,key,value,即每个query单独配一个key和value
- 右边是其他研究者提出的Multi-query Attention,它在多个query共享同一个key和value
- 中间则是折中的Grouped-query Attention,它将query进行了分组,仅在组内共享同一个key和value
具体而言,Llama2使用了8组KV映射,即GQA-8,实测效果上接近MHA,推理速度上接近MQA,尽可能做到了效果和速度兼得。