Backends d'attention

Toutes les implémentations d'attention effectuent le même calcul. Chaque jeton est comparé à tous les autres jetons. La différence réside dans la manière dont le calcul est effectué. L'attention de base échoue à l'échelle car elle matérialise la matrice d'attention complète en mémoire, créant des goulets d'étranglement qui ralentissent l'inférence. Les implémentations optimisées réorganisent les mathématiques pour réduire le trafic mémoire et permettre une inférence plus rapide et plus abordable.

AttentionInterface fournit des implémentations d'attention optimisées. Il découple l'implémentation de l'attention de l'implémentation du modèle afin de simplifier les expérimentations avec différentes fonctions. Ajoutez facilement de nouveaux backends grâce à cette interface cohérente.

Définir un backend d'attention

Utilisez l'argument attn_implementation dans from_pretrained() pour instancier un modèle avec une fonction d'attention spécifique.

import torch
from transformers import AutoModelForCausalLM

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

Basculez entre les backends d'attention à l'exécution sans recharger le modèle en utilisant set_attn_implementation().

model.set_attn_implementation("sdpa")

Kernels

Téléchargez et chargez les noyaux de calcul compilés directement depuis le Hub à l'exécution avec la bibliothèque Kernels. Cela évite les problèmes d'empaquetage liés à des versions de PyTorch ou de CUDA incompatibles.

Les noyaux s’enregistrent automatiquement dans AttentionInterface lors de leur détection. Il n’est pas nécessaire d’installer explicitement le package FlashAttention. Demander FlashAttention par son nom bascule également sur le noyau du Hub, voir FlashAttention fallback.

import torch
from transformers import AutoModelForCausalLM

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

Gestionnaire de contexte SDPA

La fonction d’attention à produit scalaire mis à l’échelle (SDPA) de PyTorch sélectionne automatiquement la fonction d’attention la plus rapide pour les backends CUDA. Pour les autres backends, elle utilise par défaut l’implémentation C++ de PyTorch.

Forcer SDPA à utiliser une implémentation spécifique avec le gestionnaire de contexte 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)

Attention spécifique au backbone

Les modèles multimodaux utilisent des backbones différents pour chaque modalité. Optimisez les performances en affectant des fonctions d'attention spécifiques à chaque backbone. Certains backbones de vision fonctionnent mieux en fp32, par exemple, ce que FlashAttention ne prend pas en charge.

Associez les squelettes de vision à différentes fonctions d'attention à l'aide d'un dictionnaire, tandis que le squelette de texte continue d'utiliser FlashAttention. Les clés de l'implémentation de l'attention doivent correspondre aux noms des sous-configurations.

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
)

Omettez certains backbones du dictionnaire pour utiliser la fonction d'attention par défaut (SDPA).

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

Définissez la même fonction d’attention pour tous les backbones avec une seule chaîne de caractères.

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

Définissez la fonction d'attention globalement avec une clé vide.

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

Créer une nouvelle fonction d'attention

Personnalisez ou créez de nouvelles fonctions d'attention en les ajoutant au registre d'attention via AttentionInterface.register(). Les modèles utilisent ces fonctions à travers l'argument attn_implementation.

Enregistrez une fonction de masque d'attention correspondante lorsque vous enregistrez une fonction d'attention personnalisée. Si le nom de attn_implementation personnalisé n'est pas enregistré dans AttentionMaskInterface, Transformers ignore la création du masque et transmet attention_mask=None aux couches d'attention. Votre fonction d'attention doit gérer elle-même les contraintes causales, de bourrage, de regroupement ou de fenêtre glissante, sinon ces contraintes peuvent être silencieusement ignorées.

Cet exemple personnalise la fonction d'attention pour afficher un message pour chaque couche. Il conserve le masque de l'implémentation d'origine en enregistrant masking_utils.sdpa_mask en tant que fonction de masque d'attention.

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

Vous pouvez également ajouter de nouveaux arguments à la fonction d'attention. Les modèles prenant en charge AttentionInterface propagent les kwargs aux couches d'attention et à la fonction d'attention. Passez les arguments en tant que kwargs dans la fonction forward du modèle. Les fonctions d'attention personnalisées doivent respecter cette signature et ce format de retour.

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

Consultez le code de modélisation d’un modèle pour vérifier les arguments et kwargs qu’il transmet à la fonction d’attention.

AttentionMaskInterface

AttentionMaskInterface est le registre que les fonctions create_*_mask consultent pour convertir un masque dans le format attendu par le backend d'attention actif. FlexAttention a besoin d'un BlockMask, SDPA a besoin d'un tenseur 4D et FlashAttention a besoin du masque de padding de base 2D. Enregistrez un backend personnalisé, ou écrasez le formateur d'un backend existant, avec 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)

En l’absence d’un formateur enregistré pour le attn_implementation actif, la création du masque est ignorée et attention_mask=None est transmis aux couches d’attention.

Les fonctions enregistrées doivent correspondre à cette signature.

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

L'argument mask_function est un Callable qui imite les fonctions mask_mod de PyTorch. Il prend 4 indices (batch_idx, head_idx, q_idx, kv_idx) et renvoie un booléen indiquant si cette position contribue au calcul de l'attention. C'est la même forme primitive utilisée par or_mask_function et and_mask_function dans Construire un masque d'attention.

Utilisez ce contournement pour torch.export si mask_function échoue à créer un masque.

Construire un masque d'attention

Construisez les masques d’attention à l’aide des fonctions create_*_mask de transformers.masking_utils. Chaque fonction lit le backend d’attention actif depuis la configuration du modèle, recherche le formateur de masque du backend dans AttentionMaskInterface et renvoie le format attendu par ce backend. Vous n’avez pas besoin d’inverser, d’étendre ou de convertir le masque vous-même.

Choisissez la fonction qui correspond au motif d'attention.

Les utilitaires de masque appelables hérités - get_extended_attention_mask, create_extended_attention_mask_for_decoder, invert_attention_mask - émettent un avertissement de non-recommandation et seront supprimés dans une future version. Utilisez à la place les fonctions 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,
)

Ajoutez des contraintes supplémentaires au masque de base à l'aide des arguments or_mask_function et and_mask_function. Utilisez or_mask_function pour autoriser l'attention sur des positions supplémentaires, et and_mask_function pour restreindre davantage le motif de base. Les deux suivent la signature à 4 indices de mask_function décrite dans AttentionMaskInterface. Elles prennent (batch_idx, head_idx, q_idx, kv_idx) en entrée et renvoient un booléen.

or_mask_function et and_mask_function peuvent exprimer n’importe quel motif d’attention, mais ils sont plus lents que les motifs intégrés et ne sont pas compatibles avec ExecuTorch. Le surcoût est le plus notable sur les petits modèles (~200M paramètres), où la création du masque représente une part plus importante du temps de passe avant. N’y recourrez que lorsque les fonctions standard create_*_mask ne peuvent pas exprimer ce dont vous avez besoin.

Par exemple, superposez une fonction qui retourne True partout sur un masque causal pour le transformer en un masque entièrement bidirectionnel. L’union avec le motif causal permet à chaque jeton de prêter attention à tous les autres jetons.

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)

Pendant la génération, generate() construit les masques via create_masks_for_generate, qui achemine l'appel vers le bon create_*_mask en fonction de la configuration du modèle. Surchargez-le dans une classe de modèle pour intégrer une stratégie de masquage personnalisée pour la génération.

Passer un masque d'attention 4D personnalisé

Fournissez votre propre masque 4D lorsque vous avez besoin d’un motif d’attention que les fonctions create_*_mask ne peuvent pas exprimer. Un masque 4D a la forme (batch_size, 1, query_length, kv_length), où 1 diffuse le même masque sur toutes les têtes d’attention. Transformers détecte et utilise le masque tel quel, en ignorant create_*_mask.

Un masque 4D utilise l'une des deux conventions de valeurs.

La convention float ajoute le masque aux scores d’attention avant le softmax. Un score plus 0.0 reste inchangé, donc la position contribue. Un score plus -inf tombe à zéro après le softmax, donc la position est exclue.

La convention acceptée dépend du backend d’attention. sdpa accepte un masque booléen ou flottant. eager ajoute le masque aux scores, il n’accepte donc qu’un masque flottant. flash_attention_2 et flex_attention consomment leurs propres formats (un masque de padding 2D et un BlockMask) et n’acceptent pas de masque 4D brut.

Une erreur courante consiste à réutiliser la convention 1/0 d'un masque de padding 2D dans un masque 4D en float. Comme le masque est ajouté aux scores, 0.0 conserve une position et 1.0 n'ajoute qu'un petit biais.

L'exemple ci-dessous met en contraste un masque incorrect avec un masque correct. Les deux partent du même motif 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]

Le mauvais masque conserve toutes les positions car 0.0 est la valeur qui masque, il n’exclut donc jamais quoi que ce soit.

        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

Attention bidirectionnelle

Les modèles uniquement décodeurs utilisent par défaut une attention causale (unidirectionnelle), où chaque jeton ne porte attention qu’à lui-même et aux jetons précédents. Définissez is_causal=False pour passer à une attention bidirectionnelle, où chaque jeton porte attention à tous les autres jetons. Cela permet d’utiliser des modèles uniquement décodeurs comme encodeurs de texte, par exemple pour générer des plongements.

Cela ne fonctionne que pour les modèles causaux (décodeurs). Cela ne transforme pas les modèles encodeurs en modèles décodeurs.

Définissez is_causal=False dans la configuration du modèle pour que l'attention bidirectionnelle soit la valeur par défaut pour chaque passe avant.

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)

Passez is_causal dans l'appel forward plutôt que dans la configuration du modèle pour basculer entre une attention causale et une attention bidirectionnelle sans charger le modèle deux fois. Ce paramètre mot-clé écrase temporairement la configuration et est restauré après l'appel.

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)