Các backend chú ý

Tất cả các triển khai attention đều thực hiện cùng một phép tính. Mỗi token được so sánh với mọi token khác. Sự khác biệt nằm ở cách phép tính được thực hiện. Attention cơ bản có hiệu suất kém khi quy mô tăng vì nó vật chất hóa ma trận attention đầy đủ trong bộ nhớ, tạo ra các nút thắt cổ chai làm chậm suy luận. Các triển khai tối ưu sắp xếp lại các phép toán để giảm lưu lượng bộ nhớ, giúp suy luận nhanh hơn và tiết kiệm chi phí hơn.

AttentionInterface cung cấp các triển khai attention được tối ưu hóa. Nó tách rời triển khai attention khỏi triển khai mô hình để đơn giản hóa việc thử nghiệm với các hàm khác nhau. Dễ dàng thêm các backend mới với giao diện nhất quán này.

Đặt backend chú ý

Sử dụng đối số attn_implementation trong from_pretrained() để khởi tạo một mô hình với một hàm chú ý cụ thể.

import torch
from transformers import AutoModelForCausalLM

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

Chuyển đổi giữa các backend attention tại thời điểm chạy mà không cần tải lại mô hình bằng cách sử dụng set_attn_implementation().

model.set_attn_implementation("sdpa")

Kernels

Tải xuống và nạp các kernel tính toán đã biên dịch trực tiếp từ Hub vào thời điểm chạy bằng thư viện Kernels. Điều này giúp tránh các vấn đề đóng gói do sự không tương thích giữa các phiên bản PyTorch hoặc CUDA.

Kernels tự động đăng ký vào AttentionInterface khi được phát hiện. Bạn không cần cài đặt gói FlashAttention một cách tường minh. Việc yêu cầu FlashAttention theo tên cũng sẽ chuyển sang kernel Hub, xem FlashAttention fallback.

import torch
from transformers import AutoModelForCausalLM

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

Trình quản lý ngữ cảnh SDPA

Scaled dot product attention (SDPA) của PyTorch tự động chọn hàm attention nhanh nhất cho các backend CUDA. Đối với các backend khác, mặc định sẽ sử dụng bản triển khai C++ của PyTorch.

Buộc SDPA sử dụng một triển khai cụ thể bằng trình quản lý ngữ cảnh 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)

Chú ý cụ thể cho từng backbone

Mô hình đa phương thức sử dụng các backbone khác nhau cho từng phương thức. Tối ưu hóa hiệu suất bằng cách gán các hàm attention cụ thể cho từng backbone. Ví dụ, một số backbone thị giác hoạt động tốt hơn ở fp32, điều mà FlashAttention không hỗ trợ.

Ánh xạ các bộ khung thị giác (vision backbones) đến các hàm chú ý (attention) khác nhau bằng một dict trong khi bộ khung văn bản (text backbone) tiếp tục sử dụng FlashAttention. Các khóa trong cài đặt chú ý phải khớp với tên của các cấu hình con (sub-config).

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
)

Bỏ qua một số backbone nhất định khỏi dict để sử dụng hàm attention mặc định (SDPA).

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

Đặt cùng một hàm chú ý cho tất cả các bộ khung bằng một chuỗi duy nhất.

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

Đặt hàm chú ý toàn cục bằng một khóa rỗng.

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

Tạo một hàm chú ý mới

Tùy chỉnh hoặc tạo các hàm chú ý mới bằng cách thêm chúng vào registry chú ý với AttentionInterface.register(). Các mô hình sử dụng những hàm này thông qua đối số attn_implementation.

Đăng ký một hàm mask chú ý tương ứng khi bạn đăng ký một hàm chú ý tùy chỉnh. Nếu tên attn_implementation tùy chỉnh không được đăng ký trong AttentionMaskInterface, Transformers sẽ bỏ qua việc tạo mask và truyền attention_mask=None vào các lớp chú ý. Hàm chú ý của bạn phải tự xử lý các ràng buộc causal, padding, packing hoặc sliding-window, nếu không các ràng buộc đó có thể bị loại bỏ mà không có cảnh báo.

Ví dụ này tùy chỉnh hàm attention để in ra một câu cho mỗi lớp. Nó giữ nguyên mask trong cài đặt gốc bằng cách đăng ký masking_utils.sdpa_mask làm hàm attention 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))

Bạn cũng có thể thêm các đối số mới vào hàm attention. Các model hỗ trợ AttentionInterface sẽ truyền kwargs đến các lớp attention và hàm attention. Hãy truyền các đối số dưới dạng kwargs trong hàm forward của model. Các hàm attention tùy chỉnh phải tuân theo chữ ký và định dạng trả về này.

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

Kiểm tra mã modeling của một model để xác nhận các arguments và kwargs mà nó gửi đến hàm attention.

AttentionMaskInterface

AttentionMaskInterface là registry mà các hàm create_*_mask tham chiếu để chuyển đổi mask sang định dạng mà backend attention đang hoạt động mong đợi. FlexAttention cần một BlockMask, SDPA cần một tensor 4D, và FlashAttention cần mask padding cơ bản 2D. Hãy đăng ký một backend tùy chỉnh, hoặc ghi đè formatter cho một backend hiện có, bằng 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)

Nếu không có formatter được đăng ký cho attn_implementation đang hoạt động, việc tạo mask sẽ bị bỏ qua và attention_mask=None sẽ được truyền đến các lớp attention.

Các hàm đã đăng ký phải khớp với chữ ký này.

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

Tham số mask_function là một Callable mô phỏng các hàm mask_mod của PyTorch. Nó nhận 4 chỉ số (batch_idx, head_idx, q_idx, kv_idx) và trả về một giá trị boolean cho biết vị trí đó có đóng góp vào phép tính attention hay không. Đây là cùng một hình dạng nguyên thủy được sử dụng bởi or_mask_function và and_mask_function trong Xây dựng mặt nạ attention.

Sử dụng giải pháp thay thế này cho torch.export nếu mask_function không tạo được mask.

Xây dựng mặt nạ chú ý

Xây dựng attention mask bằng các hàm create_*_mask trong transformers.masking_utils. Mỗi hàm sẽ đọc backend attention đang hoạt động từ cấu hình mô hình, tra cứu bộ định dạng mask của backend đó trong AttentionMaskInterface, và trả về định dạng mà backend đó yêu cầu. Bạn không cần phải đảo ngược, mở rộng hoặc chuyển kiểu mask một cách thủ công.

Chọn hàm phù hợp với mẫu chú ý.

Các hàm trợ giúp mask dạng callable cũ - get_extended_attention_mask, create_extended_attention_mask_for_decoder, invert_attention_mask - sẽ phát ra cảnh báo không còn được khuyến khích sử dụng và sẽ bị loại bỏ trong một bản phát hành tương lai. Hãy sử dụng các hàm create_*_mask thay thế.

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

Thêm các ràng buộc bổ sung lên trên mặt nạ cơ bản bằng các đối số or_mask_function và and_mask_function. Sử dụng or_mask_function để cho phép các vị trí bổ sung tham gia chú ý, và and_mask_function để hạn chế thêm mẫu cơ bản. Cả hai đều tuân theo chữ ký mask_function có 4 chỉ số được mô tả trong AttentionMaskInterface. Chúng nhận (batch_idx, head_idx, q_idx, kv_idx) và trả về một giá trị boolean.

or_mask_function và and_mask_function có thể biểu diễn bất kỳ mẫu chú ý nào, nhưng chúng chậm hơn các mẫu tích hợp sẵn và không tương thích với ExecuTorch. Chi phí phát sinh rõ rệt nhất trên các mô hình nhỏ hơn (~200M tham số), nơi việc tạo mask chiếm tỷ trọng lớn hơn trong thời gian của bước truyền thẳng. Chỉ nên dùng chúng khi các hàm create_*_mask chuẩn không thể biểu diễn được điều bạn cần.

Ví dụ, hãy phủ một hàm trả về True ở mọi vị trí lên một mặt nạ nhân quả (causal mask) để biến nó thành một mặt nạ hai chiều hoàn toàn. Việc hợp nhất với mẫu nhân quả cho phép mọi token chú ý đến mọi token khác.

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)

Trong quá trình sinh, generate() xây dựng các mặt nạ thông qua create_masks_for_generate, hàm này sẽ điều hướng đến create_*_mask phù hợp dựa trên cấu hình của mô hình. Hãy ghi đè hàm này trên một lớp mô hình để tích hợp chiến lược che mặt nạ tùy chỉnh cho quá trình sinh.

Truyền vào mặt nạ chú ý 4D tùy chỉnh

Truyền mask 4D của riêng bạn khi cần một mẫu attention mà các hàm create_*_mask không thể biểu diễn. Mask 4D có hình dạng (batch_size, 1, query_length, kv_length), trong đó 1 broadcast cùng một mask cho mọi attention head. Transformers phát hiện và sử dụng mask như vậy, bỏ qua create_*_mask.

Mặt nạ 4D sử dụng một trong hai quy ước giá trị.

Quy ước float cộng mask vào điểm chú ý trước softmax. Một điểm cộng 0.0 giữ nguyên, nên vị trí đó được tính đến. Một điểm cộng -inf giảm về không sau softmax, nên vị trí đó bị loại.

Quy ước được chấp nhận phụ thuộc vào backend chú ý. sdpa nhận một mask boolean hoặc float. eager cộng mask vào các điểm số, do đó chỉ nhận mask float. flash_attention_2 và flex_attention sử dụng các định dạng riêng (một mask padding 2D và một BlockMask) và không chấp nhận mask 4D thô.

Một sai lầm phổ biến là tái sử dụng quy ước 1/0 của mặt nạ đệm 2D trong mặt nạ float 4D. Vì mặt nạ được cộng vào điểm số, 0.0 giữ nguyên một vị trí còn 1.0 chỉ thêm một độ lệch nhỏ.

Ví dụ dưới đây so sánh một mask sai với một mask đúng. Cả hai đều bắt đầu từ cùng một mẫu nhân quả 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]

Mặt nạ sai giữ lại mọi vị trí vì 0.0 là giá trị dùng để che, nên nó không bao giờ loại trừ bất kỳ thứ gì.

        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

Chú ý hai chiều

Mô hình chỉ có decoder sử dụng chú ý nhân quả (một chiều) theo mặc định, trong đó mỗi token chỉ chú ý đến chính nó và các token trước đó. Đặt is_causal=False để chuyển sang chú ý hai chiều, trong đó mỗi token chú ý đến mọi token khác. Điều này cho phép bạn sử dụng các mô hình chỉ có decoder như bộ mã hóa văn bản, chẳng hạn để tạo ra các vector nhúng (embeddings).

Điều này chỉ hoạt động với các mô hình nhân quả (decoder). Nó không biến các mô hình encoder thành mô hình decoder.

Đặt is_causal=False trong cấu hình mô hình để chú ý hai chiều trở thành mặc định cho mọi lần truyền tiến.

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)

Truyền is_causal trong lời gọi forward thay vì trong cấu hình mô hình để chuyển đổi giữa chú ý nhân quả và chú ý hai chiều mà không cần tải mô hình hai lần. Tham số kwarg này tạm thời ghi đè lên cấu hình và được khôi phục sau khi gọi.

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)