Backend perhatian

Semua implementasi perhatian melakukan komputasi yang sama. Setiap token dibandingkan dengan setiap token lainnya. Perbedaannya terletak pada bagaimana komputasi tersebut dilakukan. Perhatian dasar memiliki skalabilitas yang buruk karena mematerialisasi matriks perhatian penuh di memori, menciptakan hambatan yang memperlambat inferensi. Implementasi yang dioptimalkan menyusun ulang matematisnya untuk mengurangi lalu lintas memori demi inferensi yang lebih cepat dan lebih terjangkau.

AttentionInterface menyediakan implementasi perhatian yang dioptimalkan. Antarmuka ini memisahkan implementasi perhatian dari implementasi model untuk menyederhanakan eksperimen dengan fungsi yang berbeda. Tambahkan backend baru dengan mudah melalui antarmuka yang konsisten ini.

Atur backend perhatian

Gunakan argumen attn_implementation di from_pretrained() untuk membuat model dengan fungsi perhatian tertentu.

import torch
from transformers import AutoModelForCausalLM

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

Beralih antara backend perhatian pada waktu runtime tanpa memuat ulang model menggunakan set_attn_implementation().

model.set_attn_implementation("sdpa")

Kernels

Unduh dan muat kernel komputasi yang telah dikompilasi langsung dari Hub saat runtime dengan pustaka Kernels. Hal ini menghindari masalah pengemasan akibat ketidakcocokan versi PyTorch atau CUDA.

Kernels secara otomatis terdaftar ke AttentionInterface saat terdeteksi. Anda tidak perlu menginstal paket FlashAttention secara eksplisit. Meminta FlashAttention berdasarkan nama juga akan kembali ke kernel Hub, lihat FlashAttention fallback.

import torch
from transformers import AutoModelForCausalLM

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

Manajer konteks SDPA

Scaled dot product attention (SDPA) dari PyTorch secara otomatis memilih fungsi attention tercepat untuk backend CUDA. Untuk backend lainnya, default-nya adalah implementasi C++ PyTorch.

Paksa SDPA untuk menggunakan implementasi tertentu dengan context manager 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)

Perhatian spesifik backbone

Model multimodal menggunakan backbone yang berbeda untuk setiap modalitas. Optimalkan kinerja dengan menugaskan fungsi perhatian tertentu ke setiap backbone. Beberapa backbone visi, misalnya, memberikan kinerja yang lebih baik dalam fp32, yang tidak didukung oleh FlashAttention.

Petakan tulang punggung visi ke fungsi perhatian yang berbeda menggunakan dict, sementara tulang punggung teks tetap menggunakan FlashAttention. Kunci dalam implementasi perhatian harus cocok dengan nama sub-konfigurasi.

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
)

Hapus beberapa backbone dari dict untuk menggunakan fungsi perhatian default (SDPA).

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

Atur fungsi perhatian yang sama untuk semua backbone dengan satu string.

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

Atur fungsi perhatian secara global dengan kunci kosong.

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

Buat fungsi perhatian baru

Sesuaikan atau buat fungsi perhatian baru dengan menambahkannya ke registri perhatian menggunakan AttentionInterface.register(). Model menggunakan fungsi-fungsi ini melalui argumen attn_implementation.

Daftarkan fungsi mask perhatian yang sesuai saat Anda mendaftarkan fungsi perhatian kustom. Jika nama attn_implementation kustom tidak terdaftar di AttentionMaskInterface, Transformers akan melewati pembuatan mask dan meneruskan attention_mask=None ke lapisan perhatian. Fungsi perhatian Anda harus menangani kendala kausal, padding, packing, atau sliding-window sendiri, atau kendala tersebut dapat diabaikan secara diam-diam.

Contoh ini menyesuaikan fungsi perhatian untuk mencetak pernyataan untuk setiap lapisan. Mask dipertahankan dalam implementasi asli dengan mendaftarkan masking_utils.sdpa_mask sebagai fungsi mask perhatian.

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

Anda juga dapat menambahkan argumen baru ke fungsi perhatian. Model yang mendukung AttentionInterface meneruskan kwargs ke lapisan perhatian dan fungsi perhatian. Lewatkan argumen sebagai kwargs dalam fungsi forward model. Fungsi perhatian kustom harus mengikuti tanda tangan dan format pengembalian ini.

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

Periksa kode pemodelan model untuk mengonfirmasi argumen dan kwargs apa yang dikirimkannya ke fungsi perhatian.

AttentionMaskInterface

AttentionMaskInterface adalah registry yang dirujuk oleh fungsi create_*_mask untuk mengonversi mask menjadi format yang diharapkan oleh backend attention yang aktif. FlexAttention membutuhkan BlockMask, SDPA membutuhkan tensor 4D, dan FlashAttention membutuhkan mask padding dasar 2D. Daftarkan backend kustom, atau timpa formatter untuk backend yang sudah ada, dengan 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)

Tanpa formatter yang terdaftar untuk attn_implementation yang aktif, pembuatan mask dilewati dan attention_mask=None diteruskan ke lapisan attention.

Fungsi yang terdaftar harus cocok dengan tanda tangan ini.

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

Argumen mask_function adalah Callable yang meniru fungsi mask_mod PyTorch. Argumen ini menerima 4 indeks (batch_idx, head_idx, q_idx, kv_idx) dan mengembalikan nilai boolean yang menunjukkan apakah posisi tersebut berkontribusi pada komputasi perhatian. Ini adalah bentuk primitif yang sama yang digunakan oleh or_mask_function dan and_mask_function di Membangun mask perhatian.

Gunakan workaround ini untuk torch.export jika mask_function gagal membuat mask.

Membangun attention mask

Bangun attention mask dengan fungsi create_*_mask di transformers.masking_utils. Setiap fungsi membaca backend attention aktif dari konfigurasi model, mencari formatter mask backend di AttentionMaskInterface, dan mengembalikan format yang diharapkan oleh backend tersebut. Anda tidak perlu membalik, memperluas, atau meng-cast mask secara manual.

Pilih fungsi yang sesuai dengan pola perhatian.

Helper mask callable lama - get_extended_attention_mask, create_extended_attention_mask_for_decoder, invert_attention_mask - mengeluarkan peringatan deprekasi dan akan dihapus pada rilis mendatang. Gunakan fungsi create_*_mask sebagai gantinya.

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

Tambahkan kendala tambahan di atas mask dasar dengan argumen or_mask_function dan and_mask_function. Gunakan or_mask_function untuk memungkinkan posisi tambahan melakukan attend, dan and_mask_function untuk membatasi pola dasar lebih lanjut. Keduanya mengikuti tanda tangan mask_function 4-indeks yang dijelaskan dalam AttentionMaskInterface. Mereka menerima (batch_idx, head_idx, q_idx, kv_idx) dan mengembalikan nilai boolean.

or_mask_function dan and_mask_function dapat mengekspresikan pola perhatian apa pun, tetapi lebih lambat daripada pola bawaan dan tidak kompatibel dengan ExecuTorch. Overhead paling terasa pada model yang lebih kecil (~200M parameter), di mana pembuatan mask mengambil porsi yang lebih besar dari waktu forward-pass. Gunakan fungsi-fungsi ini hanya ketika fungsi create_*_mask standar tidak dapat mengekspresikan apa yang Anda butuhkan.

Sebagai contoh, tambahkan fungsi yang mengembalikan True di mana saja pada mask kausal untuk mengubahnya menjadi mask bidireksional penuh. Gabungan dengan pola kausal memungkinkan setiap token memperhatikan setiap token lainnya.

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)

Saat generasi, generate() membangun mask melalui create_masks_for_generate, yang meneruskan ke create_*_mask yang tepat berdasarkan konfigurasi model. Timpa metode ini pada kelas model untuk menerapkan strategi masking kustom untuk generasi.

Lewati mask perhatian 4D kustom

Lewati mask 4D Anda sendiri ketika Anda membutuhkan pola perhatian yang tidak dapat diekspresikan oleh fungsi create_*_mask. Mask 4D memiliki bentuk (batch_size, 1, query_length, kv_length), di mana 1 melakukan broadcast mask yang sama di setiap kepala perhatian. Transformers mendeteksi dan menggunakan mask apa adanya, melewati create_*_mask.

Mask 4D menggunakan salah satu dari dua konvensi nilai.

Konvensi float menambahkan mask ke skor perhatian sebelum softmax. Skor ditambah 0.0 tetap tidak berubah, sehingga posisi tersebut berkontribusi. Skor ditambah -inf menjadi nol setelah softmax, sehingga posisi tersebut dikecualikan.

Konvensi yang diterima bergantung pada backend perhatian. sdpa menerima mask boolean atau float. eager menambahkan mask ke skor, sehingga hanya menerima mask float. flash_attention_2 dan flex_attention menggunakan format mereka sendiri (mask padding 2D dan BlockMask) dan tidak menerima mask 4D mentah.

Kesalahan umum adalah menggunakan kembali konvensi 1/0 dari mask padding 2D pada mask 4D float. Karena mask ditambahkan ke skor, 0.0 mempertahankan posisi dan 1.0 hanya menambahkan bias kecil.

Contoh di bawah ini membandingkan mask yang salah dengan yang benar. Keduanya dimulai dari pola kausal 1/0 yang sama.

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]

Mask yang salah mempertahankan semua posisi karena 0.0 adalah nilai yang digunakan untuk masking, sehingga tidak pernah mengecualikan apa pun.

        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

Perhatian bidireksional

Model decoder-only menggunakan perhatian kausal (unidireksional) secara default, di mana setiap token hanya memperhatikan dirinya sendiri dan token-token sebelumnya. Atur is_causal=False untuk beralih ke perhatian bidireksional, di mana setiap token memperhatikan semua token lainnya. Hal ini memungkinkan Anda menggunakan model decoder-only sebagai encoder teks, misalnya untuk menghasilkan embedding.

Ini hanya berfungsi untuk model kausal (decoder). Hal ini tidak mengubah model encoder menjadi model decoder.

Setel is_causal=False dalam konfigurasi model untuk menjadikan perhatian bidireksional sebagai default untuk setiap forward pass.

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)

Lewati is_causal dalam panggilan forward alih-alih melalui konfigurasi model untuk beralih antara perhatian kausal dan bidireksional tanpa memuat model dua kali. Argumen kwarg ini secara sementara menimpa konfigurasi dan dipulihkan setelah panggilan selesai.

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)