TransformerSelfAttentionLayer

class torchtune.modules.TransformerSelfAttentionLayer(attn: MultiHeadAttention, mlp: Module, *, sa_norm: Optional[Module] = None, mlp_norm: Optional[Module] = None, sa_scale: Optional[Module] = None, mlp_scale: Optional[Module] = None)

源自 Llama2 模型的 Transformer 层。在注意力**和** FF 层之前应用归一化。

参数：

• attn (MultiHeadAttention) – 注意力模块。

• mlp (nn.Module) – 前馈模块。

• sa_norm (Optional[nn.Module]) – 在自注意力之前应用的归一化。

• mlp_norm (Optional[nn.Module]) – 在前馈层之前应用的归一化。

• sa_scale (Optional[nn.Module]) – 用于缩放自注意力输出的模块。

• mlp_scale (Optional[nn.Module]) – 用于缩放前馈输出的模块。

property cache_enabled: bool

检查键值缓存是否已设置。

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

参数：

• x (torch.Tensor) – 形状为 [batch_size x seq_length x embed_dim] 的输入张量

• mask (Optional[_MaskType]) –

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

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

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

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

• **kwargs (Dict) – 与自注意力无关的 Transformer 层输入。

返回值:

与输入形状相同的输出张量

[batch_size x seq_length x embed_dim]

返回类型:

torch.Tensor

reset_cache()

重置键值缓存。

setup_cache(batch_size: int, dtype: dtype, *, encoder_max_seq_len: int, decoder_max_seq_len: int) → None

为注意力计算设置键值缓存。

参数：

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

• dtype (torch.dpython:type) – 缓存的数据类型。

• encoder_max_seq_len (int) – 此参数在此层中被忽略。

• decoder_max_seq_len (int) – 最大缓存序列长度。