GQA分组注意力机制

一、目录

  1. 定义
  2. demo

二、实现

  1. 定义
    grouped query attention(GQA)
    1 GQA 原理与优点:将query 进行分组,每组query 参数共享一份key,value, 从而使key, value 矩阵变小。
    2. 优点: 降低内存读取模型权重的时间开销:由于Key矩阵和Value矩阵数量变少了,因此权重参数量也减少了,需要读取到内存的数量量少了,因此减少了读取权重的等待时间。
    3. 效果(并未降低模型性能):GQA通过设置合适的分组大小,可以和MQA的推理性能几乎相等,同时逼近MHA的模型性能。

  2. llama3 分组数为4, chatglm2 分组数为2 .


    参考:https://zhuanlan.zhihu.com/p/693928854
    demo

    import torch
    import torch.nn as nn
    import math

    #GQA
    bs=3
    seq_len =5
    hidden_size= 32
    n_heads=4
    n_kv_heads = 2
    head_dim = hidden_size//n_heads #
    groups = n_heads//n_kv_heads # 4/2
    print("groups=",groups)
    x=torch.randn((bs,seq_len,hidden_size))
    print("x:", x.shape)
    wq = nn.Linear(hidden_size,n_heads*head_dim,bias=False)
    wk = nn.Linear(hidden_size, n_kv_heads * head_dim, bias=False)
    wv = nn.Linear(hidden_size, n_kv_heads * head_dim, bias=False)
    xq,xk,xv=wq(x),wk(x),wv(x)
    xq = xq.view(bs,seq_len, n_heads, head_dim).transpose(1, 2)
    xk = xk.view(bs,seq_len, n_kv_heads, head_dim).transpose(1, 2)
    xv = xv.view(bs,seq_len, n_kv_heads, head_dim).transpose(1, 2)
    print("xq:",xq.shape) #[bs,n_heads,seq_len, head_dim]
    print("xk:", xk.shape)#[bs,n_kv_heads,seq_len, head_dim]
    print("xv:", xv.shape)#[bs,n_kv_heads,seq_len, head_dim]
    def repeat_kv(keys: torch.Tensor, values: torch.Tensor, repeats: int, dim: int):
    keys = torch.repeat_interleave(keys, repeats=repeats, dim=dim)
    values = torch.repeat_interleave(values, repeats=repeats, dim=dim)
    return keys, values
    #复制kv head
    key,val = repeat_kv(xk,xv, groups,dim=1)
    print("key:", key.shape)
    print("val:", val.shape)
    attn_weights = torch.matmul(xq, key.transpose(2, 3)) / math.sqrt(head_dim)
    print("attn_weights:", attn_weights.shape) #[bs,n_heads,seq_len,seq_len]
    attn_output = torch.matmul(attn_weights, val)
    print("attn_output:", attn_output.shape) # [bs,n_heads,seq_len,head_dim]

相关推荐
cver1231 小时前
CSGO 训练数据集介绍-2,427 张图片 AI 游戏助手 游戏数据分析
人工智能·深度学习·yolo·目标检测·游戏·计算机视觉
FreeBuf_1 小时前
新型BERT勒索软件肆虐:多线程攻击同时针对Windows、Linux及ESXi系统
人工智能·深度学习·bert
强哥之神2 小时前
Meta AI 推出 Multi - SpatialMLLM:借助多模态大语言模型实现多帧空间理解
人工智能·深度学习·计算机视觉·语言模型·自然语言处理·llama
慕婉03073 小时前
Tensor自动微分
人工智能·pytorch·python
神经星星4 小时前
专治AI审稿?论文暗藏好评提示词,谢赛宁呼吁关注AI时代科研伦理的演变
人工智能·深度学习·机器学习
lucky_lyovo4 小时前
深度学习--tensor(创建、属性)
人工智能·深度学习
李加号pluuuus4 小时前
【论文阅读】CogVideoX: Text-to-Video Diffusion Models with An Expert Transformer
论文阅读·深度学习·transformer
盼小辉丶5 小时前
TensorFlow深度学习实战——基于自编码器构建句子向量
人工智能·深度学习·tensorflow
冬天给予的预感8 小时前
DAY 54 Inception网络及其思考
网络·python·深度学习
董厂长11 小时前
langchain :记忆组件混淆概念澄清 & 创建Conversational ReAct后显示指定 记忆组件
人工智能·深度学习·langchain·llm