融合层

class torchtune.modules.model_fusion.FusionLayer(layer: Module, fusion_layer: Module, fusion_first: bool = True)

融合层，如Flamingo: a Visual Language Model for Few-Shot Learning中介绍的。

深度融合模型架构通过将编码器输出注入到LLM的中间层，将预训练的编码器模型与预训练的语言模型结合起来。这使得语言模型能够将编码器输出解释为文本，并“理解”任何可以训练编码器的模态。为了使语言模型适应编码器输出，FusionLayer将一个新的可学习层融合到现有的解码器（语言模型）层中。这个额外的层可以接收编码器嵌入，并学习将它们与解码器的标记嵌入结合起来。该模块支持在原始层之前或之后融合新层，在Flamingo中，新层在原始层之前融合。

原始层被包裹在FusionLayer中，以便保持其原始的state_dict键，并且预训练的检查点不会被破坏。新层的参数通过fusion_params可用，以分别控制它们是否可训练。

示例

>>> # Original decoder style transformer
>>> layer = nn.TransformerSelfAttentionLayer(...)
>>> model = TransformerDecoder(layers=layer, num_layers=32, ...)
>>>
>>> # Fuse a cross attention layer to each self attention layer to adapt for the encoder
>>> fusion_layer = nn.TransformerCrossAttentionLayer(...)
>>> fused_layer = FusionLayer(layer, fusion_layer)
>>> model = TransformerDecoder(layers=fused_layer, num_layers=32, ...)
>>>
>>> # Original decoder state_dict still works
>>> model.load_state_dict(..., strict=False)

Parameters:

• layer (nn.Module) – 原始解码器层

• fusion_layer (nn.Module) – 新的融合层

• fusion_first (bool) – 布尔值，用于确定是否在解码器层之前或之后插入融合层。

caches_are_enabled() → bool

检查self.layer上的键值缓存是否启用。 参见 :func:~torchtune.modules.TransformerDecoder.caches_are_enabled`。

caches_are_setup() → bool

检查键值缓存是否在self.layer上设置。 参见 :func:~torchtune.modules.TransformerDecoder.caches_are_setup`。

forward(x: Tensor, **kwargs: Dict) → Tensor

Parameters:

• x (torch.Tensor) – 输入张量，形状为 [batch_size x seq_length x embed_dim]

• **kwargs (Dict) – 所有额外的层参数

Returns:

输出张量与输入形状相同

[batch_size x seq_length x embed_dim]`

Return type:

张量

fusion_params() → List[str]

返回融合层的参数。

reset_cache()

重置两个层的键值缓存。

setup_caches(batch_size: int, dtype: dtype, *, encoder_max_seq_len: int, decoder_max_seq_len: int) → None

为两个层设置键值缓存。

Parameters:

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

• dtype (torch.dpython:type) – 缓存的dtype。

• encoder_max_seq_len (int) – 交叉注意力层的最大缓存序列长度。

• decoder_max_seq_len (int) – 自注意力层的最大缓存序列长度。