devbox
你好赵伟
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

LLMs之llama3-from-scratch：llama3-from-scratch(从头开始利用pytorch来实现并解读LLaMA-3模型的每层代码)的简介、核心思路梳理_llama3 scratch

作者：你好赵伟 | 2024-05-31 19:08:04

赞

踩

llama3 scratch

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模型底层工作机制，是一篇值得推荐的论文。

目录

llama3-from-scratch的简介

llama3-from-scratch的核心思路梳理

0、前置

0.1、加载tokenizer对文本进行tokenize：将文本转换为模型可以理解的数字序列（即词元或tokens）+并在生成模型输出后能将tokens转换回可读文本（解码）

0.2、从模型文件加载各层参数：从预训练模型文件中读取权重参数

1、输入嵌入层：采用tiktoken库的BPE算法将输入文本的tokens转换为嵌入向量，为模型处理做好准备

1.1、文本分词并编码：将一段文本(输入)转换为token ID(输出)。将输入的自然语言经过分词技术转换为token ID，如下所示，一句话包含17个token(可以简单理解为单词)

1.2、离散token进行嵌入向量化：将离散的token ID(输入，形状为【17,1】)转换为其连续的嵌入向量表示(输出，形状为【17,4096】)，以便进行后续的运算

1.3、采用RMSNorm对嵌入向量进行均方根归一化：仅对嵌入进行归一化处理，前后形状未变，只是确保数值稳定，避免后续计算过程中出现数值0问题

2、位置编码层：为词元嵌入添加位置信息，使模型能理解词元的顺序

2.1、背景——构建Transformer的第一个层：准备嵌入以输入到Transformer层进行处理，实现Transformer网络结构的前向传播

2.2、背景——计算每个token在Transformer第一层第一个头的query向量

2.3、真正实现RoPE：此处以Query向量进行RoPE为例讲解，但Key与Query几乎相同

3、LLaMA的核心模块

3.1、MHA模块—从零开始实现注意力机制：计算不同词元之间的关联度，提取重要信息。包括Q、K、V矩阵的生成、位置编码、掩码、softmax计算注意力权重等

3.1.1、计算查询(Query)、键(Key)：生成用于注意力机制中每一个token的query和key；

3.1.2、计算注意力分数矩阵并掩码：计算注意力分数并进行掩码，确保模型不会关注未来的词元

3.1.3、值（Value）的计算：根据注意力分数和value向量计算最终的注意力输出

3.1.4、多头注意力的合并：将多个注意力头的输出合并为一个向量

3.1.5、注意力输出权重矩阵：将合并后的注意力输出转换为嵌入向量的修正

3.1.6、嵌入向量的更新：将注意力机制的输出加到原始嵌入向量上，更新嵌入信息

3.2、FFNN层：在嵌入向量上应用前馈网络，增加模型的非线性

3.3、循环迭代

3.3.1、所有层的迭代：将上述过程在每一层中进行迭代处理，逐层增强表示能力

3.3.2、最后的归一化：对最终嵌入向量进行归一化处理，用于预测下一个token

4、输出层

4.1、将嵌入向量转换为预测的token：确认logits张量的形状

4.2、解码：将模型输出的词元转换为可读的文本

5、模型推理：应用训练好的输出层参数进行预测(验证整个代码的正确性)，模型预测下一个token的编号为2983，这是42的token ID吗？

llama3-from-scratch的简介

2024年5月20日，Nishant Aklecha正式发布了该项目，在这个文件中，我从头实现了llama3，一次一个张量和矩阵乘法。此外，我将直接从Meta为llama3提供的模型文件中加载张量，在运行此文件之前需要下载权重。以下是下载权重的官方链接：https://llama.meta.com/llama-downloads/

GitHub地址：GitHub - naklecha/llama3-from-scratch: llama3 implementation one matrix multiplication at a time

llama3-from-scratch的核心思路梳理

注意：当前文章仍处于持续更新和梳理中……

0、前置

0.1、加载tokenizer对文本进行tokenize：将文本转换为模型可以理解的数字序列（即词元或tokens）+并在生成模型输出后能将tokens转换回可读文本（解码）

简介

我不会实现一个BPE分词器（但是Andrej Karpathy有一个非常简洁的实现）

链接到他的实现：

https://github.com/karpathy/minbpe

该项目给出了一个针对中文分词任务的最小化、清晰而可读的Python实现Byte Pair Encoding(BPE)算法的代码库。它实现了两个Tokenizer对象用于文本到词元和词元到文本的转换：BasicTokenizer和RegexTokenizer。
BasicTokenizer直接在文本上运行BPE算法，RegexTokenizer在BPE之前通过正则表达式将文本切分为不同类型(如字母、数字、标点等)以防止词汇跨类型合并。

思路步骤

从Meta下载权重文件。

使用tiktoken库加载BPE分词器，并定义特殊token。

核心技术点

BPE分词：字节对编码（Byte Pair Encoding）是一种常见的分词方法，用于处理词汇表。

代码解读

目的是设置和测试一个自定义的tokenizer，它能够处理文本数据

  • 设置路径和特殊标记:

    • tokenizer_path变量指定了tokenizer.model文件的位置。
    • special_tokens列表定义了一系列特殊标记，这些标记在文本编码过程中有特定的含义，例如表示文本的开始和结束。

  • 加载BPE（Byte Pair Encoding）模型:

    • 使用tiktoken.load.load_tiktoken_bpe函数从指定的tokenizer_path加载BPE模型，这是一种用于文本编码的算法，通过合并常见的字节对来减少词汇量。

  • 创建tokenizer对象:

    • tokenizer是tiktoken.Encoding的一个实例，它使用上面加载的BPE模型、正则表达式模式（pat_str）、合并等级（mergeable_ranks）和特殊标记（special_tokens）来创建。

  • 解码和编码测试:

    • 最后，代码测试了tokenizer的encode和decode方法。首先使用encode方法将字符串"hello world!"编码为数字表示，然后使用decode方法将这些数字解码回文本表示。
  1. from pathlib import Path
  2. import tiktoken
  3. from tiktoken.load import load_tiktoken_bpe
  4. import torch
  5. import json
  6. import matplotlib.pyplot as plt
  7.  
  8. tokenizer_path = "Meta-Llama-3-8B/tokenizer.model"
  9. special_tokens = [
  10.             "<|begin_of_text|>",
  11.             "<|end_of_text|>",
  12.             "<|reserved_special_token_0|>",
  13.             "<|reserved_special_token_1|>",
  14.             "<|reserved_special_token_2|>",
  15.             "<|reserved_special_token_3|>",
  16.             "<|start_header_id|>",
  17.             "<|end_header_id|>",
  18.             "<|reserved_special_token_4|>",
  19.             "<|eot_id|>",  # end of turn
  20.         ] + [f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5)]
  21. mergeable_ranks = load_tiktoken_bpe(tokenizer_path)
  22. tokenizer = tiktoken.Encoding(
  23.     name=Path(tokenizer_path).name,
  24.     pat_str=r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+",
  25.     mergeable_ranks=mergeable_ranks,
  26.     special_tokens={token: len(mergeable_ranks) + i for i, token in enumerate(special_tokens)},
  27. )
  28.  
  29. tokenizer.decode(tokenizer.encode("hello world!"))

0.2、从模型文件加载各层参数：从预训练模型文件中读取权重参数

简介

通常情况下，读取模型文件取决于模型类的编写方式和其中的变量名称。但由于我们要从头实现llama3，所以我们将逐个张量读取文件。

我们使用此配置来推断模型的细节，模型有32个Transformer层，每个多头注意力块有32个头，n_kv_heads，词汇表大小为128256等等。

下面是各个参数的含义：

  • 'dim': 4096：表示模型的隐藏层维度，即模型中嵌入向量的维度或者是注意力层的尺寸。
  • 'n_layers': 32：表示模型的层数，即 Transformer 的层数。
  • 'n_heads': 32：表示多头注意力机制中的注意力头数。多头注意力允许模型在不同的表示空间中并行处理信息。
  • 'n_kv_heads': 8：表示键和值（Key 和 Value）在多头注意力中的头数。这指的是键（Key）和值（Value）头数，这是一种注意力机制的变体，其中键和值的头数可以与查询头数不同。
  • 'vocab_size': 128256：表示模型词汇表的大小，即可以处理的不同词的数量。
  • 'multiple_of': 1024：表示一些尺寸（比如前馈神经网络的维度）应当是这个数值的倍数，通常用于确保计算效率。这可能是一个设计上的考虑，用于确保某些层的尺寸是某个数值的倍数，这有助于提高模型在硬件上的运行效率。
  • 'ffn_dim_multiplier': 1.3：前馈神经网络（Feed-Forward Network）的维度倍率，表示前馈网络的维度是隐藏层维度的1.3倍。
  • 'norm_eps': 1e-05：Layer Normalization中的一个非常小的值，防止除以零的情况，用于避免除以零的情况。
  • 'rope_theta': 500000.0：这是相对位置编码（rotary positional embeddings）的一个参数，通常用于提高模型在处理长序列时的性能。

思路步骤

加载模型文件，查看模型参数。

读取模型配置文件，获取模型的详细参数。

核心技术点

模型文件加载：使用torch.load加载模型权重文件。

配置解析：解析配置文件以提取模型超参数。

  1. model = torch.load("Meta-Llama-3-8B/consolidated.00.pth")
  2. print(json.dumps(list(model.keys())[:20], indent=4))
  3.  
  4. [
  5.     "tok_embeddings.weight",
  6.     "layers.0.attention.wq.weight",
  7.     "layers.0.attention.wk.weight",
  8.     "layers.0.attention.wv.weight",
  9.     "layers.0.attention.wo.weight",
  10.     "layers.0.feed_forward.w1.weight",
  11.     "layers.0.feed_forward.w3.weight",
  12.     "layers.0.feed_forward.w2.weight",
  13.     "layers.0.attention_norm.weight",
  14.     "layers.0.ffn_norm.weight",
  15.     "layers.1.attention.wq.weight",
  16.     "layers.1.attention.wk.weight",
  17.     "layers.1.attention.wv.weight",
  18.     "layers.1.attention.wo.weight",
  19.     "layers.1.feed_forward.w1.weight",
  20.     "layers.1.feed_forward.w3.weight",
  21.     "layers.1.feed_forward.w2.weight",
  22.     "layers.1.attention_norm.weight",
  23.     "layers.1.ffn_norm.weight",
  24.     "layers.2.attention.wq.weight"
  25. ]
  26.  
  27.  
  28. with open("Meta-Llama-3-8B/params.json", "r") as f:
  29.     config = json.load(f)
  30. config
  31.  
  32. {'dim': 4096,
  33.  'n_layers': 32,
  34.  'n_heads': 32,
  35.  'n_kv_heads': 8,
  36.  'vocab_size': 128256,
  37.  'multiple_of': 1024,
  38.  'ffn_dim_multiplier': 1.3,
  39.  'norm_eps': 1e-05,
  40.  'rope_theta': 500000.0}
  41.  
  42.  
  43.  
  44.  
  45. dim = config["dim"]
  46. n_layers = config["n_layers"]
  47. n_heads = config["n_heads"]
  48. n_kv_heads = config["n_kv_heads"]
  49. vocab_size = config["vocab_size"]
  50. multiple_of = config["multiple_of"]
  51. ffn_dim_multiplier = config["ffn_dim_multiplier"]
  52. norm_eps = config["norm_eps"]
  53. rope_theta = torch.tensor(config["rope_theta"])
  54.  
  55.

1、输入嵌入层：采用tiktoken库的BPE算法将输入文本的tokens转换为嵌入向量，为模型处理做好准备

1.1、文本分词并编码：将一段文本(输入)转换为token ID(输出)。将输入的自然语言经过分词技术转换为token ID，如下所示，一句话包含17个token(可以简单理解为单词)

简介

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

思路步骤

输入一个提示文本，并使用分词器将其转换为tokens。

打印tokens及其对应的文本表示。

核心技术点

分词：使用分词器将文本转换为token ID。

文本表示：解码token ID为文本以验证分词过程。

  1. prompt = "the answer to the ultimate question of life, the universe, and everything is "
  2. tokens = [128000] + tokenizer.encode(prompt)
  3. print(tokens)
  4. tokens = torch.tensor(tokens)
  5. prompt_split_as_tokens = [tokenizer.decode([token.item()]) for token in tokens]
  6. print(prompt_split_as_tokens)
  7.  
  8.  
  9.  
  10.  
  11. [128000, 1820, 4320, 311, 279, 17139, 3488, 315, 2324, 11, 279, 15861, 11, 323, 4395, 374, 220]
  12. ['<|begin_of_text|>', 'the', ' answer', ' to', ' the', ' ultimate', ' question', ' of', ' life', ',', ' the', ' universe', ',', ' and', ' everything', ' is', ' ']

1.2、离散token进行嵌入向量化：将离散的token ID(输入，形状为【17,1】)转换为其连续的嵌入向量表示(输出，形状为【17,4096】)，以便进行后续的运算

简介

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

我们的[17x1] tokens现在是[17x4096]，即17个长度为4096的嵌入（每个token一个嵌入）。

注意：跟踪形状，这使它更容易理解一切

思路步骤

定义一个嵌入层，并加载预训练的权重。

将tokens转换为嵌入表示。

核心技术点

嵌入层：使用torch.nn.Embedding定义嵌入层，并加载预训练权重。

  1. embedding_layer = torch.nn.Embedding(vocab_size, dim)
  2. embedding_layer.weight.data.copy_(model["tok_embeddings.weight"])
  3. token_embeddings_unnormalized = embedding_layer(tokens).to(torch.bfloat16)
  4. token_embeddings_unnormalized.shape
  5. torch.Size([17, 4096])

1.3、采用RMSNorm对嵌入向量进行均方根归一化：仅对嵌入进行归一化处理，前后形状未变，只是确保数值稳定，避免后续计算过程中出现数值0问题

简介

我们然后使用RMS归一化嵌入，请注意此步骤后形状不会改变，值只是被归一化了。需要注意的是我们需要一个norm_eps（从配置中获取），因为我们不想意外地将rms设为0并除以0。

公式如下：

归一化后我们的形状仍然是[17x4096]，与嵌入相同但已归一化。

思路步骤

对嵌入进行RMS归一化。

核心技术点

RMS归一化：对嵌入向量进行均方根归一化（RMS Norm）。

  1. # def rms_norm(tensor, norm_weights):
  2. #     rms = (tensor.pow(2).mean(-1, keepdim=True) + norm_eps)**0.5
  3. #     return tensor * (norm_weights / rms)
  4. def rms_norm(tensor, norm_weights):
  5.     return (tensor * torch.rsqrt(tensor.pow(2).mean(-1, keepdim=True) + norm_eps)) * norm_weights

2、位置编码层：为词元嵌入添加位置信息，使模型能理解词元的顺序

2.1、背景——构建Transformer的第一个层：准备嵌入以输入到Transformer层进行处理，实现Transformer网络结构的前向传播

简介

您会看到我从这个模型字典中访问layer.0（这是第一层） 无论如何，所以归一化后，我们的形状仍然是【17x4096】，与嵌入相同但已归一化

思路步骤

构建Transformer的第一层：

核心技术点

多头注意力机制（Multi-Head Attention）与残差连接

Transformer层：应用注意力机制和前馈神经网络进行特征提取和变换。

  1. token_embeddings = rms_norm(token_embeddings_unnormalized, model["layers.0.attention_norm.weight"])
  2. token_embeddings.shape
  3.  
  4. torch.Size([17, 4096])
  5.  
  6.

2.2、背景——计算每个token在Transformer第一层第一个头的query向量

简介

让我们加载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）的计算。

  1.  
  2.  
  3. print(
  4.     model["layers.0.attention.wq.weight"].shape,
  5.     model["layers.0.attention.wk.weight"].shape,
  6.     model["layers.0.attention.wv.weight"].shape,
  7.     model["layers.0.attention.wo.weight"].shape
  8. )
  9. torch.Size([4096, 4096]) torch.Size([1024, 4096]) torch.Size([1024, 4096]) torch.Size([4096, 4096])
  10.  
  11.  
  12.  
  13. q_layer0 = model["layers.0.attention.wq.weight"]
  14. head_dim = q_layer0.shape[0] // n_heads
  15. q_layer0 = q_layer0.view(n_heads, head_dim, dim)
  16. q_layer0.shape
  17. torch.Size([32, 128, 4096])
  18.  
  19.  
  20.  
  21. q_layer0_head0 = q_layer0[0]
  22. q_layer0_head0.shape
  23. torch.Size([128, 4096])

2.3、真正实现RoPE：此处以Query向量进行RoPE为例讲解，但Key与Query几乎相同

背景

痛点：现在，已经为prompt中的每个token生成了query向量，但每个单独的query向量并不知道它在prompt中的具体位置。我们现在处于每个token在我们的提示中有一个查询向量的阶段，但如果你想一想——单个查询向量对提示中的位置没有任何了解。例如：

  • query：“the answer to the ultimate question of life, the universe, and everything is____”
  • 查询：“生命、宇宙以及一切终极问题的答案是____”

期望：在我们的提示中我们使用了三次“the”，我们需要所有3个“the”token的查询向量（每个大小为[1x128]）根据它们在查询中的位置具有不同的query向量。

解决方案：采用RoPE（旋转位置嵌入）执行这些旋转操作。

RoPE：观看这个视频（这是我看的）来理解数学：https://www.youtube.com/watch?v=o29P0Kpobz0&t=530s

简介

在上述步骤中，我们将查询向量分成了若干对，并对每对应用一个旋转角度偏移！

现在我们有一个大小为[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函数对查询向量进行旋转，结合位置编码信息。

  1. zero_to_one_split_into_64_parts = torch.tensor(range(64))/64
  2. zero_to_one_split_into_64_parts
  3. tensor([0.0000, 0.0156, 0.0312, 0.0469, 0.0625, 0.0781, 0.0938, 0.1094, 0.1250,
  4.         0.1406, 0.1562, 0.1719, 0.1875, 0.2031, 0.2188, 0.2344, 0.2500, 0.2656,
  5.         0.2812, 0.2969, 0.3125, 0.3281, 0.3438, 0.3594, 0.3750, 0.3906, 0.4062,
  6.         0.4219, 0.4375, 0.4531, 0.4688, 0.4844, 0.5000, 0.5156, 0.5312, 0.5469,
  7.         0.5625, 0.5781, 0.5938, 0.6094, 0.6250, 0.6406, 0.6562, 0.6719, 0.6875,
  8.         0.7031, 0.7188, 0.7344, 0.7500, 0.7656, 0.7812, 0.7969, 0.8125, 0.8281,
  9.         0.8438, 0.8594, 0.8750, 0.8906, 0.9062, 0.9219, 0.9375, 0.9531, 0.9688,
  10.         0.9844])
  11. freqs = 1.0 / (rope_theta ** zero_to_one_split_into_64_parts)
  12. freqs
  13. tensor([1.0000e+00, 8.1462e-01, 6.6360e-01, 5.4058e-01, 4.4037e-01, 3.5873e-01,
  14.         2.9223e-01, 2.3805e-01, 1.9392e-01, 1.5797e-01, 1.2869e-01, 1.0483e-01,
  15.         8.5397e-02, 6.9566e-02, 5.6670e-02, 4.6164e-02, 3.7606e-02, 3.0635e-02,
  16.         2.4955e-02, 2.0329e-02, 1.6560e-02, 1.3490e-02, 1.0990e-02, 8.9523e-03,
  17.         7.2927e-03, 5.9407e-03, 4.8394e-03, 3.9423e-03, 3.2114e-03, 2.6161e-03,
  18.         2.1311e-03, 1.7360e-03, 1.4142e-03, 1.1520e-03, 9.3847e-04, 7.6450e-04,
  19.         6.2277e-04, 5.0732e-04, 4.1327e-04, 3.3666e-04, 2.7425e-04, 2.2341e-04,
  20.         1.8199e-04, 1.4825e-04, 1.2077e-04, 9.8381e-05, 8.0143e-05, 6.5286e-05,
  21.         5.3183e-05, 4.3324e-05, 3.5292e-05, 2.8750e-05, 2.3420e-05, 1.9078e-05,
  22.         1.5542e-05, 1.2660e-05, 1.0313e-05, 8.4015e-06, 6.8440e-06, 5.5752e-06,
  23.         4.5417e-06, 3.6997e-06, 3.0139e-06, 2.4551e-06])
  24. freqs_for_each_token = torch.outer(torch.arange(17), freqs)
  25. freqs_cis = torch.polar(torch.ones_like(freqs_for_each_token), freqs_for_each_token)
  26. freqs_cis.shape
  27.  
  28. # viewing tjhe third row of freqs_cis
  29. value = freqs_cis[3]
  30. plt.figure()
  31. for i, element in enumerate(value[:17]):
  32.     plt.plot([0, element.real], [0, element.imag], color='blue', linewidth=1, label=f"Index: {i}")
  33.     plt.annotate(f"{i}", xy=(element.real, element.imag), color='red')
  34. plt.xlabel('Real')
  35. plt.ylabel('Imaginary')
  36. plt.title('Plot of one row of freqs_cis')
  37. plt.show()
  38.  
  39.  
  40.  
  41.  
  42.  
  43. q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)
  44. q_per_token_as_complex_numbers.shape
  45. torch.Size([17, 64])
  46. q_per_token_as_complex_numbers_rotated = q_per_token_as_complex_numbers * freqs_cis
  47. q_per_token_as_complex_numbers_rotated.shape
  48. torch.Size([17, 64])
  49.  
  50.  
  51. q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers_rotated)
  52. q_per_token_split_into_pairs_rotated.shape
  53.  
  54. torch.Size([17, 64, 2])
  55.  
  56.  
  57. q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)
  58. q_per_token_rotated.shape
  59. torch.Size([17, 128])

3、LLaMA的核心模块

3.1、MHA模块—从零开始实现注意力机制：计算不同词元之间的关联度，提取重要信息。包括Q、K、V矩阵的生成、位置编码、掩码、softmax计算注意力权重等

3.1.1、计算查询(Query)、键(Key)：生成用于注意力机制中每一个token的query和key；

简介

查询Query、键Key、值Value和输出向量

思路步骤

Key键向量生成

>> key向量维度为128。

>> key的权重共享机制，key的权重只有query的四分之一,，这是因为每4个头共享一次，以减少计算量。

>> key向量也进行旋转以加入位置信息，原因与query相同。

Query向量生成及旋转，增加位置信息。

查询向量与键向量的维度匹配。

旋转操作用于引入位置编码。

在这个阶段，我们现在有了每个token旋转后的query和key，其中每一个query和key现在的形状是[17x128]。

核心技术点

矩阵乘法与旋转位置编码

  1. k_layer0 = model["layers.0.attention.wk.weight"]
  2. k_layer0 = k_layer0.view(n_kv_heads, k_layer0.shape[0] // n_kv_heads, dim)
  3. k_layer0.shape
  4. torch.Size([8, 128, 4096])
  5. k_layer0_head0 = k_layer0[0]
  6. k_layer0_head0.shape
  7. torch.Size([128, 4096])
  8. k_per_token = torch.matmul(token_embeddings, k_layer0_head0.T)
  9. k_per_token.shape
  10. torch.Size([17, 128])
  11. k_per_token_split_into_pairs = k_per_token.float().view(k_per_token.shape[0], -1, 2)
  12. k_per_token_split_into_pairs.shape
  13. torch.Size([17, 64, 2])
  14. k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token_split_into_pairs)
  15. k_per_token_as_complex_numbers.shape
  16. torch.Size([17, 64])
  17. k_per_token_split_into_pairs_rotated = torch.view_as_real(k_per_token_as_complex_numbers * freqs_cis)
  18. k_per_token_split_into_pairs_rotated.shape
  19. torch.Size([17, 64, 2])
  20. k_per_token_rotated = k_per_token_split_into_pairs_rotated.view(k_per_token.shape)
  21. k_per_token_rotated.shape
  22. torch.Size([17, 128])

3.1.2、计算注意力分数矩阵并掩码：计算注意力分数并进行掩码，确保模型不会关注未来的词元

简介

查询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。

核心技术点

分数计算与未来词元掩码

  1. def display_qk_heatmap(qk_per_token):
  2.     _, ax = plt.subplots()
  3.     im = ax.imshow(qk_per_token.to(float).detach(), cmap='viridis')
  4.     ax.set_xticks(range(len(prompt_split_as_tokens)))
  5.     ax.set_yticks(range(len(prompt_split_as_tokens)))
  6.     ax.set_xticklabels(prompt_split_as_tokens)
  7.     ax.set_yticklabels(prompt_split_as_tokens)
  8.     ax.figure.colorbar(im, ax=ax)
  9.     
  10. display_qk_heatmap(qk_per_token)
  11.  
  12.  
  13.  
  14. mask = torch.full((len(tokens), len(tokens)), float("-inf"), device=tokens.device)
  15. mask = torch.triu(mask, diagonal=1)
  16. mask
  17. tensor([[0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
  18.         [0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
  19.         [0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
  20.         [0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
  21.         [0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
  22.         [0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
  23.         [0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
  24.         [0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
  25.         [0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
  26.         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf],
  27.         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf],
  28.         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf],
  29.         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf],
  30.         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf],
  31.         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf],
  32.         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf],
  33.         [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
  34. qk_per_token_after_masking = qk_per_token + mask
  35. display_qk_heatmap(qk_per_token_after_masking)
  36.  
  37. qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_after_masking, dim=1).to(torch.bfloat16)
  38. display_qk_heatmap(qk_per_token_after_masking_after_softmax)

3.1.3、值（Value）的计算：根据注意力分数和value向量计算最终的注意力输出

简介

查询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]。

核心技术点

分数与值的权重合并

值向量的矩阵乘法

  1. qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T)/(head_dim)**0.5
  2. qk_per_token.shape
  3. torch.Size([17, 17])

3.1.4、多头注意力的合并：将多个注意力头的输出合并为一个向量

简介

我们现在有了第一层中第一个头的注意力值，然后，将运行一个循环，并对第一层中的每一个头执行与上述单元格完全相同的数学运算。

然后，得到了第一层所有32个头的qkv_attention矩阵，接下来我将把所有的注意力分数合并成一个大的矩阵，其大小为[17x4096]

我们快要结束了!

思路步骤

多头注意力机制：对每个注意力头分别计算注意力输出，并将其拼接。

多头注意力的并行计算和拼接。

线性层变换：通过线性变换将拼接后的多头注意力输出进行进一步处理。

线性层的矩阵乘法。

核心技术点

向量的拼接

  1. qkv_attention_store = []
  2.  
  3. for head in range(n_heads):
  4.     q_layer0_head = q_layer0[head]
  5.     k_layer0_head = k_layer0[head//4] # key weights are shared across 4 heads
  6.     v_layer0_head = v_layer0[head//4] # value weights are shared across 4 heads
  7.     q_per_token = torch.matmul(token_embeddings, q_layer0_head.T)
  8.     k_per_token = torch.matmul(token_embeddings, k_layer0_head.T)
  9.     v_per_token = torch.matmul(token_embeddings, v_layer0_head.T)
  10.  
  11.     q_per_token_split_into_pairs = q_per_token.float().view(q_per_token.shape[0], -1, 2)
  12.     q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)
  13.     q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers * freqs_cis[:len(tokens)])
  14.     q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)
  15.  
  16.     k_per_token_split_into_pairs = k_per_token.float().view(k_per_token.shape[0], -1, 2)
  17.     k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token_split_into_pairs)
  18.     k_per_token_split_into_pairs_rotated = torch.view_as_real(k_per_token_as_complex_numbers * freqs_cis[:len(tokens)])
  19.     k_per_token_rotated = k_per_token_split_into_pairs_rotated.view(k_per_token.shape)
  20.  
  21.     qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T)/(128)**0.5
  22.     mask = torch.full((len(tokens), len(tokens)), float("-inf"), device=tokens.device)
  23.     mask = torch.triu(mask, diagonal=1)
  24.     qk_per_token_after_masking = qk_per_token + mask
  25.     qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_after_masking, dim=1).to(torch.bfloat16)
  26.     qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
  27.     qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
  28.     qkv_attention_store.append(qkv_attention)
  29.  
  30. len(qkv_attention_store)
  31.  
  32.  
  33.  
  34.  
  35. stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)
  36. stacked_qkv_attention.shape
  37. torch.Size([17, 4096])

3.1.5、注意力输出权重矩阵：将合并后的注意力输出转换为嵌入向量的修正

简介

查询Query、键Key、值Value和输出向量

思路步骤

对于第0层注意力机制的最后步骤，其一是将注意力得分矩阵与权重矩阵相乘。

这是一个简单的线性层，所以我们只需进行矩阵乘法

核心技术点

线性层（全连接层）的应用。

  1. stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)
  2. stacked_qkv_attention.shape
  3. torch.Size([17, 4096])

3.1.6、嵌入向量的更新：将注意力机制的输出加到原始嵌入向量上，更新嵌入信息

简介

查询Query、键Key、值Value和输出向量

现在，我们有了注意力后嵌入Value的变化，应该加到原始的token嵌入中；

然后，对嵌入增量进行归一化，然后通过嵌入增量运行一个前馈神经网络；

思路步骤

核心技术点

残差连接

3.2、FFNN层：在嵌入向量上应用前馈网络，增加模型的非线性

简介

在Llama3中，加载前馈权重并实现前馈网络。使用了一种名为SwiGLU的前馈网络，这种网络结构在模型需要的时候，能够有效地增加非线性。 在当今的LLMs中，使用这种前馈网络架构已经相当标准。

思路步骤

归一化与前馈网络：对注意力输出进行归一化，并通过前馈神经网络进行非线性变换。

RMSNorm归一化。

SwiGLU前馈网络，用于增强非线性表达能力。

核心技术点

SwiGLU激活函数与线性层

  1. embedding_after_edit_normalized = rms_norm(embedding_after_edit, model["layers.0.ffn_norm.weight"])
  2. embedding_after_edit_normalized.shape
  3. torch.Size([17, 4096])
  4.  
  5.  
  6.  
  7.  
  8. w1 = model["layers.0.feed_forward.w1.weight"]
  9. w2 = model["layers.0.feed_forward.w2.weight"]
  10. w3 = model["layers.0.feed_forward.w3.weight"]
  11. output_after_feedforward = torch.matmul(torch.functional.F.silu(torch.matmul(embedding_after_edit_normalized, w1.T)) * torch.matmul(embedding_after_edit_normalized, w3.T), w2.T)
  12. output_after_feedforward.shape
  13. torch.Size([17, 4096])

3.3、循环迭代

3.3.1、所有层的迭代：将上述过程在每一层中进行迭代处理，逐层增强表示能力
背景现在完成了第一层之后每个token的新嵌入。现在只剩下31层了，只需通过一个循环来完成。

简介

查询Query、键Key、值Value和输出向量

我们终于在第1层之后为每个词元获得了新的编辑后的嵌入。我们还有31层才能完成（还需要一个循环）。

你可以将这个编辑后的嵌入，想象成包含了第一层中所有Query信息的嵌入。随着层数的增加，每一层都会对输入的信息进行越来越复杂的处理，直到最终得到一个，能够全面了解下一个需要预测的token的嵌入。

天哪，一切都集中在一起了

是的，就是这样。我们之前所做的所有事情，现在一次性为每个层都做了

思路步骤

重复以上步骤，处理所有Transformer子层

核心技术点

循环与层堆叠

循环迭代处理每一层。

  1. layer_0_embedding = embedding_after_edit+output_after_feedforward
  2. layer_0_embedding.shape
  3. torch.Size([17, 4096])

3.3.2、最后的归一化：对最终嵌入向量进行归一化处理，用于预测下一个token

简介

思路步骤

核心技术点

矩阵乘法，线性变换。

最终归一化处理

  1. final_embedding = token_embeddings_unnormalized
  2. for layer in range(n_layers):
  3.     qkv_attention_store = []
  4.     layer_embedding_norm = rms_norm(final_embedding, model[f"layers.{layer}.attention_norm.weight"])
  5.     q_layer = model[f"layers.{layer}.attention.wq.weight"]
  6.     q_layer = q_layer.view(n_heads, q_layer.shape[0] // n_heads, dim)
  7.     k_layer = model[f"layers.{layer}.attention.wk.weight"]
  8.     k_layer = k_layer.view(n_kv_heads, k_layer.shape[0] // n_kv_heads, dim)
  9.     v_layer = model[f"layers.{layer}.attention.wv.weight"]
  10.     v_layer = v_layer.view(n_kv_heads, v_layer.shape[0] // n_kv_heads, dim)
  11.     w_layer = model[f"layers.{layer}.attention.wo.weight"]
  12.     for head in range(n_heads):
  13.         q_layer_head = q_layer[head]
  14.         k_layer_head = k_layer[head//4]
  15.         v_layer_head = v_layer[head//4]
  16.         q_per_token = torch.matmul(layer_embedding_norm, q_layer_head.T)
  17.         k_per_token = torch.matmul(layer_embedding_norm, k_layer_head.T)
  18.         v_per_token = torch.matmul(layer_embedding_norm, v_layer_head.T)
  19.         q_per_token_split_into_pairs = q_per_token.float().view(q_per_token.shape[0], -1, 2)
  20.         q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)
  21.         q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers * freqs_cis)
  22.         q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)
  23.         k_per_token_split_into_pairs = k_per_token.float().view(k_per_token.shape[0], -1, 2)
  24.         k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token_split_into_pairs)
  25.         k_per_token_split_into_pairs_rotated = torch.view_as_real(k_per_token_as_complex_numbers * freqs_cis)
  26.         k_per_token_rotated = k_per_token_split_into_pairs_rotated.view(k_per_token.shape)
  27.         qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T)/(128)**0.5
  28.         mask = torch.full((len(token_embeddings_unnormalized), len(token_embeddings_unnormalized)), float("-inf"))
  29.         mask = torch.triu(mask, diagonal=1)
  30.         qk_per_token_after_masking = qk_per_token + mask
  31.         qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_after_masking, dim=1).to(torch.bfloat16)
  32.         qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
  33.         qkv_attention_store.append(qkv_attention)
  34.  
  35.     stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)
  36.     w_layer = model[f"layers.{layer}.attention.wo.weight"]
  37.     embedding_delta = torch.matmul(stacked_qkv_attention, w_layer.T)
  38.     embedding_after_edit = final_embedding + embedding_delta
  39.     embedding_after_edit_normalized = rms_norm(embedding_after_edit, model[f"layers.{layer}.ffn_norm.weight"])
  40.     w1 = model[f"layers.{layer}.feed_forward.w1.weight"]
  41.     w2 = model[f"layers.{layer}.feed_forward.w2.weight"]
  42.     w3 = model[f"layers.{layer}.feed_forward.w3.weight"]
  43.     output_after_feedforward = torch.matmul(torch.functional.F.silu(torch.matmul(embedding_after_edit_normalized, w1.T)) * torch.matmul(embedding_after_edit_normalized, w3.T), w2.T)
  44.     final_embedding = embedding_after_edit+output_after_feedforward

4、输出层

4.1、将嵌入向量转换为预测的token：确认logits张量的形状

简介

生成最终的嵌入向量，用于预测下一个token的最优预测。这个嵌入的形状与常规的token嵌入相同，为[17x4096]，其中17是token的数量，4096是嵌入的维度。

思路步骤

确认logits张量的形状：确保计算得到的logits具有正确的维度，以便后续处理。

核心技术点

线性层与softmax激活函数

  1. final_embedding = rms_norm(final_embedding, model["norm.weight"])
  2. final_embedding.shape
  3. torch.Size([17, 4096])

4.2、解码：将模型输出的词元转换为可读的文本

简介

我们将使用输出解码器，将最终嵌入转换为token值，使用输出解码器将最终的嵌入转换成一个token。确认输出解码器的尺寸是否正确，确保矩阵乘法可以顺利进行。

使用最后一个词元的嵌入通过矩阵乘法预测下一个令牌的概率分布。

通过最后一个令牌的嵌入和输出解码器权重矩阵的乘法，计算出每个可能令牌的logits。

思路步骤

获取预测的下一个令牌的索引，通过argmax函数从logits中选择概率最高的令牌，来确定下一个令牌，得到最终的预测结果。

核心技术点

分词器的逆过程

argmax函数，分类问题中的预测。

  1. model["output.weight"].shape
  2. torch.Size([128256, 4096])

5、模型推理：应用训练好的输出层参数进行预测(验证整个代码的正确性)，模型预测下一个token的编号为2983，这是42的token ID吗？

背景

期望：使用最后一个token的嵌入来预测下一个value，希望预测的结果是42。加载的llama3模型在样本问题"答案是什么"上正确预测答案为42，验证了从零开始实现的整个流程是正确的。

希望在我们这个例子中，是42 :)

方法：查看预测的下一个令牌是否为“42”

注意：因为根据《银河系漫游指南》一书中的说法，42是“生命、宇宙及一切的终极问题的答案”。大多数现代LLM在这里都会回答42，这将验证整个代码的正确性。

简介

模型预测下一个词元是token编号2983，这是42的token ID吗？

我正在让你兴奋，这是最后一行的代码，希望你能享受阅读！

结束！

思路步骤

输出预测结果

打印或返回预测的令牌索引，完成整个模型推理过程。

  1. logits = torch.matmul(final_embedding[-1], model["output.weight"].T)
  2. logits.shape
  3. torch.Size([128256])
  4.  
  5. next_token = torch.argmax(logits, dim=-1)
  6. next_token
  7. tensor(2983)
  8.  
  9.  
  10. tokenizer.decode([next_token.item()])
  11. '42'

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/你好赵伟/article/detail/653926
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号