手撕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 = ["<|begin_of_text|>","<|end_of_text|>","<|reserved_special_token_0|>","<|reserved_special_token_1|>","<|reserved_special_token_2|>","<|reserved_special_token_3|>","<|start_header_id|>","<|end_header_id|>","<|reserved_special_token_4|>","<|eot_id|>",# end of turn] + [f"<|reserved_special_token_{i}|>" 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)|[^\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+",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]['<|begin_of_text|>', 'the', ' answer', ' to', ' the', ' ultimate', ' question', ' of', ' life', ',', ' the', ' universe', ',', ' and', ' everything', ' is', ' ']

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

截止到目前，我们的[17x1]令牌现在变成了[17x4096]，即长度为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，得到归一化后的张量。

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

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 )。

额外说明

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

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

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

乍一看这很奇怪，因为理想情况下我们希望每个头的每个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）权重矩阵，这个查询权重矩阵的大小是 [128x4096]。

q_layer0_head0 = q_layer0[0]q_layer0_head0.shape

图片

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

在这里，你可以看到结果形状是 [17x128]，这是因为我们有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的詳細內容。更多資訊請關注PHP中文網其他相關文章！