GQA

GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints

Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai · EMNLP 2023

Read paper · arXiv:2305.13245

Share KV heads across contiguous query groups and convert multi-head projections by group-wise averaging.

What you will implement

Implement KV projection mean pooling and grouped attention; verify shared gradients and cache capacity.

• 62. GQA and KV Head Conversion

  Implement group-wise KV projection mean pooling and attention with shared KV heads; verify routing and accumulated gradients.

  Complete two core functions; supplied experiments compare MHA/GQA/MQA cache capacity and synthetic conversion error without checkpoint downloads or uptraining.