探秘Transformer系列之（27）--- MQA & GQA

探秘Transformer系列之（27）--- MQA & GQA


目录

• 探秘Transformer系列之（27）--- MQA & GQA
  
  
  
  • 0x00 概述
    
  • 0x01 MHA
    
    
    
    • 1.1 概念
      
    • 1.2 实现
      
      
      
      • 1.2.1 哈佛
        
      • 1.2.2 llm-foundry
        
      
      
    • 1.3 资源占用
      
    
    
  • 0x02 MQA
    
    
    
    • 2.1 概念
      
    • 2.2 实现
      
      
      
      • 1.2.1 精简版
        
      • 1.2.2 完整版
        
      
      
    • 2.3 效果
      
      
      
      • 2.3.1 内存
        
      • 2.3.2 速度
        
      • 2.3.3 表征能力
        
      • 2.3.3 通信
        
      
      
    
    
  • 0x03 GQA
    
    
    
    • 3.1 概念
      
    • 3.2 架构比对
      
    • 3.3 实现
      
      
      
      • 3.3.1 精简版
        
      • 3.3.2 完整版
        
      
      
    • 3.4 效果
      
      
      
      • 3.4.1 内存
        
      • 3.4.2 速度
        
      • 3.4.3 表征能力
        
      
      
    • 3.5 转换
      
      
      
      • 3.5.1 平均池化
        
      • 3.5.2 基于掩码
        
        
        
        • 网络转换
          
        • 找到更好的分组方法
          
        • 剪枝训练
          
        
        
      
      
    • 3.6 优化
      
    
    
  • 0xFF 参考
    
  
  

0x00 概述

在前文“优化KV Cache"中我们提到过，在”减少注意力头的数量“这个维度上，目前主要的相关工作有 MQA和GQA。MQA 和 GQA 是在缓存多少数量KV的思路上进行优化：直觉是如果缓存的KV个数少一些，显存就占用少一些，大模型能力的降低可以通过进一步的训练或者增加FFN/GLU的规模来弥补。
因为MQA和GQA是基于MHA进行改进，所以我们用下图展示了三者的区别。可以看到，通过缩减注意力头数目，MQA/GQA会降低KV Cache存储，让不同的注意力头或者同一组的注意力头共享一个K和V的集合，因为只单独保留了一份（或者几份）查询参数。因此K和V的矩阵仅有一份（或者几份），这大幅度减少了显存占用，使其更高效。另外，传统的基于MHA的Attention算子过于卡访存带宽，MQA和GQA，乃至后续的MLA都可以提计算访存比，这样也是对性能的极大提升。

注：

• 全部文章列表在这里，估计最终在35篇左右，后续每发一篇文章，会修改此文章列表。cnblogs 探秘Transformer系列之文章列表
  
• 本系列是对论文、博客和代码的学习和解读，借鉴了很多网上朋友的文章，在此表示感谢，并且会在参考中列出。因为本系列参考文章太多，可能有漏给出处的现象。如果原作者发现，还请指出，我在参考文献中进行增补。
  

0x01 MHA

因为MQA，GQA是基于MHA进行修改，所以我们有必要先回顾下MHA。
1.1 概念

MHA（即多头注意力机制）在2017年就随着Transformer原始论文"Attention Is All You Need"一起提出，其主要工作是：把原来一个注意力计算拆成多个小份的注意力头，即把Q、K、V分别拆分成多份，每个注意力头使用独立的Q、K、V进行计算。而多个头可以并行计算，分别得出结果，最后再合回原来的维度。

我们通过下图来看看MHA的流程，这里设
来源：程序园用户自行投稿发布，如果侵权，请联系站长删除
免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！