Llama模型源码调试笔记
Llama基础
Llama是什么？
Llama2是Meta开源发布的大型语言模型。其训练涵盖了庞大的数据集——规模达到2万亿token。Llama1可以处理2048个token，Llama2可以处理4096个token的文本。Llama1和Llama2有同样的模型结构。该系列包含7B、13B及70B规模的不同模型，均在多种基准测试集上展现了卓越的性能。Llama2可用于学术和商业用途。
Llama模型的源码位置
如果使用transformers库在调用llama模型，llama的源码位于：
transformers/src/transformers/models/llama/modeling_llama.py
本文内容
本文包括两部分：调试LlamaModel源代码和调试LlamaForCausalLM源代码。
在Hugging Face的Transformers库中，模型类的命名通常遵循一个标准模式：
AutoModel: 这个前缀指的是一个自动化的类，能够根据预训练模型的名称自动识别并加载与之对应的模型。
ForCausalLM: 这个后缀指定了模型的用途，即用于因果语言建模（causal language modeling），意味着模型被训练用于生成文本，每个时间点上的输出仅依赖于之前的输出和输入，不依赖于后续的输出。
对于"LlamaModel"和"LlamaForCausalLM"：
"LlamaModel"是一个泛指任何基于llama架构的模型的类名，这个名称不指定说明模型的具体应用或行为（例如，它可能用于分类、生成、回归等多种任务）。
“LlamaForCausalLM"则更明确地暗示了模型的用途，即用于因果语言建模。这意呩着这个特定的类被训练用来继承自"LlamaModel”来生成文本，通常是按照单向的、因果的顺序。
调用LlamaModel类
用于debug的代码
想要调用llama模型，可以先从huggingface或者modelscope把llama模型下载下来，然后通过model = AutoModel.from_pretrained(
from transformers.models.llama import LlamaModel,LlamaConfig
import torch
llamaconfig = LlamaConfig(vocab_size=32000,
hidden_size=4096//2,
intermediate_size=11008//2,
num_hidden_layers=32//2,
num_attention_heads=32//2,
max_position_embeddings=4096//2)
llamamodel = LlamaModel(config=llamaconfig)
input_ids = torch.randint(low=0,high=llamaconfig.vocab_size, size=(4,30))
res = llamamodel(input_ids)
llama-7B模型的参数量是6,607,343,616，一些信息是：hidden_size=4096,intermediate_size=11008,num_hidden_layers=32
由于目的仅仅是debug一遍llama模型的流程，因此模型的大小可以设置得小一些。
模型输入
input_ids是模型的输入，shape是torch.Size([4, 30])，代表着一个批次内有4条数据，每条数据的length是30.
LlamaModel类的forward的流程
接下来具体分析这个类里的forward方法。
模型的输入是input_ids（torch.Size([4, 30])）
第（1）步，经过embedding层
inputs_embeds = self.embed_tokens(input_ids)
embed_tokens是nn.Embedding层。经过Embedding后，输出为inputs_embeds，shape=torch.Size([4, 30, 2048])
第（2）步，生成attention_mask矩阵
在注意力机制中构造的attention_mask矩阵发挥着关键作用。在本例中，我们得到的attention_mask具有形状torch.Size([4, 1, 30, 30])，其中包含了批次（batch）维度。若撇除批次维度，attention_mask本质上应呈现为方阵，这是因为它描绘了每个token对于其他token的注意力情况。自回归模型中，生成文本的过程遵循先后顺序，即每个新词的产生依赖于前置的token。因此，任何先行的token都不会对尚未出现的后续token有注意力。下面展示了attention_mask的左上角小段，即方阵中前3行与前3列：
print(attention_mask[0,0,:3,:3])
tensor([[ 0.0000e+00, -3.4028e+38, -3.4028e+38],
[ 0.0000e+00, 0.0000e+00, -3.4028e+38],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00]])
在上述输出中，-3.4028e+38代表了float32数据类型中可能的最小数值，通常被视作负无穷大。在注意力计算中，这个值用于有效地屏蔽（或阻断）不相关的token，故而注意力模型只会专注于合理的、先行出现的token序列。
第（3）步，经过16层decoder layer
第（1）步的输出的inputs_embeds，其shape=torch.Size([4, 30, 2048])，依次经过16个解码器层。每个解码器层的输出有两个，一个是hidden_state，shape=torch.Size([4, 30, 2048])。另一个是kvcache。
第（4）步，经过RMSNorm层。
其实在每一层decoder layer的内部，都有RMSNorm。在经过了所有16层decoder layer后，再次应用一个RMSNorm层作为额外的规范化步骤，确保在进行最终的输出之前，模型的状态被有效地规范化，为输出流程提供了均衡的特征表示。
RMSNorm的源码如下：
class LlamaRMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
"""
LlamaRMSNorm is equivalent to T5LayerNorm
"""
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states):
input_dtype = hidden_states.dtype
hidden_states = hidden_states.to(torch.float32)
variance = hidden_states.pow(2).mean(-1, keepdim=True)
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
return self.weight * hidden_states.to(input_dtype)
举个例子，当前hidden_states.shape=4302048。在代码中，variance.shape=1301，表示的是每个token所有隐藏维度的“平方平均数”。然后hidden_states中的每个数除以“平方平均数”的平方根。经过这一操作后，hidden_states的形状仍然保持为4 x 30 x 2048。经过了这一标准化处理，hidden_states的维度仍然保持不变，但其数值分布已经得到了调整和对齐，可以用来提高模型训练的稳定性和性能。
注意，self.weight是一个可学习的参数，初始值为1
第（5）步，模型的输出
模型的输出有两个：odict_keys(['last_hidden_state', 'past_key_values'])
其中last_hidden_state.shape=torch.Size([4, 30, 2048])
past_key_values.shape可以理解为tupple[][]格式，打印其形状如下：
print(len(res[1]), len(res[1][0]), res[1][0][0].shape)
16 2 torch.Size([4, 16, 30, 128])
这里的16指的是有16层，2指的是key和value的cache，4指的是batch_size，16指的是注意力的头数，30指的是seq_length，128指的是每个头的hidden_size。
总结
LlamaModel类接收文本输入（经过编码的input_ids），其在本例中的形状为torch.Size([4, 30])。文本数据首先通过嵌入层（embedding layer），然后进入由16个Transformer层构成的核心处理单元，在这一系列复杂的层内部进行表示学习。模型的输出主要包括last_hidden_state和past_key_values，其中last_hidden_state的形状为torch.Size([4, 30, 2048])，包含了编码了输入文本语义信息的高维特征表征。
作为基础模型，LlamaModel为众多下游任务提供了初步处理能力。若要在此基础上执行特定的任务，比如因果关系的语言建模，就需要借助配套设计的LlamaForCausalLM这一专用类别。LlamaForCausalLM不仅继承了LlamaModel的功能，还引入了特有的结构来支持文本生成任务，使得模型能以连贯的方式预测文本序列。
调用LlamaForCausalLM类
用于debug的代码
from transformers.models.llama import LlamaModel,LlamaConfig,LlamaForCausalLM
import torch
llamaconfig = LlamaConfig(vocab_size=32000,
hidden_size=4096//2,
intermediate_size=11008//2,
num_hidden_layers=32//2,
num_attention_heads=32//2,
max_position_embeddings=4096//2)
llamamodel = LlamaForCausalLM(config=llamaconfig)
input_ids = torch.randint(low=0,high=llamaconfig.vocab_size, size=(4,30))
res = llamamodel(input_ids)
模型输入
input_ids是模型的输入，shape是torch.Size([4, 30])，代表着一个批次内有4条数据，每条数据的length是30.
LlamaForCausalLM类的forward的流程
LlamaForCausalLM类的forward方法负责控制数据在因果语言模型中的流动和处理。下面是该方法的具体流程分析：
模型以input_ids作为输入，其形状为torch.Size([4, 30])，反映了一批大小为4、序列长度为30的文本数据。
第（1）步，编码阶段。
输入被传递给LlamaModel，这个模型进行了一连串复杂的编码操作，并输出了两个组件：last_hidden_state和past_key_values，其中last_hidden_state的shape=torch.Size([4, 30, 2048])
第（2）步，head层处理。
last_hidden_state会经过一个head层，具体而言是线性层（linear layer），实现从隐藏状态空间到词汇空间的映射。这一层将维度从2048个特征（对应每个隐藏状态）扩展到32000个特征（对应词汇表中的潜在token）。因此，映射后的输出形状为torch.Size([4, 30, 32000])。
第（3）步，计算loss
如果我将label输入到网络中，那么在这里会计算loss。此时debug的代码改成：
res = llamamodel(input_ids=input_ids, labels=input_ids)
以下是计算loss的源码：
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous() #torch.Size([4, 29, 32000])
shift_labels = labels[..., 1:].contiguous() #torch.Size([4, 29])
# Flatten the tokens
loss_fct = CrossEntropyLoss()
shift_logits = shift_logits.view(-1, self.config.vocab_size) #torch.Size([116, 32000])
shift_labels = shift_labels.view(-1) ##torch.Size([116, 32000])
# Enable model parallelism
shift_labels = shift_labels.to(shift_logits.device)
loss = loss_fct(shift_logits, shift_labels)
在因果语言模型中，标签（labels）通常是输入序列（inputs）的一个副本。举例来说，假设给定一个包含以下token的序列：[START, 词1, 词2, 词3, ..., 词N]，其长度为 N+1，这作为标签。模型生成的预测输出将形如：[词1的预测, 词2的预测, ..., 词N-1的预测, 词N的预测, END]，长度同为 N+1。
在计算损失的过程中，logits[..., :-1, :]操作将logits切片，得到从第一个到倒数第二个token的预测输出，即为：[词1的预测, 词2的预测, ..., 词N的预测]；同时，labels[..., 1:]操作将从第二个token开始的真实标签切片，得到 [词1, 词2, ..., 词N]，从而实现了预测和真实标签之间的对齐。经过交叉熵损失函数（CrossEntropyLoss）的计算，这两个序列将被训练得尽可能相似。
第（4）步，模型的输出
第四步涉及到模型的输出，它通常由三个主要部分组成，分别是loss、logits以及past_key_values，可以通过输出字典的键（keys）访问。
loss：这是一个标量值，代表了批次中所有样本的平均损失。在本例中，由于使用的CrossEntropyLoss的reduction参数被设置为'mean'，该损失值是对整个批量的损失计算出的平均值。此参数设置确保了对不同批量大小的统一处理，允许模型的损失不受批量大小的影响。
logits：这个张量的形状为torch.Size([4, 30, 32000])，其包含了模型预测结果的对数概率分布。对于给定的批次中的每个输入序列的每个位置，logits都提供了一个词汇表大小的向量，表示模型预测每个可能下一个token的概率。
past_key_values：这组张量与LlamaModel输出的结构一致，包括了所有解码器层（decoder layers）的键（keys）和值（values）。这些用于在多步解码过程中维持和利用前面的注意力状态，有助于生成过程中的性能和效率。
以上组成的模型输出为后续步骤提供了全面的信息，无论是评估模型的性能（通过损失），还是继续进行文本生成和处理（通过logits和past_key_values）。
model.generate方法
上面介绍了forward流程，以下是LlamaForCausalLM类中generate方法的详细介绍，以及该方法如何用于文本生成。
调用generate方法的代码是：
res = llamamodel.generate(input_ids=input_ids)
在本例中，generate方法的实现依靠transformers库内的greedy_search函数，该函数的源代码位于transformers/src/transformers/generation/utils.py文件中。（还有更复杂的生成方法，本文暂时不做介绍）
第（1）步，调用forward方法
model通过forward方法处理模型的输入，从而产生logits。这些logits的形状为torch.Size([4, 30, 32000])，每个数字都代表在给定的位置产生特定词汇的对数概率。
第（2）步，预测下一个词
取出最后一个（最新）token的向量：
next_token_logits = logits[:, -1, :]
这里的next_token_logits将具有形状torch.Size([4, 32000])，其中包含批量中每个序列最后一个token位置的预测分数。然后找出分数最大的词：
next_tokens = torch.argmax(next_token_logits , dim=-1)
结果next_tokens的形状为torch.Size([4])，包含了这一步骤中选取的token索引。这作为预测出来的下一个词。
第（4）步，将新预测出来的token连接到原始的input_ids上，构成新的输入，然后准备进行下一轮生成。
input_ids = torch.cat((input_ids, new_tokens.unsqueeze(-1)), dim=-1)
初始的input_ids形状为4 x 30，表示有30个词的序列。预测新的token并拼接后，更新的input_ids形状变成4 x 31。
第（5）步，生成循环
重复上述步骤（1）至（4），直到模型生成了结束符（eos token）或达到了预设的最大长度限制。
generate方法通过上述迭代流程为文本生成提供了一个简单而直观的机制，使得模型可以连贯地扩展文本输出，直到完成整个生成任务。这种贪婪的策略总是选择概率最高的下一个token，导致生成的是高置信度的文本。
总结
LlamaForCausalLM类作为模型的主体，接收经过编码的文本input_ids作为其输入，在本例中input_ids的维度为torch.Size([4, 30])。处理的第一阶段涉及通过LlamaModel类的forward方法，其输出是last_hidden_state，具有维度torch.Size([4, 30, 2048])。接下来，经过（head layer）的处理，其作用是将隐藏层状态从隐空间投影到词汇空间，从而产生维度为torch.Size([4, 30, 32000])的logits。
在此步骤完成后，有两种方法可用于词汇预测。一种是手动检索logits中32000个概率分量，找到最大值对应的分量；对应的词汇表索引token即为模型预测的下一个词。另一种更自动化的方式是使用模型内置的model.generate()方法，它能便捷地直接预测出接下来的词。这个方法封装了预测的复杂细节，为快捷生成文本提供了极大的便利。
