Backends de atenção

Todas as implementações de atenção executam o mesmo cálculo. Cada token é comparado com todos os outros tokens. A diferença está em como o cálculo é realizado. A atenção básica escala mal porque materializa a matriz de atenção completa na memória, criando gargalos que atrasam a inferência. As implementações otimizadas reorganizam a matemática para reduzir o tráfego de memória, resultando em inferências mais rápidas e acessíveis.

A AttentionInterface fornece implementações de atenção otimizadas. Ela desacopla a implementação da atenção da implementação do modelo para simplificar a experimentação com diferentes funções. Adicione novos backends facilmente com essa interface consistente.

Definir um backend de atenção

Use o argumento attn_implementation em from_pretrained() para instanciar um modelo com uma função de atenção específica.

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B", attn_implementation="flash_attention_2"
)

Alterne entre backends de atenção em tempo de execução sem recarregar o modelo usando set_attn_implementation().

model.set_attn_implementation("sdpa")

Kernels

Baixe e carregue kernels de computação compilados diretamente do Hub em tempo de execução com a biblioteca Kernels. Isso evita problemas de empacotamento causados por versões incompatíveis de PyTorch ou CUDA.

Os kernels se registram automaticamente no AttentionInterface ao serem detectados. Não é necessário instalar o pacote FlashAttention explicitamente. Solicitar o FlashAttention pelo nome também recorre ao kernel do Hub, consulte FlashAttention fallback.

import torch
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B", attn_implementation="kernels-community/flash-attn2"
)

Gerenciador de contexto SDPA

O scaled dot product attention (SDPA) do PyTorch seleciona automaticamente a função de atenção mais rápida para backends CUDA. Para outros backends, ele usa por padrão a implementação em C++ do PyTorch.

Force o SDPA a usar uma implementação específica com o gerenciador de contexto torch.nn.attention.sdpa_kernel.

import torch
from torch.nn.attention import SDPBackend, sdpa_kernel
from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-1B", attn_implementation="sdpa"
)

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    outputs = model.generate(**inputs)

Atenção específica do backbone

Modelos multimodais usam backbones diferentes para cada modalidade. Otimize o desempenho atribuindo funções de atenção específicas a cada backbone. Alguns backbones de visão, por exemplo, têm melhor desempenho em fp32, o que o FlashAttention não suporta.

Mapeie os backbones de visão para diferentes funções de atenção com um dicionário, enquanto o backbone de texto continua a usar FlashAttention. As chaves na implementação de atenção devem corresponder aos nomes das subconfigurações.

from transformers import AutoModelForImageTextToText

attention_implementation_per_backbone = {"vision_config": "sdpa", "text_config": "flash_attention_2"}

for key in attention_implementation_per_backbone:
    assert key in model.config.sub_configs, f"Invalid key in `attention_implementation`"

model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation=attention_implementation_per_backbone
)

Omita determinados backbones do dict para usar a função de atenção padrão (SDPA).

model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation={"text_config": "flash_attention_2"}
)

Defina a mesma função de atenção para todos os backbones com uma única string.

model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation="eager"
)

Defina a função de atenção globalmente com uma chave vazia.

model = AutoModelForImageTextToText.from_pretrained(
    "facebook/chameleon-7b", attn_implementation={"": "eager"}
)

Crie uma nova função de atenção

Personalize ou crie novas funções de atenção adicionando-as ao registro de atenção com AttentionInterface.register(). Os modelos utilizam essas funções por meio do argumento attn_implementation.

Registre uma função de máscara de atenção correspondente ao registrar uma função de atenção personalizada. Se o nome da attn_implementation personalizada não estiver registrado em AttentionMaskInterface, o Transformers ignora a criação da máscara e passa attention_mask=None para as camadas de atenção. Sua função de atenção deve lidar com as restrições causais, de preenchimento, de empacotamento ou de janela deslizante por conta própria, caso contrário, essas restrições podem ser descartadas silenciosamente.

Este exemplo personaliza a função de atenção para imprimir uma mensagem para cada camada. Ele mantém a máscara na implementação original, registrando masking_utils.sdpa_mask como a função de máscara de atenção.

import torch
from transformers import AutoModelForCausalLM, AttentionInterface, AttentionMaskInterface
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.masking_utils import sdpa_mask

def my_new_sdpa(*args, **kwargs):
    print("I just entered the attention computation")
    return sdpa_attention_forward(*args, **kwargs)

AttentionInterface.register("my_new_sdpa", my_new_sdpa)
AttentionMaskInterface.register("my_new_sdpa", sdpa_mask)  # must have the same name as the registered attention function

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B", attn_implementation="my_new_sdpa")
model(torch.ones(1, 5, dtype=int))

Você também pode adicionar novos argumentos à função de atenção. Os modelos que suportam AttentionInterface propagam kwargs para as camadas de atenção e para a função de atenção. Passe os argumentos como kwargs na função forward do modelo. Funções de atenção personalizadas devem seguir essa assinatura e formato de retorno.

import torch
from transformers import AutoModelForCausalLM, AttentionInterface, AttentionMaskInterface
from transformers.integrations.sdpa_attention import sdpa_attention_forward
from transformers.masking_utils import sdpa_mask

def custom_attention(
    module: torch.nn.Module,  # required arg
    query: torch.Tensor,  # required arg
    key: torch.Tensor,  # required arg
    value: torch.Tensor,  # required arg
    attention_mask: Optional[torch.Tensor],  # required arg
    a_new_kwargs = None,  # You can now add as many kwargs as you need
    another_new_kwargs = None,  # You can now add as many kwargs as you need
    **kwargs,  # You need to accept **kwargs as models will pass other args
) -> tuple[torch.Tensor, Optional[torch.Tensor]]
    ...  # do your magic!
    return attn_output, attn_weights  # attn_weights are optional here

AttentionInterface.register("custom", custom_attention)
AttentionMaskInterface.register("custom", sdpa_mask)  # to leave the existing mask untouched

model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="custom")
model(torch.ones(1, 5, dtype=int), a_new_kwargs=..., another_new_kwargs=...)

Verifique o código de modelagem de um modelo para confirmar quais argumentos e kwargs ele envia para a função de atenção.

AttentionMaskInterface

AttentionMaskInterface é o registro consultado pelas funções create_*_mask para converter uma máscara no formato esperado pelo backend de atenção ativo. O FlexAttention requer uma BlockMask, o SDPA requer um tensor 4D e o FlashAttention requer a máscara de padding base 2D. Registre um backend personalizado ou substitua o formatador de um backend existente com AttentionMaskInterface.register().

import torch
from transformers import AttentionMaskInterface
from transformers.masking_utils import sdpa_mask

def my_new_sdpa_mask(*args, **kwargs):
    print("I just entered the attention mask computation")
    return sdpa_mask(*args, **kwargs)

AttentionMaskInterface.register("my_new_sdpa_mask", my_new_sdpa_mask)

Sem um formatador registrado para o attn_implementation ativo, a criação da máscara é ignorada e attention_mask=None é passado para as camadas de atenção.

Funções registradas devem corresponder a esta assinatura.

def custom_attention_mask(
    batch_size: int,  # required arg
    q_length: int,  # required arg
    kv_length: int,  # required arg
    q_offset: int = 0,  # required arg
    kv_offset: int = 0,  # required arg
    mask_function: Callable = causal_mask_function,  # required arg
    attention_mask: Optional[torch.Tensor] = None,  # required arg
    **kwargs,  # a few additional args may be passed as kwargs, especially the model's config is always passed
) -> Optional[torch.Tensor]:

O argumento mask_function é um Callable que imita as funções mask_mod do PyTorch. Ele recebe 4 índices (batch_idx, head_idx, q_idx, kv_idx) e retorna um booleano indicando se aquela posição contribui para o cálculo de atenção. Essa é a mesma forma primitiva usada por or_mask_function e and_mask_function em Build an attention mask.

Use este workaround para o torch.export caso o mask_function falhe ao criar uma máscara.

Construir uma máscara de atenção

Crie máscaras de atenção com as funções create_*_mask em transformers.masking_utils. Cada função lê o backend de atenção ativo da configuração do modelo, consulta o formatador de máscara do backend em AttentionMaskInterface e retorna o formato que o backend espera. Você não precisa inverter, expandir ou converter o tipo da máscara por conta própria.

Escolha a função que corresponde ao padrão de atenção.

Os helpers legados de máscara chamáveis - get_extended_attention_mask, create_extended_attention_mask_for_decoder, invert_attention_mask - emitem um aviso de deprecação e serão removidos em uma versão futura. Use as funções create_*_mask em vez disso.

from transformers.masking_utils import create_causal_mask

attention_mask = create_causal_mask(
    config=self.config,
    inputs_embeds=inputs_embeds,
    attention_mask=attention_mask,
    past_key_values=past_key_values,
)

from transformers.masking_utils import create_bidirectional_mask

attention_mask = create_bidirectional_mask(
    config=self.config,
    inputs_embeds=embedding_output,
    attention_mask=attention_mask,
)

encoder_attention_mask = create_bidirectional_mask(
    config=self.config,
    inputs_embeds=embedding_output,
    attention_mask=encoder_attention_mask,
    encoder_hidden_states=encoder_hidden_states,
)

Adicione restrições extras por cima da máscara base com os argumentos or_mask_function e and_mask_function. Use or_mask_function para permitir que posições adicionais participem da atenção e and_mask_function para restringir ainda mais o padrão base. Ambos seguem a assinatura de 4 índices de mask_function descrita em AttentionMaskInterface. Eles recebem (batch_idx, head_idx, q_idx, kv_idx) e retornam um booleano.

or_mask_function e and_mask_function podem expressar qualquer padrão de atenção, mas são mais lentos que os padrões integrados e não são compatíveis com ExecuTorch. O overhead é mais perceptível em modelos menores (~200M parâmetros), onde a criação da máscara ocupa uma proporção maior do tempo do forward pass. Utilize-os apenas quando as funções padrão create_*_mask não conseguem expressar o que você precisa.

Por exemplo, sobreponha uma função que retorna True em todos os pontos a uma máscara causal para transformá-la em uma totalmente bidirecional. A união com o padrão causal permite que cada token preste atenção a todos os outros tokens.

mask_kwargs = {
    "config": self.config,
    "inputs_embeds": inputs_embeds,
    "attention_mask": attention_mask,
    "past_key_values": past_key_values,
    "position_ids": position_ids,
    "or_mask_function": lambda *args: torch.tensor(True, dtype=torch.bool),
}

attention_mask = create_causal_mask(**mask_kwargs)

Durante a geração, generate() constrói máscaras por meio de create_masks_for_generate, que direciona para o create_*_mask adequado com base na configuração do modelo. Sobrescreva-o em uma classe de modelo para incorporar uma estratégia de mascaramento personalizada para a geração.

Passe uma máscara de atenção 4D personalizada

Forneça sua própria máscara 4D quando precisar de um padrão de atenção que as funções create_*_mask não conseguem expressar. Uma máscara 4D tem a forma (batch_size, 1, query_length, kv_length), em que 1 propaga a mesma máscara para todas as cabeças de atenção. O Transformers detecta e usa a máscara como está, ignorando create_*_mask.

Uma máscara 4D usa uma de duas convenções de valores.

A convenção float adiciona a máscara aos scores de atenção antes do softmax. Um score somado a 0.0 permanece inalterado, de modo que a posição contribui. Um score somado a -inf cai para zero após o softmax, de modo que a posição é excluída.

A convenção aceita depende do backend de atenção. sdpa aceita uma máscara booleana ou de ponto flutuante. eager adiciona a máscara aos scores, portanto, aceita apenas uma máscara de ponto flutuante. flash_attention_2 e flex_attention consomem seus próprios formatos (uma máscara de padding 2D e um BlockMask) e não aceitam uma máscara 4D bruta.

Um erro comum é reutilizar a convenção 1/0 de uma máscara de padding 2D em uma máscara 4D em ponto flutuante. Como a máscara é somada às pontuações, 0.0 preserva uma posição e 1.0 apenas adiciona um pequeno viés.

O exemplo abaixo contrasta uma máscara incorreta com uma correta. Ambas partem do mesmo padrão causal 1/0.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0", attn_implementation="sdpa")
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")

input_ids = tokenizer("my favorite condiment on a", return_tensors="pt").input_ids
seq_len = input_ids.shape[1]

# 1 attends, 0 masks
causal = torch.tril(torch.ones(seq_len, seq_len))

# wrong: 1.0/0.0 floats are added to the scores, so 0.0 keeps a token and 1.0 barely changes it
wrong_mask = causal[None, None]

# correct: 0.0 attends, -inf masks
correct_mask = torch.where(causal.bool(), 0.0, float("-inf"))[None, None]

A máscara incorreta mantém todas as posições, pois 0.0 é o valor que mascara, de modo que nada é excluído.

        wrong_mask                          correct_mask
   (1 attends, 0 masks)              (0 attends, -inf masks)

      k0 k1 k2 k3 k4                     k0   k1   k2   k3   k4
   q0  1  0  0  0  0                  q0  0  -inf -inf -inf -inf
   q1  1  1  0  0  0                  q1  0   0   -inf -inf -inf
   q2  1  1  1  0  0                  q2  0   0    0   -inf -inf
   q3  1  1  1  1  0                  q3  0   0    0    0   -inf
   q4  1  1  1  1  1                  q4  0   0    0    0    0

Atenção bidirecional

Modelos apenas de decodificador usam atenção causal (unidirecional) por padrão, na qual cada token atende apenas a si mesmo e aos tokens anteriores. Defina is_causal=False para alternar para atenção bidirecional, na qual cada token atende a todos os outros tokens. Isso permite usar modelos apenas de decodificador como codificadores de texto, por exemplo, para gerar embeddings.

Isso só funciona para modelos causais (decodificadores). Não transforma modelos codificadores em modelos decodificadores.

Defina is_causal=False na configuração do modelo para tornar a atenção bidirecional o padrão em cada passagem de frente.

from transformers import AutoModel, AutoConfig

config = AutoConfig.from_pretrained("meta-llama/Llama-3.2-1B")
config.is_causal = False

model = AutoModel.from_pretrained("meta-llama/Llama-3.2-1B", config=config)

# all forward passes now use bidirectional attention
outputs = model(**inputs)

Passe is_causal na chamada de forward em vez de na configuração do modelo para alternar entre atenção causal e bidirecional sem carregar o modelo duas vezes. O kwarg sobrescreve temporariamente a configuração e é restaurado após a chamada.

from transformers import AutoModel

model = AutoModel.from_pretrained("meta-llama/Llama-3.2-1B")

# run with bidirectional attention
outputs = model(**inputs, is_causal=False)

# run with default causal attention
outputs = model(**inputs)