Pixiv - おむたつ/omutatsu
GQA 详解
878 字
4 分钟
GQA 详解
Note
LLM 推理优化系列博客
GQA 来源于论文 GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints，是一种注意力机制变体。
Important
根据论文所述，设计 GQA 的原因，是为了优化 MQA (multi-query attention)，MQA 可以大幅减少内存开销，但是会导致模型性能大幅下降以及训练极其不稳定。GQA 可以在减少内存开销的同时，保证一定的模型性能。
MHA 回顾
先回顾一下 MHA (multi-head attention)，我们每个 Query 头各自拥有独立的 K、V 头，因此 Q 的头数 num_heads 与 K、V 的头数 num_kv_heads 是相等的。
class MultiHeadAttention(nn.Module): def __init__(self, d_model, num_heads, dropout=0.0, bias=True): super().__init__() assert d_model % num_heads == 0, "d_model must be divisible by num_heads" self.d_model = d_model self.num_heads = num_heads self.head_dim = d_model // num_heads
self.q_proj = nn.Linear(d_model, d_model, bias=bias) self.k_proj = nn.Linear(d_model, d_model, bias=bias) self.v_proj = nn.Linear(d_model, d_model, bias=bias) self.out_proj = nn.Linear(d_model, d_model, bias=bias) self.dropout = nn.Dropout(dropout)
def forward(self, query, key, value, mask=None): batch, seq_len_q, _ = query.shape seq_len_kv = key.size(1)
# 线性投影并重塑为多头格式 (batch, seq_len, num_heads, head_dim) q = self.q_proj(query).view(batch, seq_len_q, self.num_heads, self.head_dim) k = self.k_proj(key).view(batch, seq_len_kv, self.num_heads, self.head_dim) v = self.v_proj(value).view(batch, seq_len_kv, self.num_heads, self.head_dim)
# 转置为 (batch, num_heads, seq_len, head_dim) q = q.transpose(1, 2) k = k.transpose(1, 2) v = v.transpose(1, 2)
# 缩放点积注意力 attn_weights = torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5) if mask is not None: attn_weights = attn_weights + mask attn_weights = F.softmax(attn_weights, dim=-1) attn_weights = self.dropout(attn_weights)
attn_output = torch.matmul(attn_weights, v) # (batch, num_heads, seq_len_q, head_dim) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(batch, seq_len_q, -1) return self.out_proj(attn_output)MQA
对于 MQA (multi-query attention)，所有的 Query 头共享唯一一组 K、V，即 num_kv_heads = 1。计算时需要将这一组 K、V 重复使用 num_heads 次，以匹配 Query 头的数量。
class MultiQueryAttention(nn.Module): def __init__(self, d_model, num_heads, dropout=0.0, bias=True): super().__init__() assert d_model % num_heads == 0, "d_model must be divisible by num_heads" self.d_model = d_model self.num_heads = num_heads self.head_dim = d_model // num_heads
self.q_proj = nn.Linear(d_model, d_model, bias=bias) self.k_proj = nn.Linear(d_model, self.head_dim, bias=bias) self.v_proj = nn.Linear(d_model, self.head_dim, bias=bias) self.out_proj = nn.Linear(d_model, d_model, bias=bias) self.dropout = nn.Dropout(dropout)
def forward(self, query, key, value, mask=None):
batch, seq_len_q, _ = query.shape seq_len_kv = key.size(1)
q = self.q_proj(query).view(batch, seq_len_q, self.num_heads, self.head_dim) k = self.k_proj(key).view(batch, seq_len_kv, self.head_dim) v = self.v_proj(value).view(batch, seq_len_kv, self.head_dim)
q = q.transpose(1, 2)
outputs = [] for h in range(num_heads): q_h = query[:, h, :, :] # (batch, seq_len_q, head_dim)
# 计算缩放点积 attn_weights = torch.matmul(q_h, k.transpose(-2, -1)) / (head_dim ** 0.5) if mask is not None: attn_weights = attn_weights + mask attn_weights = F.softmax(attn_weights, dim=-1) attn_weights = self.dropout(attn_weights)
attn_output_h = torch.matmul(attn_weights, v) # (batch, seq_len_q, head_dim) outputs.append(attn_output_h)
attn_output = torch.stack(outputs, dim=1) attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(batch, seq_len_q, -1) return self.out_proj(attn_output)GQA
现在来介绍 GQA，它的 Q 头数是 num_heads，K、V 头数是 num_kv_heads，将 Q 进行分组，每 group_size = num_heads // num_kv_heads 一组共享一组 K、V 头。
class GroupedQueryAttention(nn.Module): def __init__(self, d_model, num_heads, num_kv_heads, dropout=0.0, bias=True): super().__init__() assert d_model % num_heads == 0, "d_model must be divisible by num_heads" self.d_model = d_model self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = d_model // num_heads self.group_size = num_heads // num_kv_heads
self.q_proj = nn.Linear(d_model, d_model, bias=bias) self.k_proj = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=bias) self.v_proj = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=bias) self.out_proj = nn.Linear(d_model, d_model, bias=bias) self.dropout = nn.Dropout(dropout)
def forward(self, query, key, value, mask=None): batch, seq_len_q, _ = query.shape seq_len_kv = key.size(1)
q = self.q_proj(query).view(batch, seq_len_q, self.num_heads, self.head_dim) k = self.k_proj(key).view(batch, seq_len_kv self.num_kv_heads, self.head_dim) v = self.v_proj(value).view(batch, seq_len_kv, self.num_kv_heads, self.head_dim)
q = q.transpose(1, 2) # (batch, num_heads, seq_len_q, head_dim) k = k.transpose(1, 2) # (batch, num_kv_heads, seq_len_kv, head_dim) v = v.transpose(1, 2) # (batch, num_kv_heads, seq_len_kv, head_dim)
scale = 1.0 / math.sqrt(self.head_dim)
outputs = [] for kv_head in range(self.num_kv_heads): q_start = kv_head * self.group_size q_end = q_start + self.group_size
# (batch, group_size, seq_len_q, head_dim) q_group = q[:, q_start:q_end, :, :]
# (batch, 1, seq_len_kv, head_dim) k_head = k[:, kv_head:kv_head + 1, :, :] v_head = v[:, kv_head:kv_head + 1, :, :]
# (batch, seq_len_kv, head_dim) k_head = k_head.squeeze(1) v_head = v_head.squeeze(1)
# (batch, group_size, seq_len_q, seq_len_kv) attn_weights = torch.matmul(q_group, k.transpose(-1, -2)) * scale if mask is not None: attn_weights = attn_weights + mask attn_weights = F.softmax(attn_weights, dim=-1) attn_weights = self.dropout(attn_weights)
# (batch, group_size, seq_len_q, head_dim) attn_output_group = torch.matmul(attn_weights, v_head) outputs.append(attn_output_group)
# (batch, num_heads, seq_len_q, head_dim) out = torch.cat(outputs, dim=1) out = out.transpose(1, 2).contiguous() out = out.view(batch, seq_len_q, -1) return self.out_proj(out)Note
支持与分享
如果这篇文章对你有帮助，欢迎分享给更多人或赞助支持！
这里是我的个人博客，分享 AI Infra、LLM 等技术内容。欢迎关注交流！
19
8
16
55,114
0 天
0 天前