Karpathy稱贊，從零實(shí)現(xiàn)LLaMa3項(xiàng)目爆火，半天1.5k star
作者：機(jī)器之心
十幾個小時前，有位名為「Nishant Aklecha」的開發(fā)者發(fā)布了一個從零開始實(shí)現(xiàn) llama3 的存儲庫，包括跨多個頭的注意力矩陣乘法、位置編碼和每個層在內(nèi)都有非常詳細(xì)的解釋。
一個月前，Meta 發(fā)布了開源大模型 llama3 系列，在多個關(guān)鍵基準(zhǔn)測試中優(yōu)于業(yè)界 SOTA 模型，并在代碼生成任務(wù)上全面領(lǐng)先。
此后，開發(fā)者們便開始了本地部署和實(shí)現(xiàn)，比如 llama3 的中文實(shí)現(xiàn)、llama3 的純 NumPy 實(shí)現(xiàn)等。
十幾個小時前，有位名為「Nishant Aklecha」的開發(fā)者發(fā)布了一個從零開始實(shí)現(xiàn) llama3 的存儲庫，包括跨多個頭的注意力矩陣乘法、位置編碼和每個層在內(nèi)都有非常詳細(xì)的解釋。
該項(xiàng)目得到了大神 Karpathy 的稱贊，他表示項(xiàng)目看起來不錯，完全展開后，通過模塊嵌套和相互調(diào)用，可以更容易看到實(shí)際的情況。
上傳半天的時間，該項(xiàng)目已在 GitHub 上收獲了 1.5k 的 star，足可見其含金量。
從零開始實(shí)現(xiàn) llama3
接下來項(xiàng)目作者手把手教你如何從頭開始實(shí)現(xiàn) llama3。
項(xiàng)目地址：https://github.com/naklecha/llama3-from-scratch
首先從 Meta 提供的 llama3 模型文件中加載張量。
下載地址：https://llama.meta.com/llama-downloads/
接著是分詞器（tokenizer），作者表示沒打算自己實(shí)現(xiàn)分詞器，因而借用了 Andrej Karpathy 的實(shí)現(xiàn)方式：
分詞器的實(shí)現(xiàn)鏈接：https://github.com/karpathy/minbpe
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!"))'hello world!'上述步驟完成后，就是讀取模型文件了。由于該研究是從頭開始實(shí)現(xiàn) llama3，因此代碼一次只讀取一個張量文件。
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}項(xiàng)目作者使用以下配置來推斷模型細(xì)節(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"])接下來的操作是將文本裝換為 token，這里作者使用的是 tiktoken 庫（一個用于 OpenAI 模型的 BPE tokeniser）。
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', ' ']然后將 token 轉(zhuǎ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.shapetorch.Size ([17, 4096])將嵌入進(jìn)行歸一化。該研究使用均方根 RMS 算法進(jìn)行歸一化。不過，在這一步之后，張量形狀不會改變，只是值進(jìn)行了歸一化。
# 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òu)建 transformer 第一層。完成上述準(zhǔn)備后，接著是構(gòu)建 transformer 第一層：從模型文件中訪問 layer.0（即第一層），歸一化后嵌入維度仍然是 [17x4096] 。
token_embeddings = rms_norm (token_embeddings_unnormalized, model ["layers.0.attention_norm.weight"])
token_embeddings.shapetorch.Size ([17, 4096])從頭開始實(shí)現(xiàn)注意力。加載第一層 transformer 的注意力頭：
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])展開查詢。展開來自多個注意力頭的查詢，得到的形狀是 [32x128x4096]，這里，32 是 llama3 中注意力頭的數(shù)量，128 是查詢向量的大小，4096 是 token 嵌入的大小。
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.shapetorch.Size ([32, 128, 4096])從頭實(shí)現(xiàn)第一層的第一個頭。訪問第一層的查詢權(quán)重矩陣，大小是 [128x4096]。
q_layer0_head0 = q_layer0 [0]
q_layer0_head0.shapetorch.Size ([128, 4096])將查詢權(quán)重與 token 嵌入相乘，從而得到 token 的查詢，在這里你可以看到結(jié)果大小是 [17x128]。
q_per_token = torch.matmul (token_embeddings, q_layer0_head0.T)
q_per_token.shapetorch.Size ([17, 128])定位編碼?，F(xiàn)在處于這樣一個階段，即對提示符中的每個 token 都有一個查詢向量，但是考慮單個查詢向量，我們不知道其提示符中的位置。作者使用了 RoPE（旋轉(zhuǎn)位置嵌入）來解決。
q_per_token_split_into_pairs = q_per_token.float ().view (q_per_token.shape [0], -1, 2)
q_per_token_split_into_pairs.shapetorch.Size ([17, 64, 2])在上面的步驟中，該研究將查詢向量分成對，并對每對應(yīng)用旋轉(zhuǎn)角度移位。
使用復(fù)數(shù)點(diǎn)積來旋轉(zhuǎn)向量。
zero_to_one_split_into_64_parts = torch.tensor (range (64))/64
zero_to_one_split_into_64_partstensor ([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)
freqstensor ([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 ()
現(xiàn)在每個 token 查詢都有了復(fù)數(shù)。
q_per_token_as_complex_numbers = torch.view_as_complex (q_per_token_split_into_pairs)
q_per_token_as_complex_numbers.shapetorch.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.shapetorch.Size ([17, 64])旋轉(zhuǎn)后的向量。
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.shapetorch.Size ([17, 64, 2])現(xiàn)在有了一個新的查詢向量 (旋轉(zhuǎn)查詢向量)，形狀為 [17x128]，其中 17 是 token 數(shù)量，128 是查詢向量的維度。
q_per_token_rotated = q_per_token_split_into_pairs_rotated.view (q_per_token.shape)
q_per_token_rotated.shapetorch.Size ([17, 128])鍵（幾乎和查詢一樣），鍵也生成維度為 128 的鍵向量。鍵的權(quán)重只有查詢的 1/4，這是因?yàn)殒I的權(quán)重在 4 個頭之間共享，以減少所需的計算量，鍵也會被旋轉(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.shapetorch.Size ([8, 128, 4096])k_layer0_head0 = k_layer0 [0]
k_layer0_head0.shapetorch.Size ([128, 4096])k_per_token = torch.matmul (token_embeddings, k_layer0_head0.T)
k_per_token.shapetorch.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.shapetorch.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.shapetorch.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.shapetorch.Size ([17, 64, 2])k_per_token_rotated = k_per_token_split_into_pairs_rotated.view (k_per_token.shape)
k_per_token_rotated.shapetorch.Size ([17, 128])每個 token 查詢和鍵的旋轉(zhuǎn)值如下，每個查詢和鍵現(xiàn)在的形狀都是 [17x128]。
接下來一步是將查詢和鍵矩陣相乘。注意力得分矩陣 (qk_per_token) 的形狀為 [17x17]，其中 17 是提示中 token 的數(shù)量。
qk_per_token = torch.matmul (q_per_token_rotated, k_per_token_rotated.T)/(head_dim)**0.5
qk_per_token.shapetorch.Size ([17, 17])現(xiàn)在必須掩蔽查詢鍵分?jǐn)?shù)。
在 llama3 的訓(xùn)練過程中，未來 token 的 qk 分?jǐn)?shù)被掩蔽。這是因?yàn)樵谟?xùn)練期間，只學(xué)習(xí)使用過去的 token 來預(yù)測未來的 token。因此在推理過程中，將未來的 token 標(biāo)記為零。
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, diagnotallow=1) masktensor ([[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)
值（幾乎在注意力結(jié)束時）
這些分?jǐn)?shù) (0-1) 被用于確定每個 token 使用了多少值矩陣。
- 就像鍵一樣，值權(quán)重也在 4 個注意力頭之間共享（以節(jié)省計算量）
- 結(jié)果，下面的值權(quán)重矩陣形狀為 [8x128x4096]
v_layer0 = model ["layers.0.attention.wv.weight"] v_layer0 = v_layer0.view (n_kv_heads, v_layer0.shape [0] //n_kv_heads, dim) v_layer0.shapetorch.Size ([8, 128, 4096])第一層和第一個頭的值權(quán)重矩陣如下所示。
v_layer0_head0 = v_layer0 [0] v_layer0_head0.shapetorch.Size ([128, 4096])值向量如下圖所示。
現(xiàn)在使用值權(quán)重來獲取每個 token 的注意力值，其大小為 [17x128]，其中 17 為提示中的 token 數(shù)，128 為每個 token 的值向量維數(shù)。
v_per_token = torch.matmul (token_embeddings, v_layer0_head0.T)v_per_token.shapetorch.Size ([17, 128])注意力如下圖所示。
與每個 token 的值相乘后得到的注意力向量的形狀為 [17*128]。
qkv_attention = torch.matmul (qk_per_token_after_masking_after_softmax, v_per_token) qkv_attention.shapetorch.Size ([17, 128])多頭注意力與單頭注意力如下圖所示。
現(xiàn)在有了第一層和第一個頭的注意力值。
接下來運(yùn)行一個循環(huán)并執(zhí)行與上面單元完全相同的數(shù)學(xué)運(yù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, diagnotallow=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)32
現(xiàn)在第一層上的所有 32 個頭都有了 qkv_attention 矩陣，并在快結(jié)束的時候?qū)⑺凶⒁饬Ψ謹(jǐn)?shù)合并為一個大小為 [17x4096] 的大矩陣。
stacked_qkv_attention = torch.cat (qkv_attention_store, dim=-1) stacked_qkv_attention.shapetorch.Size ([17, 4096])權(quán)重矩陣是最后的步驟之一。
第 0 層注意力要做的最后一件事是，對以下的權(quán)重矩陣進(jìn)行乘法操作。
w_layer0 = model ["layers.0.attention.wo.weight"] w_layer0.shapetorch.Size ([4096, 4096])這是一個簡單的線性層，所以只做矩陣乘法（matmul）。
embedding_delta = torch.matmul (stacked_qkv_attention, w_layer0.T) embedding_delta.shapetorch.Size ([17, 4096])
現(xiàn)在，注意力之后的嵌入值有了變化，并應(yīng)該被添加到原始 token 嵌入中。
embedding_after_edit = token_embeddings_unnormalized + embedding_delta
embedding_after_edit.shapetorch.Size ([17, 4096])歸一化并在嵌入 delta 過程中運(yùn)行一個前饋神經(jīng)網(wǎng)絡(luò)。
embedding_after_edit_normalized = rms_norm (embedding_after_edit, model ["layers.0.ffn_norm.weight"]) embedding_after_edit_normalized.shapetorch.Size ([17, 4096])加載 ff 權(quán)重，并實(shí)現(xiàn)前饋網(wǎng)絡(luò)。
llama3 使用 SwiGLU 前饋網(wǎng)絡(luò)，該網(wǎng)絡(luò)架構(gòu)非常擅長在模型需要時添加非線性。當(dāng)前，在 LLMs 中使用這一前饋網(wǎng)絡(luò)是非常標(biāo)準(zhǔn)的做法。
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.shapetorch.Size ([17, 4096])現(xiàn)在終于在第一層之后為每個 token 提供了新的編輯后的嵌入，并且在完成之前只剩下 31 層需要處理（one for loop away）。
你可以想象這個編輯后的嵌入擁有在第一層上所有查詢的信息?，F(xiàn)在每一層將在所問問題上編碼越來越復(fù)雜的查詢，直到得到的嵌入了解所需的下一個 token 的一切。
layer_0_embedding = embedding_after_edit+output_after_feedforward
layer_0_embedding.shapetorch.Size ([17, 4096])之前為每一層做的所有事情，都可以一次性完成。
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, diagnotallow=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現(xiàn)在有了最終的嵌入，即該模型對下一個 token 的最佳猜測。該嵌入的形狀與常見的 token 嵌入 [17x4096] 相同，其中 17 為 token 數(shù)，4096 為嵌入維數(shù)。
final_embedding = rms_norm (final_embedding, model ["norm.weight"]) final_embedding.shapetorch.Size ([17, 4096])將該嵌入解碼為 token 值。
使用該輸入解碼器將最終的嵌入轉(zhuǎn)換為一個 token。
model ["output.weight"].shapetorch.Size ([128256, 4096])使用最后 token 的嵌入來預(yù)測下一個值。在示例中，42 是「生命、宇宙和萬物終極問題的答案是什么」的答案，根據(jù)《銀河系漫游指南》一書，大多數(shù)現(xiàn)代 LLMs 都會回答 42，應(yīng)該驗(yàn)證了整個代碼。
logits = torch.matmul (final_embedding [-1], model ["output.weight"].T) logits.shapetorch.Size ([128256])模型預(yù)測 token 數(shù) 2983 為下一個 token，這是 42 的 token 數(shù)嗎？以下是最后的代碼單元。
next_token = torch.argmax (logits, dim=-1) next_tokentensor (2983)最后，啟動。
tokenizer.decode ([next_token.item ()])'42'完結(jié)撒花
責(zé)任編輯：張燕妮
來源：
機(jī)器之心
