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)

Transformer 解码器，源自 Llama2 架构。

参数:

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

• 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 解码器层的数量，仅当 layers 不是列表时定义。

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

Raises:

• AssertionError – 已设置 num_layers 且 layer 是列表

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

Note

参数值将检查正确性（例如：attn_dropout 属于 [0,1]），这在它们使用的模块中进行。这有助于减少代码中的 raise 语句数量并提高可读性。

caches_are_enabled() → bool

检查是否启用了键值缓存。一旦设置了 KV 缓存，相关的注意力模块将被“启用”，并且所有前向传播都将更新缓存。此行为可以禁用，而无需更改 KV 缓存的状态，方法是使用 torchtune.modules.common_utils.disable_kv_cache()“禁用” KV 缓存，此时 caches_are_enabled 将返回 False。

caches_are_setup() → bool

检查是否已设置键值缓存。这意味着已调用 setup_caches，并且模型中相关的注意力模块已创建其 KVCache。

chunked_output(last_hidden_state: Tensor) → List[Tensor]

分块应用输出投影。这应与 CEWithChunkedOutputLoss 结合使用，因为在其中完成了到 fp32 的向上转型。

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

参数:

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

Returns:

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

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

Return type:

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 值表示 token i 注意 token j。False 值表示 token i 不注意 token 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]) – 定义 token 和编码器嵌入之间关系矩阵的布尔张量。位置 i,j 处的 True 值表示 token i 可以注意解码器中的嵌入 j。掩码的形状为 [b x s x s_e]。默认为 None，但如果模型已设置为使用编码器嵌入并且已设置缓存，则在推理期间需要此参数。

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

Returns:

形状为 [b x s x v] 的输出张量或层列表

由 output_hidden_states 定义的输出张量，最终输出张量附加到列表。

Return type:

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

Note

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

形状表示法

• b: 批次大小

• s: token 序列长度

• s_e: 编码器序列长度

• v: 词汇表大小

• d: token 嵌入维度

• d_e: 编码器嵌入维度

• m_s: 最大序列长度

reset_caches()

将相关注意力模块上的 KV 缓存缓冲区重置为零，并将缓存位置重置为零，而无需删除或重新分配缓存张量。

Raises:

RuntimeError – 如果未设置 KV 缓存。使用 setup_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]) – 最大解码器缓存序列长度。