深度融合模型¶
- class torchtune.modules.model_fusion.DeepFusionModel(decoder: TransformerDecoder, encoder: Module, *, decoder_trainable: bool = False, encoder_trainable: bool = False, fusion_trainable: bool = True)[source]¶
DeepFusion 是一种融合模型架构，其中预训练的编码器与预训练的解码器（LLM）在内部解码器层中结合。这是多模态模型的一种流行架构，完整的概述可以在The Evolution of Multimodal Model Architectures中找到。一种常见的深度融合架构是将编码器输入通过交叉注意力层融合到解码器中。该模块不对编码器和解码器的融合方式做任何假设；它只是将编码器嵌入传递给解码器，并让解码器处理任何融合。
该模块具有与
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)
- Parameters:
decoder (TransformerDecoder) – 解码器模块
encoder (nn.Module) – 编码器模块
decoder_trainable (bool) – 是否训练或冻结解码器。默认值为False。
encoder_trainable (bool) – 是否训练或冻结编码器。默认值为False。
fusion_trainable (bool) – 是否训练融合参数。默认值为 True。
- caches_are_enabled() bool[source]¶
检查键值缓存是否已启用。一旦KV缓存设置完成，相关的注意力模块将被“启用”，并且所有前向传递都将更新缓存。这种行为可以通过使用
disable_kv_cache()“禁用”KV缓存来停止，而不会改变KV缓存的状态，此时caches_are_enabled将返回False。
- 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]][source]¶
- Parameters:
tokens (torch.Tensor) – 输入张量，形状为
[b x s]mask (可选[torch.Tensor]) – 可选的布尔张量，包含形状为
[b x s x s]的注意力掩码。这是在查询-键乘法之后和softmax之前应用的。第i行和第j列的值为True表示第i个token关注第j个token。值为False表示第i个token不关注第j个token。如果未指定掩码，默认使用因果掩码。默认值为None。encoder_input (Optional[Dict]) – 编码器的可选输入。
encoder_mask (Optional[torch.Tensor]) – 布尔张量定义了标记和编码器嵌入之间的关系矩阵。位置i,j处的True值表示标记i可以关注解码器中的嵌入j。掩码的形状为
[b x s x s_e]。默认值为None。input_pos (可选[torch.Tensor]) – 可选的张量，包含每个标记的位置ID。在训练期间，这用于指示每个标记相对于其样本的位置，形状为
[b x s]。在推理期间，这表示当前标记的位置。如果未提供，则假定标记的索引为其位置ID。默认值为None。
注意：在推理的第一步，当模型接收到提示时，
input_pos将包含提示中所有标记的位置 （例如：torch.arange(prompt_length)）。这是因为我们需要计算每个位置的KV值。- Returns:
输出形状为
[b x s x v]的张量，或由output_hidden_states定义的层输出张量列表，并将最终输出张量附加到列表中。- Return type:
张量
- Notation used for tensor shapes:
b: 批量大小
s: 令牌序列长度
s_e: 编码器序列长度
v: 词汇大小
d: 令牌嵌入维度
d_e: 编码器嵌入维度
m_s: 最大序列长度
- set_num_output_chunks(num_output_chunks: int) None[source]¶
用于与
CEWithChunkedOutputLoss结合以节省内存。这应该在第一次前向传递之前调用，在配方中。
- setup_caches(batch_size: int, dtype: dtype, *, encoder_max_seq_len: Optional[int] = None, decoder_max_seq_len: Optional[int] = None)[source]¶
为
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。
