FlashAttention
if is_flash_attn_available(): # 检查flashattention的可用性
from flash_attn import flash_attn_func, flash_attn_varlen_func
from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
FlashAttention是Tranformer模型中用于改进注意力机制的技术，主要目的是减少计算复杂度和内存占用。
- flash_attn_func用于标准的flashattention计算。
- flash_attn_varlen_func用于处理变长序列（长度未能确定）的flashattention计算。
- index_first_axis用于处理第一个索引轴。
- pad_input将数据进行填充处理，从而确定长度。
- unpad_input将填充后的输入还原为原始形态。
Logging模块
logger = logging.get_logger(__name__)
_CONFIG_FOR_DOC = "LlamaConfig"
创建了名为logger的日志记录器对象，__name__用于保存模块的名称，确保每个模块都有自己的日志记录器。
_CONFIG_FOR_DOC前面带有下划线，因此可以看出其代表一个模块的内部变量。
get_unpad_data模块
def _get_unpad_data(padding_mask):
seqlens_in_batch = padding_mask.sum(dim=-1, dtype=torch.int32)
indices = torch.nonzero(padding_mask.flatten(), as_tuple=False).flatten()
max_seqlen_in_batch = seqlens_in_batch.max().item()
cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))
return (
indices,
cu_seqlens,
max_seqlen_in_batch,
)
该模块的作用是padding_mask提取非填充的数据，分为以下几步：
- seqlens_in_batch计算量每个张量的有效长度，sum()函数计算每个张量的有效长度。dim等于-1意味着以最按照最后一个维度进行求和，如果是二维，就可以理解为对跨列操作，即计算了每一行非填充元素的个数。
- indices获取了非填充元素的索引。flatten()函数将张量展开成一维，as_tuple为flase意味着返回不是元组形式而是二维矩阵形式，由于返回的是二维矩阵，因此我们需要flatten()再次展平成一维。
——为什么不返回元组呢？
如果返回元组，那么返回的格式是包含一个一维张量的元组，然后还需要从元组中取出这个一维张量，类似：
torch.nonzero(padding_mask.flatten(), as_tuple=True)[0]
这样比较麻烦，不如直接返回二维数组再展平。
3. max_seqlen_in_batch获取了在seqlens_in_batch中的最大值并返回（即长度最长的那一个），然后 item()函数的作用是将一个元素的张量转换为python对应的标量，即一个数。
4. cu_seqlens计算累计长度并进行填充。cumsum()函数用于计算指定维度的累计和，(1,0)意味着只在左边添加一个元素，右边不添加。F.pad()是为张量进行填充的函数。这对于处理变长序列非常有用，因为即获得了每个序列的开始索引，容易确定起始和结束位置。
5. 最终返回的包括：非零元素的索引，左边填充过了的累计长度，最长序列的长度。
从而，达到了提取非填充数据的目的。
_make_causal_mask模块
def _make_causal_mask(
input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
):
"""
Make causal mask used for bi-directional self-attention.
"""
bsz, tgt_len = input_ids_shape
mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
mask_cond = torch.arange(mask.size(-1), device=device)
mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
mask = mask.to(dtype)
if past_key_values_length > 0:
mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
该模块用于生成因果掩码，通常用于双向自注意力机制。具体来说，该模块保证在计算注意力时，只能看到当前时间步之前的信息，而看不到未来的，来保持因果关系。
输入参数
input_ids_shape: torch.Size：输入张量的形状，通常为(batch_size, target_length)。
dtype: torch.dtype:用于生成掩码的张量类型。
device: torch.device：指定设备是GPU还是CPU。
past_key_values_length: int = 0：过去的键值对长度，用于增量计算。
步骤
- 获取批大小(bsz)和目标长度(tgt_len)。
- 创建初始掩码：形状为(tgt_len,tgt_len),所有的dtpye设置为最小，通常为负无穷。torch.full()函数创建一个均为min的矩阵
- 设置掩码条件：mask_cond生成一个mask最后一个维度大小-1长度的序列，并放置在指定设备
- mask.masked_fill:将下三角矩阵设置为0。这里用到了pytorch的广播机制，将一个行为1和列为1的向量扩充进行比较，从而将下三角都变为了0。
- 将掩码转换为指定的数据类型。
- 如果past_key_values_length>0,那么就在最后一个维度拼接上一个（tgt_len, past_key_values_length）的张量，这是为了在处理增量计算时，能够考虑过去的键值对。
其中，zeros()函数创建了（tgt_len,past_key_value_length)的全零矩阵，用cat()在mask前添加了一个全0块。 - 最终将遮罩的维度扩展为四维形状(bsz, 1, tgt_len, tgt_len + past_key_values_length)，并返回。
