手撕Llama3第1層：從零開始實現(xiàn)Llama3

作者：魏新宇 2024-05-27 09:00:00

人工智能

通過矩陣乘法將每個令牌的嵌入向量轉(zhuǎn)換為查詢向量，為實現(xiàn)注意力機制的下一步做準備。每個令牌現(xiàn)在都有了一個與之對應的查詢向量，這些查詢向量將用于計算與其他令牌的注意力得分。?

一、Llama3的架構(gòu)

在本系列文章中，我們從頭開始實現(xiàn)llama3。

Llama3的整體架構(gòu)：

圖片

Llama3的模型參數(shù)：

讓我們來看看這些參數(shù)在LlaMa 3模型中的實際數(shù)值。

圖片

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

在實例化LlaMa類時，變量max_seq_len定義了context-window。類中還有其他參數(shù)，但這個參數(shù)與transformer模型的關系最為直接。這里的max_seq_len是8K。

圖片

[2] 詞匯量（Vocabulary-size）和注意力層（Attention Layers）

接下來是Transformer類，它定義了詞匯量和層數(shù)。這里的詞匯量是指模型能夠識別和處理的單詞（和tokens）集。Attention layers指的是模型中使用的transformer block（attention和feed-forward layers的組合）。

圖片

根據(jù)這些數(shù)字，LlaMa 3的詞匯量為128K，這是相當大的。此外，它有32個transformer block。

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

特征維度和attention-heads被引入到Self-Attention模塊中。Feature dimension指的是嵌入空間中tokens的向量大小（特征維度是指輸入數(shù)據(jù)或嵌入向量的維度大小），而attention-heads包括驅(qū)動transformers中self-attention機制的QK-module。

圖片

[4] 隱藏維度（Hidden Dimensions）

隱藏維度是指在前饋神經(jīng)網(wǎng)絡（Feed Forward）中，隱藏層的維度大小。前饋神經(jīng)網(wǎng)絡通常包含一個或多個隱藏層，這些隱藏層的維度決定了網(wǎng)絡的容量和復雜度。在Transformer模型中，前饋神經(jīng)網(wǎng)絡的隱藏層維度通常是特征維度的某個倍數(shù)，以增加模型的表示能力。LLama3中，隱藏維度是特征維度的1.3倍。需要注意的是，隱藏層和隱藏維度是兩個概念。

更多的隱藏層數(shù)量允許網(wǎng)絡在將它們投射回較小的輸出維度之前，內(nèi)部創(chuàng)建和操縱更豐富的表示。

圖片

[5] 將上述參數(shù)組合成Transformer

第一個矩陣是輸入特征矩陣，通過Attention layer處理生成Attention Weighted features。在這幅圖像中，輸入特征矩陣只有5 x 3的大小，但在真實的Llama 3模型中，它增長到了8K x 4096，這是巨大的。

接下來是Feed-Forward Network中的隱藏層，增長到5325，然后在最后一層回落到4096。

圖片

[6] Transformer block的多層

LlaMa 3結(jié)合了上述32個transformer block，輸出從一個block傳遞到下一個block，直到達到最后一個。

圖片

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

一旦我們啟動了所有上述部分，就是時候把它們整合在一起，看看它們是如何產(chǎn)生LlaMa效果的。

圖片

步驟1：首先我們有我們的輸入矩陣，大小為8K（context-window）x 128K（vocabulary-size）。這個矩陣經(jīng)過嵌入處理，將這個高維矩陣轉(zhuǎn)換為低維。

步驟2：在這種情況下，這個低維結(jié)果變?yōu)?096，這是我們之前看到的LlaMa模型中特征的指定維度。

在神經(jīng)網(wǎng)絡中，升維和降維都是常見的操作，它們各自有不同的目的和效果。

升維通常是為了增加模型的容量，使其能夠捕捉更復雜的特征和模式。當輸入數(shù)據(jù)被映射到一個更高維度的空間時，不同的特征組合可以被模型更容易地區(qū)分。這在處理非線性問題時尤其有用，因為它可以幫助模型學習到更復雜的決策邊界。

降維則是為了減少模型的復雜性和過擬合的風險。通過減少特征空間的維度，模型可以被迫學習更加精煉和泛化的特征表示。此外，降維可以作為一種正則化手段，有助于提高模型的泛化能力。在某些情況下，降維還可以減少計算成本和提高模型的運行效率。

在實際應用中，升維后再降維的策略可以被視為一種特征提取和變換的過程。在這個過程中，模型首先通過增加維度來探索數(shù)據(jù)的內(nèi)在結(jié)構(gòu)，然后通過降維來提取最有用的特征和模式。這種方法可以幫助模型在保持足夠復雜性的同時，避免過度擬合訓練數(shù)據(jù) 。

步驟3：這個特征通過Transformer block進行處理，首先由Attention layer處理，然后是FFN layer。Attention layer橫向跨特征處理，而FFN layer則縱向跨維度處理。

步驟4：步驟3為Transformer block的32層重復。最終，結(jié)果矩陣的維度與用于特征維度的維度相同。

步驟5：最后，這個矩陣被轉(zhuǎn)換回原始的詞匯矩陣大小，即128K，以便模型可以選擇并映射詞匯中可用的單詞。

這就是LlaMa 3在那些基準測試中取得高分并創(chuàng)造LlaMa 3效應的方式。

我們將容易搞混的幾個術(shù)語用簡短的語言總結(jié)一下：

1. max_seq_len (最大序列長度)

這是模型在單次處理時能夠接受的最大token數(shù)。

在LlaMa 3-8B模型中，這個參數(shù)設定為8,000個tokens，即Context Window Size = 8K。這意味著模型在單次處理時可以考慮的最大token數(shù)量為8,000。這對于理解長文本或保持長期對話上下文非常關鍵。

2. Vocabulary-size (詞匯量)

這是模型能識別的所有不同token的數(shù)量。這包括所有可能的單詞、標點符號和特殊字符。模型的詞匯量是128,000，表示為Vocabulary-size = 128K。這意味著模型能夠識別和處理128,000種不同的tokens，這些tokens包括各種單詞、標點符號和特殊字符。

3. Attention Layers (注意力層)

Transformer模型中的一個主要組件。它主要負責通過學習輸入數(shù)據(jù)中哪些部分最重要（即“注意”哪些token）來處理輸入數(shù)據(jù)。一個模型可能有多個這樣的層，每層都試圖從不同的角度理解輸入數(shù)據(jù)。

LlaMa 3-8B模型包含32個處理層，即Number of Layers = 32。這些層包括多個Attention Layers及其他類型的網(wǎng)絡層，每層都從不同角度處理和理解輸入數(shù)據(jù)。

4. transformer block

包含多個不同層的模塊，通常至少包括一個Attention Layer和一個Feed-Forward Network（前饋網(wǎng)絡）。一個模型可以有多個transformer block，這些block順序連接，每個block的輸出都是下一個block的輸入。也可以稱transformer block為decoder layer。

在Transformer模型的語境中，通常我們說模型有“32層”，這可以等同于說模型有“32個Transformer blocks”。每個Transformer block通常包含一個自注意力層和一個前饋神經(jīng)網(wǎng)絡層，這兩個子層共同構(gòu)成了一個完整的處理單元或“層”。

因此，當我們說模型有32個Transformer blocks時，實際上是在描述這個模型由32個這樣的處理單元組成，每個單元都有能力進行數(shù)據(jù)的自注意力處理和前饋網(wǎng)絡處理。這種表述方式強調(diào)了模型的層級結(jié)構(gòu)和其在每個層級上的處理能力。

總結(jié)來說，"32層"和"32個Transformer blocks"在描述Transformer模型結(jié)構(gòu)時基本是同義的，都指模型包含32次獨立的數(shù)據(jù)處理周期，每個周期都包括自注意力和前饋網(wǎng)絡操作。

5. Feature-dimension (特征維度)

這是輸入token在模型中表示為向量時，每個向量的維度。

每個token在模型中被轉(zhuǎn)換成一個含4096個特征的向量，即Feature-dimension = 4096。這個高維度使得模型能夠捕捉更豐富的語義信息和上下文關系。

6. Attention-Heads (注意力頭)

在每個Attention Layer中，可以有多個Attention-Heads，每個head獨立地從不同的視角分析輸入數(shù)據(jù)。

每個Attention Layer包含32個獨立的Attention Heads，即Number of Attention Heads = 32。這些heads分別從不同的方面分析輸入數(shù)據(jù)，共同提供更全面的數(shù)據(jù)解析能力。

7. Hidden Dimensions (隱藏維度)

這通常指的是在Feed-Forward Network中的層的寬度，即每層的神經(jīng)元數(shù)量。通常，Hidden Dimensions會大于Feature-dimension，這允許模型在內(nèi)部創(chuàng)建更豐富的數(shù)據(jù)表示。

在Feed-Forward Networks中，隱藏層的維度為5325，即Hidden Dimensions = 5325。這比特征維度大，允許模型在內(nèi)部層之間進行更深層次的特征轉(zhuǎn)換和學習。

關系和數(shù)值：

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

數(shù)值關系：一個模型可能有多個transformer blocks，每個block包含一個Attention Layer和一個或多個其他層。每個Attention Layer可能有多個Attention-Heads。這樣，整個模型就在不同層和heads中進行復雜的數(shù)據(jù)處理。

下載Llama3模型的官方鏈接腳本：https://llama.meta.com/llama-downloads/

二、查看模型

下面這段代碼展示了如何使用tiktoken庫來加載和使用一個基于Byte Pair Encoding (BPE) 的分詞器。這個分詞器是為了處理文本數(shù)據(jù)，特別是在自然語言處理和機器學習模型中使用。

我們輸入hello world,看分詞器如何進行分詞。

from pathlib import Path
import tiktoken
from tiktoken.load import load_tiktoken_bpe
import torch
import json
import matplotlib.pyplot as plt




tokenizer_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個參數(shù)或權(quán)重的名稱。

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

圖片

"tok_embeddings.weight"：這表示模型有一個詞嵌入層，用于將輸入的單詞（或者更一般的，token）轉(zhuǎn)換為固定維度的向量。這是大多數(shù)自然語言處理模型的第一步。
"layers.0.attention..." 和 "layers.1.attention..."：這些參數(shù)表示多個層中，每層都包含一個注意力機制模塊。在這個模塊中，wq、wk、wv、wo分別代表查詢（Query）、鍵（Key）、值（Value）和輸出（Output）的權(quán)重矩陣。這是Transformer模型的核心組成部分，用于捕捉輸入序列中不同部分之間的關系。
"layers.0.feed_forward..." 和 "layers.1.feed_forward..."：這些參數(shù)表示每個層還包含一個前饋網(wǎng)絡（Feed Forward Network），它通常由兩個線性變換組成，中間有一個非線性激活函數(shù)。w1、w2、w3可能代表這個前饋網(wǎng)絡中的不同線性層的權(quán)重。
"layers.0.attention_norm.weight" 和 "layers.1.attention_norm.weight"：這些參數(shù)表示每個層中的注意力模塊后面有一個歸一化層（可能是Layer Normalization），用于穩(wěn)定訓練過程。
"layers.0.ffn_norm.weight" 和 "layers.1.ffn_norm.weight"：這些參數(shù)表示前饋網(wǎng)絡后面也有一個歸一化層。上面代碼輸出內(nèi)容，與下圖相同，也就是Llama3中的一個transformer block。

圖片

總的來說，這個輸出結(jié)果揭示了一個基于Transformer架構(gòu)的深度學習模型的關鍵組成部分。這種模型廣泛用于自然語言處理任務，如文本分類、機器翻譯、問答系統(tǒng)等。每一層的結(jié)構(gòu)幾乎相同，包括注意力機制、前饋網(wǎng)絡和歸一化層，這有助于模型捕捉復雜的輸入序列特征。

查看Llama3模型的參數(shù)配置：

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

圖片

'dim': 4096 - 表示模型中的隱藏層維度或特征維度。這是模型處理數(shù)據(jù)時每個向量的大小。
'n_layers': 32 - 表示模型中層的數(shù)量。在基于Transformer的模型中，這通常指的是編碼器和解碼器中的層的數(shù)量。
'n_heads': 32 - 表示在自注意力（Self-Attention）機制中，頭（head）的數(shù)量。多頭注意力機制是Transformer模型的關鍵特性之一，它允許模型在不同的表示子空間中并行捕獲信息。
'n_kv_heads': 8 - 這個參數(shù)不是標準Transformer模型的常見配置，可能指的是在某些特定的注意力機制中，用于鍵（Key）和值（Value）的頭的數(shù)量。
'vocab_size': 128256 - 表示模型使用的詞匯表大小。這是模型能夠識別的不同單詞或標記的總數(shù)。
'multiple_of': 1024 - 這可能是指模型的某些維度需要是1024的倍數(shù)，以確保模型結(jié)構(gòu)的對齊或優(yōu)化。
'ffn_dim_multiplier': 1.3 - 表示前饋網(wǎng)絡（Feed-Forward Network, FFN）的維度乘數(shù)。在Transformer模型中，F(xiàn)FN是每個注意力層后的一個網(wǎng)絡，這個乘數(shù)可能用于調(diào)整FFN的大小。
'norm_eps': 1e-05 - 表示在歸一化層（如Layer Normalization）中使用的epsilon值，用于防止除以零的錯誤。這是數(shù)值穩(wěn)定性的一個小技巧。
'rope_theta': 500000.0 - 這個參數(shù)不是標準Transformer模型的常見配置，可能是指某種特定于模型的技術(shù)或優(yōu)化的參數(shù)。它可能與位置編碼或某種正則化技術(shù)有關。

我們使用這個配置來推斷模型的細節(jié)，比如：

模型有32個Transformer層
每個多頭注意力塊有32個頭
詞匯表的大小等等

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轉(zhuǎn)化為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', ' ']

將令牌轉(zhuǎn)換為它們的嵌入表示

截止到目前，我們的[17x1]令牌現(xiàn)在變成了[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

圖片

三、構(gòu)建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 的函數(shù)，它實現(xiàn)了對輸入張量（tensor）的RMS（Root Mean Square，均方根）歸一化處理。這個函數(shù)接受兩個參數(shù)：tensor 和 norm_weights。tensor 是需要進行歸一化處理的輸入張量，而 norm_weights 是歸一化時使用的權(quán)重。

函數(shù)的工作原理如下：

首先，計算輸入張量每個元素的平方（tensor.pow(2)）。
然后，對平方后的張量沿著最后一個維度（-1）計算均值（mean），并保持維度不變（keepdim=True），這樣得到每個元素的均方值。
接著，將均方值加上一個很小的正數(shù) norm_eps（為了避免除以零的情況），然后計算其平方根的倒數(shù)（torch.rsqrt），得到RMS的倒數(shù)。
最后，將輸入張量與RMS的倒數(shù)相乘，再乘以歸一化權(quán)重 norm_weights，得到歸一化后的張量。

在進行歸一化處理后，我們的數(shù)據(jù)形狀仍然保持為 [17x4096]，這與嵌入層的形狀相同，只不過數(shù)據(jù)已經(jīng)過歸一化。

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

圖片

接下來，我們介紹注意力機制的實現(xiàn)，也就是下圖中的紅框標注的位置：

圖片

我們一步一步地解釋這張圖，詳細說明每個步驟。

1. 輸入句子

描述：這是我們的輸入句子。
解釋：輸入句子被表示為一個矩陣 ( X )，其中每一行代表一個詞的嵌入向量。

2. 嵌入每個詞

描述：我們對每個詞進行嵌入。
解釋：輸入句子中的每個詞被轉(zhuǎn)換為一個高維向量，這些向量組成了矩陣 ( X )。

3. 分成8個頭

描述：將矩陣 ( X ) 分成8個頭。我們用權(quán)重矩陣 ( W^Q )、( W^K ) 和 ( W^V ) 分別乘以 ( X )。
解釋：多頭注意力機制將輸入矩陣 ( X ) 分成多個頭（這里是8個），每個頭有自己的查詢（Query）、鍵（Key）和值（Value）矩陣。具體來說，輸入矩陣 ( X ) 分別與查詢權(quán)重矩陣 ( W^Q )、鍵權(quán)重矩陣 ( W^K ) 和值權(quán)重矩陣 ( W^V ) 相乘，得到查詢矩陣 ( Q )、鍵矩陣 ( K ) 和值矩陣 ( V )。

4. 計算注意力

描述：使用得到的查詢、鍵和值矩陣計算注意力。
解釋：對于每個頭，使用查詢矩陣 ( Q )、鍵矩陣 ( K ) 和值矩陣 ( V ) 計算注意力分數(shù)。具體步驟包括：

計算 ( Q ) 和 ( K ) 的點積。

對點積結(jié)果進行縮放。

應用softmax函數(shù)得到注意力權(quán)重。

用注意力權(quán)重乘以值矩陣 ( V ) 得到輸出矩陣 ( Z )。

5. 拼接結(jié)果矩陣

描述：將得到的 ( Z ) 矩陣拼接起來，然后用權(quán)重矩陣 ( W^O ) 乘以拼接后的矩陣，得到層的輸出。
解釋：將所有頭的輸出矩陣 ( Z ) 拼接成一個矩陣，然后用輸出權(quán)重矩陣 ( W^O ) 乘以這個拼接后的矩陣，得到最終的輸出矩陣 ( Z )。

額外說明

查詢、鍵、值和輸出向量的形狀：在加載查詢、鍵、值和輸出向量時，注意到它們的形狀分別是 [4096x4096]、[1024x4096]、[1024x4096]、[1024x4096] 和 [4096x4096]。
并行化注意力頭的乘法：將它們捆綁在一起有助于并行化注意力頭的乘法。

這張圖展示了Transformer模型中多頭注意力機制的實現(xiàn)過程，從輸入句子的嵌入開始，經(jīng)過多頭分割、注意力計算，最后拼接結(jié)果并生成輸出。每個步驟都詳細說明了如何從輸入矩陣 ( 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）權(quán)重矩陣 (wq.weight) 的形狀是 [4096, 4096]。鍵（Key）權(quán)重矩陣 (wk.weight) 的形狀是 [1024, 4096]。值（Value）權(quán)重矩陣 (wv.weight) 的形狀是 [1024, 4096]。輸出（Output）權(quán)重矩陣 (wo.weight) 的形狀是 [4096, 4096]。輸出結(jié)果表明：查詢（Q）和輸出（O）權(quán)重矩陣的形狀是相同的，都是[4096, 4096]。這意味著對于查詢和輸出，輸入特征和輸出特征的維度都是4096。鍵（K）和值（V）權(quán)重矩陣的形狀也是相同的，都是[1024, 4096]。這表明鍵和值的輸入特征維度為4096，但輸出特征維度被壓縮到了1024。這些權(quán)重矩陣的形狀反映了模型設計者如何設置注意力機制中不同部分的維度。特別是，鍵和值的維度被減小可能是為了減少計算復雜度和內(nèi)存消耗，而保持查詢和輸出的較高維度可能是為了保留更多的信息。這種設計選擇依賴于特定的模型架構(gòu)和應用場景。

讓我們用“我欣賞李鴻章”這個句子作為例子，來簡化解釋這個圖中的注意力機制的實現(xiàn)過程。

輸入句子：首先，我們有句子“我欣賞李鴻章”。在處理這個句子之前，我們需要將句子中的每個詞轉(zhuǎn)換成數(shù)學上可以處理的形式，即詞向量。這個過程叫做詞嵌入（embedding）。
詞嵌入：每個詞，比如“我”、“欣賞”、“李鴻章”，都會被轉(zhuǎn)換成一個固定大小的向量。這些向量包含了詞的語義信息。
分割成多個頭：為了讓模型能夠從不同的角度理解句子，我們將每個詞的向量分割成多個部分，這里是8個頭。每個頭都會關注句子的不同方面。
計算注意力：對于每個頭，我們都會計算一個叫做注意力的東西。這個過程涉及到三個步驟：以“我欣賞李鴻章”為例，如果我們想要關注“欣賞”這個詞，那么“欣賞”就是查詢，而其他詞比如“我”和“李鴻章”就是鍵，它們的向量就是值。

查詢（Q）：這是我們想要尋找信息的部分。
鍵（K）：這是包含信息的部分。
值（V）：這是實際的信息內(nèi)容。

拼接和輸出：計算完每個頭的注意力之后，我們將這些結(jié)果拼接起來，并通過一個權(quán)重矩陣Wo來生成最終的輸出。這個輸出將被用于下一層的處理或者作為最終結(jié)果的一部分。

在圖中的注釋中提到的形狀問題，是關于如何在計算機中有效地存儲和處理這些向量的問題。在實際的代碼實現(xiàn)中，為了提高效率，開發(fā)者可能會將多個頭的查詢、鍵、值向量打包在一起處理，而不是單獨處理每個頭。這樣可以利用現(xiàn)代計算機的并行處理能力，加快計算速度。

查詢（Query）權(quán)重矩陣 (wq.weight) 的形狀是 [4096, 4096]。
鍵（Key）權(quán)重矩陣 (wk.weight) 的形狀是 [1024, 4096]。
值（Value）權(quán)重矩陣 (wv.weight) 的形狀是 [1024, 4096]。
輸出（Output）權(quán)重矩陣 (wo.weight) 的形狀是 [4096, 4096]。

輸出結(jié)果表明：

查詢（Q）和輸出（O）權(quán)重矩陣的形狀是相同的，都是[4096, 4096]。這意味著對于查詢和輸出，輸入特征和輸出特征的維度都是4096。
鍵（K）和值（V）權(quán)重矩陣的形狀也是相同的，都是[1024, 4096]。這表明鍵和值的輸入特征維度為4096，但輸出特征維度被壓縮到了1024。

這些權(quán)重矩陣的形狀反映了模型設計者如何設置注意力機制中不同部分的維度。特別是，鍵和值的維度被減小可能是為了減少計算復雜度和內(nèi)存消耗，而保持查詢和輸出的較高維度可能是為了保留更多的信息。這種設計選擇依賴于特定的模型架構(gòu)和應用場景

讓我們用“我欣賞李鴻章”這個句子作為例子，來簡化解釋這個圖中的注意力機制的實現(xiàn)過程。

輸入句子：首先，我們有句子“我欣賞李鴻章”。在處理這個句子之前，我們需要將句子中的每個詞轉(zhuǎn)換成數(shù)學上可以處理的形式，即詞向量。這個過程叫做詞嵌入（embedding）。
詞嵌入：每個詞，比如“我”、“欣賞”、“李鴻章”，都會被轉(zhuǎn)換成一個固定大小的向量。這些向量包含了詞的語義信息。
分割成多個頭：為了讓模型能夠從不同的角度理解句子，我們將每個詞的向量分割成多個部分，這里是8個頭。每個頭都會關注句子的不同方面。

計算注意力：對于每個頭，我們都會計算一個叫做注意力的東西。這個過程涉及到三個步驟：以“我欣賞李鴻章”為例，如果我們想要關注“欣賞”這個詞，那么“欣賞”就是查詢，而其他詞比如“我”和“李鴻章”就是鍵，它們的向量就是值。查詢（Q）：這是我們想要尋找信息的部分。鍵（K）：這是包含信息的部分。值（V）：這是實際的信息內(nèi)容。

拼接和輸出：計算完每個頭的注意力之后，我們將這些結(jié)果拼接起來，并通過一個權(quán)重矩陣Wo來生成最終的輸出。這個輸出將被用于下一層的處理或者作為最終結(jié)果的一部分。

在圖中的注釋中提到的形狀問題，是關于如何在計算機中有效地存儲和處理這些向量的問題。在實際的代碼實現(xiàn)中，為了提高效率，開發(fā)者可能會將多個頭的查詢、鍵、值向量打包在一起處理，而不是單獨處理每個頭。這樣可以利用現(xiàn)代計算機的并行處理能力，加快計算速度。

我們繼續(xù)使用句子“我欣賞李鴻章”來解釋WQ、WK、WV和WO這些權(quán)重矩陣的作用。

在Transformer模型中，每個詞都會通過詞嵌入轉(zhuǎn)換成一個向量。這些向量接下來會通過一系列的線性變換來計算注意力分數(shù)。這些線性變換就是通過權(quán)重矩陣WQ、WK、WV和WO來實現(xiàn)的。

WQ（權(quán)重矩陣Q）：這個矩陣用于將每個詞的向量轉(zhuǎn)換成“查詢（Query）”向量。在我們的例子中，如果我們想要關注“欣賞”這個詞，我們會將“欣賞”的向量乘以WQ來得到查詢向量。
WK（權(quán)重矩陣K）：這個矩陣用于將每個詞的向量轉(zhuǎn)換成“鍵（Key）”向量。同樣地，我們會將每個詞，包括“我”和“李鴻章”，的向量乘以WK來得到鍵向量。
WV（權(quán)重矩陣V）：這個矩陣用于將每個詞的向量轉(zhuǎn)換成“值（Value）”向量。每個詞的向量乘以WV后，我們得到的是值向量。這三個矩陣（WQ、WK、WV）是用來為每個頭生成不同的查詢、鍵和值向量的。這樣做可以讓每個頭關注句子的不同方面。
WQ（權(quán)重矩陣Q）、WK（權(quán)重矩陣K）、WV（權(quán)重矩陣V）和WO（權(quán)重矩陣O）這些矩陣是Transformer模型中的參數(shù)，它們是在模型訓練過程中通過反向傳播算法和梯度下降等優(yōu)化方法學習得到的。

在整個過程中，WQ、WK、WV和WO是通過訓練學習得到的，它們決定了模型如何將輸入的詞向量轉(zhuǎn)換成不同的表示，以及如何組合這些表示來得到最終的輸出。這些矩陣是Transformer模型中注意力機制的核心部分，它們使得模型能夠捕捉到句子中不同詞之間的關系。

WQ（權(quán)重矩陣Q）、WK（權(quán)重矩陣K）、WV（權(quán)重矩陣V）和WO（權(quán)重矩陣O）這些矩陣是Transformer模型中的參數(shù)，它們是在模型訓練過程中通過反向傳播算法和梯度下降等優(yōu)化方法學習得到的。

讓我們來看看這個學習過程是如何進行的：

初始化：在訓練開始之前，這些矩陣通常會被隨機初始化。這意味著它們的初始值是隨機選取的，這樣可以打破對稱性并開始學習過程。
前向傳播：在模型的訓練過程中，輸入數(shù)據(jù)（如句子“我欣賞李鴻章”）會通過模型的各個層進行前向傳播。在注意力機制中，輸入的詞向量會與WQ、WK、WV矩陣相乘，以生成查詢、鍵和值向量。
計算損失：模型的輸出會與期望的輸出（通常是訓練數(shù)據(jù)中的標簽）進行比較，計算出一個損失值。這個損失值衡量了模型的預測與實際情況的差距。
反向傳播：損失值會通過反向傳播算法傳回模型，計算每個參數(shù)（包括WQ、WK、WV和WO）對損失的影響，即它們的梯度。
參數(shù)更新：根據(jù)計算出的梯度，使用梯度下降或其他優(yōu)化算法來更新這些矩陣的值。這個過程會逐漸減小損失值，使模型的預測更加準確。
迭代過程：這個前向傳播、損失計算、反向傳播和參數(shù)更新的過程會在訓練數(shù)據(jù)上多次迭代進行，直到模型的性能達到一定的標準或者不再顯著提升。

通過這個訓練過程，WQ、WK、WV和WO這些矩陣會逐漸調(diào)整它們的值，以便模型能夠更好地理解和處理輸入數(shù)據(jù)。在訓練完成后，這些矩陣將固定下來，用于模型的推理階段，即對新的輸入數(shù)據(jù)進行預測。

四、展開查詢向量

在本小節(jié)中，我們將從多個注意力頭中展開查詢向量，得到的形狀是 [32x128x4096] 這里，32 是 llama3 中注意力頭的數(shù)量，128 是查詢向量的大小，而 4096 是令牌嵌入的大小。

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

圖片

這段代碼通過對模型中第一層的查詢（Q）權(quán)重矩陣進行重塑（reshape），將其分解為多個注意力頭的形式，從而揭示了32和128這兩個維度。

q_layer0 = model["layers.0.attention.wq.weight"]：這行代碼從模型中提取第一層的查詢（Q）權(quán)重矩陣。
head_dim = q_layer0.shape[0] // n_heads：這行代碼計算每個注意力頭的維度大小。它通過將查詢權(quán)重矩陣的第一個維度（原本是4096）除以注意力頭的數(shù)量（n_heads），得到每個頭的維度。如果n_heads是32（即模型設計為有32個注意力頭），那么head_dim就是4096 // 32 = 128。
q_layer0 = q_layer0.view(n_heads, head_dim, dim)：這行代碼使用.view()方法重塑查詢權(quán)重矩陣，使其形狀變?yōu)閇n_heads, head_dim, dim]。這里dim很可能是原始特征維度4096，n_heads是32，head_dim是128，因此重塑后的形狀是[32, 128, 4096]。
q_layer0.shape 輸出：torch.Size([32, 128, 4096])：這行代碼打印重塑后的查詢權(quán)重矩陣的形狀，確認了其形狀為[32, 128, 4096]。

之所以在這段代碼中出現(xiàn)了32和128這兩個維度，而在之前的代碼段中沒有，是因為這段代碼通過重塑操作明確地將查詢權(quán)重矩陣分解為多個注意力頭，每個頭具有自己的維度。32代表了模型中注意力頭的數(shù)量，而128代表了分配給每個頭的特征維度大小。這種分解是為了實現(xiàn)多頭注意力機制，其中每個頭可以獨立地關注輸入的不同部分，最終通過組合這些頭的輸出來提高模型的表達能力。

實現(xiàn)第一層的第一個頭

訪問了第一層第一個頭的查詢（query）權(quán)重矩陣，這個查詢權(quán)重矩陣的大小是 [128x4096]。

q_layer0_head0 = q_layer0[0]
q_layer0_head0.shape

圖片

我們現(xiàn)在將查詢權(quán)重與令牌嵌入相乘，以獲得令牌的查詢

在這里，你可以看到結(jié)果形狀是 [17x128]，這是因為我們有17個令牌，每個令牌都有一個長度為128的查詢（每個令牌在一個頭上方的查詢）。

br

圖片

這段代碼執(zhí)行了一個矩陣乘法操作，將令牌嵌入（token_embeddings）與第一層第一個頭的查詢（query）權(quán)重矩陣（q_layer0_head0）的轉(zhuǎn)置（.T）相乘，以生成每個令牌的查詢向量（q_per_token）。

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

torch.matmul 是PyTorch中的矩陣乘法函數(shù)，它可以處理兩個張量的乘法。

token_embeddings 應該是一個形狀為 [17, 4096] 的張量，表示有17個令牌，每個令牌由4096維的嵌入向量表示。

q_layer0_head0 是第一層第一個頭的查詢權(quán)重矩陣，其原始形狀為 [128, 4096]。.T 是PyTorch中的轉(zhuǎn)置操作，將 q_layer0_head0 的形狀轉(zhuǎn)置為 [4096, 128]。

這樣，token_embeddings 和 q_layer0_head0.T 的矩陣乘法就是 [17, 4096] 和 [4096, 128] 的乘法，結(jié)果是一個形狀為 [17, 128] 的張量。

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

這行代碼打印出 q_per_token 張量的形狀，確認其為 [17, 128]。

這意味著對于輸入的每個令牌（共17個），我們現(xiàn)在都有了一個128維的查詢向量。這128維的查詢向量是通過將令牌嵌入與查詢權(quán)重矩陣相乘得到的，可以用于后續(xù)的注意力機制計算。

總之，這段代碼通過矩陣乘法將每個令牌的嵌入向量轉(zhuǎn)換為查詢向量，為實現(xiàn)注意力機制的下一步做準備。每個令牌現(xiàn)在都有了一個與之對應的查詢向量，這些查詢向量將用于計算與其他令牌的注意力得分。

責任編輯：武曉燕來源：大魏分享

Llama3 向量架構(gòu)