LLMs之llama3-from-scratch:llama3-from-scratch(从头开始利用pytorch来实现并解读LLaMA-3模型的每层代码)的简介、核心思路梳理
导读:这篇论文实现了transformer网络的llama3模型,从头开始利用pytorch来实现该模型。
背景:目前机器学习语言模型内容的复杂性不断增强,但是大多模型都是基于高度抽象和封装的框架来实现,对模型内部工作机制的理解不是很深入。这篇论文采用从零开始逐步实现的方式,帮助读者更好地理解transformer和llama3模型是如何工作的。
具体解决方案:作者加载了llama3预训练模型的参数,并按照每个算子依次实现。
流程和关键步骤:
>> 加载tokenizer对文本进行tokenize
>> 从模型文件加载各层参数
>> 对输入文本的tokens生成嵌入向量
>> 对嵌入向量进行RMSNorm标准化
>> 实现单头注意力机制,包括Q、K、V矩阵的生成、位置编码、掩码、softmax计算注意力权重等
>> 循环实现多头注意力
>> 加载FFN层的参数实现前向传播
>> 循环上述步骤实现整个模型
>> 应用训练好的输出层参数进行预测
关键技术点
>> 利用RoPE位置编码为Q、K向量添加位置信息
>> 多头注意力分头计算并拼接
>> SwiGLU结构的FFN层
>> 应用RMSNorm进行标准化
>> 循环实现每一层的计算
>> 代码运行效果
>> 最后利用加载的llama3模型在样本问题"答案是什么"上正确预测答案为42,验证了从零开始实现的整个流程是正确的。
总之,这篇论文采用详尽清晰的流程,帮助读者通过实际实现深入理解transformer模型底层工作机制,是一篇值得推荐的论文。
GitHub地址:GitHub - naklecha/llama3-from-scratch: llama3 implementation one matrix multiplication at a time
简介
我不会实现一个BPE分词器(但是Andrej Karpathy有一个非常简洁的实现)
链接到他的实现:
该项目给出了一个针对中文分词任务的最小化、清晰而可读的Python实现Byte Pair Encoding(BPE)算法的代码库。它实现了两个Tokenizer对象用于文本到词元和词元到文本的转换:BasicTokenizer和RegexTokenizer。 BasicTokenizer直接在文本上运行BPE算法,RegexTokenizer在BPE之前通过正则表达式将文本切分为不同类型(如字母、数字、标点等)以防止词汇跨类型合并。
思路步骤
从Meta下载权重文件。
使用tiktoken库加载BPE分词器,并定义特殊token。
核心技术点
BPE分词:字节对编码(Byte Pair Encoding)是一种常见的分词方法,用于处理词汇表。
目的是设置和测试一个自定义的tokenizer,它能够处理文本数据
设置路径和特殊标记:
加载BPE(Byte Pair Encoding)模型:
创建tokenizer对象:
解码和编码测试:
简介
通常情况下,读取模型文件取决于模型类的编写方式和其中的变量名称。但由于我们要从头实现llama3,所以我们将逐个张量读取文件。
我们使用此配置来推断模型的细节,模型有32个Transformer层,每个多头注意力块有32个头,n_kv_heads,词汇表大小为128256等等。
下面是各个参数的含义:
思路步骤
加载模型文件,查看模型参数。
读取模型配置文件,获取模型的详细参数。
核心技术点
配置解析:解析配置文件以提取模型超参数。
简介
这里我们使用tiktoken(我认为是OpenAI的库)作为分词器。
思路步骤
输入一个提示文本,并使用分词器将其转换为tokens。
打印tokens及其对应的文本表示。
核心技术点
分词:使用分词器将文本转换为token ID。
文本表示:解码token ID为文本以验证分词过程。
简介
这是代码库中唯一使用内置神经网络模块的部分。
我们的[17x1] tokens现在是[17x4096],即17个长度为4096的嵌入(每个token一个嵌入)。
注意:跟踪形状,这使它更容易理解一切
思路步骤
定义一个嵌入层,并加载预训练的权重。
将tokens转换为嵌入表示。
核心技术点
简介
我们然后使用RMS归一化嵌入,请注意此步骤后形状不会改变,值只是被归一化了。需要注意的是我们需要一个norm_eps(从配置中获取),因为我们不想意外地将rms设为0并除以0。
公式如下:
归一化后我们的形状仍然是[17x4096],与嵌入相同但已归一化。
思路步骤
对嵌入进行RMS归一化。
核心技术点
RMS归一化:对嵌入向量进行均方根归一化(RMS Norm)。
简介
您会看到我从这个模型字典中访问layer.0(这是第一层) 无论如何,所以归一化后,我们的形状仍然是【17x4096】,与嵌入相同但已归一化
思路步骤
构建Transformer的第一层:
核心技术点
多头注意力机制(Multi-Head Attention)与残差连接
Transformer层:应用注意力机制和前馈神经网络进行特征提取和变换。
简介
让我们加载Transformer第一层的注意力头。
>> 加载查询Query、键Key、值Value和输出向量时,我们注意到它们的权重矩阵形状分别为[4096x4096]、[1024x4096]、[1024x4096]、[4096x4096]。
>> 乍一看这有点奇怪,因为我们理想情况下希望每个Query、Key、Value和输出分别对应每个头。而代码作者将它们捆绑在一起,因为这样做,有助于并行化注意力头的矩阵乘法。
接下来,我将展开所有内容...
下一步,我将展开多个注意力头的查询,结果形状为[32x128x4096]。
这里,32是llama3的注意力头数,128是Query向量的大小,4096是token嵌入的大小。
这里我访问第一层第一个头的Query权重矩阵,Query权重矩阵的大小为【128x4096】。
我们现在,将Query权重【128x4096】与token嵌入【17x4096】相乘,以接收token的查询,即获得每个token的query。
结果形状为[17x128],因为我们有17个token,每个token有一个128长度的query向量。
思路步骤
通过查询、键和值向量计算注意力分数,提取输入中重要的信息。
加载第一层的查询、键、值和输出向量权重。
展开注意力头的查询权重矩阵。
计算token的查询向量。
核心技术点
自注意力(Self-Attention)的计算。
痛点:现在,已经为prompt中的每个token生成了query向量,但每个单独的query向量并不知道它在prompt中的具体位置。我们现在处于每个token在我们的提示中有一个查询向量的阶段,但如果你想一想——单个查询向量对提示中的位置没有任何了解。例如:
期望:在我们的提示中我们使用了三次“the”,我们需要所有3个“the”token的查询向量(每个大小为[1x128])根据它们在查询中的位置具有不同的query向量。
解决方案:采用RoPE(旋转位置嵌入)执行这些旋转操作。
简介
在上述步骤中,我们将查询向量分成了若干对,并对每对应用一个旋转角度偏移!
现在我们有一个大小为[17x64x2]的向量,这是将128长度的query向量,对每个prompt中的每个token分成的64对,这64对中的每一对都会按照m*(theta) 进行旋转,其中m是要为其旋转query的token的位置!
使用复数的点积来旋转向量
现在我们为每个token的query元素得到了一个复数(角度变化向量),我们可以将我们的query(即我们之前分成的对)转换为复数,然后使用点积根据位置旋转query向量。
老实说,想到这一点真是太美妙了 :)
在获得旋转后的向量后,我们可以通过将复数,重新视为实数,来得到(或来还原)成对的query向量。
旋转后的对现在被合并,我们现在有了一个新的query向量(旋转后的query向量),其形状为 [17x128],其中17是token的数量,128是query向量的维度。
准备基向量:
基向量定义为一个由二维位置(p)索引的表情形(对于大小为128的嵌入,生成的形状为[64])。此基向量本质上是一个在0到θ之间等距离的线性间隔。
为提示生成旋转向量
为了生成旋转向量,我们将基向量乘以0到16之间的整数,这些整数是我们提示中token的位置。此乘法后的结果形状为[17x64],其中每行表示token的旋转向量。
应用RoPE到查询向量
最后,我们可以将我们的查询向量与旋转向量结合,生成位置编码的查询向量。每个查询向量现在被旋转以反映它们在提示中的位置。
最终位置编码的查询向量
我们现在有一个大小为[17x128]的查询向量,每个token都有它在提示中位置的旋转编码。
思路步骤
生成RoPE(旋转位置嵌入)以编码token的位置信息。
将查询向量与位置编码向量结合,生成位置编码的查询向量。
核心技术点
旋转位置编码(RoPE):使用旋转位置嵌入(RoPE)将位置信息编码到查询向量中。
向量旋转:使用cos和sin函数对查询向量进行旋转,结合位置编码信息。
简介
查询Query、键Key、值Value和输出向量
思路步骤
Key键向量生成
>> key向量维度为128。
>> key的权重共享机制,key的权重只有query的四分之一,,这是因为每4个头共享一次,以减少计算量。
>> key向量也进行旋转以加入位置信息,原因与query相同。
Query向量生成及旋转,增加位置信息。
查询向量与键向量的维度匹配。
旋转操作用于引入位置编码。
在这个阶段,我们现在有了每个token旋转后的query和key,其中每一个query和key现在的形状是[17x128]。
核心技术点
矩阵乘法与旋转位置编码
简介
查询Query、键Key、值Value和输出向量
1、计算注意力分数矩阵:在下一步中,我们将Query和Key矩阵相乘,这样做会给我们一个分数矩阵,将每个token与另一个token映射或关联起来。这个分数矩阵描述了每个token的query与每个token的key之间的关系有多好(相关性),这就是自注意力机制(Self Attention)。其中,分数矩阵(qk_per_token)的形状是[17x17],其中17是prompt中的token数量。
2、对注意力分数矩阵中的未来token进行掩码处理:然后,我们现在必须对Query-Key分数进行掩码处理,在llama3的训练过程中,未来token的qk分数会被屏蔽。 为什么?因为在训练过程中,我们只学习使用过去的token来预测未来token。 因此,在推理时,我们将未来token的评分设置为0。
思路步骤
计算Query向量与Key向量之间的匹配得分,即自注意力得分。
使用矩阵乘法和归一化。
自注意力的核心机制:查询-键匹配得分。
训练过程中掩盖未来的token,确保只使用过去的token进行预测。
上三角掩码矩阵,屏蔽未来token。
核心技术点
分数计算与未来词元掩码
简介
查询Query、键Key、值Value和输出向量
接下来是value,接近注意力机制的最后一步。
思路步骤
这些分数(0-1)用于确定每个token使用多少value矩阵。就和key一样,value权重也在每4个注意力头之间共享(以节省计算)。因此,下面的Value权重矩阵的形状是[8x128x4096] 。
第一层,第一个注意力头的Value权重矩阵如下所示
Softmax 归一化得到注意力权重:对掩码后的得分进行Softmax归一化,得到注意力权重。
使用Softmax函数将得分转换为概率分布。
Values(值向量)计算:使用Value向量和注意力权重计算最终的注意力输出。Value向量与Query和Key向量的共享机制。
注意力计算:现在使用Value权重来获取每个token的注意力值,矩阵的大小是[17x128],其中17是prompt中的token数量,128是每个token的value向量的维度。
注意力:与每个token的value相乘后得到的注意力向量的形状为[17x128]。
核心技术点
分数与值的权重合并
值向量的矩阵乘法
简介
我们现在有了第一层中第一个头的注意力值,然后,将运行一个循环,并对第一层中的每一个头执行与上述单元格完全相同的数学运算。
然后,得到了第一层所有32个头的qkv_attention矩阵,接下来我将把所有的注意力分数合并成一个大的矩阵,其大小为[17x4096]
我们快要结束了!
思路步骤
多头注意力机制:对每个注意力头分别计算注意力输出,并将其拼接。
多头注意力的并行计算和拼接。
线性层变换:通过线性变换将拼接后的多头注意力输出进行进一步处理。
线性层的矩阵乘法。
核心技术点
向量的拼接
简介
查询Query、键Key、值Value和输出向量
思路步骤
对于第0层注意力机制的最后步骤,其一是将注意力得分矩阵与权重矩阵相乘。
这是一个简单的线性层,所以我们只需进行矩阵乘法
核心技术点
线性层(全连接层)的应用。
简介
查询Query、键Key、值Value和输出向量
现在,我们有了注意力后嵌入Value的变化,应该加到原始的token嵌入中;
然后,对嵌入增量进行归一化,然后通过嵌入增量运行一个前馈神经网络;
思路步骤
核心技术点
残差连接
简介
在Llama3中,加载前馈权重并实现前馈网络。使用了一种名为SwiGLU的前馈网络,这种网络结构在模型需要的时候,能够有效地增加非线性。 在当今的LLMs中,使用这种前馈网络架构已经相当标准。
思路步骤
归一化与前馈网络:对注意力输出进行归一化,并通过前馈神经网络进行非线性变换。
RMSNorm归一化。
SwiGLU前馈网络,用于增强非线性表达能力。
核心技术点
SwiGLU激活函数与线性层
简介
查询Query、键Key、值Value和输出向量
我们终于在第1层之后为每个词元获得了新的编辑后的嵌入。我们还有31层才能完成(还需要一个循环)。
你可以将这个编辑后的嵌入,想象成包含了第一层中所有Query信息的嵌入。随着层数的增加,每一层都会对输入的信息进行越来越复杂的处理,直到最终得到一个,能够全面了解下一个需要预测的token的嵌入。
天哪,一切都集中在一起了
是的,就是这样。我们之前所做的所有事情,现在一次性为每个层都做了
思路步骤
重复以上步骤,处理所有Transformer子层
核心技术点
循环与层堆叠
循环迭代处理每一层。
简介
思路步骤
核心技术点
矩阵乘法,线性变换。
最终归一化处理
简介
生成最终的嵌入向量,用于预测下一个token的最优预测。这个嵌入的形状与常规的token嵌入相同,为[17x4096],其中17是token的数量,4096是嵌入的维度。
思路步骤
确认logits张量的形状:确保计算得到的logits具有正确的维度,以便后续处理。
核心技术点
线性层与softmax激活函数
简介
我们将使用输出解码器,将最终嵌入转换为token值,使用输出解码器将最终的嵌入转换成一个token。确认输出解码器的尺寸是否正确,确保矩阵乘法可以顺利进行。
使用最后一个词元的嵌入通过矩阵乘法预测下一个令牌的概率分布。
通过最后一个令牌的嵌入和输出解码器权重矩阵的乘法,计算出每个可能令牌的logits。
思路步骤
获取预测的下一个令牌的索引,通过argmax函数从logits中选择概率最高的令牌,来确定下一个令牌,得到最终的预测结果。
核心技术点
分词器的逆过程
argmax函数,分类问题中的预测。
期望:使用最后一个token的嵌入来预测下一个value,希望预测的结果是42。加载的llama3模型在样本问题"答案是什么"上正确预测答案为42,验证了从零开始实现的整个流程是正确的。
希望在我们这个例子中,是42 :)
方法:查看预测的下一个令牌是否为“42”
注意:因为根据《银河系漫游指南》一书中的说法,42是“生命、宇宙及一切的终极问题的答案”。大多数现代LLM在这里都会回答42,这将验证整个代码的正确性。
简介
模型预测下一个词元是token编号2983,这是42的token ID吗?