LlamaDecoderLayer源码解析
- 1. LlamaDecoderLayer 介绍
- 2. 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