BertSdpaSelfAttention源码解析
1. BertSdpaSelfAttention类 介绍
BertSdpaSelfAttention类是 BERT 模型自注意力层的实现,继承 BertSelfAttention 类。BertSdpaSelfAttention 模块是在 Hugging Face 的 transformers 库的 4.31.0 版本后引入的。这一模块是为了增强 BERT 模型中自注意力机制的效率和性能,利用了 PyTorch 中的 scaled_dot_product_attention
函数。
2. BertSdpaSelfAttention类 源码解析
源码地址:transformers/src/transformers/models/bert/modeling_bert.py
# -*- coding: utf-8 -*-
# @time: 2024/7/15 14:30
import torch
from typing import Optional, Tuple
from packaging import version
from transformers.models.bert.modeling_bert import BertSelfAttention
from transformers.utils import get_torch_version, logging
logger = logging.get_logger(__name__)
class BertSdpaSelfAttention(BertSelfAttention):
def __init__(self, config, position_embedding_type=None):
super().__init__(config, position_embedding_type=position_embedding_type) # 初始化父类,继承其配置和属性
self.dropout_prob = config.attention_probs_dropout_prob # 存储注意力机制中的dropout概率
# 检查当前的PyTorch版本,决定是否需要连续的qkv输入,低于2.2.0版本时设为True
self.require_contiguous_qkv = version.parse(get_torch_version()) < version.parse("2.2.0")
# Adapted from BertSelfAttention
def forward(
self,
hidden_states: torch.Tensor, # 输入的隐藏状态张量
attention_mask: Optional[torch.Tensor] = None, # 注意力掩码,可选
head_mask: Optional[torch.FloatTensor] = None, # 头部掩码,可选
encoder_hidden_states: Optional[torch.FloatTensor] = None, # 编码器的隐藏状态,仅在交叉注意力中使用
encoder_attention_mask: Optional[torch.FloatTensor] = None, # 编码器的注意力掩码,仅在交叉注意力中使用
past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None, # 过去的键和值,用于缓存
output_attentions: Optional[bool] = False, # 是否输出注意力权重
) -> Tuple[torch.Tensor]:
# 1. 如果位置嵌入类型不是 "absolute" 或者需要输出注意力权重或使用头部掩码,则调用BertSelfAttention的forward方法来处理
if self.position_embedding_type != "absolute" or output_attentions or head_mask is not None:
# TODO: Improve this warning with e.g. `model.config._attn_implementation = "manual"` once implemented.
# 记录一个警告,提示用户在未来版本中需要手动指定注意力实现
logger.warning_once(
"BertSdpaSelfAttention is used but `torch.nn.functional.scaled_dot_product_attention` does not support "
"non-absolute `position_embedding_type` or `output_attentions=True` or `head_mask`. Falling back to "
"the manual attention implementation, but specifying the manual implementation will be required from "
"Transformers version v5.0.0 onwards. This warning can be removed using the argument "
'`attn_implementation="eager"` when loading the model.'
)
return super().forward(
hidden_states,
attention_mask,
head_mask,
encoder_hidden_states,
encoder_attention_mask,
past_key_value,
output_attentions,
)
# -------------------- 2. 位置嵌入类型position_embedding_type是 "absolute" -------------------------
# ----------------- 2.1 获取输入的批次大小(bsz)、目标序列长度(tgt_len), 后面会用于shape的调整-----------------
bsz, tgt_len, _ = hidden_states.size()
# ---------------- 2.2 获取query_layer, key_layer, value_layer, attention_mask, is_causal, 用于注意力的计算-------
# 将hidden_states投影到查询向量(query_layer),并使用transpose_for_scores方法调整其形状
query_layer = self.transpose_for_scores(self.query(hidden_states))
# If this is instantiated as a cross-attention module, the keys and values come from an encoder; the attention
# mask needs to be such that the encoder's padding tokens are not attended to.
""" 如果这是一个跨注意力模块实例化的情况,键和值来自编码器;注意力掩码需要确保编码器中的填充标记不被关注。 """
# 判断是否是交叉注意力,并为current_states和attention_mask赋值
is_cross_attention = encoder_hidden_states is not None
current_states = encoder_hidden_states if is_cross_attention else hidden_states
attention_mask = encoder_attention_mask if is_cross_attention else attention_mask
# Check `seq_length` of `past_key_value` == `len(current_states)` to support prefix tuning
# 如果是交叉注意力且有past_key_value,并且它的序列长度与current_states一致,直接使用缓存的键和值
if is_cross_attention and past_key_value and past_key_value[0].shape[2] == current_states.shape[1]:
key_layer, value_layer = past_key_value
else:
# 否则,计算新的键值对,并在非交叉注意力情况下,将它们与过去的键值对拼接
key_layer = self.transpose_for_scores(self.key(current_states))
value_layer = self.transpose_for_scores(self.value(current_states))
if past_key_value is not None and not is_cross_attention:
key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
# 如果是解码器,则缓存当前的键和值,以便在后续步骤中使用
if self.is_decoder:
# if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
# Further calls to cross_attention layer can then reuse all cross-attention
# key/value_states (first "if" case)
# if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
# all previous decoder key/value_states. Further calls to uni-directional self-attention
# can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
# if encoder bi-directional self-attention `past_key_value` is always `None`
"""
如果是交叉注意力,将所有交叉注意力的 key/value 状态保存为一个包含两个 torch.Tensor 的元组。
后续对交叉注意力层的调用可以重用所有交叉注意力的 key/value 状态(即第一个 "if" 情况)。
如果是单向自注意力(解码器),则保存所有先前解码器的 key/value 状态为一个包含两个 torch.Tensor 的元组。
后续对单向自注意力层的调用可以将先前解码器的 key/value 状态与当前投影的 key/value 状态拼接起来(即第三个 "elif" 情况)。
如果是编码器的双向自注意力,则 `past_key_value` 始终为 `None`。
"""
past_key_value = (key_layer, value_layer)
# SDPA with memory-efficient backend is broken in torch==2.1.2 when using non-contiguous inputs and a custom
# attn_mask, so we need to call `.contiguous()` here. This was fixed in torch==2.2.0.
# Reference: https://github.com/pytorch/pytorch/issues/112577
"""
在 torch==2.1.2 中,当使用非连续的输入和自定义的注意力掩码时,带有内存高效后端的 SDPA(缩放点积注意力)是有问题的,
因此我们需要在这里调用 `.contiguous()` 方法来确保输入是连续的。这个问题在 torch==2.2.0 中已被修复。'
参考:https://github.com/pytorch/pytorch/issues/112577
"""
# 如果 PyTorch 版本低于 2.2.0 且设备类型为 CUDA 且 attention_mask 不为空,则确保 qkv 输入是连续的
if self.require_contiguous_qkv and query_layer.device.type == "cuda" and attention_mask is not None:
query_layer = query_layer.contiguous()
key_layer = key_layer.contiguous()
value_layer = value_layer.contiguous()
# The tgt_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal
# mask in case tgt_len == 1.
"""
需要保证 tgt_len > 1,以匹配 AttentionMaskConverter.to_causal_4d 的行为,
因为当 tgt_len == 1 时,它不会创建因果掩码。
"""
# 如果是解码器且没有注意力掩码且目标序列长度大于1,则启用因果注意力
is_causal = self.is_decoder and attention_mask is None and tgt_len > 1
# ----------- 2.3 使用 torch.nn.functional.scaled_dot_product_attention 计算注意力 -------------
# 使用 PyTorch 的 scaled_dot_product_attention 函数计算注意力输出,传入查询、键、值、注意力掩码和dropout概率
attn_output = torch.nn.functional.scaled_dot_product_attention(
query_layer,
key_layer,
value_layer,
attn_mask=attention_mask,
dropout_p=self.dropout_prob if self.training else 0.0,
is_causal=is_causal,
)
# ----------- 2.4 调整 attn_output 的形状, 用于返回计算后的注意力输出, 以及在解码器模式下缓存的键值 -----------------
# 将注意力输出张量转置并调整形状以匹配原始输入的形状
attn_output = attn_output.transpose(1, 2)
attn_output = attn_output.reshape(bsz, tgt_len, self.all_head_size)
# 返回包含注意力输出的元组,如果是解码器,还会返回缓存的键和值
outputs = (attn_output,)
if self.is_decoder:
outputs = outputs + (past_key_value,)
return outputs