Agregar un modelo con transformers modulares

Modular transformers reduce el código necesario para añadir un modelo al permitir importaciones y herencia, a diferencia de la política de un solo modelo, un solo archivo. En lugar de repetir componentes del modelo en varios archivos, añade un archivo modular a la carpeta de tu modelo y hereda de clases existentes.

Un conversor genera archivos independientes a partir del archivo modular. Los usuarios obtienen la misma interfaz de archivo único que ya conocen.

Modular transformers no está diseñado para reemplazar el código de modelado legacy. Si tu modelo no se basa en un modelo existente, añade manualmente un archivo modeling.py. Lo mismo aplica a los archivos de configuración, tokenización o procesamiento que no pueden heredar limpiamente de un archivo similar.

No existe un orden correcto único. Algunos colaboradores escriben primero el archivo modular y lo generan a partir de él. Otros comienzan con un modeling.py escrito a mano y lo refactorizan en un archivo modular más adelante. Ambos enfoques funcionan.

Implementar un archivo modular

Comience buscando un modelo en Transformers similar al suyo. Buenos puntos de partida son Mistral, Qwen2, Cohere y Cohere2, y Llama. La tabla siguiente asocia los componentes comunes con modelos de los que puede heredar.

Utilice la herramienta modular-detector-v2 para encontrar implementaciones existentes de las que heredar. Pegue un fragmento de código y le devolverá los métodos más similares que ya existen en Transformers, de modo que pueda identificar la mejor clase padre antes de empezar a escribir.

No modifiques un modelo existente solo para que la herencia funcione con el tuyo. Si renombrar o subclasear una clase padre resulta demasiado incómodo, copia directamente el código relevante.

Cree src/transformers/models/<name>/modular_<name>.py, donde <name> coincide con el nombre del directorio del modelo en snake_case. En esta sección se le guía para implementar Olmo2 a partir de Olmo con el enfoque modular (consulte el archivo original modular_olmo2.py).

Config

Hay dos puntos en los que Olmo2Config difiere de OlmoConfig.

1. Hay un nuevo argumento, rms_norm_eps.
2. El argumento clip_qkv ya no se utiliza.

Declará nuevos argumentos como anotaciones de tipo a nivel de clase con un valor predeterminado. Para los argumentos eliminados, asigná AttributeError() para suprimir el atributo heredado en el archivo generado (ver Eliminar atributos).

- @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 genera automáticamente la documentación estándar de los argumentos (consulte la guía de @auto_docstring). @strict rechaza los kwargs desconocidos en el momento de la instanciación, lo que permite detectar errores tipográficos y argumentos obsoletos de forma temprana. Añada ambos a cada clase de configuración, ya que los decoradores no se heredan de la clase padre. Declárelos explícitamente incluso si la configuración padre ya los tiene.

Para establecer un atributo derivado o manejar la lógica de compatibilidad hacia atrás, use __post_init__ en lugar de __init__. Por ejemplo, Cohere2 calcula head_dim y deriva layer_types en el momento de la inicialización.

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)

Para los modelos con soporte de paralelismo de tensor o de pipeline, defina base_model_tp_plan y base_model_pp_plan como diccionarios a nivel de clase en la configuración. Ambos diccionarios definen cómo dividir el modelo entre dispositivos. Consulte configuraciones existentes como Olmo2 o Cohere2 para ver ejemplos.

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"]),
    }

Norma

Para copiar una clase padre sin cambios, hereda con pass. El linter copia el contenido de la clase padre y renombra todas las referencias para que coincidan con el nuevo modelo.

from ..olmo.modeling_olmo import OlmoRotaryEmbedding

class Olmo2RotaryEmbedding(OlmoRotaryEmbedding):
    pass

Para cambiar un comportamiento específico, herede y sobrescriba solo lo que difiere. Olmo2RMSNorm difiere de LlamaRMSNorm en una línea. La multiplicación se realiza antes de volver a convertir al dtype de entrada, no después.

  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)

Atención

La atención de Olmo2 es idéntica a la de Olmo, excepto que aplica RMSNorm a las consultas y claves, y elimina el recorte de qkv. super().__init__(...) copia el cuerpo del padre y añade las dos nuevas líneas de normalización. El forward se redefine por completo porque las consultas y claves ahora pasan por normalizaciones antes de la proyección. El linter también incorpora en el archivo generado cualquier función importada, incluyendo apply_rotary_pos_emb, eager_attention_forward y sus dependencias.

  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

Después de super().__init__(...), sobrescribe los atributos de normalización con instancias de Olmo2RMSNorm y reasigna self.self_attn a la nueva clase Olmo2Attention. La instrucción del self.input_layernorm elimina la asignación de input_layernorm del padre, ya que Olmo2 aplica la normalización después de la atención y no antes. Consulta Eliminar atributos para ver qué elimina y qué no elimina del.

El forward se reescribe para reflejar la ubicación de la normalización posterior a la atención. Solo se necesita reescribir forward cuando se cambia el nombre de un atributo, no cuando solo cambia su tipo.

  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

Modelo

Aquí solo cambia el tipo de self.norm. El método forward es idéntico al de la clase padre, por lo que el linter lo hereda automáticamente.

  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)]
+         )

Cabeza del modelo

La lógica es idéntica a la de OlmoForCausalLM, por lo que no se necesitan cambios.

from ..olmo.modeling_olmo import OlmoForCausalLM

class Olmo2ForCausalLM(OlmoForCausalLM):
    pass

Otras clases

El modeling_olmo2.py generado por el linter también contiene clases (Olmo2MLP, Olmo2RotaryEmbedding, Olmo2PreTrainedModel) que no se definieron explícitamente en modular_olmo2.py.

El linter incorpora cualquier clase de la que dependa una clase heredada, a menos que la redefinas explícitamente. Las funciones importadas, como apply_rotary_pos_emb, siguen la misma regla.

Por ejemplo, OlmoDecoderLayer tiene self.mlp = OlmoMLP(config). Olmo2MLP nunca se definió en el archivo modular, por lo que el linter lo crea automáticamente, equivalente a usar pass.

from ..olmo.modeling_olmo import OlmoMLP

class Olmo2MLP(OlmoMLP):
    pass

Si quieres que Olmo2MLP herede de un modelo diferente, sé explícito.

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

class Olmo2MLP(MistralMLP):
    pass

Finalizar el archivo

Cada archivo modular debe declarar un logger y una lista __all__ a nivel de módulo.

logger = logging.get_logger(__name__)

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

__all__ debe listar todas las clases públicas del archivo. El conversor y las importaciones posteriores dependen de él. Una clase que falte en __all__ no se exportará correctamente.

Generar los archivos de modelado

El script modular_model_converter.py genera archivos independientes como modeling.py, configuration.py y otros a partir de su archivo modular. Para cada clase heredada, copia el cuerpo de la clase padre en la hija, renombra todas las referencias para que coincidan con el nuevo modelo e incorpora las funciones auxiliares o clases de las que dependen esos padres.

Los archivos de salida no contienen importaciones entre modelos ni herencia de otros directorios de modelos. El linter aplanar la herencia a un solo nivel. Si Olmo2Attention hereda de OlmoAttention, el Olmo2Attention generado es completamente autónomo. Pero si OlmoAttention a su vez heredaba de otra cosa, el linter no incorpora ese abuelo.

Ejecuta el siguiente comando para generar archivos a partir de un archivo modular.

python utils/modular_model_converter.py your_model

No edites directamente los archivos generados, ya que cualquier cambio se sobrescribirá en la siguiente ejecución.

Patrones para archivos modulares

Las secciones siguientes documentan los patrones de uso habituales, como eliminar atributos o sobrescribir métodos decorados, al trabajar con un archivo modular.

Eliminar atributos

Eliminar un atributo heredado depende de si estás trabajando con una clase de configuración o con una subclase de nn.Module.

Para una clase de configuración, asigne AttributeError() al atributo a nivel de clase.

class MyNewConfig(ParentConfig):
    removed_attr = AttributeError()

El linter elimina por completo la declaración del atributo del archivo de configuración generado. Las clases de configuración utilizan un diseño de tipo dataclass sin __init__, por lo que asignar AttributeError() a nivel de clase es el enfoque correcto.

Para una subclase de nn.Module, usa del self.attribute después de super().__init__(...).

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

del self.attribute elimina únicamente la línea de asignación self.attribute = ... del cuerpo del padre copiado. No elimina otras líneas que hagan referencia a self.attribute. Si el método forward u otros métodos del padre también hacen referencia al atributo, debe sobrescribir esos métodos también.

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.

Trabajar con super()

super().__init__(config) le indica al conversor que copie el cuerpo del padre en el hijo. Dos patrones permiten sobrescribir este comportamiento.

• Llama directamente a una clase padre específica cuando necesites que la salida generada llame a un abuelo (nn.Module.__init__) en lugar de a la clase padre modular.
• Use **super_kwargs para heredar la firma completa de un método padre mientras se añade una docstring personalizada o se sustituye un decorador.

Llamar directamente a una clase abuela

Sé explícito sobre qué clase estás llamando cuando necesitas que super() apunte al padre de la clase generada en lugar del padre modular. El ejemplo siguiente llama directamente a nn.Module.__init__(self). DummyModule es en sí mismo un nn.Module, por lo que el conversor lo escribe como super().__init__() en el MyNewDummyModule generado.

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

Use **super_kwargs para heredar la firma completa de un método padre mientras se añade una docstring personalizada o se sustituye un decorador. En la firma sobrescrita, indica al linter que expanda todos los argumentos del padre en la salida generada.

El uso más común es añadir una docstring específica del modelo, como documentar el argumento labels, sin reescribir la firma completa.

# 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)

El GemmaForCausalLM.forward generado tiene la firma completa de LlamaForCausalLM sin necesidad de copiar nada manualmente.

**super_kwargs es un atajo para casos de nicho. Si vas a cambiar el comportamiento, escribe la firma completa en su lugar.

Eliminación de métodos no utilizados

Elimine un método de la clase padre sobrescribiéndolo con una declaración raise AttributeError(""). El linter elimina el método del archivo generado.

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")

Sobrescribir métodos decorados

Cuando se sobrescribe un método padre decorado, el decorador del padre se hereda automáticamente. Si se añade un decorador propio, este reemplaza al del padre.

En toda la biblioteca aparecen dos decoradores, uno para capturar las salidas intermedias del modelo y otro para generar automáticamente las docstrings.

En el ejemplo siguiente, una subclase sobrescribe un método decorado de la clase padre sin añadir su propio decorador. El decorador de la clase padre se hereda.

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

Si agregas un nuevo decorador, tu decorador reemplaza al del padre.

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

Nomenclatura especial

El linter renombra automáticamente todo al heredar de una clase. Use el mismo prefijo de nombre de clase en todas las clases del mismo archivo.

Evite mezclar prefijos, como en el ejemplo siguiente. MyModelIncredibleMLP incumple las convenciones de nomenclatura, y el linter no sabrá si debe usar MyModelIncredible o MyModel al renombrar dependencias de orden superior.

class MyModelIncredibleMLP(LlamaMLP):
    ...

class MyModelDecoderLayer(LlamaDecoderLayer):
    ...

Al no tener dependencias implícitas, puedes renombrar una sola clase localmente. Redefine explícitamente cada otra mención de esa clase con el nuevo patrón de nombre. De lo contrario, el linter añadirá una clase no deseada MyModelMLP junto a MyModelIncredibleMLP.

El linter emite una advertencia cuando detecta un prefijo ambiguo.

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).

Los prefijos ambiguos son más comunes en los modelos multimodales, donde los nombres de clase incluyen un calificador de modalidad como Text. Para asignar un prefijo específico a una dependencia, renómbrala explícitamente con un pass.

class Emu3TextMLP(LlamaMLP):
    pass

Documentación de docstrings

El linter aún no admite la herencia parcial de docstrings. Al agregar o eliminar atributos de configuración, agregue el docstring completo directamente en el archivo modular bajo la definición de la clase.

Conversión de puntos de control

Una vez que hayas generado tus archivos de modelado, verifica que los pesos reales se carguen correctamente. Escribe un script de conversión para traducir el formato del checkpoint de origen a uno compatible con Transformers y, a continuación, guárdalo en el Hub.

Escribir un script de conversión

Agrega un archivo convert_<model>_to_hf.py en src/transformers/models/<model>/. El script carga los pesos de la fuente, renombra y redimensiona las claves para que coincidan con los nombres de parámetros de tu módulo y guarda el resultado con save_pretrained().

Busca un script existente para copiar y adaptar. Los modelos en src/transformers/models/ incluyen un convert_*_to_hf.py que puedes usar como punto de partida.

Después de ejecutar el script, carga el punto de control guardado con from_pretrained() y confirma que cada peso esperado se haya cargado correctamente. Las claves del punto de control sin usar indican nombres que no coinciden, por lo que imprímelas para detectar problemas a tiempo.

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

Comprueba que la forma y el nombre coincidan al iterar sobre las claves. Las discrepancias de forma suelen indicar que un parámetro en tu configuración es incorrecto, que la arquitectura difiere de la original o que un peso necesita ser transpuesto.

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}"
    )

Corrige cualquier problema iterando entre tu archivo modular, el archivo de modelado generado y el script de conversión hasta que todos los pesos se carguen correctamente.

Una vez que el punto de control se carga correctamente, súbelo al Hub mediante push_to_hub(). Consulta la guía de compartir modelos para obtener más detalles.

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

Mapeo de conversión en tiempo de ejecución

Añade un mapeo de tiempo de ejecución a src/transformers/conversion_mapping.py cuando los pesos publicados no coincidan con la disposición de parámetros de tu módulo. Los casos habituales incluyen pesos fusionados almacenados por separado y tensores de expertos MoE que requieren apilamiento. El mapeo permite que from_pretrained() cargue el punto de control del Hub sin una etapa de exportación separada.

Consulte la guía de carga dinámica de pesos para saber cómo escribir reglas de WeightRenaming y WeightConverter y registrarlas para su model_type.

Próximos pasos

• Las reglas de estructura del modelo son reglas estáticas que se aplican a todos los archivos modeling_*.py, modular_*.py y configuration_*.py. Ejecute make typing para verificarlas antes de abrir una PR.
• Añadir componentes de procesamiento de visión explica cómo añadir un procesador de imágenes, un procesador de vídeo y un procesador para un modelo multimodal.
• Generación automática de docstrings muestra cómo usar @auto_docstring para no tener que escribir a mano la documentación de los argumentos de las APIs de modelos compartidas.
• Escritura de pruebas de modelos explica cómo escribir pruebas de integración para tu nuevo modelo y ejecutarlo localmente.
• Controles de la solicitud de extracción explica los controles de CI que debe superar su PR antes de poder fusionarse, y cómo reproducirlos y corregirlos localmente.