LLMs之llama3-from-scratch：llama3-from-scratch(從頭開始利用pytorch來(lái)實(shí)現(xiàn)并解讀LLaMA-3模型的每層代碼)的簡(jiǎn)介,、核心思路梳理
導(dǎo)讀：這篇論文實(shí)現(xiàn)了transformer網(wǎng)絡(luò)的llama3模型，從頭開始利用pytorch來(lái)實(shí)現(xiàn)該模型,。
背景：目前機(jī)器學(xué)習(xí)語(yǔ)言模型內(nèi)容的復(fù)雜性不斷增強(qiáng)，但是大多模型都是基于高度抽象和封裝的框架來(lái)實(shí)現(xiàn)，對(duì)模型內(nèi)部工作機(jī)制的理解不是很深入。這篇論文采用從零開始逐步實(shí)現(xiàn)的方式，幫助讀者更好地理解transformer和llama3模型是如何工作的,。
具體解決方案：作者加載了llama3預(yù)訓(xùn)練模型的參數(shù)，并按照每個(gè)算子依次實(shí)現(xiàn),。
流程和關(guān)鍵步驟：
>> 加載tokenizer對(duì)文本進(jìn)行tokenize
>> 從模型文件加載各層參數(shù)
>> 對(duì)輸入文本的tokens生成嵌入向量
>> 對(duì)嵌入向量進(jìn)行RMSNorm標(biāo)準(zhǔn)化
>> 實(shí)現(xiàn)單頭注意力機(jī)制，包括Q、K,、V矩陣的生成,、位置編碼、掩碼,、softmax計(jì)算注意力權(quán)重等
>> 循環(huán)實(shí)現(xiàn)多頭注意力
>> 加載FFN層的參數(shù)實(shí)現(xiàn)前向傳播
>> 循環(huán)上述步驟實(shí)現(xiàn)整個(gè)模型
>> 應(yīng)用訓(xùn)練好的輸出層參數(shù)進(jìn)行預(yù)測(cè)
關(guān)鍵技術(shù)點(diǎn)
>> 利用RoPE位置編碼為Q,、K向量添加位置信息
>> 多頭注意力分頭計(jì)算并拼接
>> SwiGLU結(jié)構(gòu)的FFN層
>> 應(yīng)用RMSNorm進(jìn)行標(biāo)準(zhǔn)化
>> 循環(huán)實(shí)現(xiàn)每一層的計(jì)算
>> 代碼運(yùn)行效果
>> 最后利用加載的llama3模型在樣本問(wèn)題"答案是什么"上正確預(yù)測(cè)答案為42，驗(yàn)證了從零開始實(shí)現(xiàn)的整個(gè)流程是正確的。
總之，這篇論文采用詳盡清晰的流程，幫助讀者通過(guò)實(shí)際實(shí)現(xiàn)深入理解transformer模型底層工作機(jī)制，是一篇值得推薦的論文。
llama3-from-scratch的簡(jiǎn)介
2024年5月20日，Nishant Aklecha正式發(fā)布了該項(xiàng)目，在這個(gè)文件中，我從頭實(shí)現(xiàn)了llama3，一次一個(gè)張量和矩陣乘法,。此外，我將直接從Meta為llama3提供的模型文件中加載張量，在運(yùn)行此文件之前需要下載權(quán)重,。以下是下載權(quán)重的官方鏈接：https://llama./llama-downloads/
GitHub地址：GitHub - naklecha/llama3-from-scratch: llama3 implementation one matrix multiplication at a time
llama3-from-scratch的核心思路梳理
注意：當(dāng)前文章仍處于持續(xù)更新和梳理中……
0、前置
0.1,、加載tokenizer對(duì)文本進(jìn)行tokenize：將文本轉(zhuǎn)換為模型可以理解的數(shù)字序列（即詞元或tokens）+并在生成模型輸出后能將tokens轉(zhuǎn)換回可讀文本（解碼）
| 簡(jiǎn)介 | 我不會(huì)實(shí)現(xiàn)一個(gè)BPE分詞器（但是Andrej Karpathy有一個(gè)非常簡(jiǎn)潔的實(shí)現(xiàn)） 鏈接到他的實(shí)現(xiàn)： https://github.com/karpathy/minbpe 該項(xiàng)目給出了一個(gè)針對(duì)中文分詞任務(wù)的最小化,、清晰而可讀的Python實(shí)現(xiàn)Byte Pair Encoding(BPE)算法的代碼庫(kù)。它實(shí)現(xiàn)了兩個(gè)Tokenizer對(duì)象用于文本到詞元和詞元到文本的轉(zhuǎn)換：BasicTokenizer和RegexTokenizer,。
BasicTokenizer直接在文本上運(yùn)行BPE算法，RegexTokenizer在BPE之前通過(guò)正則表達(dá)式將文本切分為不同類型(如字母,、數(shù)字,、標(biāo)點(diǎn)等)以防止詞匯跨類型合并,。 | | 思路步驟 | 從Meta下載權(quán)重文件,。 使用tiktoken庫(kù)加載BPE分詞器，并定義特殊token。 | | 核心技術(shù)點(diǎn) | BPE分詞：字節(jié)對(duì)編碼（Byte Pair Encoding）是一種常見的分詞方法，用于處理詞匯表,。 | | 代碼解讀 | 目的是設(shè)置和測(cè)試一個(gè)自定義的tokenizer，它能夠處理文本數(shù)據(jù)
|
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!"))
0.2,、從模型文件加載各層參數(shù)：從預(yù)訓(xùn)練模型文件中讀取權(quán)重參數(shù)
| 簡(jiǎn)介 | 通常情況下，讀取模型文件取決于模型類的編寫方式和其中的變量名稱,。但由于我們要從頭實(shí)現(xiàn)llama3，所以我們將逐個(gè)張量讀取文件。 我們使用此配置來(lái)推斷模型的細(xì)節(jié)，模型有32個(gè)Transformer層，每個(gè)多頭注意力塊有32個(gè)頭，n_kv_heads，詞匯表大小為128256等等,。 下面是各個(gè)參數(shù)的含義：
'dim': 4096：表示模型的隱藏層維度，即模型中嵌入向量的維度或者是注意力層的尺寸,。'n_layers': 32：表示模型的層數(shù)，即 Transformer 的層數(shù),。'n_heads': 32：表示多頭注意力機(jī)制中的注意力頭數(shù),。多頭注意力允許模型在不同的表示空間中并行處理信息。'n_kv_heads': 8：表示鍵和值（Key 和 Value）在多頭注意力中的頭數(shù),。這指的是鍵（Key）和值（Value）頭數(shù)，這是一種注意力機(jī)制的變體，其中鍵和值的頭數(shù)可以與查詢頭數(shù)不同,。'vocab_size': 128256：表示模型詞匯表的大小，即可以處理的不同詞的數(shù)量。'multiple_of': 1024：表示一些尺寸（比如前饋神經(jīng)網(wǎng)絡(luò)的維度）應(yīng)當(dāng)是這個(gè)數(shù)值的倍數(shù)，通常用于確保計(jì)算效率,。這可能是一個(gè)設(shè)計(jì)上的考慮，用于確保某些層的尺寸是某個(gè)數(shù)值的倍數(shù)，這有助于提高模型在硬件上的運(yùn)行效率,。'ffn_dim_multiplier': 1.3：前饋神經(jīng)網(wǎng)絡(luò)（Feed-Forward Network）的維度倍率，表示前饋網(wǎng)絡(luò)的維度是隱藏層維度的1.3倍。'norm_eps': 1e-05：Layer Normalization中的一個(gè)非常小的值，防止除以零的情況，用于避免除以零的情況,。'rope_theta': 500000.0：這是相對(duì)位置編碼（rotary positional embeddings）的一個(gè)參數(shù)，通常用于提高模型在處理長(zhǎng)序列時(shí)的性能,。
| | 思路步驟 | 加載模型文件，查看模型參數(shù)。 讀取模型配置文件，獲取模型的詳細(xì)參數(shù),。 | | 核心技術(shù)點(diǎn) | 模型文件加載：使用torch.load加載模型權(quán)重文件,。 配置解析：解析配置文件以提取模型超參數(shù),。 |
model = torch.load("Meta-Llama-3-8B/consolidated.00.pth")
print(json.dumps(list(model.keys())[:20], indent=4))
[
"tok_embeddings.weight",
"layers.0.attention.wq.weight",
"layers.0.attention.wk.weight",
"layers.0.attention.wv.weight",
"layers.0.attention.wo.weight",
"layers.0.feed_forward.w1.weight",
"layers.0.feed_forward.w3.weight",
"layers.0.feed_forward.w2.weight",
"layers.0.attention_norm.weight",
"layers.0.ffn_norm.weight",
"layers.1.attention.wq.weight",
"layers.1.attention.wk.weight",
"layers.1.attention.wv.weight",
"layers.1.attention.wo.weight",
"layers.1.feed_forward.w1.weight",
"layers.1.feed_forward.w3.weight",
"layers.1.feed_forward.w2.weight",
"layers.1.attention_norm.weight",
"layers.1.ffn_norm.weight",
"layers.2.attention.wq.weight"
]
with open("Meta-Llama-3-8B/params.json", "r") as f:
config = json.load(f)
config
{'dim': 4096,
'n_layers': 32,
'n_heads': 32,
'n_kv_heads': 8,
'vocab_size': 128256,
'multiple_of': 1024,
'ffn_dim_multiplier': 1.3,
'norm_eps': 1e-05,
'rope_theta': 500000.0}
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"])
1、輸入嵌入層：采用tiktoken庫(kù)的BPE算法將輸入文本的tokens轉(zhuǎn)換為嵌入向量，為模型處理做好準(zhǔn)備
1.1,、文本分詞并編碼：將一段文本(輸入)轉(zhuǎn)換為token ID(輸出),。將輸入的自然語(yǔ)言經(jīng)過(guò)分詞技術(shù)轉(zhuǎn)換為token ID，如下所示，一句話包含17個(gè)token(可以簡(jiǎn)單理解為單詞)
| 簡(jiǎn)介 | 這里我們使用tiktoken（我認(rèn)為是OpenAI的庫(kù)）作為分詞器。 | | 思路步驟 | 輸入一個(gè)提示文本，并使用分詞器將其轉(zhuǎn)換為tokens,。 打印tokens及其對(duì)應(yīng)的文本表示,。 | | 核心技術(shù)點(diǎn) | 分詞：使用分詞器將文本轉(zhuǎn)換為token ID。 文本表示：解碼token ID為文本以驗(yàn)證分詞過(guò)程,。 |
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', ' ']
1.2,、離散token進(jìn)行嵌入向量化：將離散的token ID(輸入，形狀為【17,1】)轉(zhuǎn)換為其連續(xù)的嵌入向量表示(輸出，形狀為【17,4096】)，以便進(jìn)行后續(xù)的運(yùn)算
| 簡(jiǎn)介 | 這是代碼庫(kù)中唯一使用內(nèi)置神經(jīng)網(wǎng)絡(luò)模塊的部分。 我們的[17x1] tokens現(xiàn)在是[17x4096]，即17個(gè)長(zhǎng)度為4096的嵌入（每個(gè)token一個(gè)嵌入）,。 注意：跟蹤形狀，這使它更容易理解一切 | | 思路步驟 | 定義一個(gè)嵌入層，并加載預(yù)訓(xùn)練的權(quán)重,。 將tokens轉(zhuǎn)換為嵌入表示。 | | 核心技術(shù)點(diǎn) | 嵌入層：使用torch.nn.Embedding定義嵌入層，并加載預(yù)訓(xùn)練權(quán)重,。 |
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
torch.Size([17, 4096])
1.3,、采用RMSNorm對(duì)嵌入向量進(jìn)行均方根歸一化：僅對(duì)嵌入進(jìn)行歸一化處理，前后形狀未變，只是確保數(shù)值穩(wěn)定，避免后續(xù)計(jì)算過(guò)程中出現(xiàn)數(shù)值0問(wèn)題
| 簡(jiǎn)介 | 我們?nèi)缓笫褂肦MS歸一化嵌入，請(qǐng)注意此步驟后形狀不會(huì)改變，值只是被歸一化了。需要注意的是我們需要一個(gè)norm_eps（從配置中獲取），因?yàn)槲覀儾幌胍馔獾貙ms設(shè)為0并除以0,。 公式如下： 歸一化后我們的形狀仍然是[17x4096]，與嵌入相同但已歸一化,。 | | 思路步驟 | 對(duì)嵌入進(jìn)行RMS歸一化。 | | 核心技術(shù)點(diǎn) | RMS歸一化：對(duì)嵌入向量進(jìn)行均方根歸一化（RMS Norm）,。 |
# 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
2,、位置編碼層：為詞元嵌入添加位置信息，使模型能理解詞元的順序
2.1、背景——構(gòu)建Transformer的第一個(gè)層：準(zhǔn)備嵌入以輸入到Transformer層進(jìn)行處理，實(shí)現(xiàn)Transformer網(wǎng)絡(luò)結(jié)構(gòu)的前向傳播
| 簡(jiǎn)介 | 您會(huì)看到我從這個(gè)模型字典中訪問(wèn)layer.0（這是第一層） 無(wú)論如何，所以歸一化后，我們的形狀仍然是【17x4096】，與嵌入相同但已歸一化 | | 思路步驟 | 構(gòu)建Transformer的第一層： | | 核心技術(shù)點(diǎn) | 多頭注意力機(jī)制（Multi-Head Attention）與殘差連接 Transformer層：應(yīng)用注意力機(jī)制和前饋神經(jīng)網(wǎng)絡(luò)進(jìn)行特征提取和變換,。 |
token_embeddings = rms_norm(token_embeddings_unnormalized, model["layers.0.attention_norm.weight"])
token_embeddings.shape
torch.Size([17, 4096])
2.2,、背景——計(jì)算每個(gè)token在Transformer第一層第一個(gè)頭的query向量
| 簡(jiǎn)介 | 讓我們加載Transformer第一層的注意力頭。 >> 加載查詢Query,、鍵Key,、值Value和輸出向量時(shí)，我們注意到它們的權(quán)重矩陣形狀分別為[4096x4096]、[1024x4096],、[1024x4096],、[4096x4096]。 >> 乍一看這有點(diǎn)奇怪，因?yàn)槲覀兝硐肭闆r下希望每個(gè)Query,、Key,、Value和輸出分別對(duì)應(yīng)每個(gè)頭。而代碼作者將它們捆綁在一起，因?yàn)檫@樣做，有助于并行化注意力頭的矩陣乘法,。 接下來(lái)，我將展開所有內(nèi)容... | | 下一步，我將展開多個(gè)注意力頭的查詢，結(jié)果形狀為[32x128x4096],。 這里，32是llama3的注意力頭數(shù)，128是Query向量的大小，4096是token嵌入的大小。 | | 這里我訪問(wèn)第一層第一個(gè)頭的Query權(quán)重矩陣，Query權(quán)重矩陣的大小為【128x4096】。 | | 我們現(xiàn)在，將Query權(quán)重【128x4096】與token嵌入【17x4096】相乘，以接收token的查詢，即獲得每個(gè)token的query,。 結(jié)果形狀為[17x128]，因?yàn)槲覀冇?7個(gè)token，每個(gè)token有一個(gè)128長(zhǎng)度的query向量,。 | | 思路步驟 | 通過(guò)查詢、鍵和值向量計(jì)算注意力分?jǐn)?shù)，提取輸入中重要的信息,。 加載第一層的查詢,、鍵、值和輸出向量權(quán)重,。 展開注意力頭的查詢權(quán)重矩陣,。 計(jì)算token的查詢向量。 | | 核心技術(shù)點(diǎn) | 自注意力（Self-Attention）的計(jì)算,。 |
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
)
torch.Size([4096, 4096]) torch.Size([1024, 4096]) torch.Size([1024, 4096]) torch.Size([4096, 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
torch.Size([32, 128, 4096])
q_layer0_head0 = q_layer0[0]
q_layer0_head0.shape
torch.Size([128, 4096])
2.3,、真正實(shí)現(xiàn)RoPE：此處以Query向量進(jìn)行RoPE為例講解，但Key與Query幾乎相同
  
| 背景 | | | 痛點(diǎn)：現(xiàn)在，已經(jīng)為prompt中的每個(gè)token生成了query向量，但每個(gè)單獨(dú)的query向量并不知道它在prompt中的具體位置。我們現(xiàn)在處于每個(gè)token在我們的提示中有一個(gè)查詢向量的階段，但如果你想一想——單個(gè)查詢向量對(duì)提示中的位置沒(méi)有任何了解,。例如：
- query：“the answer to the ultimate question of life, the universe, and everything is____”
- 查詢：“生命,、宇宙以及一切終極問(wèn)題的答案是____”
期望：在我們的提示中我們使用了三次“the”，我們需要所有3個(gè)“the”token的查詢向量（每個(gè)大小為[1x128]）根據(jù)它們?cè)诓樵冎械奈恢?span style="color:#fe2c24;">具有不同的query向量。 解決方案：采用RoPE（旋轉(zhuǎn)位置嵌入）執(zhí)行這些旋轉(zhuǎn)操作,。 RoPE：觀看這個(gè)視頻（這是我看的）來(lái)理解數(shù)學(xué)：https://www./watch?v=o29P0Kpobz0&t=530s | | 簡(jiǎn)介 | 在上述步驟中，我們將查詢向量分成了若干對(duì)，并對(duì)每對(duì)應(yīng)用一個(gè)旋轉(zhuǎn)角度偏移！ 現(xiàn)在我們有一個(gè)大小為[17x64x2]的向量，這是將128長(zhǎng)度的query向量，對(duì)每個(gè)prompt中的每個(gè)token分成的64對(duì)，這64對(duì)中的每一對(duì)都會(huì)按照m*(theta) 進(jìn)行旋轉(zhuǎn)，其中m是要為其旋轉(zhuǎn)query的token的位置！ 使用復(fù)數(shù)的點(diǎn)積來(lái)旋轉(zhuǎn)向量 | | 現(xiàn)在我們?yōu)槊總€(gè)token的query元素得到了一個(gè)復(fù)數(shù)（角度變化向量），我們可以將我們的query（即我們之前分成的對(duì)）轉(zhuǎn)換為復(fù)數(shù)，然后使用點(diǎn)積根據(jù)位置旋轉(zhuǎn)query向量,。 老實(shí)說(shuō)，想到這一點(diǎn)真是太美妙了 :) | | 在獲得旋轉(zhuǎn)后的向量后，我們可以通過(guò)將復(fù)數(shù)，重新視為實(shí)數(shù)，來(lái)得到(或來(lái)還原)成對(duì)的query向量。 旋轉(zhuǎn)后的對(duì)現(xiàn)在被合并，我們現(xiàn)在有了一個(gè)新的query向量（旋轉(zhuǎn)后的query向量），其形狀為 [17x128]，其中17是token的數(shù)量，128是query向量的維度,。 | | 準(zhǔn)備基向量: 基向量定義為一個(gè)由二維位置（p）索引的表情形（對(duì)于大小為128的嵌入，生成的形狀為[64]）,。此基向量本質(zhì)上是一個(gè)在0到θ之間等距離的線性間隔。 為提示生成旋轉(zhuǎn)向量 為了生成旋轉(zhuǎn)向量，我們將基向量乘以0到16之間的整數(shù)，這些整數(shù)是我們提示中token的位置,。此乘法后的結(jié)果形狀為[17x64]，其中每行表示token的旋轉(zhuǎn)向量,。 應(yīng)用RoPE到查詢向量 最后，我們可以將我們的查詢向量與旋轉(zhuǎn)向量結(jié)合，生成位置編碼的查詢向量。每個(gè)查詢向量現(xiàn)在被旋轉(zhuǎn)以反映它們?cè)谔崾局械奈恢谩?/span> 最終位置編碼的查詢向量 我們現(xiàn)在有一個(gè)大小為[17x128]的查詢向量，每個(gè)token都有它在提示中位置的旋轉(zhuǎn)編碼,。 | | 思路步驟 | 生成RoPE（旋轉(zhuǎn)位置嵌入）以編碼token的位置信息,。 將查詢向量與位置編碼向量結(jié)合，生成位置編碼的查詢向量。 | | 核心技術(shù)點(diǎn) | 旋轉(zhuǎn)位置編碼（RoPE）：使用旋轉(zhuǎn)位置嵌入（RoPE）將位置信息編碼到查詢向量中,。 向量旋轉(zhuǎn)：使用cos和sin函數(shù)對(duì)查詢向量進(jìn)行旋轉(zhuǎn)，結(jié)合位置編碼信息,。 |
zero_to_one_split_into_64_parts = torch.tensor(range(64))/64
zero_to_one_split_into_64_parts
tensor([0.0000, 0.0156, 0.0312, 0.0469, 0.0625, 0.0781, 0.0938, 0.1094, 0.1250,
0.1406, 0.1562, 0.1719, 0.1875, 0.2031, 0.2188, 0.2344, 0.2500, 0.2656,
0.2812, 0.2969, 0.3125, 0.3281, 0.3438, 0.3594, 0.3750, 0.3906, 0.4062,
0.4219, 0.4375, 0.4531, 0.4688, 0.4844, 0.5000, 0.5156, 0.5312, 0.5469,
0.5625, 0.5781, 0.5938, 0.6094, 0.6250, 0.6406, 0.6562, 0.6719, 0.6875,
0.7031, 0.7188, 0.7344, 0.7500, 0.7656, 0.7812, 0.7969, 0.8125, 0.8281,
0.8438, 0.8594, 0.8750, 0.8906, 0.9062, 0.9219, 0.9375, 0.9531, 0.9688,
0.9844])
freqs = 1.0 / (rope_theta ** zero_to_one_split_into_64_parts)
freqs
tensor([1.0000e+00, 8.1462e-01, 6.6360e-01, 5.4058e-01, 4.4037e-01, 3.5873e-01,
2.9223e-01, 2.3805e-01, 1.9392e-01, 1.5797e-01, 1.2869e-01, 1.0483e-01,
8.5397e-02, 6.9566e-02, 5.6670e-02, 4.6164e-02, 3.7606e-02, 3.0635e-02,
2.4955e-02, 2.0329e-02, 1.6560e-02, 1.3490e-02, 1.0990e-02, 8.9523e-03,
7.2927e-03, 5.9407e-03, 4.8394e-03, 3.9423e-03, 3.2114e-03, 2.6161e-03,
2.1311e-03, 1.7360e-03, 1.4142e-03, 1.1520e-03, 9.3847e-04, 7.6450e-04,
6.2277e-04, 5.0732e-04, 4.1327e-04, 3.3666e-04, 2.7425e-04, 2.2341e-04,
1.8199e-04, 1.4825e-04, 1.2077e-04, 9.8381e-05, 8.0143e-05, 6.5286e-05,
5.3183e-05, 4.3324e-05, 3.5292e-05, 2.8750e-05, 2.3420e-05, 1.9078e-05,
1.5542e-05, 1.2660e-05, 1.0313e-05, 8.4015e-06, 6.8440e-06, 5.5752e-06,
4.5417e-06, 3.6997e-06, 3.0139e-06, 2.4551e-06])
freqs_for_each_token = torch.outer(torch.arange(17), freqs)
freqs_cis = torch.polar(torch.ones_like(freqs_for_each_token), freqs_for_each_token)
freqs_cis.shape
# viewing tjhe third row of freqs_cis
value = freqs_cis[3]
plt.figure()
for i, element in enumerate(value[:17]):
plt.plot([0, element.real], [0, element.imag], color='blue', linewidth=1, label=f"Index: {i}")
plt.annotate(f"{i}", xy=(element.real, element.imag), color='red')
plt.xlabel('Real')
plt.ylabel('Imaginary')
plt.title('Plot of one row of freqs_cis')
plt.show()
q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)
q_per_token_as_complex_numbers.shape
torch.Size([17, 64])
q_per_token_as_complex_numbers_rotated = q_per_token_as_complex_numbers * freqs_cis
q_per_token_as_complex_numbers_rotated.shape
torch.Size([17, 64])
q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers_rotated)
q_per_token_split_into_pairs_rotated.shape
torch.Size([17, 64, 2])
q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)
q_per_token_rotated.shape
torch.Size([17, 128])
3、LLaMA的核心模塊
3.1,、MHA模塊—從零開始實(shí)現(xiàn)注意力機(jī)制：計(jì)算不同詞元之間的關(guān)聯(lián)度，提取重要信息,。包括Q,、K,、V矩陣的生成、位置編碼,、掩碼,、softmax計(jì)算注意力權(quán)重等
3.1.1、計(jì)算查詢(Query)、鍵(Key)：生成用于注意力機(jī)制中每一個(gè)token的query和key；
| 簡(jiǎn)介 | 查詢Query,、鍵Key,、值Value和輸出向量 | | 思路步驟 | Key鍵向量生成 >> key向量維度為128。 >> key的權(quán)重共享機(jī)制，key的權(quán)重只有query的四分之一,，這是因?yàn)槊?個(gè)頭共享一次，以減少計(jì)算量,。 >> key向量也進(jìn)行旋轉(zhuǎn)以加入位置信息，原因與query相同,。 | | Query向量生成及旋轉(zhuǎn)，增加位置信息。 查詢向量與鍵向量的維度匹配,。 旋轉(zhuǎn)操作用于引入位置編碼,。 | | 在這個(gè)階段，我們現(xiàn)在有了每個(gè)token旋轉(zhuǎn)后的query和key，其中每一個(gè)query和key現(xiàn)在的形狀是[17x128]。 | | 核心技術(shù)點(diǎn) | 矩陣乘法與旋轉(zhuǎn)位置編碼 |
k_layer0 = model["layers.0.attention.wk.weight"]
k_layer0 = k_layer0.view(n_kv_heads, k_layer0.shape[0] // n_kv_heads, dim)
k_layer0.shape
torch.Size([8, 128, 4096])
k_layer0_head0 = k_layer0[0]
k_layer0_head0.shape
torch.Size([128, 4096])
k_per_token = torch.matmul(token_embeddings, k_layer0_head0.T)
k_per_token.shape
torch.Size([17, 128])
k_per_token_split_into_pairs = k_per_token.float().view(k_per_token.shape[0], -1, 2)
k_per_token_split_into_pairs.shape
torch.Size([17, 64, 2])
k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token_split_into_pairs)
k_per_token_as_complex_numbers.shape
torch.Size([17, 64])
k_per_token_split_into_pairs_rotated = torch.view_as_real(k_per_token_as_complex_numbers * freqs_cis)
k_per_token_split_into_pairs_rotated.shape
torch.Size([17, 64, 2])
k_per_token_rotated = k_per_token_split_into_pairs_rotated.view(k_per_token.shape)
k_per_token_rotated.shape
torch.Size([17, 128])
3.1.2,、計(jì)算注意力分?jǐn)?shù)矩陣并掩碼：計(jì)算注意力分?jǐn)?shù)并進(jìn)行掩碼，確保模型不會(huì)關(guān)注未來(lái)的詞元
| 簡(jiǎn)介 | 查詢Query,、鍵Key、值Value和輸出向量 1,、計(jì)算注意力分?jǐn)?shù)矩陣：在下一步中，我們將Query和Key矩陣相乘，這樣做會(huì)給我們一個(gè)分?jǐn)?shù)矩陣，將每個(gè)token與另一個(gè)token映射或關(guān)聯(lián)起來(lái),。這個(gè)分?jǐn)?shù)矩陣描述了每個(gè)token的query與每個(gè)token的key之間的關(guān)系有多好(相關(guān)性)，這就是自注意力機(jī)制（Self Attention）。其中，分?jǐn)?shù)矩陣（qk_per_token）的形狀是[17x17]，其中17是prompt中的token數(shù)量,。
2,、對(duì)注意力分?jǐn)?shù)矩陣中的未來(lái)token進(jìn)行掩碼處理：然后，我們現(xiàn)在必須對(duì)Query-Key分?jǐn)?shù)進(jìn)行掩碼處理，在llama3的訓(xùn)練過(guò)程中，未來(lái)token的qk分?jǐn)?shù)會(huì)被屏蔽。 為什么？因?yàn)樵谟?xùn)練過(guò)程中，我們只學(xué)習(xí)使用過(guò)去的token來(lái)預(yù)測(cè)未來(lái)token,。 因此，在推理時(shí)，我們將未來(lái)token的評(píng)分設(shè)置為0,。
| | 思路步驟 | 計(jì)算Query向量與Key向量之間的匹配得分，即自注意力得分。 使用矩陣乘法和歸一化,。 自注意力的核心機(jī)制：查詢-鍵匹配得分,。 | | 訓(xùn)練過(guò)程中掩蓋未來(lái)的token，確保只使用過(guò)去的token進(jìn)行預(yù)測(cè)。 上三角掩碼矩陣，屏蔽未來(lái)token,。 | | 核心技術(shù)點(diǎn) | 分?jǐn)?shù)計(jì)算與未來(lái)詞元掩碼 |
def display_qk_heatmap(qk_per_token):
_, ax = plt.subplots()
im = ax.imshow(qk_per_token.to(float).detach(), cmap='viridis')
ax.set_xticks(range(len(prompt_split_as_tokens)))
ax.set_yticks(range(len(prompt_split_as_tokens)))
ax.set_xticklabels(prompt_split_as_tokens)
ax.set_yticklabels(prompt_split_as_tokens)
ax.figure.colorbar(im, ax=ax)
display_qk_heatmap(qk_per_token)
mask = torch.full((len(tokens), len(tokens)), float("-inf"), device=tokens.device)
mask = torch.triu(mask, diagonal=1)
mask
tensor([[0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf, -inf],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf, -inf],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf, -inf],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf, -inf],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., -inf],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
qk_per_token_after_masking = qk_per_token + mask
display_qk_heatmap(qk_per_token_after_masking)
qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_after_masking, dim=1).to(torch.bfloat16)
display_qk_heatmap(qk_per_token_after_masking_after_softmax)
3.1.3,、值（Value）的計(jì)算：根據(jù)注意力分?jǐn)?shù)和value向量計(jì)算最終的注意力輸出
| 簡(jiǎn)介 | 查詢Query、鍵Key,、值Value和輸出向量 接下來(lái)是value，接近注意力機(jī)制的最后一步,。 | | 思路步驟 | 這些分?jǐn)?shù)（0-1）用于確定每個(gè)token使用多少value矩陣。就和key一樣，value權(quán)重也在每4個(gè)注意力頭之間共享（以節(jié)省計(jì)算）,。因此，下面的Value權(quán)重矩陣的形狀是[8x128x4096] ,。 第一層，第一個(gè)注意力頭的Value權(quán)重矩陣如下所示 | | Softmax 歸一化得到注意力權(quán)重：對(duì)掩碼后的得分進(jìn)行Softmax歸一化，得到注意力權(quán)重。 使用Softmax函數(shù)將得分轉(zhuǎn)換為概率分布,。 | | Values（值向量）計(jì)算：使用Value向量和注意力權(quán)重計(jì)算最終的注意力輸出,。Value向量與Query和Key向量的共享機(jī)制。 | | 注意力計(jì)算：現(xiàn)在使用Value權(quán)重來(lái)獲取每個(gè)token的注意力值，矩陣的大小是[17x128]，其中17是prompt中的token數(shù)量，128是每個(gè)token的value向量的維度,。 注意力：與每個(gè)token的value相乘后得到的注意力向量的形狀為[17x128],。 | | 核心技術(shù)點(diǎn) | 分?jǐn)?shù)與值的權(quán)重合并 值向量的矩陣乘法 |
qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T)/(head_dim)**0.5
qk_per_token.shape
torch.Size([17, 17])
3.1.4,、多頭注意力的合并：將多個(gè)注意力頭的輸出合并為一個(gè)向量
| 簡(jiǎn)介 | 我們現(xiàn)在有了第一層中第一個(gè)頭的注意力值，然后，將運(yùn)行一個(gè)循環(huán)，并對(duì)第一層中的每一個(gè)頭執(zhí)行與上述單元格完全相同的數(shù)學(xué)運(yùn)算。 然后，得到了第一層所有32個(gè)頭的qkv_attention矩陣，接下來(lái)我將把所有的注意力分?jǐn)?shù)合并成一個(gè)大的矩陣，其大小為[17x4096] 我們快要結(jié)束了! | | 思路步驟 | 多頭注意力機(jī)制：對(duì)每個(gè)注意力頭分別計(jì)算注意力輸出，并將其拼接,。 多頭注意力的并行計(jì)算和拼接,。 | | 線性層變換：通過(guò)線性變換將拼接后的多頭注意力輸出進(jìn)行進(jìn)一步處理。 線性層的矩陣乘法,。 | | 核心技術(shù)點(diǎn) | 向量的拼接 |
qkv_attention_store = []
for head in range(n_heads):
q_layer0_head = q_layer0[head]
k_layer0_head = k_layer0[head//4] # key weights are shared across 4 heads
v_layer0_head = v_layer0[head//4] # value weights are shared across 4 heads
q_per_token = torch.matmul(token_embeddings, q_layer0_head.T)
k_per_token = torch.matmul(token_embeddings, k_layer0_head.T)
v_per_token = torch.matmul(token_embeddings, v_layer0_head.T)
q_per_token_split_into_pairs = q_per_token.float().view(q_per_token.shape[0], -1, 2)
q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)
q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers * freqs_cis[:len(tokens)])
q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)
k_per_token_split_into_pairs = k_per_token.float().view(k_per_token.shape[0], -1, 2)
k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token_split_into_pairs)
k_per_token_split_into_pairs_rotated = torch.view_as_real(k_per_token_as_complex_numbers * freqs_cis[:len(tokens)])
k_per_token_rotated = k_per_token_split_into_pairs_rotated.view(k_per_token.shape)
qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T)/(128)**0.5
mask = torch.full((len(tokens), len(tokens)), float("-inf"), device=tokens.device)
mask = torch.triu(mask, diagonal=1)
qk_per_token_after_masking = qk_per_token + mask
qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_after_masking, dim=1).to(torch.bfloat16)
qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
qkv_attention_store.append(qkv_attention)
len(qkv_attention_store)
stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)
stacked_qkv_attention.shape
torch.Size([17, 4096])
3.1.5,、注意力輸出權(quán)重矩陣：將合并后的注意力輸出轉(zhuǎn)換為嵌入向量的修正
| 簡(jiǎn)介 | 查詢Query、鍵Key,、值Value和輸出向量 | | 思路步驟 | 對(duì)于第0層注意力機(jī)制的最后步驟，其一是將注意力得分矩陣與權(quán)重矩陣相乘,。 這是一個(gè)簡(jiǎn)單的線性層，所以我們只需進(jìn)行矩陣乘法 | | 核心技術(shù)點(diǎn) | 線性層（全連接層）的應(yīng)用。 |
stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)
stacked_qkv_attention.shape
torch.Size([17, 4096])
3.1.6,、嵌入向量的更新：將注意力機(jī)制的輸出加到原始嵌入向量上，更新嵌入信息
| 簡(jiǎn)介 | 查詢Query,、鍵Key、值Value和輸出向量 現(xiàn)在，我們有了注意力后嵌入Value的變化，應(yīng)該加到原始的token嵌入中； 然后，對(duì)嵌入增量進(jìn)行歸一化，然后通過(guò)嵌入增量運(yùn)行一個(gè)前饋神經(jīng)網(wǎng)絡(luò)； | | 思路步驟 | | | 核心技術(shù)點(diǎn) | 殘差連接 |
3.2,、FFNN層：在嵌入向量上應(yīng)用前饋網(wǎng)絡(luò)，增加模型的非線性
| 簡(jiǎn)介 | 在Llama3中，加載前饋權(quán)重并實(shí)現(xiàn)前饋網(wǎng)絡(luò),。使用了一種名為SwiGLU的前饋網(wǎng)絡(luò)，這種網(wǎng)絡(luò)結(jié)構(gòu)在模型需要的時(shí)候，能夠有效地增加非線性。 在當(dāng)今的LLMs中，使用這種前饋網(wǎng)絡(luò)架構(gòu)已經(jīng)相當(dāng)標(biāo)準(zhǔn),。 | | 思路步驟 | 歸一化與前饋網(wǎng)絡(luò)：對(duì)注意力輸出進(jìn)行歸一化，并通過(guò)前饋神經(jīng)網(wǎng)絡(luò)進(jìn)行非線性變換,。 RMSNorm歸一化。 SwiGLU前饋網(wǎng)絡(luò)，用于增強(qiáng)非線性表達(dá)能力,。 | | | | | 核心技術(shù)點(diǎn) | SwiGLU激活函數(shù)與線性層 |
embedding_after_edit_normalized = rms_norm(embedding_after_edit, model["layers.0.ffn_norm.weight"])
embedding_after_edit_normalized.shape
torch.Size([17, 4096])
w1 = model["layers.0.feed_forward.w1.weight"]
w2 = model["layers.0.feed_forward.w2.weight"]
w3 = model["layers.0.feed_forward.w3.weight"]
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)
output_after_feedforward.shape
torch.Size([17, 4096])
3.3,、循環(huán)迭代
3.3.1、所有層的迭代：將上述過(guò)程在每一層中進(jìn)行迭代處理，逐層增強(qiáng)表示能力
| 背景 | 現(xiàn)在完成了第一層之后每個(gè)token的新嵌入?，F(xiàn)在只剩下31層了，只需通過(guò)一個(gè)循環(huán)來(lái)完成,。 | | 簡(jiǎn)介 | 查詢Query、鍵Key,、值Value和輸出向量 我們終于在第1層之后為每個(gè)詞元獲得了新的編輯后的嵌入,。我們還有31層才能完成（還需要一個(gè)循環(huán)）。 你可以將這個(gè)編輯后的嵌入，想象成包含了第一層中所有Query信息的嵌入,。隨著層數(shù)的增加，每一層都會(huì)對(duì)輸入的信息進(jìn)行越來(lái)越復(fù)雜的處理，直到最終得到一個(gè)，能夠全面了解下一個(gè)需要預(yù)測(cè)的token的嵌入,。 天哪，一切都集中在一起了 是的，就是這樣。我們之前所做的所有事情，現(xiàn)在一次性為每個(gè)層都做了 | | 思路步驟 | 重復(fù)以上步驟，處理所有Transformer子層 | | 核心技術(shù)點(diǎn) | 循環(huán)與層堆疊 循環(huán)迭代處理每一層,。 |
layer_0_embedding = embedding_after_edit+output_after_feedforward
layer_0_embedding.shape
torch.Size([17, 4096])
3.3.2,、最后的歸一化：對(duì)最終嵌入向量進(jìn)行歸一化處理，用于預(yù)測(cè)下一個(gè)token
| 簡(jiǎn)介 | | | 思路步驟 | | | 核心技術(shù)點(diǎn) | 矩陣乘法，線性變換。 最終歸一化處理 |
final_embedding = token_embeddings_unnormalized
for layer in range(n_layers):
qkv_attention_store = []
layer_embedding_norm = rms_norm(final_embedding, model[f"layers.{layer}.attention_norm.weight"])
q_layer = model[f"layers.{layer}.attention.wq.weight"]
q_layer = q_layer.view(n_heads, q_layer.shape[0] // n_heads, dim)
k_layer = model[f"layers.{layer}.attention.wk.weight"]
k_layer = k_layer.view(n_kv_heads, k_layer.shape[0] // n_kv_heads, dim)
v_layer = model[f"layers.{layer}.attention.wv.weight"]
v_layer = v_layer.view(n_kv_heads, v_layer.shape[0] // n_kv_heads, dim)
w_layer = model[f"layers.{layer}.attention.wo.weight"]
for head in range(n_heads):
q_layer_head = q_layer[head]
k_layer_head = k_layer[head//4]
v_layer_head = v_layer[head//4]
q_per_token = torch.matmul(layer_embedding_norm, q_layer_head.T)
k_per_token = torch.matmul(layer_embedding_norm, k_layer_head.T)
v_per_token = torch.matmul(layer_embedding_norm, v_layer_head.T)
q_per_token_split_into_pairs = q_per_token.float().view(q_per_token.shape[0], -1, 2)
q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)
q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers * freqs_cis)
q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)
k_per_token_split_into_pairs = k_per_token.float().view(k_per_token.shape[0], -1, 2)
k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token_split_into_pairs)
k_per_token_split_into_pairs_rotated = torch.view_as_real(k_per_token_as_complex_numbers * freqs_cis)
k_per_token_rotated = k_per_token_split_into_pairs_rotated.view(k_per_token.shape)
qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T)/(128)**0.5
mask = torch.full((len(token_embeddings_unnormalized), len(token_embeddings_unnormalized)), float("-inf"))
mask = torch.triu(mask, diagonal=1)
qk_per_token_after_masking = qk_per_token + mask
qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_after_masking, dim=1).to(torch.bfloat16)
qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
qkv_attention_store.append(qkv_attention)
stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)
w_layer = model[f"layers.{layer}.attention.wo.weight"]
embedding_delta = torch.matmul(stacked_qkv_attention, w_layer.T)
embedding_after_edit = final_embedding + embedding_delta
embedding_after_edit_normalized = rms_norm(embedding_after_edit, model[f"layers.{layer}.ffn_norm.weight"])
w1 = model[f"layers.{layer}.feed_forward.w1.weight"]
w2 = model[f"layers.{layer}.feed_forward.w2.weight"]
w3 = model[f"layers.{layer}.feed_forward.w3.weight"]
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)
final_embedding = embedding_after_edit+output_after_feedforward
4,、輸出層
4.1,、將嵌入向量轉(zhuǎn)換為預(yù)測(cè)的token：確認(rèn)logits張量的形狀
| 簡(jiǎn)介 | 生成最終的嵌入向量，用于預(yù)測(cè)下一個(gè)token的最優(yōu)預(yù)測(cè)。這個(gè)嵌入的形狀與常規(guī)的token嵌入相同，為[17x4096]，其中17是token的數(shù)量，4096是嵌入的維度,。 | | 思路步驟 | 確認(rèn)logits張量的形狀：確保計(jì)算得到的logits具有正確的維度，以便后續(xù)處理,。 | | 核心技術(shù)點(diǎn) | 線性層與softmax激活函數(shù) |
final_embedding = rms_norm(final_embedding, model["norm.weight"])
final_embedding.shape
torch.Size([17, 4096])
4.2,、解碼：將模型輸出的詞元轉(zhuǎn)換為可讀的文本
| 簡(jiǎn)介 | 我們將使用輸出解碼器，將最終嵌入轉(zhuǎn)換為token值，使用輸出解碼器將最終的嵌入轉(zhuǎn)換成一個(gè)token,。確認(rèn)輸出解碼器的尺寸是否正確，確保矩陣乘法可以順利進(jìn)行,。 使用最后一個(gè)詞元的嵌入通過(guò)矩陣乘法預(yù)測(cè)下一個(gè)令牌的概率分布。 通過(guò)最后一個(gè)令牌的嵌入和輸出解碼器權(quán)重矩陣的乘法，計(jì)算出每個(gè)可能令牌的logits,。 | | 思路步驟 | 獲取預(yù)測(cè)的下一個(gè)令牌的索引，通過(guò)argmax函數(shù)從logits中選擇概率最高的令牌，來(lái)確定下一個(gè)令牌，得到最終的預(yù)測(cè)結(jié)果,。 | | 核心技術(shù)點(diǎn) | 分詞器的逆過(guò)程 argmax函數(shù)，分類問(wèn)題中的預(yù)測(cè)。 |
model["output.weight"].shape
torch.Size([128256, 4096])
5,、模型推理：應(yīng)用訓(xùn)練好的輸出層參數(shù)進(jìn)行預(yù)測(cè)(驗(yàn)證整個(gè)代碼的正確性)，模型預(yù)測(cè)下一個(gè)token的編號(hào)為2983，這是42的token ID嗎？
| 背景 | 期望：使用最后一個(gè)token的嵌入來(lái)預(yù)測(cè)下一個(gè)value，希望預(yù)測(cè)的結(jié)果是42,。加載的llama3模型在樣本問(wèn)題"答案是什么"上正確預(yù)測(cè)答案為42，驗(yàn)證了從零開始實(shí)現(xiàn)的整個(gè)流程是正確的。 希望在我們這個(gè)例子中，是42 :) 方法：查看預(yù)測(cè)的下一個(gè)令牌是否為“42” 注意：因?yàn)楦鶕?jù)《銀河系漫游指南》一書中的說(shuō)法，42是“生命,、宇宙及一切的終極問(wèn)題的答案”,。大多數(shù)現(xiàn)代LLM在這里都會(huì)回答42，這將驗(yàn)證整個(gè)代碼的正確性。 | | 簡(jiǎn)介 | 模型預(yù)測(cè)下一個(gè)詞元是token編號(hào)2983，這是42的token ID嗎？ 我正在讓你興奮，這是最后一行的代碼，希望你能享受閱讀！ 結(jié)束！ | | 思路步驟 | 輸出預(yù)測(cè)結(jié)果 打印或返回預(yù)測(cè)的令牌索引，完成整個(gè)模型推理過(guò)程,。 |
logits = torch.matmul(final_embedding[-1], model["output.weight"].T)
logits.shape
torch.Size([128256])
next_token = torch.argmax(logits, dim=-1)
next_token
tensor(2983)
tokenizer.decode([next_token.item()])
'42'
|
