DeepFusionModel

class torchtune.modules.model_fusion.DeepFusionModel(decoder: TransformerDecoder, encoder: Module, *, decoder_trainable: bool = False, encoder_trainable: bool = False, fusion_trainable: bool = True)

DeepFusion 是一种融合模型架构，其中预训练的编码器与预训练的解码器 (LLM) 相结合。这是一种用于多模态模型的流行架构，在 多模态模型架构的演变 中提供了全面概述。

此模块具有与 TransformerDecoder 相同的方法和前向签名，并且可以在使用 TransformerDecoder 的任何地方互换使用。它将编码器与解码器组合成一个模块，用于检查点和微调。预计编码器和解码器已经定义了任何额外的可学习 fusion_params：可学习的参数，以帮助将预训练的编码器适应预训练的解码器。

示例

>>> # decoder is a TransformerDecoder (e.g. llama3_8b) with fused cross attention layers
>>> embed = FusionEmbedding(...)
>>> layer = FusionLayer(
...     layer=TransformerSelfAttentionLayer(...),
...     fusion_layer=TransformerCrossAttentionLayer(...),
... )
>>> decoder = TransformerDecoder(tok_embeddings=embed, layers=layer, num_layers=32, ...)
>>>
>>> # encoder is pre-trained encoder (e.g. clip_vit_224) with an added projection head
>>> projection_head = FeedForward(...)
>>> register_fusion_module(projection_head))
>>> encoder = nn.Sequential(clip_vit_224(), projection_head)
>>>
>>> # DeepFusionModel combines the encoder and decoder
>>> model = DeepFusionModel(decoder, encoder)
>>>
>>> # Load full fused checkpoints (e.g. a Flamingo checkpoint)
>>> model.load_state_dict(...)
>>>
>>> # Or load pretrained individual models (fusion_params are not loaded)
>>> model.encoder.load_state_dict(..., strict=False)
>>> model.decoder.load_state_dict(..., strict=False)
>>>
>>> # Forward pass
>>> output = model(tokens, mask, encoder_input, encoder_mask, input_pos)

参数：

• decoder (TransformerDecoder) – 解码器模块

• encoder (nn.Module) – 编码器模块

• decoder_trainable (bool) – 是否训练或冻结解码器。默认为 False。

• encoder_trainable (bool) – 是否训练或冻结编码器。默认为 False。

• fusion_trainable (bool) – 是否训练融合参数。默认为 True。

caches_are_enabled() → bool

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

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

参数：

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

• mask (可选[torch.Tensor]) – 可选的布尔张量，形状为 [b x s x s]，包含注意力掩码。它应用于查询-键乘法之后，softmax之前。第 i 行第 j 列为 True 表示 token i 注意到 token j。为 False 表示 token i 不注意 token j。如果没有指定掩码，则默认使用因果掩码。默认为 None。

• encoder_input (可选[Dict]) – 编码器的可选输入。

• encoder_mask (可选[torch.Tensor]) – 布尔张量，定义 token 与编码器嵌入之间的关系矩阵。位置 i,j 为 True 表示 token i 可以注意到解码器中的嵌入 j。掩码形状为 [b x s x s_e]。默认为 None。

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

注意：在推理的第一步，当模型收到提示时，input_pos 将包含提示中所有 token 的位置（例如：torch.arange(prompt_length)）。这是因为我们需要为每个位置计算 KV 值。

返回:

形状为 [b x s x v] 的输出张量，或者由 output_hidden_states 定义的层输出张量列表，并将最终输出张量附加到列表中。

返回类型:

张量

张量形状表示法

• b: 批次大小

• s: token 序列长度

• s_e: 编码器序列长度

• v: 词汇表大小

• d: token 嵌入维度

• 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: int = None, decoder_max_seq_len: int = None)

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

参数：

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

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

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

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