为新的架构添加 BetterTransformer 支持

您想为Better Transformer（PyTorch Transformer API 的快速路径）添加一个新的模型？请查看此指南！

应该支持的模型

理论上，任何具有 Transformer 编码器层的模型，类似于“注意力就是你所需要的一切”论文中描述的经典编码器，都应该得到支持。更具体地说，一个模型，如果其编码器块具有多头注意力模块（具有注意力前或注意力后的层归一化），则可以转换为其BetterTransformer等效项。条件可以概括如下

使用经典的多头注意力模块（例如，DeBERTa 不支持）
使用gelu或relu激活函数
具有偶数个注意力头
不要使用任何注意力偏差（例如，T5 使用注意力偏差，因此不支持）
每层的第一个和第二个层归一化的eps必须相等

如何将模型转换为 BetterTransformer 格式？

步骤 1：识别要更改的源层

首先，转到optimum/bettertransformer/__init__.py，您将看到字典BetterTransformerManager.MODEL_MAPPING。这应该包含模型类型与其Tuple[str, BetterTransformerBaseLayer]之间的映射，该映射由可以转换为其BetterTransformer等效项的nn.Module的名称以及有效的等效BetterTransformer层类组成。

让我们尝试一步一步地为Bert执行此操作，首先我们需要识别需要替换的层

>>> from transformers import AutoModel

>>> model = AutoModel.from_pretrained("bert-base-uncased")
>>> print(model)
...
          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
      (11): BertLayer(
        (attention): BertAttention(
          (self): BertSelfAttention(
            (query): Linear(in_features=768, out_features=768, bias=True)
            (key): Linear(in_features=768, out_features=768, bias=True)
            (value): Linear(in_features=768, out_features=768, bias=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
          (output): BertSelfOutput(
            (dense): Linear(in_features=768, out_features=768, bias=True)
            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
            (dropout): Dropout(p=0.1, inplace=False)
          )
        )
        (intermediate): BertIntermediate(
          (dense): Linear(in_features=768, out_features=3072, bias=True)
          (intermediate_act_fn): GELUActivation()
        )
        (output): BertOutput(
          (dense): Linear(in_features=3072, out_features=768, bias=True)
          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)
          (dropout): Dropout(p=0.1, inplace=False)
        )
      )
    )
  )
  (pooler): BertPooler(
    (dense): Linear(in_features=768, out_features=768, bias=True)
    (activation): Tanh()
  )
)

您可以清楚地看到，需要替换的层是BertLayer模块，因为它们包含整个编码器层模块。

步骤 2：构建 xxxLayerBetterTransformer 模块

检查已识别的模块是否尚未从另一个模块复制（通过检查transformers中的源代码并检查类定义是否以# Copied from ...开头） - 如果没有，请在bettertransformer/models/encoder_model.py中创建一个类。以以下几行开头

import torch
import torch.nn as nn

from ..base import BetterTransformerBaseLayer


class BertLayerBetterTransformer(BetterTransformerBaseLayer):
    def __init__(self, bert_layer, config):
...

现在，确保填写所有必要的属性，属性列表为

in_proj_weight
in_proj_bias
out_proj_weight
out_proj_bias
linear1_weight
linear1_bias
linear2_weight
linear2_bias
norm1_eps
norm1_weight
norm1_bias
norm2_weight
norm2_bias
num_heads
embed_dim

请注意，这些属性对应于运行 Transformer 编码器模块所需的所有组件，请查看“注意力就是你所需要的一切”论文上的图 1。

填写完所有这些属性后（有时query、key和value层需要“连续化”，请检查modeling_encoder.py文件以了解更多信息。）

还要确保添加以下几行

self.is_last_layer = False
self.validate_bettertransformer()

步骤 3：构建前向传递

首先，以super().forward_checker()开头，这是必要的，以便父类可以在之前运行所有安全检查器。

在第一次前向传递后，隐藏状态需要使用注意力掩码进行嵌套。一旦它们被嵌套，注意力掩码就不再需要了，因此可以设置为None。这就是Bert的前向传递构建方式，这些行在不同模型之间应该保持大致相同，但有时不同模型的注意力掩码的形状是不同的。

super().forward_checker()

if hidden_states.is_nested:
    attention_mask = None

if attention_mask is not None:
    # attention mask comes in with values 0 and -inf. we convert to torch.nn.TransformerEncoder style bool mask
    # 0->false->keep this token -inf->true->mask this token
    attention_mask = attention_mask.bool()
    attention_mask = torch.reshape(attention_mask, (attention_mask.shape[0], attention_mask.shape[-1]))
    seqlen = attention_mask.shape[1]
    lengths = torch.sum(~attention_mask, 1)
    if not all([l == seqlen for l in lengths]):
        hidden_states = torch._nested_tensor_from_mask(hidden_states, ~attention_mask)
    attention_mask = None

一旦hidden_states被嵌套，使用正确的参数调用torch._transformer_encoder_layer_fwd，如下所示

hidden_states = torch._transformer_encoder_layer_fwd(
    hidden_states,
    self.embed_dim,
    self.num_heads,
    self.in_proj_weight,
    self.in_proj_bias,
    self.out_proj_weight,
    self.out_proj_bias,
    self.use_gelu,
    self.norm_first,
    self.norm1_eps,
    self.norm1_weight,
    self.norm1_bias,
    self.norm2_weight,
    self.norm2_bias,
    self.linear1_weight,
    self.linear1_bias,
    self.linear2_weight,
    self.linear2_bias,
    attention_mask,
)

在最后一层，重要的是要“取消嵌套”隐藏状态，以便它可以被下一个模块处理，这在以下几行中完成

if hidden_states.is_nested and self.is_last_layer:
    hidden_states = hidden_states.to_padded_tensor(0.0)
return (hidden_states,)

还要确保返回一个tuple，以遵循transformers的约定。

在您自己的模型上重现此实验的最佳方法是尝试从提供的建模脚本中获得一些灵感。当然，如果您在optimum上打开一个问题或拉取请求，我们将很乐意帮助您转换您的模型！

步骤 4：健全性检查！

作为最后一步，请确保在optimum/bettertransformer/__init__.py中更新BetterTransformerManager.MODEL_MAPPING字典，并使用正确的名称，然后您就可以准备转换您的模型了。例如，对于Bert，将是

MODEL_MAPPING = {
  ...
  "bert": ("BertLayer", BertLayerBetterTransformer),
  ...
}

使用在教程部分中介绍的转换方法试用一下！

< > 更新在GitHub上