Backends de atención

Todas las implementaciones de atención realizan el mismo cálculo. Cada token se compara con todos los demás. La diferencia es cómo se realiza el cálculo. La atención básica escala mal porque materializa la matriz de atención completa en memoria, creando cuellos de botella que ralentizan la inferencia. Las implementaciones optimizadas reorganizan la matemática para reducir el tráfico de memoria y lograr una inferencia más rápida y asequible.

La AttentionInterface proporciona implementaciones de atención optimizadas. Desacopla la implementación de la atención de la implementación del modelo para simplificar la experimentación con diferentes funciones. Añade nuevos backends fácilmente con esta interfaz coherente.

Establecer un backend de atención

Use el argumento attn_implementation en from_pretrained() para instanciar un modelo con una función de atención específica.

import torch
from transformers import AutoModelForCausalLM

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

Cambia entre backends de atención en tiempo de ejecución sin recargar el modelo usando set_attn_implementation().

model.set_attn_implementation("sdpa")

Kernels

Descargue y cargue kernels de cálculo compilados directamente desde el Hub en tiempo de ejecución con la biblioteca Kernels. Esto evita problemas de empaquetado derivados de versiones incompatibles de PyTorch o CUDA.

Los kernels se registran automáticamente en AttentionInterface al detectarse. No es necesario instalar el paquete FlashAttention explícitamente. Solicitar FlashAttention por nombre también recurre al kernel del 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"
)

Gestor de contexto de SDPA

El scaled dot product attention (SDPA) de PyTorch selecciona automáticamente la función de atención más rápida para los backends de CUDA. Para otros backends, utiliza por defecto la implementación de C++ de PyTorch.

Fuerza a SDPA a usar una implementación específica con el gestor 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)

Atención específica del backbone

Los modelos multimodales utilizan backbones diferentes para cada modalidad. Optimiza el rendimiento asignando funciones de atención específicas a cada backbone. Algunos backbones de visión funcionan mejor en fp32, por ejemplo, algo que FlashAttention no admite.

Mapee los backbones de visión a diferentes funciones de atención mediante un diccionario, mientras que el backbone de texto continúa utilizando FlashAttention. Las claves en la implementación de atención deben coincidir con los nombres de las subconfiguraciones.

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 ciertos backbones del diccionario para usar la función de atención predeterminada (SDPA).

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

Establezca la misma función de atención para todos los backbones con una sola cadena.

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

Establezca la función de atención globalmente con una clave vacía.

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

Crear una nueva función de atención

Personalice o cree nuevas funciones de atención añadiéndolas al registro de atención con AttentionInterface.register(). Los modelos utilizan estas funciones a través del argumento attn_implementation.

Registre una función de máscara de atención correspondiente cuando registre una función de atención personalizada. Si el nombre de attn_implementation personalizado no está registrado en AttentionMaskInterface, Transformers omite la creación de la máscara y pasa attention_mask=None a las capas de atención. Su función de atención debe manejar por sí misma las restricciones causales, de relleno, de empaquetado o de ventana deslizante, o bien dichas restricciones pueden omitirse silenciosamente.

Este ejemplo personaliza la función de atención para imprimir un mensaje en cada capa. Mantiene la máscara de la implementación original registrando masking_utils.sdpa_mask como la función de máscara de atención.

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

También puede agregar nuevos argumentos a la función de atención. Los modelos que admiten AttentionInterface propagan los kwargs a las capas de atención y a la función de atención. Pase los argumentos como kwargs en la función forward del modelo. Las funciones de atención personalizadas deben seguir esta firma y 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=...)

Comprueba el código de modelado de un modelo para confirmar qué argumentos y kwargs envía a la función de atención.

AttentionMaskInterface

AttentionMaskInterface es el registro que consultan las funciones create_*_mask para convertir una máscara al formato que espera el backend de atención activo. FlexAttention necesita una BlockMask, SDPA necesita un tensor 4D y FlashAttention necesita la máscara de relleno base 2D. Registra un backend personalizado o sobrescribe el formateador de uno existente con 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)

Sin un formateador registrado para el attn_implementation activo, se omite la creación de la máscara y attention_mask=None se pasa a las capas de atención.

Las funciones registradas deben coincidir con esta firma.

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]:

El argumento mask_function es un Callable que imita a las funciones mask_mod de PyTorch. Recibe 4 índices (batch_idx, head_idx, q_idx, kv_idx) y devuelve un booleano que indica si esa posición contribuye al cálculo de la atención. Esta es la misma forma primitiva utilizada por or_mask_function y and_mask_function en Build an attention mask.

Usa este solución temporal para torch.export si mask_function no logra crear una máscara.

Construir una máscara de atención

Construya máscaras de atención con las funciones create_*_mask en transformers.masking_utils. Cada función lee el backend de atención activo desde la configuración del modelo, busca el formateador de máscaras del backend en AttentionMaskInterface y devuelve el formato que ese backend espera. No necesita invertir, expandir ni convertir la máscara usted mismo.

Elija la función que coincide con el patrón de atención.

Los helpers de máscara legibles legados - get_extended_attention_mask, create_extended_attention_mask_for_decoder, invert_attention_mask - emiten una advertencia de obsolescencia y se eliminarán en una versión futura. Utilice en su lugar las funciones create_*_mask.

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

Añada restricciones adicionales sobre la máscara base con los argumentos or_mask_function y and_mask_function. Use or_mask_function para permitir que posiciones adicionales atiendan, y and_mask_function para restringir aún más el patrón base. Ambos siguen la firma de mask_function de 4 índices descrita en AttentionMaskInterface. Reciben (batch_idx, head_idx, q_idx, kv_idx) y devuelven un booleano.

or_mask_function y and_mask_function pueden expresar cualquier patrón de atención, pero son más lentos que los patrones integrados y no son compatibles con ExecuTorch. El sobrecoste es más notable en modelos más pequeños (~200M parámetros), donde la creación de la máscara ocupa una mayor proporción del tiempo de la pasada hacia adelante. Recúrrase a ellos solo cuando las funciones estándar create_*_mask no puedan expresar lo que se necesita.

Por ejemplo, superponga una función que devuelve True en todas partes sobre una máscara causal para convertirla en una completamente bidireccional. La unión con el patrón causal permite que cada token atienda a cualquier otro token.

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 la generación, generate() construye las máscaras a través de create_masks_for_generate, que deriva a la función create_*_mask adecuada según la configuración del modelo. Sobrescríbela en una clase de modelo para incorporar una estrategia de enmascarado personalizada para la generación.

Pasar una máscara de atención 4D personalizada

Pasa tu propia máscara 4D cuando necesites un patrón de atención que las funciones create_*_mask no puedan expresar. Una máscara 4D tiene la forma (batch_size, 1, query_length, kv_length), donde 1 difunde la misma máscara a través de todas las cabezas de atención. Transformers detecta y utiliza la máscara tal cual, omitiendo create_*_mask.

Una máscara 4D utiliza una de dos convenciones de valores.

La convención float añade la máscara a las puntuaciones de atención antes del softmax. Una puntuación más 0.0 permanece sin cambios, por lo que la posición contribuye. Una puntuación más -inf cae a cero después del softmax, por lo que la posición se excluye.

La convención aceptada depende del backend de atención. sdpa acepta una máscara booleana o de punto flotante. eager suma la máscara a las puntuaciones, por lo que solo acepta una máscara de punto flotante. flash_attention_2 y flex_attention consumen sus propios formatos (una máscara de relleno 2D y una BlockMask) y no aceptan una máscara 4D sin procesar.

Un error común es reutilizar la convención 1/0 de una máscara de relleno 2D en una máscara 4D de tipo float. Dado que la máscara se suma a las puntuaciones, 0.0 conserva una posición y 1.0 solo añade un pequeño sesgo.

El ejemplo siguiente contrasta una máscara incorrecta con una correcta. Ambas parten del mismo patrón 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]

La máscara incorrecta conserva todas las posiciones porque 0.0 es el valor que enmascara, por lo que nunca excluye nada.

        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

Atención bidireccional

Los modelos solo decodificador utilizan atención causal (unidireccional) de forma predeterminada, donde cada token solo atiende a sí mismo y a los tokens anteriores. Establezca is_causal=False para cambiar a atención bidireccional, donde cada token atiende a todos los demás tokens. Esto le permite utilizar modelos solo decodificador como codificadores de texto, por ejemplo, para generar embeddings.

Esto solo funciona para modelos causales (decodificadores). No convierte los modelos codificadores en modelos decodificadores.

Establezca is_causal=False en la configuración del modelo para que la atención bidireccional sea el valor predeterminado en cada paso de propagación hacia adelante.

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)

Pase is_causal en la llamada forward en lugar de en la configuración del modelo para alternar entre atención causal y bidireccional sin cargar el modelo dos veces. Este parámetro de palabra clave sobrescribe temporalmente la configuración y se restaura después de la llamada.

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)