GQA (group query attention)

什么是GQA?

多个head的Query共用一组K和V。llama模型就用到该技术。

需要明确几点:

1.group有几组

2.每个group对应几个head

3.q以head为单位 k,v以group为单位 每个head/group特征维度都是head_dim

代码实现

python 复制代码
import torch.nn as nn
import torch
import math

# 自注意力
class GroupQueryAttention(nn.Module):
  def __init__(self, d_model, n_heads, n_groups):
    super().__init__()
    self.d_model = d_model
    self.n_heads = n_heads
    self.n_groups = n_groups

    assert d_model % n_heads == 0
    self.head_dim = d_model // n_heads
    # 每个group对应几个head
    self.heads_per_group = n_heads // n_groups

    # 以head为单位
    self.w_q = nn.Linear(d_model, n_heads*self.head_dim)  # n_heads*head_dim=d_model
    # 以group为单位
    self.w_k = nn.Linear(d_model, n_groups*self.head_dim)
    self.w_v = nn.Linear(d_model, n_groups*self.head_dim)

    self.w_combine = nn.Linear(d_model, d_model)
    self.softmax = nn.Softmax(dim=-1)
  
  # 给k,v进行复制,假设每个组对应3个head,那就要把每个组的数据复制3遍
  def expand(self, data): # data:k/v [b, group, seq_len, head_dim]
    b,_,seq_len,_ = data.shape
    data = data[:,:,None,:,:].expand(b, self.n_groups, self.heads_per_group, seq_len, self.head_dim)
    data = data.contiguous().view(b, -1, seq_len, self.head_dim)
    return data  # [b, group*heads_per_group, seq_len, head_dim]

  def forward(self, x, use_mask=False): # x: [b, seq_len, d_model]
    b, seq_len, _ = x.shape
    q,k,v = self.w_q(x), self.w_k(x), self.w_v(x)
    q = q.view(b, seq_len, self.n_heads, self.head_dim).permute(0,2,1,3)  # 以head为单位
    k = k.view(b, seq_len, self.n_groups, self.head_dim).permute(0,2,1,3) # 以group为单位
    v = v.view(b, seq_len, self.n_groups, self.head_dim).permute(0,2,1,3)

    # 复制
    k,v = self.expand(k), self.expand(v)
    score = q @ k.transpose(-1,-2) / math.sqrt(self.head_dim)
    if use_mask:
      mask = torch.tril(torch.ones(seq_len, seq_len))
      score = score.masked_fill(mask==0, float('-inf'))
    
    score = self.softmax(score) @ v
    score = score.permute(0,2,1,3).contiguous().view(b, seq_len, self.n_heads*self.head_dim)

    out = self.w_combine(score)
    return out

x = torch.rand(2,100,384)
model = GroupQueryAttention(384, 4, 2)
out = model(x)
print(out.shape)
相关推荐
九章云极AladdinEdu5 小时前
超参数自动化调优指南:Optuna vs. Ray Tune 对比评测
运维·人工智能·深度学习·ai·自动化·gpu算力
酷飞飞6 小时前
Python网络与多任务编程:TCP/UDP实战指南
网络·python·tcp/ip
研梦非凡7 小时前
ICCV 2025|从粗到细:用于高效3D高斯溅射的可学习离散小波变换
人工智能·深度学习·学习·3d
数字化顾问8 小时前
Python:OpenCV 教程——从传统视觉到深度学习:YOLOv8 与 OpenCV DNN 模块协同实现工业缺陷检测
python
学生信的大叔8 小时前
【Python自动化】Ubuntu24.04配置Selenium并测试
python·selenium·自动化
诗句藏于尽头10 小时前
Django模型与数据库表映射的两种方式
数据库·python·django
通街市密人有10 小时前
IDF: Iterative Dynamic Filtering Networks for Generalizable Image Denoising
人工智能·深度学习·计算机视觉
智数研析社10 小时前
9120 部 TMDb 高分电影数据集 | 7 列全维度指标 (评分 / 热度 / 剧情)+API 权威源 | 电影趋势分析 / 推荐系统 / NLP 建模用
大数据·人工智能·python·深度学习·数据分析·数据集·数据清洗
扯淡的闲人10 小时前
多语言编码Agent解决方案(5)-IntelliJ插件实现
开发语言·python
moxiaoran575310 小时前
Flask学习笔记(一)
后端·python·flask