每次操作一个向量和矩阵乘法
西风 发自 凹非寺
量子位 | 公众号 QbitAI
让大神 Andrej Karpathy 一键三连❤️（点赞 + 转发 + 评论），一个教你从头开始实现 Llama3 的代码库爆火。
X 上转赞收藏量超 6.8k，GitHub 揽星 2k+。
火就火在，它教你从头用 Meta 开源的权重进行推理，详细解释和展开了注意力机制中多个头的矩阵乘法、位置编码以及所有中间层。
换句话说，他解释了每行代码都在干啥。
Karpathy 看后直呼打造者 Nishant Aklecha（后文暂称 “纳哥”）是个有品的人：
完全展开后，比起模块相互嵌套和调用时，更容易理解每一步具体在做什么。
网友们对其也是赞不绝口，纷纷致敬：
话不多说，一起来看纳哥是如何手把手教的。
（量子位在不改变原意的基础上，进行了编译整理）
在运行纳哥提供的文件前，大伙儿需要预先下载 Meta 官方提供的 Llama3 模型权重。
纳哥表示自己没搞分词器，推荐用 Karpathy 的现成简洁版 BPE 代码。
PS：
“字节级（byte-level）”BPE 算法，在 UTF-8 编码的字符串上运行，广泛应用于大模型分词。Karpathy 提供的这个代码库包含两个分词器，都能在给定文本上训练分词器的词汇表和合并规则、将文本编码为 token、将 token 解码为文本。
读取模型文件的方式通常取决于 model classes 的编写方式以及 class 中变量的命名。但由于纳哥是从头开始实现 Llama3，所以将逐个张量地读取文件内容。
通过此配置可以推断出模型的结构和参数信息，例如模型包含的 Transformer 层数、多头注意力块中的头数，以及词汇表的大小等细节。
将文本转换为 token 时，纳哥使用 tiktoken 作为分词器。
接下来，纳哥展示了在代码中将 token 转换为高维的嵌入表示。这是代码库中唯一使用内置神经网络模块的部分。
[17×1]的 token 矩阵变成了 [17×4096] 的嵌入矩阵。也就是说，每个 token 被转换为一个长度为 4096 的嵌入向量，总共有 17 个这样的嵌入向量。
然后，纳哥对嵌入进行 RMS 归一化。经过这一步后，嵌入的形状不会改变，只有数值被归一化了。纳哥强调需要一个 norm_eps，避免意外将 RMS 值设为 0 导致除以 0 的错误。
以下是公式：
构建 Transformer 的第一层，进行归一化处理，从模型字典中访问 layer.0（即第一层）。归一化之后，张量的形状仍然是 [17×4096]，与嵌入时相同，但数值已被归一化。
跟着纳哥从头实现注意力机制，加载 Transformer 第一层的注意力头。
从模型中加载 query、key、value 和 output 向量时，它们的形状分别是 [4096×4096]、[1024×4096]、[1024×4096] 和 [4096×4096]。
纳哥表示乍一看有点奇怪，因为理想情况是每个注意力头的 q、k、v 和 o 向量是独立的。而代码作者将它们捆绑在一起，是为了方便并行计算注意力头的矩阵乘法。
把所有这些向量解包开来：
下一步，纳哥将从多个注意力头中解包 query，解包后的形状是 [32x128x4096]，32 是 Llama3 中的注意力头数量，128 是 query 向量的大小，4096 是 token 嵌入的大小。
在这里，纳哥访问了第一层第一个注意力头的 query 权重矩阵，query 权重矩阵的大小是 [128×4096]。
将 query 权重矩阵与 token 嵌入相乘，获得每个 token 的 query 向量。结果的形状为 [17×128]，有 17 个 token，每个 token 对应一个长度为 128 的 query 向量。
接下来需要位置编码。
现在已经为 prompt 中的每个 token 生成了 query 向量，但每个单独的 query 向量并不知道它在 prompt 中的具体位置。
例如，query：“the answer to the ultimate question of life, the universe, and everything is”（生命、宇宙和一切的终极问题的答案是）。
在这个 prompt 中，使用了三次”the”，需要根据它们在 prompt 中的位置，使这三个”the”token 的 query 向量有所不同（每个向量的大小为 [1×128]）。
通过使用 RoPE（旋转位置嵌入）来进行这些旋转操作。
上一步中，纳哥将 query 向量分成对，并对每一对应用一个旋转角度偏移。
由此，得到的向量大小为 [17x64x2]，这是将长度为 128 的 query 向量对每个 prompt 中的 token 分成 64 对。这 64 对中的每一对都会根据 m*(theta) 进行旋转，其中 m 是要旋转 query 的 token 的位置。
使用复数的点积来旋转一个向量：
现在每个 token 的 query 元素都有一个复数（角度变化向量），可以将 query 向量（之前分成的对）转换为复数，然后通过点积根据位置旋转 query 向量。
获得旋转后的向量后，可以通过将复数重新视为实数来得到成对的 query 向量。
旋转后的对现在已经合并，有一个新的 query 向量（旋转后的 query 向量），其形状为 [17×128]，其中 17 是 token 的数量，128 是 query 向量的维度。
key 与 query 几乎相同。
纳哥表示自己不会详细讲解 key 的数学原理，只需要记住以下几点：
key 生成的 key 向量维度也是 128；key 的权重只有 query 的四分之一，这是因为 key 的权重在同一时间内被 4 个头共享，来减少计算量；key 也会旋转添加位置信息，原因与 query 相同。
此时，纳哥已经为每个 token 获得了旋转后的 query 和 key。每个 query 和 key 现在的形状都是 [17×128]。
下一步，纳哥将对 query 矩阵和 key 矩阵进行相乘操作。这样做会生成一个评分矩阵，将每个 token 关联起来。这些评分描述了每个 token 的 query 与每个 token 的 key 之间的相关性，这就是自注意力机制。
注意力评分矩阵（qk_per_token）的形状为 [17×17]，其中 17 是 prompt 中的 token 数量。
接下来需要对 query key 评分进行掩码处理。在 Llama3 的训练过程中，未来 token 的 qk 评分是被掩码的，只通过过去的 token 来预测 token。
因此，在推理时，要将未来的 token 评分设置为 0。
接下来是 value，接近注意力机制的最后一步。
这些评分（0-1）用于确定每个 token 使用多少 value 矩阵。
和 key 一样，value 的权重也在每 4 个注意力头之间共享，所以下面 value 权重矩阵的形状是 [8x128x4096]。
第一层，第一个注意力头的 value 权重矩阵如下所示：
然后是 value 向量。
使用 value 权重来获取每个 token 的注意力值，矩阵的大小是 [17×128]，其中 17 是 prompt 中的 token 数量，128 是每个 token 的 value 向量的维度。
注意力：与每个 token 的 value 相乘后得到的注意力向量的形状为 [17×128]。
现在有了第一层第一个头的注意力 value。然后纳哥运行一个循环，对第一层的每个头执行与上面的计算完全相同的数学运算。
然后得到了第一层所有 32 个头的 qkv_attention 矩阵，接下来将所有注意力得分合并成一个大小为 [17×4096] 的大矩阵。
对于第 0 层注意力机制的最后步骤，其一是将注意力得分矩阵与权重矩阵相乘。
这是一个简单的线性层，所以只需进行矩阵乘法。
现在得到了注意力机制后的嵌入 value 变化，应该被添加到原始的 token 嵌入中。
对嵌入增量进行归一化处理，然后通过嵌入增量运行一个前馈神经网络。
在 Llama3 中，加载前馈权重并实现前馈网络。使用了一种名为 SwiGLU 的前馈网络，这种网络结构在模型需要的时候，能够有效地增加非线性。
现在完成了第一层之后每个 token 的新嵌入。现在只剩下 31 层了，只需通过一个循环来完成。
纳哥表示可以将这个编辑后的嵌入想象成包含了第一层中所有查询信息的嵌入。随着层数的增加，每一层都会对输入的信息进行越来越复杂的处理，直到最终得到一个能够全面了解下一个需要预测的 token 的嵌入。
之前做的所有事情，对每一层都重复一次。
然后得到了最终的嵌入，这是模型对下一个 token 的最优预测。这个嵌入的形状与常规的 token 嵌入相同，为 [17×4096]，其中 17 是 token 的数量，4096 是嵌入的维度。
最后，将嵌入解码成 token 值。
使用输出解码器将最终的嵌入转换成一个 token。
接下来看纳哥使用最后一个 token 的嵌入来预测下一个 value，希望预测的结果是 42。
因为根据《银河系漫游指南》一书中的说法，42 是 “生命、宇宙及一切的终极问题的答案”。大多数 LLM 在这里都会回答 42，这将验证整个代码的正确性。
模型预测下一个 token 的编号为 2983。这个编号对应数字 42 吗？
OK，结束。
简单介绍一下 Nishant Aklecha。
Nishant Aklecha 是构建和改进定制语言模型平台 Glaive AI 的研究员，曾任职于摩根士丹利，负责训练和微调大语言模型。
此外，他还和朋友一同创立了一个研究实验室，名为 A10（AAAAAAAAAA）。
他们的目标可以总结成一句话：让研究变得更加触手可及。
除了放出这个代码库，Nishant Aklecha 可谓好人做到底。
网友想更好地理解这个代码库的内容，Nishant 直接一个 YouTube 视频甩了过来：
之前 Nishant Aklecha 还曾写过一篇 Blog，详解了潜在一致性模型（LCM），同样收获了不少好评。
啥也不说了，感兴趣的家人们赶紧码住吧。
GitHub 链接：github.com/naklecha/ll…
参考链接：
[1]x.com/naklecha/st…
[2]naklecha.notion.site/explained-l…
[3]www.youtube.com/watch?v=o29…
[4]www.youtube.com/watch?v=eMl…
版权所有，未经授权不得以任何形式转载及使用，违者必究。
