🦙从头实现Llama3

type
status
slug
summary
tags
category
icon
password
Date
从头实现的 llama3
在这个文件中,我从头开始实现了 llama3,每次一个张量和矩阵乘法。
另外,我将直接从 meta 为 llama3 提供的模型文件中加载张量,您需要在运行此文件之前下载权重。这是下载权重的官方链接:https://llama.meta.com/llama-downloads/
notion image
标记器
我不会实现 bpe 标记器(但 andrej karpathy 有一个非常干净的实现)
notion image
读取模型文件
通常,读取它取决于模型类的编写方式以及其中的变量名称。
但由于我们从头开始实现 llama3,因此我们将一次读取一个张量的文件。
notion image
我们使用此配置来推断模型的详细信息,例如
  1. 该模型有 32 个 Transformer 层
  1. 每个多头注意力模块有 32 个头
  1. 词汇量等等
将文本转换为标记
这里我们使用 tiktoken(我认为是一个 openai 库)作为 tokenizer
notion image
将 token 转换为其嵌入
很抱歉,但这是我使用内置神经网络模块的代码库中唯一的一部分
,所以我们的 [17x1] 标记现在是 [17x4096],即 17 个嵌入(每个标记一个),长度为 4096
注意:跟踪形状,它使理解所有内容变得更容易
notion image
然后我们使用 rms 归一化对嵌入进行归一化
请注意,在此步骤之后形状不会改变,这些值只是
需要记住的标准化事物,我们需要一个 norm_eps(来自配置),因为我们不想意外地将 rms 设置为 0 并除以 0,
这里是公式:
notion image
构建 Transformer 的第一层
正常化
无论如何,你会看到我从模型字典中访问 layer.0(这是第一层)
,因此在规范化之后,我们的形状仍然是 [17x4096],与嵌入相同,但经过了规范化
notion image
从头开始实现注意力
让我们加载 Transformer 第一层的注意力头
notion image
> 当我们从模型中加载查询、键、值和输出向量时,我们注意到形状为 [4096x4096]、[1024x4096]、[1024x4096]、[4096x4096]
> 乍一看这很奇怪,因为理想情况下我们希望每个头部的每个 q、k、v 和 o 单独存在
> 代码的作者将它们捆绑在一起,因为它很容易,有助于并行化注意力头部乘法。
> 我要解开一切......
展开查询
在下一节中,我们将从多个注意力头中解开查询,
这里得到的形状是 [32x128x4096],32 是 llama3 中注意力头的数量,128 是查询向量的大小,4096 是 token 嵌入的大小
我要实现第一层的第一个头
这里我访问第一层的查询权重矩阵第一个头,这个查询权重矩阵的大小是[128x4096]
我们现在将查询权重与 token 嵌入相乘,以接收针对 token 的查询
在这里您可以看到结果的形状是 [17x128],这是因为我们有 17 个标记,并且每个标记都有一个 128 长度的查询。
notion image
定位编码
我们现在处于一个阶段,我们的提示中的每个标记都有一个查询向量,但是如果你仔细想想——单独的查询向量根本不知道提示中的位置。
查询:“生命、宇宙和一切的终极问题的答案是”
在我们的提示中,我们使用了“the”三次,我们需要所有 3 个“the”标记的查询向量根据它们在查询中的位置具有不同的查询向量(每个大小为 [1x128])。我们使用 RoPE(旋转位置嵌入)执行这些旋转。
绳索
观看此视频(这是我观看的)以了解数学。 https://www.youtube.com/watch?v=o29P0Kpobz0&t=530s
notion image
在上面的步骤中,我们将查询向量分成几对,然后对每对应用旋转角度移位!
我们现在有一个大小为 [17x64x2] 的向量,这是将提示中每个标记分成 64 对的 128 个长度查询!这 64 对中的每一对都将旋转 m*(theta),其中 m 是我们要旋转查询的标记的位置!
notion image
使用复数的点积来旋转矢量
notion image
notion image
现在我们对于每个 token 的查询元素都有一个复数(角度变化向量)
我们可以将查询(分成两对的查询)转换为复数,然后使用点积根据位置旋转查询,
老实说,这很美妙:)
获得旋转向量后
我们可以通过再次将复数视为实数来获取成对的查询
旋转后的对现在已合并,我们现在有了一个新的查询向量(旋转后的查询向量),其形状为 [17x128],其中 17 是标记的数量,128 是查询向量的尺寸
键(几乎与查询相同)
notion image
我太懒了,所以我不打算对密钥进行数学运算,你唯一需要记住的是:
> 密钥生成的密钥向量也是 128 维的
> 密钥的权重数量只有查询的 1/4,这是因为密钥的权重一次在 4 个头之间共享,以减少所需的计算次数
> 密钥也会旋转以添加位置信息,就像查询一样,原因相同
在此阶段,每个标记都有查询和键的旋转值。
notion image
每个查询和键现在的形状都是[17x128]。
下一步,我们将查询和密钥矩阵相乘
这样做会给我们一个分数,将每个 token 相互映射,
这个分数描述了每个 token 的查询与每个 token 的键的关联程度。这是自我注意力 :)
注意力得分矩阵 (qk_per_token) 的形状是 [17x17],其中 17 是提示中的 token 数量
notion image
我们现在必须屏蔽查询关键分数
在 llama3 的训练过程中,未来的 token qk 分数被掩盖了。
为什么?因为在训练期间我们只学习使用过去的 token 来预测 token。
因此,在推理期间我们将未来的 token 设置为零。
notion image
notion image
notion image
notion image
notion image
价值观(注意力几乎已经结束)
notion image
这些分数(0-1)用于确定每个 token 使用多少值矩阵
> 就像键一样,值权重也在每 4 个注意力头之间共享(以节省计算量)
> 因此,下面的值权重矩阵的形状为 [8x128x4096]
第一层,第一个头值权重矩阵如下
值向量
notion image
我们现在使用值权重来获取每个 token 的注意力值,其大小为 [17x128],其中 17 是提示中的 token 数量,128 是每个 token 的值向量的大小
注意力
notion image
与每个 token 的值相乘后得到的注意力向量形状为 [17*128]
多头注意力机制
notion image
我们现在有了第一层和第一个头部的注意力值,
现在我要运行一个循环,对第一层的每个头部执行与上面的单元格完全相同的数学运算
notion image
我们现在有了第一层所有 32 个头的 qkv_attention 矩阵,接下来我要把所有的注意力分数合并成一个大小为 [17x4096] 的大矩阵,
我们快到了最后 :)
权重矩阵,最后一步
notion image
对于第 0 层注意力,最后要做的事情之一是将权重矩阵乘以
这是一个简单的线性层,所以我们只需要 matmul
notion image
我们现在有了经过注意之后嵌入值的变化,这应该会添加到原始的 token 嵌入中
我们进行归一化,然后通过嵌入增量运行前馈神经网络
notion image
加载 ff 权重并实现前馈网络
notion image
在 llama3 中,他们使用了 SwiGLU 前馈网络,这种网络架构非常适合在模型需要时添加非线性。
如今,在 llms 中使用这种前馈网络架构相当标准
我们最终在第一层之后为每个标记提供了新的编辑嵌入
在我们完成之前只需要再经过 31 层(一个 for 循环),
您可以想象这个编辑过的嵌入具有关于第一层上提出的所有查询的信息,
现在每一层都会对所问的问题进行编码越来越复杂的查询,直到我们有一个知道我们需要的下一个标记的一切的嵌入。
天啊,一切突然间
notion image
是的,就是这样。我们之前所做的一切,都是针对每一层进行的。
祝阅读愉快:)
我们现在有了最终的嵌入,这是模型对下一个标记做出的最佳猜测
嵌入的形状与常规 token 嵌入相同 [17x4096],其中 17 是 token 的数量,4096 是嵌入尺寸
notion image
最后,让我们将嵌入解码为 token 值
notion image
我们将使用输出解码器将最终的嵌入转换为 token
我们使用最后一个标记的嵌入来预测下一个值
希望在我们的例子中是 42 :) 注意:42 是“生命、宇宙和一切的终极问题的答案”的答案,根据《银河系漫游指南》一书,大多数现代法学硕士都会在这里回答 42,这应该可以验证我们的整个代码!祝我好运 :)
该模型预测下一个标记是标记号 2983,这是 42 的标记号吗?
我正在为你加油打气,这是最后的代码单元,希望你玩得开心:)
我们走吧
notion image
谢谢你,我爱你 :)
到此结束。希望您喜欢阅读!
如果你想支持我的工作
  1. 在 Twitter 上关注我https://twitter.com/naklecha
  1. 或者,给我买杯咖啡https://www.buymeacoffee.com/naklecha
老实说,如果你能做到这一步,你已经让我很开心了 :)
我的动力是什么?
我和我的朋友们肩负着一个使命——让研究更容易获得!我们创建了一个名为 A10 的研究实验室 - AAAAAAAAAA.org
我们的论点:
notion image
Hello,算法!FaceFusion 换脸神器