Transformers documentation

使用模块化 Transformers 添加模型

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.2).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

使用模块化 Transformers 添加模型

Modular transformers 通过允许导入和继承来减少添加模型所需的代码,这与单模型、单文件策略形成对比。无需在不同文件中重复模型组件,只需在模型文件夹中添加一个modular文件,并从现有类继承即可。

转换器从模块化文件生成独立文件。用户将获得他们熟悉的单文件接口。

Modular transformers 并非旨在取代旧版建模代码。如果您的模型并非基于现有模型,请手动添加modeling.py文件。对于无法从类似文件干净地继承的配置、分词或处理文件,同样适用此规则。

也不存在唯一正确的顺序。一些贡献者先编写模块化文件,再从中生成代码;另一些人则先手写modeling.py,之后再将其重构为模块化文件。两种方法都可行。

实现模块化文件

首先,在 Transformers 中找一个与你的模型相似的模型。好的起点包括 Mistral、Qwen2、Cohere、Cohere2 和 Llama。下表将常见组件映射到你可以继承的模型。

组件模型
专家混合Mixtral 或 Qwen2-MoE
交错(和/或部分)旋转嵌入GLM、Phi
状态空间模型Jamba、Bamba、Zamba、Mamba2
循环隐藏状态Gemma2
每层的滑动窗口注意力/全注意力模式Gemma2、Cohere2
QKV 裁剪Olmo
QK 归一化Olmo2,Cohere
融合的 QKV(不推荐)Phi3

使用 modular-detector-v2 工具查找可继承的现有实现。粘贴一段代码片段,它会返回 Transformers 中已有的最相似方法,从而帮助你在开始编写代码前确定最佳的父类。

不要仅仅为了让新模型支持继承而修改现有模型。如果重命名或子类化父类过于繁琐,请直接复制相关代码。

创建 src/transformers/models/<name>/modular_<name>.py,其中 <name> 需与 snake_case 模型目录名一致。本节将引导你使用模块化方法,从 Olmo 实现 Olmo2(请参阅原始的 modular_olmo2.py 文件)。

配置

Olmo2Config 与 OlmoConfig 有两处不同。

  1. 新增了一个参数 rms_norm_eps。
  2. clip_qkv 参数不再使用。

将新参数声明为带有默认值的类级别类型注解。对于已移除的参数,请为其赋值 AttributeError(),以在生成的文件中抑制继承的属性(参见 移除属性)。

- @auto_docstring(checkpoint="allenai/OLMo-7B-hf")
+ @auto_docstring(checkpoint="allenai/Olmo2-7B-1124-hf")
+ @strict
- class OlmoConfig(PreTrainedConfig):
+ class Olmo2Config(OlmoConfig):
      ...
-     model_type = "olmo"
+     model_type = "olmo2"
      ...
+     rms_norm_eps: float = 1e-5
-     clip_qkv: float | None = None
+     clip_qkv = AttributeError()

@auto_docstring 会自动生成标准的参数文档(参见 @auto_docstring 指南)。@strict 会在实例化时拒绝未知的 kwargs,从而尽早发现拼写错误和过时的参数。请为每个配置类都添加这两个装饰器,因为装饰器不会从父类继承。即使父配置中已经声明了它们,也要显式地重新声明。

若要设置派生属性或处理向后兼容逻辑,请使用 __post_init__ 而非 __init__。例如,Cohere2 在初始化时计算 head_dim 并推导出 layer_types。

def __post_init__(self, **kwargs):
    if self.num_key_value_heads is None:
        self.num_key_value_heads = self.num_attention_heads
    self.head_dim = self.hidden_size // self.num_attention_heads
    super().__post_init__(**kwargs)

对于支持张量并行或流水线并行的模型,请在配置上以类级别字典的形式定义 base_model_tp_plan 和 base_model_pp_plan。这两个字典定义了如何在设备间分片模型。可参考 Olmo2 或 Cohere2 等现有配置作为示例。

class MyNewModelConfig(PreTrainedConfig):
    model_type = "my_new_model"

    # Tensor parallelism: maps layer name patterns to sharding strategies.
    # Use "colwise" / "rowwise" for standard sharding, or the "gather_output" /
    # "split_input" variants when an extra op (e.g. a QK norm) prevents fusing.
    base_model_tp_plan = {
        "layers.*.self_attn.q_proj": "colwise",
        "layers.*.self_attn.k_proj": "colwise",
        "layers.*.self_attn.v_proj": "colwise",
        "layers.*.self_attn.o_proj": "rowwise",
        "layers.*.mlp.gate_proj": "colwise",
        "layers.*.mlp.up_proj": "colwise",
        "layers.*.mlp.down_proj": "rowwise",
    }

    # Pipeline parallelism: maps submodule names to their (input, output) tensor names.
    base_model_pp_plan = {
        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
        "norm": (["hidden_states"], ["hidden_states"]),
    }

范数

若要原样复制父类,请使用 pass 进行继承。linter 会复制父类的内容,并将所有引用重命名以匹配新模型。

from ..olmo.modeling_olmo import OlmoRotaryEmbedding

class Olmo2RotaryEmbedding(OlmoRotaryEmbedding):
    pass

若要更改特定行为,只需继承并覆盖有差异的部分。Olmo2RMSNorm 与 LlamaRMSNorm 仅有一行不同。乘法运算发生在将数据类型转换回输入 dtype 之前,而不是之后。

  from ..llama.modeling_llama import LlamaRMSNorm

  class Olmo2RMSNorm(LlamaRMSNorm):
      def forward(self, hidden_states):
          input_dtype = hidden_states.dtype
          hidden_states = hidden_states.to(torch.float32)
          variance = hidden_states.pow(2).mean(-1, keepdim=True)
          hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
-         return self.weight * hidden_states.to(input_dtype)
+         return (self.weight * hidden_states).to(input_dtype)

注意力

Olmo2 的注意力机制与 Olmo 完全相同,区别在于它对 queries 和 keys 应用了 RMSNorm,并移除了 qkv clipping。super().__init__(...) 会复制父类的 body,并追加两行新的 norm 代码。由于 queries 和 keys 现在需要先经过 norm 再进行投影,forward 被完全重新定义。linter 还会将任何已导入的函数拉入生成的文件中,包括 apply_rotary_pos_emb、eager_attention_forward 以及它们的依赖项。

  class Olmo2Attention(OlmoAttention):
      def __init__(self, config: Olmo2Config, layer_idx: int | None = None):
          super().__init__(config, layer_idx=layer_idx)
+         self.q_norm = Olmo2RMSNorm(config.num_attention_heads * self.head_dim, config.rms_norm_eps)
+         self.k_norm = Olmo2RMSNorm(config.num_key_value_heads * self.head_dim, config.rms_norm_eps)

      def forward(self, ...):
          ...
-         query_states = self.q_proj(hidden_states)
-         key_states = self.k_proj(hidden_states)
+         query_states = self.q_norm(self.q_proj(hidden_states))
+         key_states = self.k_norm(self.k_proj(hidden_states))
          value_states = self.v_proj(hidden_states)

-         if self.config.clip_qkv is not None:
-             query_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
-             key_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
-             value_states.clamp_(min=-self.config.clip_qkv, max=self.config.clip_qkv)
-
          ...

DecoderLayer

在 super().__init__(...) 之后,使用 Olmo2RMSNorm 实例覆盖 norm 属性,并将 self.self_attn 重新赋值为新的 Olmo2Attention 类。del self.input_layernorm 会移除父类中的 input_layernorm 赋值,因为 Olmo2 是在注意力机制之后而非之前应用 norm。有关 del 会移除和不会移除哪些内容的详细信息,请参阅移除属性。

forward 已被重写,以反映注意力后归一化的位置。只有当属性被重命名时才需要重写 forward,如果仅其类型发生变化则不需要。

  class Olmo2DecoderLayer(OlmoDecoderLayer):
      def __init__(self, config: Olmo2Config, layer_idx: int):
          super().__init__(config, layer_idx=layer_idx)
-         self.self_attn = OlmoAttention(config=config, layer_idx=layer_idx)
-         self.input_layernorm = OlmoLayerNorm(config.hidden_size)
-         self.post_attention_layernorm = OlmoLayerNorm(config.hidden_size)
+         self.self_attn = Olmo2Attention(config=config, layer_idx=layer_idx)
+         del self.input_layernorm
+         self.post_attention_layernorm = Olmo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+         self.post_feedforward_layernorm = Olmo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

      def forward(self, ...):
          residual = hidden_states
-         hidden_states = self.input_layernorm(hidden_states)
          # Self Attention
          hidden_states, _ = self.self_attn(...)
-         hidden_states = residual + hidden_states
+         hidden_states = self.post_attention_layernorm(hidden_states)
+         hidden_states = residual + hidden_states

          # Fully Connected
          residual = hidden_states
-         hidden_states = self.post_attention_layernorm(hidden_states)
          hidden_states = self.mlp(hidden_states)
-         hidden_states = residual + hidden_states
+         hidden_states = self.post_feedforward_layernorm(hidden_states)
+         hidden_states = residual + hidden_states
          return hidden_states

模型

这里只有 self.norm 的类型发生了变化。forward 方法与父类完全相同,因此 linter 会自动将其继承过来。

  class Olmo2Model(OlmoModel):
      def __init__(self, config: Olmo2Config):
          super().__init__(config)
-         self.layers = nn.ModuleList(
-             [OlmoDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
-         )
-         self.norm = OlmoLayerNorm(config.hidden_size)
+         self.norm = Olmo2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+         self.layers = nn.ModuleList(
+             [Olmo2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+         )

模型头部

该逻辑与 OlmoForCausalLM 完全相同,因此无需进行任何更改。

from ..olmo.modeling_olmo import OlmoForCausalLM

class Olmo2ForCausalLM(OlmoForCausalLM):
    pass

其他类

由 linter 生成的 modeling_olmo2.py 还包含一些未在 modular_olmo2.py 中显式定义的类(Olmo2MLP、Olmo2RotaryEmbedding、Olmo2PreTrainedModel)。

除非您显式重新定义,否则 linter 会引入任何继承类所依赖的类。像 apply_rotary_pos_emb 这样的导入函数也遵循相同的规则。

例如,OlmoDecoderLayer 具有 self.mlp = OlmoMLP(config)。Olmo2MLP 从未在模块化文件中定义,因此 linter 会自动创建它,相当于使用 pass。

from ..olmo.modeling_olmo import OlmoMLP

class Olmo2MLP(OlmoMLP):
    pass

如果您希望 Olmo2MLP 从其他模型继承,请明确指定。

# switch to Mistral definition
from ..mistral.modeling_mistral import MistralMLP

class Olmo2MLP(MistralMLP):
    pass

完成文件

每个模块化文件都必须在模块级别声明一个 logger 和一个 __all__ 列表。

logger = logging.get_logger(__name__)

__all__ = [
    "Olmo2Config",
    "Olmo2ForCausalLM",
    "Olmo2Model",
    "Olmo2PreTrainedModel",
]

__all__ 必须列出文件中的所有公共类。转换器和下游导入都依赖于它。如果某个类未包含在 __all__ 中,则无法正确导出。

生成建模文件

modular_model_converter.py 脚本会根据你的模块化文件生成独立的 modeling.py、configuration.py 以及其他文件。对于每个继承类,它会将父类的主体内容复制到子类中,将所有引用重命名以匹配新模型,并引入这些父类所依赖的任何辅助函数或类。

输出文件不包含跨模型的导入,也不从其他模型目录继承。linter 会将继承展平为单一层级。如果 Olmo2Attention 继承自 OlmoAttention,则生成的 Olmo2Attention 是完全自包含的。但如果 OlmoAttention 本身又继承自其他类,linter 不会将该祖父类内联。

运行以下命令,从模块化文件生成文件。

python utils/modular_model_converter.py your_model

请勿直接编辑生成的文件,因为任何更改都会在下次运行时被覆盖。

模块化文件的模式

以下各部分记录了在使用模块化文件时的常见使用模式,例如删除属性或重写已装饰的方法。

移除属性

删除继承的属性取决于你正在处理的是 config 类还是 nn.Module 子类。

对于配置类,请在类级别将 AttributeError() 分配给该属性。

class MyNewConfig(ParentConfig):
    removed_attr = AttributeError()

linter 会完全从生成的配置文件中移除该属性声明。配置类采用 dataclass 风格的布局,没有 __init__,因此在类级别分配 AttributeError() 是正确的做法。

对于 nn.Module 的子类,请在 super().__init__(...) 之后使用 del self.attribute。

class MyNewModel(ParentModel):
    def __init__(self, config: MyNewConfig):
        super().__init__(config)
        del self.attribute

del self.attribute 仅从复制的父类主体中移除 self.attribute = ... 这一赋值行。它不会移除其他任何引用 self.attribute 的行。如果父类的 forward 或其他方法也引用了该属性,请同时重写这些方法。

class DummyModel(nn.Module):
    def __init__(self, config: DummyConfig):
        super().__init__()
        self.attribute = config.attribute
        if self.attribute:
            # do more stuff with `self.attribute` here
            ...

class MyNewDummyModel(DummyModel):
    def __init__(self, config: MyNewDummyConfig):
        super().__init__(config)
        del self.attribute
        # 'self.attribute = config.attribute' is removed, but the 'if self.attribute:' block remains.
        # Override forward() or any other method that references self.attribute.

使用 super()

super().__init__(config) 指示转换器将父类的主体复制到子类中。有两种模式可以覆盖此行为。

  • 当需要生成的输出调用祖父类(nn.Module.__init__)而非模块化父类时,直接调用特定的父类。
  • 使用 **super_kwargs 可以继承父方法的完整签名,同时添加自定义文档字符串或替换装饰器。

直接调用祖父类

当需要 super() 指向生成的类的父类而非模块化父类时,请明确指定所调用的类。下面的示例直接调用了 nn.Module.__init__(self)。DummyModule 本身是一个 nn.Module,因此转换器在生成的 MyNewDummyModule 中将其写为 super().__init__()。

class MyNewDummyModule(DummyModule):                   |     class MyNewDummyModule(nn.Module):
                                                       |
  def __init__(self):                                  |       def __init__(self):
    nn.Module.__init__(self)                           |         super().__init__()
    self.foo = config.foo                              |         self.foo = config.foo
    ...                                                |         ...

super_kwargs

使用 **super_kwargs 可以继承父方法的完整签名,同时添加自定义文档字符串或替换装饰器。在被重写的签名中,它指示 linter 在生成的输出中展开所有父方法参数。

最常见的用法是添加模型特定的文档字符串,例如记录 labels 参数,而无需重写完整的函数签名。

# modular_gemma.py
class GemmaForCausalLM(LlamaForCausalLM):
    def forward(**super_kwargs):
        r"""
        Example:

        ```python
        >>> from transformers import AutoTokenizer, GemmaForCausalLM
        >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b")
        ...
        ```"""
        return super().forward(**super_kwargs)

生成的 GemmaForCausalLM.forward 具有完整的 LlamaForCausalLM 签名,无需手动复制。

**super_kwargs 是用于特殊场景的快捷方式。如果你要修改行为,请改为编写完整的签名。

删除未使用的方法

通过用 raise AttributeError("") 语句重写父类方法来移除该方法。linter 会从生成的文件中删除该方法。

class GemmaTokenizer(LlamaTokenizer):
    ...

    def get_spm_processor(self):
        raise AttributeError("Not needed for Gemma")

    def unk_token_length(self):
        raise AttributeError("Not needed for Gemma")

重写装饰方法

当重写带有装饰器的父类方法时,父类的装饰器会自动继承。如果您添加自己的装饰器,它将替换父类的装饰器。

该库中出现了两个装饰器,一个用于捕获模型的中间输出,另一个用于自动生成文档字符串。

在下面的示例中,子类重写了一个带有装饰器的父类方法,但没有添加自己的装饰器。父类的装饰器会被继承。

class NewModel(DummyModel):       |   class NewModel(nn.Module):
  ...                             |     ...
                                  |
  def forward(...):               |     @decorator(...)
    ...                           |     def forward(...):
                                  |       ...

如果您添加了一个新的装饰器,您的装饰器将替换父类的装饰器。

class NewModel(DummyModel):       |   class NewModel(nn.Module):
  ...                             |     ...
                                  |
  @my_new_decorator(...)          |     @my_new_decorator(...)
  def forward(...):               |     def forward(...):
    ...                           |       ...

特殊命名

继承自某个类时,linter 会自动重命名所有内容。请确保同一文件中的所有类使用相同的类名前缀。

避免混用前缀,如下例所示。MyModelIncredibleMLP 违反了命名规范,当重命名高阶依赖项时,linter 将无法确定应使用 MyModelIncredible 还是 MyModel。

class MyModelIncredibleMLP(LlamaMLP):
    ...

class MyModelDecoderLayer(LlamaDecoderLayer):
    ...

在没有隐式依赖的情况下,您可以在本地重命名单个类。请使用新的名称模式显式重新定义该类的其他所有引用。否则,linter 会在 MyModelIncredibleMLP 旁边添加一个不需要的 MyModelMLP 类。

当 linter 检测到存在歧义的前缀时,会发出警告。

We detected multiple prefix names when inheriting from transformers.models.llama.modeling_llama: ('Emu3Text', 'Emu3'). We will only use the most used 'Emu3' prefix when grabbing args and dependencies. Make sure to subclass the intermediate classes with the prefix you want (if different from 'Emu3') or use a single prefix in all the modular (best).

在类名包含 Text 等模态限定词的多模态模型中,前缀歧义最为常见。若要为某个依赖项指定特定前缀,请使用 pass 显式重命名它。

class Emu3TextMLP(LlamaMLP):
    pass

配置文档字符串

linter 目前尚不支持部分 docstring 继承。在添加或移除配置属性时,请在模块化文件中类定义下方直接添加完整的 docstring。

检查点转换

生成建模文件后,请验证真实权重能否正确加载。编写一个转换脚本,将上游检查点格式转换为 Transformers 兼容的格式,然后将其保存到 Hub。

编写转换脚本

在 src/transformers/models/<model>/ 中添加一个 convert_<model>_to_hf.py 文件。该脚本会加载上游权重,重命名并重塑键以匹配你模块中的参数名称,然后使用 save_pretrained() 保存结果。

查找现有的脚本进行复制和修改。src/transformers/models/ 下的模型包含一个 convert_*_to_hf.py,您可以将其作为起点。

运行脚本后,使用 from_pretrained() 加载已保存的检查点,并确认所有预期的权重均已正确加载。未使用的检查点键表示名称不匹配,因此请打印这些键以尽早发现问题。

model = YourModelForTask.from_pretrained("path/to/output/")

在遍历键时,检查形状和名称是否匹配。形状不匹配通常意味着配置中的某个参数有误、架构与原始架构不同,或者某个权重需要转置。

for key, tensor in original_state_dict.items():
    hf_tensor = hf_model.state_dict().get(mapped_key)
    assert hf_tensor.shape == tensor.shape, (
        f"Shape mismatch for {key}: expected {tensor.shape}, got {hf_tensor.shape}"
    )

通过在您的模块化文件、生成的 modeling 文件和转换脚本之间反复迭代,修复所有问题,直到所有权重都能顺利加载。

当检查点成功加载后,使用 push_to_hub() 将其推送到 Hub。有关更多详细信息,请参阅 模型共享 指南。

model.push_to_hub("username/your-model-name")

运行时转换映射

当发布的权重与模块的参数布局不匹配时,请在 src/transformers/conversion_mapping.py 中添加一个运行时映射。常见情况包括单独存储的融合权重,以及需要堆叠的 MoE 专家张量。该映射可让 from_pretrained() 直接加载 Hub 上的检查点,而无需单独的导出步骤。

有关如何为model_type编写WeightRenaming和WeightConverter规则并将其注册,请参阅动态权重加载指南。

后续步骤

  • 模型结构规则是应用于所有 modeling_*.py、modular_*.py 和 configuration_*.py 文件的静态规则。在提交 PR 之前,请运行 make typing 进行检查。
  • 添加视觉处理组件 介绍了如何为多模态模型添加图像处理器、视频处理器和处理器。
  • 自动生成文档字符串 展示了如何使用 @auto_docstring,这样你就不必为共享模型 API 手动编写参数文档。
  • 编写模型测试介绍了如何为新模型编写集成测试并在本地运行。
  • Pull request checks 介绍了在合并 PR 之前必须通过的 CI 检查,以及如何在本地复现和修复这些问题。
Update on GitHub