【HuggingFace Transformers】LlamaDecoderLayer源码解析

1. LlamaDecoderLayer 介绍

LlamaDecoderLayer 是 LLaMA 模型中的一个关键组件，它结合了自注意力机制、全连接层和残差连接，以及对输入数据的归一化。主要流程为：

• 输入归一化：在进入自注意力模块之前，首先对输入的隐藏状态进行归一化处理
• 自注意力计算：将归一化后的隐藏状态传递给自注意力模块
• 残差连接：将自注意力处理后的隐藏状态与输入状态进行相加
• 后注意力处理：对残差连接后的隐藏状态进行后注意力归一化处理
• 全连接层处理：对隐藏状态进一步处理
• 最终残差连接：将经过全连接层处理后的隐藏状态与处理前的隐藏状态进行相加，再次形成残差连接
• 返回输出：返回最终的隐藏状态

Llama Decoder Layer 的结构如下：

2. LlamaDecoderLayer 类源码解析

源码地址：transformers/src/transformers/models/llama/modeling_llama.py

# -*- coding: utf-8 -*-
# @time: 2024/8/28 14:52
import torchfrom typing import Optional, Tuple
from torch import nn
from transformers import LlamaConfig, Cache
from transformers.models.llama.modeling_llama import LlamaAttention, LlamaFlashAttention2, LlamaSdpaAttention, LlamaMLP, LlamaRMSNorm# 定义一个字典，将不同类型的注意力实现类与字符串名称对应起来
LLAMA_ATTENTION_CLASSES = {"eager": LlamaAttention,  # 默认的注意力实现"flash_attention_2": LlamaFlashAttention2,  # 使用 Flash Attention 2 的实现"sdpa": LlamaSdpaAttention,  # 使用 SDPA 的注意力实现
}class LlamaDecoderLayer(nn.Module):def __init__(self, config: LlamaConfig, layer_idx: int):super().__init__()self.hidden_size = config.hidden_size  # 从配置中获取隐藏层的维度大小# 根据配置中的注意力实现类型，选择相应的注意力类并实例化self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx)# 实例化一个全连接层（MLP），用于后续的非线性变换self.mlp = LlamaMLP(config)# 输入归一化层，用于规范化输入的隐藏状态self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)# 注意力后的归一化层，用于规范化经过注意力计算后的隐藏状态self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)def forward(self,hidden_states: torch.Tensor,  # 输入的隐藏状态attention_mask: Optional[torch.Tensor] = None,  # 注意力掩码position_ids: Optional[torch.LongTensor] = None,  # 位置编码IDpast_key_value: Optional[Cache] = None,  # 缓存的键值对，用于加速推理output_attentions: Optional[bool] = False,  # 是否输出注意力权重use_cache: Optional[bool] = False,  # 是否使用缓存来加速解码cache_position: Optional[torch.LongTensor] = None,  # 缓存中位置的索引position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # 位置嵌入，将在v4.45版本中变为必填项**kwargs,  # 其他额外参数，主要用于特殊情况（如分布式训练）) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:"""Args:hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`attention_mask (`torch.FloatTensor`, *optional*):attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,query_sequence_length, key_sequence_length)` if default attention is used.output_attentions (`bool`, *optional*):Whether or not to return the attentions tensors of all attention layers. See `attentions` underreturned tensors for more detail.use_cache (`bool`, *optional*):If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding(see `past_key_values`).past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection statescache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):Indices depicting the position of the input sequence tokens in the sequenceposition_embeddings (`Tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,with `head_dim` being the embedding dimension of each attention head.kwargs (`dict`, *optional*):Arbitrary kwargs to be ignored, used for FSDP and other methods that injects codeinto the model"""residual = hidden_states  # 保存原始的输入状态，用于后续的残差连接hidden_states = self.input_layernorm(hidden_states)  # 对输入的隐藏状态进行归一化处理# Self Attention: 执行自注意力计算hidden_states, self_attn_weights, present_key_value = self.self_attn(hidden_states=hidden_states,attention_mask=attention_mask,position_ids=position_ids,past_key_value=past_key_value,output_attentions=output_attentions,use_cache=use_cache,cache_position=cache_position,position_embeddings=position_embeddings,**kwargs,  # 传递额外的参数)hidden_states = residual + hidden_states  # 将注意力输出与原始输入相加，形成残差连接# Fully Connectedresidual = hidden_states  # 更新残差连接的值为当前隐藏状态hidden_states = self.post_attention_layernorm(hidden_states)  # 对经过注意力后的隐藏状态进行归一化处理hidden_states = self.mlp(hidden_states)  # 通过全连接层（MLP）进行进一步的处理hidden_states = residual + hidden_states  # 再次形成残差连接outputs = (hidden_states,)  # 输出当前层的隐藏状态if output_attentions:  # 如果需要输出注意力权重，则将其添加到输出中outputs += (self_attn_weights,)if use_cache:  # 如果使用缓存，则将更新后的键值对添加到输出中outputs += (present_key_value,)# 返回输出return outputs