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)

Transformer 层，源自 Llama2 模型。归一化应用于注意力层和 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]) – 用于缩放前馈输出的模块。

caches_are_enabled() → bool

检查 self.attn 上的键值缓存是否已启用。请参阅 :func:~torchtune.modules.TransformerDecoder.caches_are_enabled`。

caches_are_setup() → bool

检查 self.attn 上的键值缓存是否已设置。请参阅 :func:~torchtune.modules.TransformerDecoder.caches_are_setup`。

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 值表示 token i 注意 token j。False 值表示 token i 不注意 token j。如果未指定掩码，则默认使用因果掩码。

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

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

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

返回:

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

[batch_size x seq_length x embed_dim]

返回类型:

torch.Tensor

reset_cache()

重置键值缓存。

setup_caches(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) – 最大缓存序列长度。