Transformer的优化

本文总结 Transformer 和 attention 机制相关的 trick。留下学习痕迹。

Multi Query Attention (MQA)

早在 2019 年就被提出，但最近才被重视。

相比 Multi Head Attention，MQA 让多头注意力层的各个 head 共享同一份 Key 和 Value 参数（Query 不参与共享，各 head 独立）。如此，以不太多的精度代价，可减少参数量，提升推理速度。

目前 PyTorch 没有 MQA 的原生实现，下文的 GQA 也是。若不使用这些 trick，使用原生的 F.scaled_dot_product_attention 能够享受速度加成。

这个仓库 使用纯 Python 的方式实现了 MQA 和 GQA。据作者实验，当 \(n\) 较小时 GQA 仍然比 PyTorch 原生 MHA 更快。

这个视频 讲述了LLaMA2 模型中用到的技巧，包括 MQA 和 KV Cache。可以参考。

Grouped Query Attention (GQA)

MQA 的做法也许太极端——所有 head 只共享一份 Key 和 Value 势必会影响模型效果。

GQA 作为折中方案，实现方法非常直接。将 head 分为 \(n\) 组，各组各自共享 Key 和 Value。

很明显，当 \(n=1\) 时就是 MQA，当 \(n\) 等于 head 数时就是 Multi Head Attention。

Attention with Linear Bias（ALiBi）

为了让 transformer 模型知道输入序列的顺序关系，需要对输入序列添加 Position Embedding。

而 ALiBi 去掉了 Position Embedding 步骤，转而在计算 Query × Key 值时添加一个偏置常量来注入位置信息。实验证明，这样的修改可以让模型训练长度小于推理长度时也能获得良好推理效果。

具体实现的话，部分代码如下（参考）。

首先要获得 bias。这个 bias 要加到 Query × Key 结果中，充当位置信息。

def get_relative_positions(seq_len: int) -> torch.tensor:
    x = torch.arange(seq_len)[None, :]
    y = torch.arange(seq_len)[:, None]
    return x - y

这个函数会返回类似以下的矩阵：

tensor([[ 0,  1,  2,  3,  4],
        [-1,  0,  1,  2,  3],
        [-2, -1,  0,  1,  2],
        [-3, -2, -1,  0,  1],
        [-4, -3, -2, -1,  0]])

针对不同的 head，有着不同的 bias 权重大小。

def get_alibi_slope(num_heads):
    x = (2 ** 8) ** (1 / num_heads)
    return (
        torch.tensor([1 / x ** (i + 1) for i in range(num_heads)])
        .unsqueeze(-1)
        .unsqueeze(-1)
    )

如此，就能获得 [head, seq_len, seq_len] 大小的 bias 了：

alibi = (
    get_relative_positions(seq_len) * get_relative_positions(seq_len)
).unsqueeze(0)

为了让运算更快，可以用上 PyTorch 原生的 F.scaled_dot_product_attention（参考）。

context_layer = F.scaled_dot_product_attention(
         query_layer, key_layer, value_layer, attn_mask=alibi, dropout_p=0.0
     )

参考来源

• 何枝，“【LLM 加速技巧】Muti Query Attention 和 Attention with Linear Bias（附源码）”，https://zhuanlan.zhihu.com/p/634236135?utm_id=0