TransformerDecoder

class torchtune.modules.TransformerDecoder(*, tok_embeddings: Embedding, layers: Union[Module, List[Module], ModuleList], max_seq_len: int, num_heads: int, head_dim: int, norm: Module, output: Union[Linear, Callable], num_layers: Optional[int] = None, output_hidden_states: Optional[List[int]] = None)

源自 Llama2 架构的 Transformer 解码器。

参数：

• tok_embeddings (nn.Embedding) – PyTorch 嵌入层，用于将令牌移动到嵌入空间。

• layers (Union[nn.Module, List[nn.Module], nn.ModuleList]) – 单个 Transformer 解码器层、nn.ModuleList 层列表或层列表。建议使用 nn.ModuleList。

• max_seq_len (int) – 模型将运行的最大序列长度，如 KVCache() 中所用

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

• head_dim (int) – 自注意力中每个头的嵌入维度。这用于设置 KVCache()

• norm (nn.Module) – 在最终 MLP 之前，将规范化应用于解码器输出的可调用对象。

• output (Union[nn.Linear, Callable]) – 将线性变换应用于解码器输出的可调用对象。

• num_layers (Optional[int]) – Transformer 解码器层的数量，仅在层不是列表时定义。

• output_hidden_states (Optional[List[int]]) – 要包含在输出中的层（索引）列表

引发：

• AssertionError – 设置了 num_layers 且层是列表

• AssertionError – 未设置 num_layers 且层是 nn.Module

注意

模块会在使用参数时检查参数的正确性（例如：attn_dropout 应该在 [0,1] 之间）。这有助于减少代码中的 raise 语句数量，提高可读性。

caches_are_enabled() → bool

检查键值缓存是否已设置。这对于高效推理很有用。

chunked_output(last_hidden_state: Tensor) → List[Tensor]

分块应用输出投影。这应该与 CEWithChunkedOutputLoss 结合使用，因为向上转换为 fp32 在那里完成。

要使用此方法，您应该首先调用 set_num_output_chunks()。

参数：

last_hidden_state (torch.Tensor) – 解码器的最后一个隐藏状态，形状为 [b, seq_len, embed_dim]。

返回值:

包含 num_chunks 个输出张量的列表，每个张量形状为

[b, seq_len/num_chunks, out_dim]，其中 out_dim 通常是词汇量大小。

返回类型:

List[torch.Tensor]

forward(tokens: Tensor, *, mask: Optional[Tensor] = None, encoder_input: Optional[Tensor] = None, encoder_mask: Optional[Tensor] = None, input_pos: Optional[Tensor] = None) → Union[Tensor, List[Tensor]]

参数：

• tokens (torch.Tensor) – 输入张量，形状为 [b x s]

• 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。如果没有指定掩码，则默认情况下使用因果掩码。

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

• encoder_input (Optional[torch.Tensor]) – 编码器的可选输入嵌入。形状为 [b x s_e x d_e]

• encoder_mask (Optional[torch.Tensor]) – 布尔张量，定义令牌和编码器嵌入之间的关系矩阵。位置 i,j 处的 True 值表示令牌 i 可以关注解码器中的嵌入 j。掩码的形状为 [b x s x s_e]。默认值为 None，但在推理期间，如果模型已使用任何使用编码器嵌入的层并已设置缓存，则这是必需的。

• input_pos (Optional[torch.Tensor]) – 包含每个令牌位置 ID 的可选张量。在训练期间，这用于指示每个令牌相对于其打包样本的位置，形状为 [b x s]。在推理期间，这指示当前令牌的位置。如果已设置缓存，此参数在推理期间是必需的。默认值为 None。

返回值:

形状为 [b x s x v] 的输出张量，或由 output_hidden_states 定义的层

输出张量列表，将最终输出张量附加到列表中。

返回类型:

Union[torch.Tensor, List[torch.Tensor]]

注意

在推理的第一步，当模型被提供提示时，input_pos 应该包含提示中所有令牌的位置。对于单批次提示，或具有相同长度的提示批次，这将是 torch.arange(prompt_length)。对于一批长度不同的提示，较短的提示会在左侧填充，位置 ID 会相应地右移，因此位置 ID 的形状应该为 [b, padded_prompt_length]。这是因为我们需要检索每个输入 ID 的位置嵌入。在后续步骤中，如果模型已使用 KV 缓存，input_pos 将包含当前令牌的位置 torch.tensor([padded_prompt_length])。否则，input_pos 将包含所有位置 ID 直到当前令牌。

形状表示法

• b: 批次大小

• s: 令牌序列长度

• s_e: 编码器序列长度

• v: 词汇量大小

• d: 令牌嵌入维度

• d_e: 编码器嵌入维度

• m_s: 最大序列长度

reset_caches()

重置键值缓存。

set_num_output_chunks(num_output_chunks: int) → None

用于与 CEWithChunkedOutputLoss 结合使用以节省内存。这应该在配方中的第一次正向传递之前调用。

setup_caches(batch_size: int, dtype: dtype, *, encoder_max_seq_len: Optional[int] = None, decoder_max_seq_len: Optional[int] = None)

为推理设置键值注意力缓存。对于 self.layers 中的每一层： - TransformerSelfAttentionLayer 将使用 decoder_max_seq_len。 - TransformerCrossAttentionLayer 将使用 encoder_max_seq_len。 - FusionLayer 将使用 decoder_max_seq_len 和 encoder_max_seq_len。

参数：

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

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

• encoder_max_seq_len (Optional[int]) – 最大编码器缓存序列长度。

• decoder_max_seq_len (Optional[int]) – 最大解码器缓存序列长度。