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：可学习参数，用于帮助预训练的编码器适应预训练的解码器。

DeepFusionModel 目前仅支持单个编码器。

示例： >>> # decoder 是具有融合交叉注意力层的 TransformerDecoder（例如 llama3_8b） >>> embed = FusionEmbedding(…) >>> layer = FusionLayer( … layer=TransformerSelfAttentionLayer(…), … fusion_layer=TransformerCrossAttentionLayer(…), … ) >>> decoder = TransformerDecoder(tok_embeddings=embed, layers=layer, num_layers=32, …) >>> >>> # encoder 是预训练的编码器（例如 clip_vit_224），带有一个添加的投影头 >>> projection_head = FeedForward(…) >>> register_fusion_module(projection_head)) >>> encoder = nn.Sequential(clip_vit_224(), projection_head) >>> >>> # DeepFusionModel 组合了编码器和解码器 >>> model = DeepFusionModel(decoder, encoder) >>> >>> # 加载完全融合的检查点（例如 Llama3.2 Vision 检查点） >>> model.load_state_dict(…) >>> >>> # 或加载预训练的单个模型（未加载 fusion_params） >>> model.encoder.load_state_dict(…, strict=False) >>> model.decoder.load_state_dict(…, strict=False) >>> >>> # 前向传递 >>> 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

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

caches_are_setup() → bool

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

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 (Optional[torch.Tensor]) – 可选的布尔张量，其中包含形状为 [b x s x s] 的注意力掩码。这在查询-键乘法之后和 softmax 之前应用。行 i 列 j 中的 True 值表示标记 i 注意标记 j。False 值表示标记 i 不注意标记 j。如果未指定掩码，则默认使用因果掩码。默认为 None。

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

• 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，则假定标记的索引是其位置 ID。默认为 None。

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

返回:

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

返回类型:

Tensor

张量形状的表示法

• b：批大小

• s：标记序列长度

• s_e：编码器序列长度

• v：词汇表大小

• d：标记嵌入维度

• d_e：编码器嵌入维度

• m_s：最大序列长度

reset_caches()

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

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.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) – 缓存的 dtype。

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

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