llama3 从头开始实现

原文参考连接：https://github.com/naklecha/llama3-from-scratch

在本篇文章中，从头开始实现了 llama3，一次一个张量和矩阵乘法。

此外，将直接从 Meta 为 llama3 提供的模型文件加载张量，您需要在运行此文件之前下载权重。以下是下载llama3权重的官方链接： https://llama.meta.com/llama-downloads/

实验是在 Windows 11 环境进行的，使用 jupyter notebook 进行调试

1. 创建 jupyter 环境

1.1 确认 python.exe 路径

1.2 确认 pip.exe 路径

1.3 确认 jupyter.exe 路径

1.4 运行 notebook

2. 安装必要的库

3. 分词器(tokenizer)

这里不打算实现 BPE 分词器（但 Andrej Karpathy 有一个非常干净的实现）

https://github.com/karpathy/minbpe

可以从 llama3 的 tokenizer.json 文件知道分词器类型是 BPE

在文件 tokenizer.json 的前面还有其他特殊的 tokens

下面的代码含有了 tokenizer.json 文件里面所有特定的字符

运行结果：

可以看到 tokenizer.n_vocab 的值为 128256，和 config.json 文件里的 vocab_size 一致

4. 读取模型文件

从头开始实现 llama3，我们将一次读取一个张量文件。

运行结果：

这些键名代表了模型中不同部件的参数。以下是对每个键的解释：

1. tok_embeddings.weight – 代表词嵌入层的权重，它负责将输入的单词（或标记）转换为固定大小的向量。
2. layers.x.attention.wq.weight, wk.weight, wv.weight, wo.weight – 这些分别表示在Transformer的第x层中，多头注意力（multi-head attention）机制的查询（Query）、键（Key）、值（Value）和输出（Output）的权重。
3. layers.x.feed_forward.w1.weight, w2.weight, w3.weight – 这些代表在第x层的前馈神经网络（feed-forward network, FFN）中的不同权重。通常在BERT或Transformer模型中，FFN包含两个线性变换，这里w1和w2可能代表这两个变换的权重，w3可能是扩展实现的一部分。
4. layers.x.attention_norm.weight, layers.x.ffn_norm.weight – 这些是应用在多头注意力和前馈神经网络输出上的层归一化（layer normalization）的权重。
5. norm.weight – 这可能是模型最后的层归一化权重，用于整个模型输出的最终归一化处理。
6. output.weight – 这是模型输出层的权重。

以上每个参数都在模型内部承担着重要的角色，确保信息能够在模型层间有效传递并进行适当的变换。

下面是查看其他模型参数：

运行结果：

这些参数是用于配置一个基于Transformer架构的神经网络模型的各种设置。每个参数的具体作用如下：

1. dim: 表示模型中每层的维度或隐藏层的大小。在这里，它被设置为4096，意味着每个Transformer层的内部向量将有4096个维度。
2. n_layers: 模型中的层数。这个参数值为32，意味着这个模型有32层Transformer层。
3. n_heads: 多头注意力机制中头的数量。在这里，它是32，表示每个注意力层会有32个不同的头并行处理信息。
4. n_kv_heads: 特定于某些模型的参数，可能表示在处理键（K）和值（V）时使用的注意力头的数量。在这个模型中，这个数值为8。
5. vocab_size: 词汇表的大小，即模型可以识别的不同单词或标记的总数。这里设置为128,256。
6. multiple_of: 这通常用于确保模型维度的一些方面（如层大小）是某个数值（1024）的倍数，这有助于优化计算效率。
7. ffn_dim_multiplier: 前馈网络维度乘数。这里是1.3，意味着前馈层的内部维度是输入/输出维度（dim）的1.3倍，即计算为 (4096 \times 1.3).
8. norm_eps: 层归一化操作中使用的一个极小数（epsilon），用于避免除以零的错误。这里设置为 (10^{-5})。
9. rope_theta: 特定于某些Transformer模型的参数，可能与旋转位置编码（RoPE）有关。在这里，它设置为500000.0，这可能影响模型处理位置信息的方式。

这些参数共同定义了一个复杂的深度学习模型的结构和行为，每个都对模型的性能和能力产生重要影响。

保存这些参数

这些行从config的字典中提取特定的模型配置参数。这包括模型的维度、层数、头的数量等，以及将rope_theta参数转换为一个PyTorch张量。

5. 将文本转换为标记

在这里，我们使用 tiktoken（我认为是 OpenAI 库）作为分词器

运行结果：

这里 128000，代表 bos_token_id，bos_token，”<|begin_of_text|>”

17 表示tokens 的长度。

这里定义了一个文本提示，并使用tokenizer.encode方法将其转换为标记。标记128000被添加到编码的标记列表的开头，可能代表某种特殊或起始标记。接着打印这些标记及其数量。

我们再从tokens 还原出来原来的内容:

运行结果：

将标记列表转换为PyTorch张量，然后迭代每个标记，将其解码回文本形式并打印。这是为了验证标记是否正确解码回其原始文本。

6. 将令牌转换为其嵌入

对不起，这是使用内置神经网络模块的代码库中唯一部分

无论如何，我们的 [17×1] 令牌现在是 [17×4096]，即 17 个长度为 4096 的嵌入（每个令牌一个）

运行结果：

这部分代码创建了一个词嵌入层，其大小为vocab_size x dim。然后将预训练模型的词嵌入权重复制到这个新层的权重中。使用这个词嵌入层将输入的标记转换为嵌入向量，并将数据类型转换为bfloat16（一种用于深度学习中节省内存的浮点格式）。最后，打印嵌入向量的形状(shapes)和内容。

7. 使用 RMS 规范化对嵌入进行规范化

请注意，在此步骤之后，形状(shapes)不会改变，值只是归一化

要记住的事情，我们需要一个norm_eps（来自 config），因为我们不想意外地将 rms 设置为 0 并除以 0

公式如下：

这个Python函数rms_norm是一个应用根均方（RMS）标准化技术来标准化张量的函数，并使用权重进行缩放。

函数rms_norm接受两个参数：

• tensor：待标准化的张量。
• norm_weights：用于缩放标准化张量的权重张量。

标准化计算:

• tensor.pow(2): 将张量中的每个元素平方。
• .mean(-1, keepdim=True): 计算这些平方值沿最后一个维度的平均值，并保持维度以便进行广播。
• + norm_eps: 在平均值上加上一个小数（norm_eps），防止除零，这对于数值稳定性是必需的。
• torch.rsqrt(...): 计算这个结果的平方根的倒数。
• tensor * torch.rsqrt(...): 将原始张量乘以这个值进行标准化。
• ... * norm_weights: 最后，用norm_weights缩放标准化后的张量。

整个操作确保了张量基于其最后一个维度上的RMS值被缩小，然后通过norm_weights进行调整，这可以用来实现一种参数化的标准化形式。

8. 构建变压器的第一层 

8.1  规范化

您将看到我从模型字典访问 layer.0（这是第一层）

无论如何，在归一化后，我们的形状仍然 [17×4096] 与嵌入相同，但归一化

运行结果：

这行代码调用了之前定义的rms_norm函数，用于对token_embeddings_unnormalized（未标准化的token嵌入）进行标准化处理。标准化使用的权重是从模型中提取的，特定于第0层的注意力层的归一化权重（model["layers.0.attention_norm.weight"]）。

这个函数将返回一个新的张量token_embeddings，其维度应与输入张量token_embeddings_unnormalized保持一致。因此，如果您问及token_embeddings.shape的值，它应该与token_embeddings_unnormalized的形状相同。具体形状取决于输入张量token_embeddings_unnormalized的维度，通常在处理NLP任务时，这个张量的形状为(batch_size, sequence_length, embedding_dim)，其中embedding_dim是在前面定义的dim。

8.2 注意头从头开始实现

让我们加载变压器第一层的注意头

当我们从模型中加载查询（query）、键（key）、值（value）和输出（output）向量时，我们注意到它们的形状分别是 [4096×4096]、[1024×4096]、[1024×4096] 和 [4096×4096]。乍一看这很奇怪，因为理想情况下我们希望每个头分别有自己的q、k、v和o。
代码的作者将它们捆绑在一起，因为这样做很方便，有助于并行化注意力头的乘法运算。
我将解开所有这些…

运行结果：

8.3 解包查询

在下一节中，我们将解包来自多个注意力头的查询，生成的形状为 [32x128x4096]

这里，32 是 llama3 中的注意力头数，128 是查询向量的大小，4096 是令牌嵌入的大小

运行结果：

9. 实现第一层的第一个头

这里我访问查询权重矩阵第一层的第一头，这个查询权重矩阵的大小是[128×4096]

运行结果：

10. 将查询权重与令牌嵌入相乘，以接收对令牌的查询

在这里，您可以看到生成的形状是 [17×128]，这是因为我们有 17 个标记，每个标记都有一个 128 长度的查询。

运行结果：

11. 定位编码

我们现在正处于一个阶段，我们在提示中为每个标记都有一个查询向量，但如果你仔细想想，个性化查询向量不知道提示中的位置。

查询：“生命、宇宙和万物的终极问题的答案是”

在我们的提示中，我们已经使用了 “the” 三次，我们需要所有 3 个 “the” 标记的查询向量根据它们在查询中的位置具有不同的查询向量（每个大小为 [1×128]）。我们使用 RoPE（旋转位置嵌入）执行这些旋转。

RoPE

观看此视频（这是我观看的）以了解数学。 https://www.youtube.com/watch?v=o29P0Kpobz0&t=530s

运行结果：

在上面的步骤中，我们将查询向量拆分为对，我们对每对应用旋转角度偏移！

我们现在有一个大小为 [17x64x2] 的向量，这是提示中每个标记的 128 个长度查询，分为 64 对！这 64 对中的每一对都将由 m*（theta） 旋转，其中 m 是我们旋转查询的令牌的位置！

关于 q_per_token_split_into_pairs.shape 的中文解释：

这段代码首先将q_per_token张量的数据类型转换为浮点型，然后调用view方法来改变张量的形状。具体的变形操作是保持第一个维度（通常是批处理大小）不变，第三个维度设置为 2，这意味着每两个元素被分为一组，而第二个维度使用 -1 自动计算，以确保元素总数保持不变。这样操作通常用于准备数据以适配某些特定的处理需求，例如在某些神经网络操作中需要成对处理数据。

最终得到的 q_per_token_split_into_pairs 的形状会是一个三维张量，其中：

• 第一维度是原张量的第一维度（如批次大小）。
• 第二维度是自动计算出的，确保所有元素都被恰当地分配到大小为2的小组中。
• 第三维度固定为2，表示分组的大小。

这样的变形操作对于进一步处理数据，如在处理双向关系或成对特征时，非常有用。

12. 使用复数的点积旋转向量

 

这段代码创建了一个从0到1均匀分成64部分的张量。具体解释如下：

1. 创建序列： range(64)生成了一个从0到63的整数序列。
2. 转换为张量： torch.tensor(range(64))将这个序列转换成了一个PyTorch张量。
3. 除以64： .../64操作将张量中的每个元素除以64，因此元素的值从0开始（0/64），以1/64为步长，最终到达63/64，不包括1（因为是从0开始的64个数，最大数为63）。

所得到的张量zero_to_one_split_into_64_parts包含64个浮点数，这些数值均匀地分布在0到1之间（不包括1）。这样的数据常用于生成需要等间距数值的场合，例如在某些图形渲染、数据标准化或机器学习模型输入预处理中。

在这段代码中：

• rope_theta 是一个预定义的标量，用于调整生成频率的尺度。
• zero_to_one_split_into_64_parts 是一个从0增加到接近1的64个值的数组，每个数值都等间隔分布。
• rope_theta ** zero_to_one_split_into_64_parts 计算了rope_theta的各个幂次，幂的底数是rope_theta，指数是zero_to_one_split_into_64_parts中的值。这样可以生成一系列经过非线性变换的尺度。
• 取倒数（1.0 / ...）将这些幂次转换成对应的频率值，这些频率通常用于某些周期性编码过程，如在某些类型的神经网络（尤其是那些处理时间序列数据的）中。

freqs张量包含了根据rope_theta和均匀间隔的指数值调整后的一系列频率。这些频率可用于进一步的数据处理或特征工程，特别是在处理与周期性或波形相关的任务时。

运行结果：

运行结果：

运行结果：

这段代码首先使用torch.outer计算17个标记每个对应的频率集合。然后，使用torch.polar函数将这些频率转换为复数表示，其中实部和虚部分别对应复平面上的横坐标和纵坐标。最后，使用matplotlib绘图库在复平面上绘制出这些复数点，并通过直线将它们与原点相连，形象地展示了这些复数的分布和相位变化。这可以用于可视化分析频率分布的特性，尤其是在处理涉及信号处理或周期性数据的场景中。

13. 有了每个令牌的查询元素的复数（角度变化向量）

我们可以将查询（我们分成对的查询）转换为复数，然后点积根据位置旋转查询

运行结果：

在这段代码中：

• torch.view_as_complex(q_per_token_split_into_pairs)：该函数将输入的实数张量视为复数张量，其中输入张量的最后一个维度大小必须是2，分别代表复数的实部和虚部。

对于张量形状的解释：

• 假设q_per_token_split_into_pairs的形状是(batch_size, sequence_length, 2)，其中2表示成对的实部和虚部。
• 使用torch.view_as_complex后，最后一个维度会被视为复数的一部分，因此新的张量q_per_token_as_complex_numbers的形状将变为(batch_size, sequence_length)，每个元素都是一个复数。

这种转换通常用于信号处理或神经网络中处理复数数据，使得后续的计算可以直接在复数域中进行。

运行结果：

• q_per_token_as_complex_numbers * freqs_cis：这个操作将两个复数张量相乘。在复数乘法中，两个复数相乘的结果是一个新的复数，其模长是原两个复数模长的乘积，其角度是原两个复数角度的和。

对于张量形状的解释：

• 假设 q_per_token_as_complex_numbers 的形状是 (batch_size, sequence_length)，其中每个元素是一个复数。
• freqs_cis 的形状应当与 q_per_token_as_complex_numbers 兼容，例如也是 (batch_size, sequence_length) 或能够广播到这个形状。
• 结果张量 q_per_token_as_complex_numbers_rotated 将保持与 q_per_token_as_complex_numbers 相同的形状 (batch_size, sequence_length)。

这样的操作常用于信号处理和深度学习中的特征变换，尤其是在处理需要调整相位或频率特性的应用中。每个复数的旋转可以看作是在复平面上的一个角度变化，这对于某些类型的算法（比如那些涉及时间序列或周期性数据处理的算法）特别有用。

14. 获得旋转向量后 

我们可以通过再次将复数视为实数来将查询作为对返回

运行结果：

这个操作的具体行为如下：

• torch.view_as_real(q_per_token_as_complex_numbers_rotated)：此函数将输入的复数张量转换为实数张量，每个复数的实部和虚部成对出现，形成最后一个维度的两个连续元素。

关于张量形状的解释：

• 假设q_per_token_as_complex_numbers_rotated的形状是(batch_size, sequence_length)，其中每个元素是一个复数。
• 转换后的张量q_per_token_split_into_pairs_rotated的形状将变为(batch_size, sequence_length, 2)。最后的2表示复数的两个组成部分：实部和虚部。

这样的转换使得复数数据的实部和虚部可以作为独立的特征在后续处理中使用，适用于需要分别处理复数实部和虚部的场景，例如在某些机器学习模型的输入处理中。

旋转后的对现在已经合并，我们现在有了一个新的查询向量（旋转后的查询向量），其形状为 [17×128]，其中 17 是标记的数量，128 是查询向量的维度。

运行结果：

在这里：

• q_per_token.shape 原本的形状可能是 (batch_size, sequence_length)，假设之前操作中每个token对应的复数（实部和虚部）被视为独立的特征或维度处理。
• view(q_per_token.shape) 方法用于将张量 q_per_token_split_into_pairs_rotated 从包含复数实虚部的三维形状 (batch_size, sequence_length, 2) 转换回二维形状 (batch_size, sequence_length)。这种操作通常意味着在某种处理（如复数到实数的转换）之后，我们需要恢复数据到某个特定的维度结构以适配后续的处理流程。

因此，最终的张量 q_per_token_rotated 将保持与原始 q_per_token 相同的形状，这也意味着每个 token 现在对应一个合并后的特征向量，而不再是分开的实部和虚部。

15. 键（几乎与查询相同）

不打算通过钥匙的数学计算，你唯一需要记住的是：

键生成键向量，也是 dimention 128

键的权重数只有查询权重的 1/4，这是因为键的权重一次在 4 个头之间共享，以减少所需的计算次数

键也会旋转以添加位置信息，就像出于相同原因的查询一样

运行结果：

• n_kv_heads 表示用于键向量的头的数量。
• k_layer0.shape[0] // n_kv_heads 计算每个头需要处理的键向量的数量。
• dim 是每个键向量的维度，通常与模型的嵌入维度相匹配。

该操作后的 k_layer0 的形状将是 (n_kv_heads, k_layer0.shape[0] // n_kv_heads, dim)。这意味着：

• 第一维 n_kv_heads 表示头的数量。
• 第二维 k_layer0.shape[0] // n_kv_heads 表示每个头处理的键向量的数量。
• 第三维 dim 表示每个键向量的维度。

这种重塑是为了便于在后续操作中，每个注意力头能够独立地对应并处理其分配的键向量，从而实现有效的并行处理和更精细的注意力机制。

运行结果：

运行结果：

运行结果：

运行结果：

运行结果：