一、Llama3的架構(gòu)

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

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

圖片

Llama3的模型參數(shù)：

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

圖片

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

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

圖片

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

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

圖片

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

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

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

圖片

[4] 隱藏維度（Hidden Dimensions）

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

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

圖片

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

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

接下來(lái)是Feed-Forward Network中的隱藏層，增長(zhǎng)到5325，然后在最后一層回落到4096。

圖片

[6] Transformer block的多層

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

圖片

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

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

圖片

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

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

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

升維通常是為了增加模型的容量，使其能夠捕捉更復(fù)雜的特征和模式。當(dāng)輸入數(shù)據(jù)被映射到一個(gè)更高維度的空間時(shí)，不同的特征組合可以被模型更容易地區(qū)分。這在處理非線性問(wèn)題時(shí)尤其有用，因?yàn)樗梢詭椭Ｐ蛯W(xué)習(xí)到更復(fù)雜的決策邊界  。

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

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

步驟3：這個(gè)特征通過(guò)Transformer block進(jìn)行處理，首先由Attention layer處理，然后是FFN layer。Attention layer橫向跨特征處理，而FFN layer則縱向跨維度處理。

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

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

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

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

1. max_seq_len (最大序列長(zhǎng)度)

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

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

2. Vocabulary-size (詞匯量)

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

3. Attention Layers (注意力層)

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

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

4. transformer block 

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

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

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

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

5. Feature-dimension (特征維度)

這是輸入token在模型中表示為向量時(shí)，每個(gè)向量的維度。

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

6. Attention-Heads (注意力頭)

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

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

7. Hidden Dimensions (隱藏維度)

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

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

關(guān)系和數(shù)值：

Attention Layers 和 Attention-Heads 的關(guān)系：每個(gè)Attention Layer可以包含多個(gè)Attention-Heads。

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

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

二、查看模型

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

我們輸入hello world,看分詞器如何進(jìn)行分詞。

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

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

圖片

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

圖片

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

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

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

圖片

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

我們使用這個(gè)配置來(lái)推斷模型的細(xì)節(jié)，比如：

1. 模型有32個(gè)Transformer層
2. 每個(gè)多頭注意力塊有32個(gè)頭
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轉(zhuǎn)化為T(mén)oken

代碼如下：

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]，即長(zhǎng)度為4096的17個(gè)嵌入（每個(gè)令牌一個(gè)）。

下圖是為了驗(yàn)證我們輸入的這句話，是17個(gè)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 歸一化對(duì)嵌入進(jìn)行歸一化，也就是圖中這個(gè)位置：

圖片

使用公式如下：

圖片

代碼如下：

# 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

這段代碼定義了一個(gè)名為 rms_norm 的函數(shù)，它實(shí)現(xiàn)了對(duì)輸入張量（tensor）的RMS（Root Mean Square，均方根）歸一化處理。這個(gè)函數(shù)接受兩個(gè)參數(shù)：tensor 和 norm_weights。tensor 是需要進(jìn)行歸一化處理的輸入張量，而 norm_weights 是歸一化時(shí)使用的權(quán)重。

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

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

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

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

圖片

圖片

接下來(lái)，我們介紹注意力機(jī)制的實(shí)現(xiàn)，也就是下圖中的紅框標(biāo)注的位置：

圖片

圖片

我們一步一步地解釋這張圖，詳細(xì)說(shuō)明每個(gè)步驟。

1. 輸入句子

• 描述：這是我們的輸入句子。
• 解釋?zhuān)狠斎刖渥颖槐硎緸橐粋€(gè)矩陣 ( X )，其中每一行代表一個(gè)詞的嵌入向量。

2. 嵌入每個(gè)詞

• 描述：我們對(duì)每個(gè)詞進(jìn)行嵌入。
• 解釋?zhuān)狠斎刖渥又械拿總€(gè)詞被轉(zhuǎn)換為一個(gè)高維向量，這些向量組成了矩陣 ( X )。

3. 分成8個(gè)頭

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

4. 計(jì)算注意力

• 描述：使用得到的查詢(xún)、鍵和值矩陣計(jì)算注意力。
• 解釋?zhuān)簩?duì)于每個(gè)頭，使用查詢(xún)矩陣 ( Q )、鍵矩陣 ( K ) 和值矩陣 ( V ) 計(jì)算注意力分?jǐn)?shù)。具體步驟包括：

計(jì)算 ( Q ) 和 ( K ) 的點(diǎn)積。

對(duì)點(diǎn)積結(jié)果進(jìn)行縮放。

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

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

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

• 描述：將得到的 ( Z ) 矩陣拼接起來(lái)，然后用權(quán)重矩陣 ( W^O ) 乘以拼接后的矩陣，得到層的輸出。
• 解釋?zhuān)簩⑺蓄^的輸出矩陣 ( Z ) 拼接成一個(gè)矩陣，然后用輸出權(quán)重矩陣 ( W^O ) 乘以這個(gè)拼接后的矩陣，得到最終的輸出矩陣 ( Z )。

額外說(shuō)明

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

這張圖展示了Transformer模型中多頭注意力機(jī)制的實(shí)現(xiàn)過(guò)程，從輸入句子的嵌入開(kāi)始，經(jīng)過(guò)多頭分割、注意力計(jì)算，最后拼接結(jié)果并生成輸出。每個(gè)步驟都詳細(xì)說(shuō)明了如何從輸入矩陣 ( X ) 生成最終的輸出矩陣 ( Z )。

當(dāng)我們從模型中加載查詢(xún)（query）、鍵（key）、值（value）和輸出（output）向量時(shí)，我們注意到它們的形狀分別是 [4096x4096]、[1024x4096]、[1024x4096]、[4096x4096]

乍一看這很奇怪，因?yàn)槔硐肭闆r下我們希望每個(gè)頭的每個(gè)q、k、v和o都是單獨(dú)的

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
)

圖片

查詢(xú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é)果表明：查詢(xún)（Q）和輸出（O）權(quán)重矩陣的形狀是相同的，都是[4096, 4096]。這意味著對(duì)于查詢(xún)和輸出，輸入特征和輸出特征的維度都是4096。鍵（K）和值（V）權(quán)重矩陣的形狀也是相同的，都是[1024, 4096]。這表明鍵和值的輸入特征維度為4096，但輸出特征維度被壓縮到了1024。這些權(quán)重矩陣的形狀反映了模型設(shè)計(jì)者如何設(shè)置注意力機(jī)制中不同部分的維度。特別是，鍵和值的維度被減小可能是為了減少計(jì)算復(fù)雜度和內(nèi)存消耗，而保持查詢(xún)和輸出的較高維度可能是為了保留更多的信息。這種設(shè)計(jì)選擇依賴(lài)于特定的模型架構(gòu)和應(yīng)用場(chǎng)景。

讓我們用“我欣賞李鴻章”這個(gè)句子作為例子，來(lái)簡(jiǎn)化解釋這個(gè)圖中的注意力機(jī)制的實(shí)現(xiàn)過(guò)程。

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

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

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

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

• 查詢(xú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é)果表明：

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

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

讓我們用“我欣賞李鴻章”這個(gè)句子作為例子，來(lái)簡(jiǎn)化解釋這個(gè)圖中的注意力機(jī)制的實(shí)現(xiàn)過(guò)程。

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

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

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

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

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

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

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

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

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

讓我們來(lái)看看這個(gè)學(xué)習(xí)過(guò)程是如何進(jìn)行的：

1. 初始化：在訓(xùn)練開(kāi)始之前，這些矩陣通常會(huì)被隨機(jī)初始化。這意味著它們的初始值是隨機(jī)選取的，這樣可以打破對(duì)稱(chēng)性并開(kāi)始學(xué)習(xí)過(guò)程。
2. 前向傳播：在模型的訓(xùn)練過(guò)程中，輸入數(shù)據(jù)（如句子“我欣賞李鴻章”）會(huì)通過(guò)模型的各個(gè)層進(jìn)行前向傳播。在注意力機(jī)制中，輸入的詞向量會(huì)與WQ、WK、WV矩陣相乘，以生成查詢(xún)、鍵和值向量。
3. 計(jì)算損失：模型的輸出會(huì)與期望的輸出（通常是訓(xùn)練數(shù)據(jù)中的標(biāo)簽）進(jìn)行比較，計(jì)算出一個(gè)損失值。這個(gè)損失值衡量了模型的預(yù)測(cè)與實(shí)際情況的差距。
4. 反向傳播：損失值會(huì)通過(guò)反向傳播算法傳回模型，計(jì)算每個(gè)參數(shù)（包括WQ、WK、WV和WO）對(duì)損失的影響，即它們的梯度。
5. 參數(shù)更新：根據(jù)計(jì)算出的梯度，使用梯度下降或其他優(yōu)化算法來(lái)更新這些矩陣的值。這個(gè)過(guò)程會(huì)逐漸減小損失值，使模型的預(yù)測(cè)更加準(zhǔn)確。
6. 迭代過(guò)程：這個(gè)前向傳播、損失計(jì)算、反向傳播和參數(shù)更新的過(guò)程會(huì)在訓(xùn)練數(shù)據(jù)上多次迭代進(jìn)行，直到模型的性能達(dá)到一定的標(biāo)準(zhǔn)或者不再顯著提升。
   
   通過(guò)這個(gè)訓(xùn)練過(guò)程，WQ、WK、WV和WO這些矩陣會(huì)逐漸調(diào)整它們的值，以便模型能夠更好地理解和處理輸入數(shù)據(jù)。在訓(xùn)練完成后，這些矩陣將固定下來(lái)，用于模型的推理階段，即對(duì)新的輸入數(shù)據(jù)進(jìn)行預(yù)測(cè)。

四、展開(kāi)查詢(xún)向量

在本小節(jié)中，我們將從多個(gè)注意力頭中展開(kāi)查詢(xún)向量，得到的形狀是 [32x128x4096] 這里，32 是 llama3 中注意力頭的數(shù)量，128 是查詢(xún)向量的大小，而 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

圖片

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

1. q_layer0 = model["layers.0.attention.wq.weight"]：這行代碼從模型中提取第一層的查詢(xún)（Q）權(quán)重矩陣。
2. head_dim = q_layer0.shape[0] // n_heads：這行代碼計(jì)算每個(gè)注意力頭的維度大小。它通過(guò)將查詢(xún)權(quán)重矩陣的第一個(gè)維度（原本是4096）除以注意力頭的數(shù)量（n_heads），得到每個(gè)頭的維度。如果n_heads是32（即模型設(shè)計(jì)為有32個(gè)注意力頭），那么head_dim就是4096 // 32 = 128。
3. q_layer0 = q_layer0.view(n_heads, head_dim, dim)：這行代碼使用.view()方法重塑查詢(xún)權(quán)重矩陣，使其形狀變?yōu)閇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])：這行代碼打印重塑后的查詢(xún)權(quán)重矩陣的形狀，確認(rèn)了其形狀為[32, 128, 4096]。

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

實(shí)現(xiàn)第一層的第一個(gè)頭

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

q_layer0_head0 = q_layer0[0]
q_layer0_head0.shape

圖片

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

在這里，你可以看到結(jié)果形狀是 [17x128]，這是因?yàn)槲覀冇?7個(gè)令牌，每個(gè)令牌都有一個(gè)長(zhǎng)度為128的查詢(xún)（每個(gè)令牌在一個(gè)頭上方的查詢(xún)）。

br

圖片

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

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

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

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

q_layer0_head0 是第一層第一個(gè)頭的查詢(xún)權(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é)果是一個(gè)形狀為 [17, 128] 的張量。

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

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

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

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