手撕Llama3第1层: 从零开始实现llama3

一、Llama3的架构

在本系列文章中，我们从头开始实现llama3。

Llama3的整体架构：

图片

Llama3的模型参数：

让我们来看看这些参数在LlaMa 3模型中的实际数值。

图片

[1] 上下文窗口（context-window）

在实例化LlaMa类时，变量max_seq_len定义了context-window。类中还有其他参数，但这个参数与transformer模型的关系最为直接。这里的max_seq_len是8K。

图片

[2] 词汇量（Vocabulary-size）和注意力层（Attention Layers）

Transformer类是一种定义了词汇量和层数的模型。这里的词汇量是指模型能够识别和处理的单词（和标记）集合。Attention layers指的是模型中使用的transformer block（attention和feed-forward layers的组合）。

图片

根据这些数字，LlaMa 3的词汇量为128K，这是相当大的。此外，它有32个transformer block。

[3] 特征维度（Feature-dimension）和注意力头（Attention-Heads）

特征维度和attention-heads被引入到Self-Attention模块中。Feature dimension指的是嵌入空间中tokens的向量大小（特征维度是指输入数据或嵌入向量的维度大小），而attention-heads包括驱动transformers中self-attention机制的QK-module。

图片

[4] 隐藏维度（Hidden Dimensions）

隐藏维度是指在前馈神经网络（Feed Forward）中，隐藏层的维度大小。前馈神经网络通常包含一个或多个隐藏层，这些隐藏层的维度决定了网络的容量和复杂度。在Transformer模型中，前馈神经网络的隐藏层维度通常是特征维度的某个倍数，以增加模型的表示能力。LLama3中，隐藏维度是特征维度的1.3倍。需要注意的是，隐藏层和隐藏维度是两个概念。

更多的隐藏层数量允许网络在将它们投射回较小的输出维度之前，内部创建和操纵更丰富的表示。

图片

[5] 将上述参数组合成Transformer

第一个矩阵是输入特征矩阵，通过Attention layer处理生成Attention Weighted features。在这幅图像中，输入特征矩阵只有5 x 3的大小，但在真实的Llama 3模型中，它增长到了8K x 4096，这是巨大的。

接下来是Feed-Forward Network中的隐藏层，增长到5325，然后在最后一层回落到4096。

图片

[6] Transformer block的多层

LlaMa 3结合了上述32个transformer block，输出从一个block传递到下一个block，直到达到最后一个。

图片

[7] 把所有这些放在一起

一旦我们启动了所有上述部分，就是时候把它们整合在一起，看看它们是如何产生LlaMa效果的。

图片

步骤1：首先我们有我们的输入矩阵，大小为8K（context-window）x 128K（vocabulary-size）。这个矩阵经过嵌入处理，将这个高维矩阵转换为低维。

步骤2：在这种情况下，这个低维结果变为4096，这是我们之前看到的LlaMa模型中特征的指定维度。

在神经网络中，升维和降维都是常见的操作，它们各自有不同的目的和效果。

升维通常是为了增加模型的容量，使其能够捕捉更复杂的特征和模式。当输入数据被映射到一个更高维度的空间时，不同的特征组合可以被模型更容易地区分。这在处理非线性问题时尤其有用，因为它可以帮助模型学习到更复杂的决策边界  。

降维则是为了减少模型的复杂性和过拟合的风险。通过减少特征空间的维度，模型可以被迫学习更加精炼和泛化的特征表示。此外，降维可以作为一种正则化手段，有助于提高模型的泛化能力。在某些情况下，降维还可以减少计算成本和提高模型的运行效率 。

在实际应用中，升维后再降维的策略可以被视为一种特征提取和变换的过程。在这个过程中，模型首先通过增加维度来探索数据的内在结构，然后通过降维来提取最有用的特征和模式。这种方法可以帮助模型在保持足够复杂性的同时，避免过度拟合训练数据  。

步骤3：这个特征通过Transformer block进行处理，首先由Attention layer处理，然后是FFN layer。Attention layer横向跨特征处理，而FFN layer则纵向跨维度处理。

步骤4：步骤3为Transformer block的32层重复。最终，结果矩阵的维度与用于特征维度的维度相同。

步骤5：最后，这个矩阵被转换回原始的词汇矩阵大小，即128K，以便模型可以选择并映射词汇中可用的单词。

这就是LlaMa 3在那些基准测试中取得高分并创造LlaMa 3效应的方式。

我们将容易搞混的几个术语用简短的语言总结一下：

1. max_seq_len (最大序列长度)

这是模型在单次处理时能够接受的最大token数。

在LlaMa 3-8B模型中，这个参数设定为8,000个tokens，即Context Window Size = 8K。这意味着模型在单次处理时可以考虑的最大token数量为8,000。这对于理解长文本或保持长期对话上下文非常关键。

2. Vocabulary-size (词汇量)

这是模型能识别的所有不同token的数量。这包括所有可能的单词、标点符号和特殊字符。模型的词汇量是128,000，表示为Vocabulary-size = 128K。这意味着模型能够识别和处理128,000种不同的tokens，这些tokens包括各种单词、标点符号和特殊字符。

3. Attention Layers (注意力层)

Transformer模型中的一个主要组件。它主要负责通过学习输入数据中哪些部分最重要（即“注意”哪些token）来处理输入数据。一个模型可能有多个这样的层，每层都试图从不同的角度理解输入数据。

LlaMa 3-8B模型包含32个处理层，即Number of Layers = 32。这些层包括多个Attention Layers及其他类型的网络层，每层都从不同角度处理和理解输入数据。

4. transformer block 

包含多个不同层的模块，通常至少包括一个Attention Layer和一个Feed-Forward Network（前馈网络）。一个模型可以有多个transformer block，这些block顺序连接，每个block的输出都是下一个block的输入。也可以称transformer block为decoder layer。 

在Transformer模型的语境中，通常我们说模型有“32层”，这可以等同于说模型有“32个Transformer blocks”。每个Transformer block通常包含一个自注意力层和一个前馈神经网络层，这两个子层共同构成了一个完整的处理单元或“层”。

因此，当我们说模型有32个Transformer blocks时，实际上是在描述这个模型由32个这样的处理单元组成，每个单元都有能力进行数据的自注意力处理和前馈网络处理。这种表述方式强调了模型的层级结构和其在每个层级上的处理能力。

总结来说，”32层”和”32个Transformer blocks”在描述Transformer模型结构时基本是同义的，都指模型包含32次独立的数据处理周期，每个周期都包括自注意力和前馈网络操作。

5. Feature-dimension (特征维度)

这是输入token在模型中表示为向量时，每个向量的维度。

每个token在模型中被转换成一个含4096个特征的向量，即Feature-dimension = 4096。这个高维度使得模型能够捕捉更丰富的语义信息和上下文关系。

6. Attention-Heads (注意力头)

在每个Attention Layer中，可以有多个Attention-Heads，每个head独立地从不同的视角分析输入数据。

每个Attention Layer包含32个独立的Attention Heads，即Number of Attention Heads = 32。这些heads分别从不同的方面分析输入数据，共同提供更全面的数据解析能力。

7. Hidden Dimensions (隐藏维度)

这通常指的是在Feed-Forward Network中的层的宽度，即每层的神经元数量。通常，Hidden Dimensions会大于Feature-dimension，这允许模型在内部创建更丰富的数据表示。

在Feed-Forward Networks中，隐藏层的维度为5325，即Hidden Dimensions = 5325。这比特征维度大，允许模型在内部层之间进行更深层次的特征转换和学习。

关系和数值：

Attention Layers 和 Attention-Heads 的关系：每个Attention Layer可以包含多个Attention-Heads。

数值关系：一个模型可能有多个transformer blocks，每个block包含一个Attention Layer和一个或多个其他层。每个Attention Layer可能有多个Attention-Heads。这样，整个模型就在不同层和heads中进行复杂的数据处理。

下载Llama3模型的官方链接脚本：https://llama.meta.com/llama-downloads/ 

二、查看模型

下面这段代码展示了如何使用tiktoken库来加载和使用一个基于Byte Pair Encoding (BPE) 的分词器。这个分词器是为了处理文本数据，特别是在自然语言处理和机器学习模型中使用。

我们输入hello world,看分词器如何进行分词。

from pathlib import Pathimport tiktokenfrom tiktoken.load import load_tiktoken_bpeimport torchimport jsonimport matplotlib.pyplot as plttokenizer_path = "Meta-Llama-3-8B/tokenizer.model"special_tokens = ["","","","","","","","","","",# end of turn] + [f"" for i in range(5, 256 - 5)]mergeable_ranks = load_tiktoken_bpe(tokenizer_path)tokenizer = tiktoken.Encoding(name=Path(tokenizer_path).name,pat_str=r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^rnp{L}p{N}]?p{L}+|p{N}{1,3}| ?[^sp{L}p{N}]+[rn]*|s*[rn]+|s+(?!S)|s+",mergeable_ranks=mergeable_ranks,special_tokens={token: len(mergeable_ranks) + i for i, token in enumerate(special_tokens)},)tokenizer.decode(tokenizer.encode("hello world!"))

登录后复制

图片

读取模型文件

查看加载的模型文件中包含的前20个参数或权重的名称。

model = torch.load("Meta-Llama-3-8B/consolidated.00.pth")print(json.dumps(list(model.keys())[:20], indent=4))

登录后复制

图片

1. “tok_embeddings.weight”：这表示模型有一个词嵌入层，用于将输入的单词（或者更一般的，token）转换为固定维度的向量。这是大多数自然语言处理模型的第一步。
2. “layers.0.attention…” 和 “layers.1.attention…”：这些参数表示多个层中，每层都包含一个注意力机制模块。在这个模块中，wq、wk、wv、wo分别代表查询（Query）、键（Key）、值（Value）和输出（Output）的权重矩阵。这是Transformer模型的核心组成部分，用于捕捉输入序列中不同部分之间的关系。
3. “layers.0.feed_forward…” 和 “layers.1.feed_forward…”：这些参数表示每个层还包含一个前馈网络（Feed Forward Network），它通常由两个线性变换组成，中间有一个非线性激活函数。w1、w2、w3可能代表这个前馈网络中的不同线性层的权重。
4. “layers.0.attention_norm.weight” 和 “layers.1.attention_norm.weight”：这些参数表示每个层中的注意力模块后面有一个归一化层（可能是Layer Normalization），用于稳定训练过程。
5. “layers.0.ffn_norm.weight” 和 “layers.1.ffn_norm.weight”：这些参数表示前馈网络后面也有一个归一化层。上面代码输出内容，与下图相同，也就是Llama3中的一个transformer block。

图片

总的来说，这个输出结果揭示了一个基于Transformer架构的深度学习模型的关键组成部分。这种模型广泛用于自然语言处理任务，如文本分类、机器翻译、问答系统等。每一层的结构几乎相同，包括注意力机制、前馈网络和归一化层，这有助于模型捕捉复杂的输入序列特征。

查看Llama3模型的参数配置：

with open("Meta-Llama-3-8B/params.json", "r") as f:config = json.load(f)config

登录后复制

图片

1. ‘dim’: 4096 – 表示模型中的隐藏层维度或特征维度。这是模型处理数据时每个向量的大小。
2. ‘n_layers’: 32 – 表示模型中层的数量。在基于Transformer的模型中，这通常指的是编码器和解码器中的层的数量。
3. ‘n_heads’: 32 – 表示在自注意力（Self-Attention）机制中，头（head）的数量。多头注意力机制是Transformer模型的关键特性之一，它允许模型在不同的表示子空间中并行捕获信息。
4. ‘n_kv_heads’: 8 – 这个参数不是标准Transformer模型的常见配置，可能指的是在某些特定的注意力机制中，用于键（Key）和值（Value）的头的数量。
5. ‘vocab_size’: 128256 – 表示模型使用的词汇表大小。这是模型能够识别的不同单词或标记的总数。
6. ‘multiple_of’: 1024 – 这可能是指模型的某些维度需要是1024的倍数，以确保模型结构的对齐或优化。
7. ‘ffn_dim_multiplier’: 1.3 – 表示前馈网络（Feed-Forward Network, FFN）的维度乘数。在Transformer模型中，FFN是每个注意力层后的一个网络，这个乘数可能用于调整FFN的大小。
8. ‘norm_eps’: 1e-05 – 表示在归一化层（如Layer Normalization）中使用的epsilon值，用于防止除以零的错误。这是数值稳定性的一个小技巧。
9. ‘rope_theta’: 500000.0 – 这个参数不是标准Transformer模型的常见配置，可能是指某种特定于模型的技术或优化的参数。它可能与位置编码或某种正则化技术有关。

我们使用这个配置来推断模型的细节，比如：

1. 模型有32个Transformer层
2. 每个多头注意力块有32个头
3. 词汇表的大小等等 

dim = config["dim"]n_layers = config["n_layers"]n_heads = config["n_heads"]n_kv_heads = config["n_kv_heads"]vocab_size = config["vocab_size"]multiple_of = config["multiple_of"]ffn_dim_multiplier = config["ffn_dim_multiplier"]norm_eps = config["norm_eps"]rope_theta = torch.tensor(config["rope_theta"])

登录后复制

图片

将Text转化为Token

代码如下：

prompt = "the answer to the ultimate question of life, the universe, and everything is "tokens = [128000] + tokenizer.encode(prompt)print(tokens)tokens = torch.tensor(tokens)prompt_split_as_tokens = [tokenizer.decode([token.item()]) for token in tokens]print(prompt_split_as_tokens)

登录后复制

[128000, 1820, 4320, 311, 279, 17139, 3488, 315, 2324, 11, 279, 15861, 11, 323, 4395, 374, 220][”, ‘the’, ‘ answer’, ‘ to’, ‘ the’, ‘ ultimate’, ‘ question’, ‘ of’, ‘ life’, ‘,’, ‘ the’, ‘ universe’, ‘,’, ‘ and’, ‘ everything’, ‘ is’, ‘ ‘]

将令牌转换为它们的嵌入表示

截止到目前，我们的[17×1]令牌现在变成了[17×4096]，即长度为4096的17个嵌入（每个令牌一个）。

下图是为了验证我们输入的这句话，是17个token。

图片

代码如下：

embedding_layer = torch.nn.Embedding(vocab_size, dim)embedding_layer.weight.data.copy_(model["tok_embeddings.weight"])token_embeddings_unnormalized = embedding_layer(tokens).to(torch.bfloat16)token_embeddings_unnormalized.shape

登录后复制

图片

三、构建Transformer的第一层

我们接着使用 RMS 归一化对嵌入进行归一化，也就是图中这个位置：

图片

使用公式如下：

图片

代码如下：

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

登录后复制

这段代码定义了一个名为 rms_norm 的函数，它实现了对输入张量（tensor）的RMS（Root Mean Square，均方根）归一化处理。这个函数接受两个参数：tensor 和 norm_weights。tensor 是需要进行归一化处理的输入张量，而 norm_weights 是归一化时使用的权重。

函数的工作原理如下：

1. 首先，计算输入张量每个元素的平方（tensor.pow(2)）。
2. 然后，对平方后的张量沿着最后一个维度（-1）计算均值（mean），并保持维度不变（keepdim=True），这样得到每个元素的均方值。
3. 接着，将均方值加上一个很小的正数 norm_eps（为了避免除以零的情况），然后计算其平方根的倒数（torch.rsqrt），得到RMS的倒数。
4. 最后，将输入张量与RMS的倒数相乘，再乘以归一化权重 norm_weights，得到归一化后的张量。

在进行归一化处理后，我们的数据形状仍然保持为 [17×4096]，这与嵌入层的形状相同，只不过数据已经过归一化。

token_embeddings = rms_norm(token_embeddings_unnormalized, model["layers.0.attention_norm.weight"])token_embeddings.shape

登录后复制

图片

图片

接下来，我们介绍注意力机制的实现，也就是下图中的红框标注的位置：

图片

图片

我们一步一步地解释这张图，详细说明每个步骤。

1. 输入句子

• 描述：这是我们的输入句子。
• 解释：输入句子被表示为一个矩阵 ( X )，其中每一行代表一个词的嵌入向量。

2. 嵌入每个词

• 描述：我们对每个词进行嵌入。
• 解释：输入句子中的每个词被转换为一个高维向量，这些向量组成了矩阵 ( X )。

3. 分成8个头

• 描述：将矩阵 ( X ) 分成8个头。我们用权重矩阵 ( W^Q )、( W^K ) 和 ( W^V ) 分别乘以 ( X )。
• 解释：多头注意力机制将输入矩阵 ( X ) 分成多个头（这里是8个），每个头有自己的查询（Query）、键（Key）和值（Value）矩阵。具体来说，输入矩阵 ( X ) 分别与查询权重矩阵 ( W^Q )、键权重矩阵 ( W^K ) 和值权重矩阵 ( W^V ) 相乘，得到查询矩阵 ( Q )、键矩阵 ( K ) 和值矩阵 ( V )。

4. 计算注意力

• 描述：使用得到的查询、键和值矩阵计算注意力。
• 解释：对于每个头，使用查询矩阵 ( Q )、键矩阵 ( K ) 和值矩阵 ( V ) 计算注意力分数。具体步骤包括：

计算 ( Q ) 和 ( K ) 的点积。

对点积结果进行缩放。

应用softmax函数得到注意力权重。

用注意力权重乘以值矩阵 ( V ) 得到输出矩阵 ( Z )。

5. 拼接结果矩阵

• 描述：将得到的 ( Z ) 矩阵拼接起来，然后用权重矩阵 ( W^O ) 乘以拼接后的矩阵，得到层的输出。
• 解释：将所有头的输出矩阵 ( Z ) 拼接成一个矩阵，然后用输出权重矩阵 ( W^O ) 乘以这个拼接后的矩阵，得到最终的输出矩阵 ( Z )。

额外说明

• 查询、键、值和输出向量的形状：在加载查询、键、值和输出向量时，注意到它们的形状分别是 [4096×4096]、[1024×4096]、[1024×4096]、[1024×4096] 和 [4096×4096]。
• 并行化注意力头的乘法：将它们捆绑在一起有助于并行化注意力头的乘法。

这张图展示了Transformer模型中多头注意力机制的实现过程，从输入句子的嵌入开始，经过多头分割、注意力计算，最后拼接结果并生成输出。每个步骤都详细说明了如何从输入矩阵 ( X ) 生成最终的输出矩阵 ( Z )。

当我们从模型中加载查询（query）、键（key）、值（value）和输出（output）向量时，我们注意到它们的形状分别是 [4096×4096]、[1024×4096]、[1024×4096]、[4096×4096]

乍一看这很奇怪，因为理想情况下我们希望每个头的每个q、k、v和o都是单独的

print(model["layers.0.attention.wq.weight"].shape,model["layers.0.attention.wk.weight"].shape,model["layers.0.attention.wv.weight"].shape,model["layers.0.attention.wo.weight"].shape)

登录后复制

图片

查询（Query）权重矩阵 (wq.weight) 的形状是 [4096, 4096]。键（Key）权重矩阵 (wk.weight) 的形状是 [1024, 4096]。值（Value）权重矩阵 (wv.weight) 的形状是 [1024, 4096]。输出（Output）权重矩阵 (wo.weight) 的形状是 [4096, 4096]。输出结果表明：查询（Q）和输出（O）权重矩阵的形状是相同的，都是[4096, 4096]。这意味着对于查询和输出，输入特征和输出特征的维度都是4096。键（K）和值（V）权重矩阵的形状也是相同的，都是[1024, 4096]。这表明键和值的输入特征维度为4096，但输出特征维度被压缩到了1024。这些权重矩阵的形状反映了模型设计者如何设置注意力机制中不同部分的维度。特别是，键和值的维度被减小可能是为了减少计算复杂度和内存消耗，而保持查询和输出的较高维度可能是为了保留更多的信息。这种设计选择依赖于特定的模型架构和应用场景
让我们用“我欣赏李鸿章”这个句子作为例子，来简化解释这个图中的注意力机制的实现过程。输入句子：首先，我们有句子“我欣赏李鸿章”。在处理这个句子之前，我们需要将句子中的每个词转换成数学上可以处理的形式，即词向量。这个过程叫做词嵌入（embedding）。
词嵌入：每个词，比如“我”、“欣赏”、“李鸿章”，都会被转换成一个固定大小的向量。这些向量包含了词的语义信息。
分割成多个头：为了让模型能够从不同的角度理解句子，我们将每个词的向量分割成多个部分，这里是8个头。每个头都会关注句子的不同方面。
计算注意力：对于每个头，我们都会计算一个叫做注意力的东西。这个过程涉及到三个步骤：以“我欣赏李鸿章”为例，如果我们想要关注“欣赏”这个词，那么“欣赏”就是查询，而其他词比如“我”和“李鸿章”就是键，它们的向量就是值。
查询（Q）：这是我们想要寻找信息的部分。键（K）：这是包含信息的部分。值（V）：这是实际的信息内容。拼接和输出：计算完每个头的注意力之后，我们将这些结果拼接起来，并通过一个权重矩阵Wo来生成最终的输出。这个输出将被用于下一层的处理或者作为最终结果的一部分。

在图中的注释中提到的形状问题，是关于如何在计算机中有效地存储和处理这些向量的问题。在实际的代码实现中，为了提高效率，开发者可能会将多个头的查询、键、值向量打包在一起处理，而不是单独处理每个头。这样可以利用现代计算机的并行处理能力，加快计算速度。

• 查询（Query）权重矩阵 (wq.weight) 的形状是 [4096, 4096]。
• 键（Key）权重矩阵 (wk.weight) 的形状是 [1024, 4096]。
• 值（Value）权重矩阵 (wv.weight) 的形状是 [1024, 4096]。
• 输出（Output）权重矩阵 (wo.weight) 的形状是 [4096, 4096]。

输出结果表明：

• 查询（Q）和输出（O）权重矩阵的形状是相同的，都是[4096, 4096]。这意味着对于查询和输出，输入特征和输出特征的维度都是4096。
• 键（K）和值（V）权重矩阵的形状也是相同的，都是[1024, 4096]。这表明键和值的输入特征维度为4096，但输出特征维度被压缩到了1024。

这些权重矩阵的形状反映了模型设计者如何设置注意力机制中不同部分的维度。特别是，键和值的维度被减小可能是为了减少计算复杂度和内存消耗，而保持查询和输出的较高维度可能是为了保留更多的信息。这种设计选择依赖于特定的模型架构和应用场景

让我们用“我欣赏李鸿章”这个句子作为例子，来简化解释这个图中的注意力机制的实现过程。

• 输入句子：首先，我们有句子“我欣赏李鸿章”。在处理这个句子之前，我们需要将句子中的每个词转换成数学上可以处理的形式，即词向量。这个过程叫做词嵌入（embedding）。
• 词嵌入：每个词，比如“我”、“欣赏”、“李鸿章”，都会被转换成一个固定大小的向量。这些向量包含了词的语义信息。
• 分割成多个头：为了让模型能够从不同的角度理解句子，我们将每个词的向量分割成多个部分，这里是8个头。每个头都会关注句子的不同方面。
• 计算注意力：对于每个头，我们都会计算一个叫做注意力的东西。这个过程涉及到三个步骤：以“我欣赏李鸿章”为例，如果我们想要关注“欣赏”这个词，那么“欣赏”就是查询，而其他词比如“我”和“李鸿章”就是键，它们的向量就是值。

      查询（Q）：这是我们想要寻找信息的部分。

      键（K）：这是包含信息的部分。

      值（V）：这是实际的信息内容。

• 拼接和输出：计算完每个头的注意力之后，我们将这些结果拼接起来，并通过一个权重矩阵Wo来生成最终的输出。这个输出将被用于下一层的处理或者作为最终结果的一部分。

在图中的注释中提到的形状问题，是关于如何在计算机中有效地存储和处理这些向量的问题。在实际的代码实现中，为了提高效率，开发者可能会将多个头的查询、键、值向量打包在一起处理，而不是单独处理每个头。这样可以利用现代计算机的并行处理能力，加快计算速度。

我们继续使用句子“我欣赏李鸿章”来解释WQ、WK、WV和WO这些权重矩阵的作用。

在Transformer模型中，每个词都会通过词嵌入转换成一个向量。这些向量接下来会通过一系列的线性变换来计算注意力分数。这些线性变换就是通过权重矩阵WQ、WK、WV和WO来实现的。

1. WQ（权重矩阵Q）：这个矩阵用于将每个词的向量转换成“查询（Query）”向量。在我们的例子中，如果我们想要关注“欣赏”这个词，我们会将“欣赏”的向量乘以WQ来得到查询向量。
2. WK（权重矩阵K）：这个矩阵用于将每个词的向量转换成“键（Key）”向量。同样地，我们会将每个词，包括“我”和“李鸿章”，的向量乘以WK来得到键向量。
3. WV（权重矩阵V）：这个矩阵用于将每个词的向量转换成“值（Value）”向量。每个词的向量乘以WV后，我们得到的是值向量。这三个矩阵（WQ、WK、WV）是用来为每个头生成不同的查询、键和值向量的。这样做可以让每个头关注句子的不同方面。
4. WQ（权重矩阵Q）、WK（权重矩阵K）、WV（权重矩阵V）和WO（权重矩阵O）这些矩阵是Transformer模型中的参数，它们是在模型训练过程中通过反向传播算法和梯度下降等优化方法学习得到的。

在整个过程中，WQ、WK、WV和WO是通过训练学习得到的，它们决定了模型如何将输入的词向量转换成不同的表示，以及如何组合这些表示来得到最终的输出。这些矩阵是Transformer模型中注意力机制的核心部分，它们使得模型能够捕捉到句子中不同词之间的关系。

WQ（权重矩阵Q）、WK（权重矩阵K）、WV（权重矩阵V）和WO（权重矩阵O）这些矩阵是Transformer模型中的参数，它们是在模型训练过程中通过反向传播算法和梯度下降等优化方法学习得到的。

让我们来看看这个学习过程是如何进行的：

1. 初始化：在训练开始之前，这些矩阵通常会被随机初始化。这意味着它们的初始值是随机选取的，这样可以打破对称性并开始学习过程。
2. 前向传播：在模型的训练过程中，输入数据（如句子“我欣赏李鸿章”）会通过模型的各个层进行前向传播。在注意力机制中，输入的词向量会与WQ、WK、WV矩阵相乘，以生成查询、键和值向量。
3. 计算损失：模型的输出会与期望的输出（通常是训练数据中的标签）进行比较，计算出一个损失值。这个损失值衡量了模型的预测与实际情况的差距。
4. 反向传播：损失值会通过反向传播算法传回模型，计算每个参数（包括WQ、WK、WV和WO）对损失的影响，即它们的梯度。
5. 参数更新：根据计算出的梯度，使用梯度下降或其他优化算法来更新这些矩阵的值。这个过程会逐渐减小损失值，使模型的预测更加准确。
6. 迭代过程：这个前向传播、损失计算、反向传播和参数更新的过程会在训练数据上多次迭代进行，直到模型的性能达到一定的标准或者不再显著提升。

   通过这个训练过程，WQ、WK、WV和WO这些矩阵会逐渐调整它们的值，以便模型能够更好地理解和处理输入数据。在训练完成后，这些矩阵将固定下来，用于模型的推理阶段，即对新的输入数据进行预测。

四、展开查询向量

在本小节中，我们将从多个注意力头中展开查询向量，得到的形状是 [32x128x4096] 这里，32 是 llama3 中注意力头的数量，128 是查询向量的大小，而 4096 是令牌嵌入的大小。

q_layer0 = model["layers.0.attention.wq.weight"]head_dim = q_layer0.shape[0] // n_headsq_layer0 = q_layer0.view(n_heads, head_dim, dim)q_layer0.shape

登录后复制

图片

这段代码通过对模型中第一层的查询（Q）权重矩阵进行重塑（reshape），将其分解为多个注意力头的形式，从而揭示了32和128这两个维度。

1. q_layer0 = model[“layers.0.attention.wq.weight”]：这行代码从模型中提取第一层的查询（Q）权重矩阵。
2. head_dim = q_layer0.shape[0] // n_heads：这行代码计算每个注意力头的维度大小。它通过将查询权重矩阵的第一个维度（原本是4096）除以注意力头的数量（n_heads），得到每个头的维度。如果n_heads是32（即模型设计为有32个注意力头），那么head_dim就是4096 // 32 = 128。
3. q_layer0 = q_layer0.view(n_heads, head_dim, dim)：这行代码使用.view()方法重塑查询权重矩阵，使其形状变为[n_heads, head_dim, dim]。这里dim很可能是原始特征维度4096，n_heads是32，head_dim是128，因此重塑后的形状是[32, 128, 4096]。
4. q_layer0.shape 输出：torch.Size([32, 128, 4096])：这行代码打印重塑后的查询权重矩阵的形状，确认了其形状为[32, 128, 4096]。

之所以在这段代码中出现了32和128这两个维度，而在之前的代码段中没有，是因为这段代码通过重塑操作明确地将查询权重矩阵分解为多个注意力头，每个头具有自己的维度。32代表了模型中注意力头的数量，而128代表了分配给每个头的特征维度大小。这种分解是为了实现多头注意力机制，其中每个头可以独立地关注输入的不同部分，最终通过组合这些头的输出来提高模型的表达能力。 

实现第一层的第一个头

访问了第一层第一个头的查询（query）权重矩阵，这个查询权重矩阵的大小是 [128×4096]。

q_layer0_head0 = q_layer0[0]q_layer0_head0.shape

登录后复制

图片

我们现在将查询权重与令牌嵌入相乘，以获得令牌的查询

在这里，你可以看到结果形状是 [17×128]，这是因为我们有17个令牌，每个令牌都有一个长度为128的查询（每个令牌在一个头上方的查询）。

br

登录后复制

图片

这段代码执行了一个矩阵乘法操作，将令牌嵌入（token_embeddings）与第一层第一个头的查询（query）权重矩阵（q_layer0_head0）的转置（.T）相乘，以生成每个令牌的查询向量（q_per_token）。

1. q_per_token = torch.matmul(token_embeddings, q_layer0_head0.T)：

torch.matmul 是PyTorch中的矩阵乘法函数，它可以处理两个张量的乘法。

token_embeddings 应该是一个形状为 [17, 4096] 的张量，表示有17个令牌，每个令牌由4096维的嵌入向量表示。

q_layer0_head0 是第一层第一个头的查询权重矩阵，其原始形状为 [128, 4096]。.T 是PyTorch中的转置操作，将 q_layer0_head0 的形状转置为 [4096, 128]。

这样，token_embeddings 和 q_layer0_head0.T 的矩阵乘法就是 [17, 4096] 和 [4096, 128] 的乘法，结果是一个形状为 [17, 128] 的张量。

2. q_per_token.shape 和输出：torch.Size([17, 128])：

这行代码打印出 q_per_token 张量的形状，确认其为 [17, 128]。

这意味着对于输入的每个令牌（共17个），我们现在都有了一个128维的查询向量。这128维的查询向量是通过将令牌嵌入与查询权重矩阵相乘得到的，可以用于后续的注意力机制计算。

总之，这段代码通过矩阵乘法将每个令牌的嵌入向量转换为查询向量，为实现注意力机制的下一步做准备。每个令牌现在都有了一个与之对应的查询向量，这些查询向量将用于计算与其他令牌的注意力得分。

以上就是手撕Llama3第1层: 从零开始实现llama3的详细内容，更多请关注叮当号网其它相关文章！