Бэкенды внимания

Все реализации внимания выполняют одинаковые вычисления. Каждый токен сравнивается с каждым другим токеном. Разница заключается в том, как выполняются вычисления. Базовое внимание масштабируется плохо, поскольку оно материализует полную матрицу внимания в памяти, создавая узкие места, которые замедляют инференс. Оптимизированные реализации перестраивают математические операции, чтобы снизить нагрузку на память и обеспечить более быстрый и доступный инференс.

AttentionInterface предоставляет оптимизированные реализации механизма внимания. Он отделяет реализацию внимания от реализации модели, упрощая эксперименты с различными функциями. С помощью этого единообразного интерфейса легко добавлять новые бэкенды.

Настройка бэкенда внимания

Используйте аргумент attn_implementation в from_pretrained(), чтобы создать модель с конкретной функцией внимания.

import torch
from transformers import AutoModelForCausalLM

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

Переключайте бэкенды внимания во время выполнения без перезагрузки модели с помощью set_attn_implementation().

model.set_attn_implementation("sdpa")

Ядра

Загружайте и подключайте скомпилированные вычислительные ядра непосредственно из Hub во время выполнения с помощью библиотеки Kernels. Это позволяет избежать проблем с упаковкой, вызванных несовпадением версий PyTorch или CUDA.

Ядра автоматически регистрируются в AttentionInterface при обнаружении. Вам не нужно явно устанавливать пакет FlashAttention. Запрос FlashAttention по имени также приводит к использованию ядра из Hub, см. FlashAttention fallback.

import torch
from transformers import AutoModelForCausalLM

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

Контекстный менеджер SDPA

Масштабированное внимание на основе скалярного произведения (SDPA) в PyTorch автоматически выбирает наиболее быструю функцию внимания для бэкендов CUDA. Для остальных бэкендов по умолчанию используется реализация на C++ в PyTorch.

Принудительно заставить SDPA использовать конкретную реализацию с помощью контекстного менеджера 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)

Внимание, специфичное для базовой модели

Мультимодальные модели используют различные базовые архитектуры (backbones) для каждого модальности. Оптимизируйте производительность, назначая специфические функции внимания для каждой базовой архитектуры. Например, некоторые визуальные базовые архитектуры лучше работают в fp32, что не поддерживается FlashAttention.

Сопоставьте визуальные базовые модели с различными функциями внимания с помощью словаря, при этом текстовая базовая модель продолжает использовать FlashAttention. Ключи в реализации внимания должны совпадать с именами подконфигураций.

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
)

Опустите определённые backbones из dict, чтобы использовать функцию внимания по умолчанию (SDPA).

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

Установите одну и ту же функцию внимания для всех базовых моделей с помощью одной строки.

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

Установите функцию внимания глобально с пустым ключом.

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

Создание новой функции внимания

Настройте или создайте новые функции внимания, добавив их в реестр внимания с помощью AttentionInterface.register(). Модели используют эти функции через аргумент attn_implementation.

При регистрации пользовательской функции внимания регистрируйте соответствующую функцию маски внимания. Если имя пользовательского attn_implementation не зарегистрировано в AttentionMaskInterface, Transformers пропускает создание маски и передаёт attention_mask=None в слои внимания. Ваша функция внимания должна самостоятельно обрабатывать ограничения каузальности, отсечения, упаковки или скользящего окна, иначе эти ограничения могут быть молча проигнорированы.

В этом примере функция внимания настраивается так, чтобы выводить сообщение для каждого слоя. Маска сохраняется из исходной реализации путём регистрации masking_utils.sdpa_mask в качестве функции маски внимания.

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

Вы также можете добавлять новые аргументы в функцию внимания. Модели, поддерживающие AttentionInterface, передают kwargs в слои внимания и функцию внимания. Передавайте аргументы как kwargs в forward-функции модели. Пользовательские функции внимания должны соответствовать этой сигнатуре и формату возврата.

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

Проверьте код моделирования модели, чтобы убедиться, какие аргументы и kwargs она передает в функцию внимания.

AttentionMaskInterface

AttentionMaskInterface — это реестр, к которому обращаются функции create_*_mask для преобразования маски в формат, ожидаемый активным бэкендом внимания. FlexAttention требует BlockMask, SDPA — тензор размерности 4, а FlashAttention — базовую 2D-маску отсечения. Зарегистрируйте собственный бэкенд или переопределите форматтер для существующего с помощью 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)

Если для активного attn_implementation не зарегистрирован форматтер, создание маски пропускается, и attention_mask=None передаётся в слои внимания.

Зарегистрированные функции должны соответствовать этой сигнатуре.

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

Аргумент mask_function представляет собой Callable, имитирующий функции mask_mod из PyTorch. Он принимает 4 индекса (batch_idx, head_idx, q_idx, kv_idx) и возвращает булево значение, указывающее, участвует ли данная позиция в вычислении внимания. Это тот же примитивный формат, который используется в or_mask_function и and_mask_function в разделе Build an attention mask.

Используйте этот обходной путь для torch.export, если mask_function не может создать маску.

Создание маски внимания

Создавайте маски внимания с помощью функций create_*_mask в модуле transformers.masking_utils. Каждая функция считывает активный бэкенд внимания из конфигурации модели, находит форматтер маски для этого бэкенда в AttentionMaskInterface и возвращает ожидаемый этим бэкендом формат. Вам не нужно самостоятельно инвертировать, расширять или приводить маску к нужному типу.

Выберите функцию, соответствующую шаблону внимания.

Устаревшие вызываемые помощники для масок - get_extended_attention_mask, create_extended_attention_mask_for_decoder, invert_attention_mask - выводят предупреждение об устаревании и будут удалены в будущей версии. Вместо них используйте функции 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,
)

Добавляйте дополнительные ограничения поверх базовой маски с помощью аргументов or_mask_function и and_mask_function. Используйте or_mask_function, чтобы разрешить внимание для дополнительных позиций, и and_mask_function, чтобы дополнительно ограничить базовый паттерн. Оба следуют сигнатуре mask_function с четырьмя индексами, описанной в AttentionMaskInterface. Они принимают (batch_idx, head_idx, q_idx, kv_idx) и возвращают булево значение.

or_mask_function и and_mask_function позволяют задать любой паттерн внимания, но они медленнее встроенных паттернов и несовместимы с ExecuTorch. Нагрузка наиболее заметна на небольших моделях (~200M параметров), где создание маски занимает значительную долю времени прямого прохода. Используйте их только в тех случаях, когда стандартные функции create_*_mask не позволяют выразить нужную логику.

Например, наложите на каузальную маску функцию, возвращающую True везде, чтобы превратить её в полностью двунаправленную. Объединение с каузальным паттерном позволяет каждому токену обращать внимание на любой другой токен.

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)

В процессе генерации generate() формирует маски с помощью create_masks_for_generate, который вызывает соответствующий create_*_mask в зависимости от конфигурации модели. Переопределите этот метод в классе модели, чтобы подключить собственную стратегию маскирования для генерации.

Передача пользовательской 4D-маски внимания

Передавайте собственную 4D-маску, когда вам нужен паттерн внимания, который не могут выразить функции create_*_mask. Форма 4D-маски — (batch_size, 1, query_length, kv_length), где 1 означает, что одна и та же маска применяется ко всем головам внимания. Transformers обнаруживает и использует маску как есть, пропуская вызов create_*_mask.

Маска 4D использует один из двух вариантов соглашения о значениях.

Конвенция float добавляет маску к оценкам внимания перед softmax. Оценка плюс 0.0 остаётся без изменений, поэтому позиция учитывается. Оценка плюс -inf обнуляется после softmax, поэтому позиция исключается.

Принятая конвенция зависит от бэкенда внимания. sdpa принимает булеву или float-маску. eager прибавляет маску к скорингу, поэтому принимает только float-маску. flash_attention_2 и flex_attention используют собственные форматы (2D-маску padding и BlockMask) и не принимают «сырую» 4D-маску.

Частая ошибка — использовать соглашение 1/0 из 2D-маски отсечки в float-маске 4D. Поскольку маска прибавляется к оценкам, 0.0 сохраняет позицию, а 1.0 лишь добавляет небольшое смещение.

Пример ниже противопоставляет неверную маску правильной. Обе начинаются с одного и того же причинного паттерна 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]

Неверная маска сохраняет все позиции, так как 0.0 — это значение, которое маскирует, поэтому оно никогда ничего не исключает.

        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

Двунаправленное внимание

Модели только с декодером по умолчанию используют каузальное (однонаправленное) внимание, при котором каждый токен обращается только к себе и предыдущим токенам. Установите is_causal=False, чтобы переключиться на двунаправленное внимание, при котором каждый токен обращается ко всем остальным токенам. Это позволяет использовать модели только с декодером в качестве текстовых энкодеров, например, для генерации эмбеддингов.

Это работает только для каузальных (декодерных) моделей. Это не превращает энкодерные модели в декодерные.

Установите is_causal=False в конфигурации модели, чтобы двунаправленное внимание стало значением по умолчанию для каждого прямого прохода.

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)

Передавайте is_causal в вызове forward вместо конфигурации модели, чтобы переключаться между каузальным и двунаправленным вниманием без повторной загрузки модели. Этот аргумент временно переопределяет конфигурацию и восстанавливается после вызова.

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)