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]) – 用于缩放前馈输出的模块。

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]) –

  用于在 query-key 相乘后、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] (如果对 encoder/decoder 层使用 KV-cacheing)。如果第 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) – 缓存的 batch size。

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

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

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