MultiHeadAttention

class torchtune.modules.MultiHeadAttention(*, embed_dim: int, num_heads: int, num_kv_heads: int, head_dim: int, q_proj: Module, k_proj: Module, v_proj: Module, output_proj: Module, pos_embeddings: Optional[Module] = None, q_norm: Optional[Module] = None, k_norm: Optional[Module] = None, kv_cache: Optional[KVCache] = None, max_seq_len: int = 4096, is_causal: bool = True, attn_dropout: float = 0.0)

多头注意力层，支持在 https://arxiv.org/abs/2305.13245v1 中引入的分组查询注意力 (GQA)。

GQA 是多头注意力 (MHA) 的一个版本，它使用的键/值头比查询头少，方法是将 n 个查询头分组到每个键和值头。多查询注意力是一个极端版本，其中所有查询头共享一个键和值头。

以下是 num_heads = 4 的 MHA、GQA 和 MQA 的示例

（文档来源：litgpt.Config）。

┌───┐┌───┐┌───┐┌───┐     ┌───┐    ┌───┐             ┌───┐
│ v ││ v ││ v ││ v │     │ v │    │ v │             │ v │
└───┘└───┘└───┘└───┘     └───┘    └───┘             └───┘
│    │    │    │         │        │                 │
┌───┐┌───┐┌───┐┌───┐     ┌───┐    ┌───┐             ┌───┐
│ k ││ k ││ k ││ k │     │ k │    │ k │             │ k │
└───┘└───┘└───┘└───┘     └───┘    └───┘             └───┘
│    │    │    │      ┌──┴──┐  ┌──┴──┐      ┌────┬──┴─┬────┐
┌───┐┌───┐┌───┐┌───┐  ┌───┐┌───┐┌───┐┌───┐  ┌───┐┌───┐┌───┐┌───┐
│ q ││ q ││ q ││ q │  │ q ││ q ││ q ││ q │  │ q ││ q ││ q ││ q │
└───┘└───┘└───┘└───┘  └───┘└───┘└───┘└───┘  └───┘└───┘└───┘└───┘
◀──────────────────▶  ◀──────────────────▶  ◀──────────────────▶
        MHA                    GQA                   MQA
n_kv_heads =4          n_kv_heads=2           n_kv_heads=1

参数:

• embed_dim (int) – 模型的嵌入维度

• num_heads (int) – 查询头的数量。对于 MHA，这也是键和值的头的数量

• num_kv_heads (int) – 键和值头的数量。用户应确保 num_heads % num_kv_heads == 0。对于标准 MHA，设置 num_kv_heads == num_heads，对于 GQA，num_kv_heads < num_heads，对于 MQA，设置 num_kv_heads == 1。

• head_dim (int) – 每个头的维度，计算方式为 embed_dim // num_heads。

• q_proj (nn.Module) – 查询的投影层。

• k_proj (nn.Module) – 键的投影层。

• v_proj (nn.Module) – 值的投影层。

• output_proj (nn.Module) – 输出的投影层。

• pos_embeddings (Optional[nn.Module]) – 位置嵌入层，例如 RotaryPositionalEmbeddings。

• q_norm (Optional[nn.Module]) – 查询的归一化层，例如 RMSNorm。对于解码，这在从 kv_cache 更新之前应用。这意味着它将仅支持令牌宽度归一化，而不支持批处理或序列宽度归一化。

• k_norm (Optional[nn.Module]) – 键的归一化层，如果设置了 q_norm，则必须设置。

• kv_cache (Optional[KVCache]) – 用于缓存键和值的 KVCache 对象

• max_seq_len (int) – 模型支持的最大序列长度。这是计算 RoPE 缓存所必需的。默认值：4096。

• is_causal (bool) – 当未提供掩码时，将默认掩码设置为因果掩码

• attn_dropout (float) – 传递到 scaled_dot_product_attention 函数的 dropout 值。默认值为 0.0。

Raises:

• ValueError – 如果 num_heads % num_kv_heads != 0

• ValueError – 如果 embed_dim % num_heads != 0

• ValueError – 如果 attn_dropout < 0 或 attn_dropout > 1

• ValueError – 如果在未设置 k_norm 的情况下定义 q_norm，反之亦然

forward(x: Tensor, y: Optional[Tensor] = None, *, mask: Optional[Tensor] = None, input_pos: Optional[Tensor] = None) → Tensor

参数:

• x (torch.Tensor) – 形状为 [b x s_x x d] 的输入张量，用于查询

• y (Optional[torch.Tensor]) – 第二个输入张量，形状为 [b x s_y x d]，是 k 和 v 的输入。对于自注意力，x=y。仅在启用 kv_cache 时为可选。

• mask (Optional[_MaskType]) –

  用于在查询-键乘法之后和 softmax 之前掩盖分数。可以是

  形状为 [b x s x s]、[b x s x self.encoder_max_cache_seq_len] 或 [b x s x self.decoder_max_cache_seq_len] 的布尔张量（如果使用带有编码器/解码器层的 KV 缓存）。行 i 和列 j 中的 True 值表示令牌 i 注意令牌 j。False 值表示令牌 i 不注意令牌 j。如果未指定掩码，则默认使用因果掩码。

  BlockMask，用于通过 create_block_mask 创建的打包序列中的文档掩码。当使用块掩码计算注意力时，我们使用 flex_attention()。默认为 None。

• input_pos (Optional[torch.Tensor]) – 可选张量，其中包含每个令牌的位置 ID。在训练期间，这用于指示每个令牌相对于其打包时的样本的位置，形状为 [b x s]。在推理期间，这表示当前令牌的位置。如果为 None，则假定令牌的索引是其位置 ID。默认为 None。

Raises:

ValueError – 如果没有 y 输入且未启用 kv_cache。

Returns:

应用注意力后的输出张量

Return type:

torch.Tensor

用于张量形状的符号

• b: 批大小

• s_x: x 的序列长度

• s_y: y 的序列长度

• n_h: 头数

• n_kv: kv 头数

• d: 嵌入维度

• h_d: 头维度

reset_cache()

重置键值缓存。

setup_cache(batch_size: int, dtype: dtype, max_seq_len: int) → None

设置用于注意力计算的键值缓存。如果在已设置 kv_cache 后调用，则将跳过此操作。

参数:

• batch_size (int) – 缓存的批大小。

• dtype (torch.dpython:type) – 缓存的 dtype。

• max_seq_len (int) – 模型将运行的最大序列长度。