Transformers documentation

Adicionar um modelo com transformers modulares

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

Adicionar um modelo com transformers modulares

O Transformers modular reduz o código necessário para adicionar um modelo ao permitir imports e herança, em contraste com a política de um modelo, um arquivo. Em vez de repetir componentes do modelo em vários arquivos, adicione um arquivo modular à pasta do seu modelo e herde de classes existentes.

Um conversor gera arquivos independentes a partir do arquivo modular. Os usuários obtêm a mesma interface de arquivo único que já conhecem.

O Modular transformers não tem a intenção de substituir o código de modelagem legado. Se o seu modelo não for baseado em um modelo existente, adicione manualmente um arquivo modeling.py. O mesmo se aplica a arquivos de configuração, tokenização ou processamento que não podem heritar limpa e diretamente de um arquivo semelhante.

Não existe uma única ordem correta. Alguns contribuidores escrevem primeiro o arquivo modular e geram a partir dele. Outros começam com um modeling.py escrito manualmente e o refatoram para um arquivo modular depois. Ambas as abordagens funcionam.

Implementando um arquivo modular

Comece encontrando um modelo no Transformers semelhante ao seu. Bons pontos de partida são Mistral, Qwen2, Cohere e Cohere2, e Llama. A tabela abaixo mapeia componentes comuns para modelos dos quais você pode herdar.

ComponenteModel
Mistura de especialistasMixtral ou Qwen2-MoE
Incorporação rotacional intercalada (e/ou parcial)GLM, Phi
Modelos de espaço de estadoJamba, Bamba, Zamba, Mamba2
Estados ocultos recorrentesGemma2
Padrões de atenção em janela deslizante/atenção completa por camadaGemma2, Cohere2
Clipping de QKVOlmo
Normalização QKOlmo2, Cohere
QKV fundido (não recomendado)Phi3

Use a ferramenta modular-detector-v2 para encontrar implementações existentes das quais herdar. Cole um trecho de código e ela retorna os métodos mais semelhantes já presentes no Transformers, permitindo que você identifique a melhor classe pai antes de começar a escrever.

Não modifique um modelo existente apenas para fazer a herança funcionar no seu novo modelo. Se renomear ou criar uma subclasse de uma classe pai for muito trabalhoso, copie o código relevante diretamente.

Crie src/transformers/models/<name>/modular_<name>.py, em que <name> corresponde ao nome do diretório do modelo em snake_case. Esta seção orienta você na implementação de Olmo2 a partir de Olmo com a abordagem modular (consulte o arquivo original modular_olmo2.py).

Config

Há dois pontos em que Olmo2Config difere de OlmoConfig.

  1. Há um novo argumento, rms_norm_eps.
  2. O argumento clip_qkv não é mais utilizado.

Declare novos argumentos como anotações de tipo em nível de classe com um valor padrão. Para argumentos removidos, atribua AttributeError() para suprimir o atributo herdado no arquivo gerado (veja Removendo 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 gera automaticamente a documentação padrão dos argumentos (veja o guia @auto_docstring). @strict rejeita kwargs desconhecidos no momento da instanciação, capturando erros de digitação e argumentos obsoletos precocemente. Adicione ambos a cada classe de configuração, pois os decoradores não são herdados da classe pai. Declare-os explicitamente, mesmo que a configuração pai já os possua.

Para definir um atributo derivado ou lidar com a lógica de compatibilidade reversa, use __post_init__ em vez de __init__. Por exemplo, o Cohere2 calcula head_dim e deriva layer_types no momento da inicialização.

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 modelos com suporte a paralelismo de tensor ou de pipeline, defina base_model_tp_plan e base_model_pp_plan como dicionários de nível de classe na configuração. Ambos os dicionários definem como particionar o modelo entre os dispositivos. Consulte configurações existentes, como Olmo2 ou Cohere2, para ver exemplos.

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 uma classe pai sem alterações, herde com pass. O linter copia o conteúdo da classe pai e renomeia todas as referências para corresponder ao novo modelo.

from ..olmo.modeling_olmo import OlmoRotaryEmbedding

class Olmo2RotaryEmbedding(OlmoRotaryEmbedding):
    pass

Para alterar um comportamento específico, herde e sobrescreva apenas o que difere. Olmo2RMSNorm difere de LlamaRMSNorm em uma linha. A multiplicação ocorre antes do retorno ao dtype de entrada, e não depois.

  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)

Atenção

A atenção do Olmo2 é idêntica à do Olmo, exceto pelo fato de aplicar RMSNorm às queries e keys, e remover o qkv clipping. super().__init__(...) copia o corpo do pai e adiciona as duas novas linhas de norm. O forward é totalmente redefinido porque as queries e keys agora passam por norms antes da projeção. O linter também inclui no arquivo gerado quaisquer funções importadas, incluindo apply_rotary_pos_emb, eager_attention_forward e suas dependências.

  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

Após super().__init__(...), sobrescreva os atributos de norm com instâncias de Olmo2RMSNorm e reatribua self.self_attn para a nova classe Olmo2Attention. O del self.input_layernorm remove a atribuição de input_layernorm do pai, pois o Olmo2 aplica a norm após a atenção, em vez de antes. Consulte Removendo atributos para detalhes sobre o que del remove e o que não remove.

O forward é reescrito para refletir a posição da normalização pós-atenção. A reescrita de um forward é necessária apenas quando um atributo é renomeado, e não quando apenas seu tipo muda.

  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

Model

Aqui, apenas o tipo de self.norm muda. O método forward é idêntico ao do pai, portanto o linter o carrega automaticamente.

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

Cabeçalho do modelo

A lógica é idêntica à de OlmoForCausalLM, portanto, nenhuma alteração é necessária.

from ..olmo.modeling_olmo import OlmoForCausalLM

class Olmo2ForCausalLM(OlmoForCausalLM):
    pass

Outras classes

O modeling_olmo2.py gerado pelo linter também contém classes (Olmo2MLP, Olmo2RotaryEmbedding, Olmo2PreTrainedModel) que não foram explicitamente definidas em modular_olmo2.py.

O linter importa qualquer classe na qual uma classe herdada depende, a menos que você a redefina explicitamente. Funções importadas, como apply_rotary_pos_emb, seguem a mesma regra.

Por exemplo, OlmoDecoderLayer tem self.mlp = OlmoMLP(config). Olmo2MLP nunca foi definido no arquivo modular, então o linter o cria automaticamente, equivalente a usar pass.

from ..olmo.modeling_olmo import OlmoMLP

class Olmo2MLP(OlmoMLP):
    pass

Se você quiser que Olmo2MLP herde de um modelo diferente, seja explícito.

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

class Olmo2MLP(MistralMLP):
    pass

Finalizando o arquivo

Todo arquivo modular deve declarar um logger e uma lista __all__ no nível do módulo.

logger = logging.get_logger(__name__)

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

__all__ deve listar todas as classes públicas do arquivo. O conversor e as importações subsequentes dependem dele. Uma classe ausente em __all__ não será exportada corretamente.

Gerar os arquivos de modelagem

O script modular_model_converter.py gera arquivos independentes, como modeling.py, configuration.py e outros, a partir do seu arquivo modular. Para cada classe herdada, ele copia o corpo da classe pai para a classe filha, renomeia todas as referências para corresponder ao novo modelo e inclui as funções auxiliares ou classes das quais esses pais dependem.

Os arquivos de saída não contêm importações entre modelos nem herança de outros diretórios de modelos. O linter aplanada a herança para um único nível. Se Olmo2Attention herda de OlmoAttention, o Olmo2Attention gerado é totalmente autossuficiente. Mas se OlmoAttention em si herda de outra coisa, o linter não incorpora esse avô.

Execute o comando abaixo para gerar arquivos a partir de um arquivo modular.

python utils/modular_model_converter.py your_model

Nunca edite os arquivos gerados diretamente, pois quaisquer alterações serão sobrescritas na próxima execução.

Padrões para arquivos modulares

As seções abaixo documentam padrões de uso comuns, como a remoção de atributos ou a sobrescrita de métodos decorados, ao trabalhar com um arquivo modular.

Removendo atributos

Remover um atributo herdado depende de você estar trabalhando com uma classe de configuração ou com uma subclasse de nn.Module.

Para uma classe de configuração, atribua AttributeError() ao atributo no nível da classe.

class MyNewConfig(ParentConfig):
    removed_attr = AttributeError()

O linter remove a declaração do atributo do arquivo de configuração gerado por completo. As classes de configuração usam um layout no estilo dataclass, sem __init__, portanto, atribuir AttributeError() no nível da classe é a abordagem correta.

Para uma subclasse de nn.Module, use del self.attribute após super().__init__(...).

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

del self.attribute remove apenas a linha de atribuição self.attribute = ... do corpo do pai copiado. Ele não remove outras linhas que referenciam self.attribute. Se o forward do pai ou outros métodos também referenciarem o atributo, sobrescreva esses métodos também.

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.

Trabalhando com super()

super().__init__(config) instrui o conversor a copiar o corpo do pai para o filho. Dois padrões permitem sobrescrever esse comportamento.

  • Chame uma classe pai específica diretamente quando a saída gerada precisar chamar um avô (nn.Module.__init__) em vez da classe pai modular.
  • Use **super_kwargs para herdar a assinatura completa de um método pai, adicionando uma docstring personalizada ou substituindo um decorador.

Chamar uma classe avó diretamente

Seja explícito sobre qual classe você está chamando quando precisar que super() aponte para a classe pai gerada, em vez da classe pai modular. O exemplo abaixo chama nn.Module.__init__(self) diretamente. DummyModule é, em si, um nn.Module, portanto o conversor o escreve como super().__init__() na MyNewDummyModule gerada.

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 herdar a assinatura completa de um método pai, adicionando uma docstring personalizada ou substituindo um decorador. Na assinatura sobrescrita, ele instrui o linter a expandir todos os argumentos do pai na saída gerada.

O uso mais comum é adicionar uma docstring específica do modelo, como documentar o argumento labels, sem reescrever a assinatura 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)

O GemmaForCausalLM.forward gerado possui a assinatura completa de LlamaForCausalLM, sem necessidade de cópia manual.

**super_kwargs é um atalho para casos específicos. Se você estiver alterando o comportamento, escreva a assinatura completa em vez disso.

Excluindo métodos não utilizados

Remova um método pai sobrescrevendo-o com a instrução raise AttributeError(""). O linter remove o método do arquivo gerado.

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

Sobrescrevendo métodos decorados

Quando você sobrescreve um método do pai decorado, o decorador do pai é herdado automaticamente. Se você adicionar seu próprio decorador, ele substitui o do pai.

Dois decoradores aparecem em toda a biblioteca, um para capturar saídas intermediárias do modelo e outro para gerar automaticamente docstrings.

No exemplo abaixo, uma subclasse sobrescreve um método do pai decorado sem adicionar seu próprio decorador. O decorador do pai é herdado.

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

Se você adicionar um novo decorador, o seu decorador substituirá o do pai.

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

Nomenclatura especial

O linter renomeia automaticamente tudo ao herdar de uma classe. Use o mesmo prefixo de nome de classe em todas as classes do mesmo arquivo.

Evite misturar prefixos, como no exemplo abaixo. MyModelIncredibleMLP viola as convenções de nomenclatura, e o linter não saberá se deve usar MyModelIncredible ou MyModel ao renomear dependências de ordem superior.

class MyModelIncredibleMLP(LlamaMLP):
    ...

class MyModelDecoderLayer(LlamaDecoderLayer):
    ...

Sem dependências implícitas, você pode renomear uma única classe localmente. Redefina explicitamente todas as outras menções àquela classe com o novo padrão de nome. Caso contrário, o linter adiciona uma classe MyModelMLP indesejada ao lado de MyModelIncredibleMLP.

O linter emite um aviso quando detecta um prefixo ambíguo.

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

Prefixos ambíguos são mais comuns em modelos multimodais, nos quais os nomes de classes incluem um qualificador de modalidade, como Text. Para atribuir um prefixo específico a uma dependência, renomeie-a explicitamente com um pass.

class Emu3TextMLP(LlamaMLP):
    pass

Docstrings de Config

O linter ainda não suporta herança parcial de docstrings. Ao adicionar ou remover atributos de configuração, adicione a docstring completa diretamente no arquivo modular, sob a definição da classe.

Conversão de checkpoints

Depois de gerar os arquivos de modelagem, verifique se os pesos reais são carregados corretamente. Escreva um script de conversão para traduzir o formato do checkpoint upstream para um formato compatível com o Transformers e, em seguida, salve-o no Hub.

Escreva um script de conversão

Adicione um arquivo convert_<model>_to_hf.py em src/transformers/models/<model>/. O script carrega os pesos upstream, renomeia e redimensiona as chaves para corresponder aos nomes dos parâmetros do seu módulo e salva o resultado com save_pretrained().

Procure um script existente para copiar e adaptar. Os modelos em src/transformers/models/ incluem um convert_*_to_hf.py que você pode usar como ponto de partida.

Após executar o script, carregue o checkpoint salvo com from_pretrained() e confirme se todos os pesos esperados foram carregados corretamente. Chaves de checkpoint não utilizadas indicam nomes incompatíveis, portanto, imprima-as para identificar problemas precocemente.

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

Verifique se as formas e os nomes correspondem ao iterar sobre as chaves. Divergências de forma geralmente indicam que um parâmetro na sua configuração está incorreto, que a arquitetura difere da original ou que um peso precisa ser transposto.

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

Corrija quaisquer problemas iterando entre o arquivo modular, o arquivo de modelagem gerado e o script de conversão até que todos os pesos sejam carregados corretamente.

Depois que o checkpoint for carregado sem problemas, envie-o para o Hub usando push_to_hub(). Consulte o guia de compartilhamento de modelos para mais detalhes.

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

Mapeamento de conversão em tempo de execução

Adicione um mapeamento de tempo de execução em src/transformers/conversion_mapping.py quando os pesos publicados não corresponderem ao layout de parâmetros do seu módulo. Casos comuns incluem pesos fundidos armazenados separadamente e tensores de especialistas MoE que precisam de empilhamento. O mapeamento permite que from_pretrained() carregue o checkpoint do Hub sem uma etapa de exportação separada.

Consulte o guia de carregamento dinâmico de pesos para saber como escrever regras de WeightRenaming e WeightConverter e registrá-las para o seu model_type.

Próximas etapas

Update on GitHub