250行代碼從頭搭建Llama 3，GitHub一天4.6k星！Karpathy大贊
Llama系列作為為數(shù)不多的優(yōu)質(zhì)開源LLM，一直受到開發(fā)者們的追捧。在Hugging Face社區(qū)的文本生成模型中，幾乎是「霸榜」的存在。
就在520這天，一位名叫Nishant Aklecha的開發(fā)者在推特上宣布了自己的一個開源項目，名為「從頭開始實現(xiàn)Llama 3」。
這個項目詳細(xì)到什么程度呢——
矩陣乘法、注意力頭、位置編碼等模塊全部都拆開解釋。
圖片
而且項目全部用Jupyter Notebook寫成，小白都可以直接上手運(yùn)行。
堪比哈佛NLP小組曾經(jīng)出品的「The Annotated Transformer」。
圖片
https://nlp.seas.harvard.edu/annotated-transformer/
才一天多的時間，小哥發(fā)表的這篇推特已經(jīng)有32萬次閱讀，甚至被Andrej Karpathy大佬親自點(diǎn)贊——
「全部拆開解釋之后，通過模塊的嵌套以及互相調(diào)用，可以更清楚地看到模型到底做了什么?！?/span>
圖片
項目也在GitHub上獲得了4.6k星。
圖片
項目地址：https://github.com/naklecha/llama3-from-scratch
那就讓我們來看看作者是如何深入拆解Llama 3的。
下載并讀取模型權(quán)重
首先需要從Meta官網(wǎng)下載模型權(quán)重文件，以便后續(xù)運(yùn)行時使用。
圖片
https://github.com/meta-llama/llama3/blob/main/README.md
下載后需要先讀取權(quán)重文件中的變量名：
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}根據(jù)以上輸出，可以推斷出模型架構(gòu)的信息——
- 32個transformer層
- 每個多頭注意力模塊有32個注意力頭
- 分詞器的詞匯量為128256
直接將模型配置信息存儲到變量中，方便使用。
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"])分詞器與編碼
那么就從語言模型的第一步——分詞器開始，但是這一步并不需要我們自己手寫。
Llama 3使用了GPT等大模型常用的BPE分詞器，karpathy大佬之前就復(fù)現(xiàn)過一個最簡版。
圖片
https://github.com/karpathy/minbpe
除了Karapthy大佬復(fù)現(xiàn)的版本，OpenAI也開源了一個運(yùn)行速度很快的分詞器tiktoken。這兩個隨便挑，估計都比自己從頭訓(xùn)練的要強(qiáng)。
圖片
https://github.com/openai/tiktoken
有了分詞器，下一步就是要把輸入的文本切分為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', ' ']再使用PyTorch內(nèi)置的神經(jīng)網(wǎng)絡(luò)模塊（torch.nn）將token轉(zhuǎn)換為embedding，[17x1]的token維度變?yōu)閇17x4096]。
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])此處應(yīng)該是整個項目中唯一使用PyTorch內(nèi)置模塊的地方。而且，作者給出了溫馨提示——記得經(jīng)常打印一下張量維度，更容易理解。
之后再使用RMS對embedding進(jìn)行歸一化處理。這一步不會改變張量形狀，只是歸一化其中的數(shù)值，公式如下：
圖片
模型配置中的norm_eps變量設(shè)置為1e-5，就是用在此處，防止rms值意外設(shè)置為0。
# 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層
每一個Transformer層都需要經(jīng)過如下步驟：
圖片
由于是從頭構(gòu)建，我們只需要訪問模型字典中第一層（layer.0）的權(quán)重。
先用剛才定義的rms_norm函數(shù)，結(jié)合模型權(quán)重，進(jìn)行embedding的歸一化處理。
token_embeddings = rms_norm(token_embeddings_unnormalized, model["layers.0.attention_norm.weight"])
token_embeddings.shapetorch.Size([17, 4096])多頭注意力
查詢向量
讓我們先用一張圖復(fù)習(xí)注意力機(jī)制的計算過程：
圖片
如果從模型直接加載查詢、鍵、值和輸出的權(quán)重，我們會得到四個二維矩陣，形狀分別為 [4096x4096]、[1024x4096]、[1024x4096]、[4096x4096]。
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])因為大模型考慮了注意力中乘法并行化的需求，壓縮了矩陣維度。但是為了更清楚地展示機(jī)制，作者決定將這些矩陣都展開。
模型有32個注意力頭，因此查詢權(quán)重矩陣應(yīng)該展開為[32x128x4096]，其中128是查詢向量的長度，4096是embedding的維度。
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])于是可以訪問第一個注意力頭的查詢權(quán)重，維度是[128x4096]。
q_layer0_head0 = q_layer0[0]
q_layer0_head0.shapetorch.Size([128, 4096])現(xiàn)在將查詢權(quán)重與embedding相乘，就得到了查詢矩陣，維度為[17x128]，表示長度為17的句子，其中每個token都有維度為128的查詢向量。
q_per_token = torch.matmul(token_embeddings, q_layer0_head0.T)
q_per_token.shapetorch.Size([17, 128])位置編碼
由于注意力機(jī)制中對每個token沒有序列「位置」的概念，第一個詞和最后一個詞在Q、K、V矩陣看來都是一樣的，因此需要在查詢向量中嵌入維度為[1x128]的位置編碼。
位置編碼有多種方法，Llama模型采用的是旋轉(zhuǎn)位置編碼（RoPE）。
圖片
首先將查詢向量兩兩分為一對，共有64對。
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])句子中在m位置的一對查詢向量，旋轉(zhuǎn)角度為m*(rope_theta)，其中rope_theta也在模型的配置信息中。
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)經(jīng)過以上操作后，我們構(gòu)建了freq_cis矩陣，存儲句子中每個位置的、對查詢向量每個值的旋轉(zhuǎn)角度。
圖片
將每對查詢向量轉(zhuǎn)換為復(fù)數(shù)，之后進(jìn)行與旋轉(zhuǎn)角度進(jìn)行點(diǎn)積操作。
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)后的查詢向量，需要再轉(zhuǎn)換回實數(shù)形式。
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])q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)
q_per_token_rotated.shapetorch.Size([17, 128])旋轉(zhuǎn)后的查詢向量，維度依舊是 [17x128]。
鍵向量
鍵向量的計算與查詢向量非常類似，也需要進(jìn)行旋轉(zhuǎn)位置編碼，只是維度有所差異。
鍵的權(quán)重數(shù)量僅為查詢的1/4，因為需要減少模型計算量，每個權(quán)重值被4個注意力頭共享。
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])因此這里第一個維度的值為8，而不是我們在查詢權(quán)重中看到的32。
k_layer0_head0 = k_layer0[0]
k_per_token = torch.matmul(token_embeddings, k_layer0_head0.T)
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)
k_per_token_rotated.shapetorch.Size([17, 128])照著前面查詢向量部分的計算流程，就可以得到句子中每個token的鍵向量了。
查詢和鍵相乘
對句子進(jìn)行「自注意力」的過程，就是將查詢向量和鍵向量相乘，得到的QK矩陣中的每個值描述了對應(yīng)位置token查詢值和鍵值的相關(guān)程度。
圖片
相乘后，我們會得到一個維度為[17x17]自注意力矩陣。
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])掩碼
語言模型的學(xué)習(xí)目標(biāo)，是根據(jù)句子之前的內(nèi)容預(yù)測下一個token，因此訓(xùn)練和推理時需要將token位置之后的QK分?jǐn)?shù)屏蔽。
圖片
值向量
值權(quán)重數(shù)量和鍵權(quán)重一樣，都是在4個注意力頭之間共享（以節(jié)省計算量）。
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)重，與句子embedding相乘，獲取值向量。
v_layer0_head0 = v_layer0[0]
v_per_token = torch.matmul(token_embeddings, v_layer0_head0.T)
v_per_token.shapetorch.Size([17, 128])注意力向量
圖片
將進(jìn)行過掩碼的QK矩陣和句子的值向量相乘，就得到了注意力矩陣，維度為[17x128]。
qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
qkv_attention.shapetorch.Size([17, 128])多頭注意力
以上得到的注意力矩陣，是第一層第一個注意力頭的計算結(jié)果。
接下來需要運(yùn)行一個循環(huán)，對第一層中所有32個注意力頭進(jìn)行上述運(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為了并行計算的方便，我們需要把上面展開的矩陣壓縮回去。
也就是將32個維度為[17x128]的注意力矩陣，壓縮成一個維度為[17x4096]的大矩陣。
stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)
stacked_qkv_attention.shapetorch.Size([17, 4096])最后，別忘了乘以輸出權(quán)重矩陣。
w_layer0 = model["layers.0.attention.wo.weight"]
w_layer0.shape
# torch.Size([4096, 4096])
embedding_delta = torch.matmul(stacked_qkv_attention, w_layer0.T)
embedding_delta.shapetorch.Size([17, 4096])至此，注意力模塊的計算就結(jié)束了。
相加與歸一化
圖片
對照這張Transformer層的架構(gòu)圖，在多頭自注意力模塊之后還需要完成一些運(yùn)算。
首先將注意力模塊的輸出與原始的embedding相加。
embedding_after_edit = token_embeddings_unnormalized + embedding_delta
embedding_after_edit.shapetorch.Size([17, 4096])之后進(jìn)行RMS歸一化。
embedding_after_edit_normalized = rms_norm(embedding_after_edit, model["layers.0.ffn_norm.weight"])
embedding_after_edit_normalized.shapetorch.Size([17, 4096])前饋神經(jīng)網(wǎng)絡(luò)層
Llama 3的Transformer層中使用了SwiGLU前饋網(wǎng)絡(luò)，這種架構(gòu)非常擅長在必要情況下為模型添加非線性，這也是當(dāng)今LLM中的常見操作。
SwiGLU與Vanilla兩種前饋神經(jīng)網(wǎng)絡(luò)架構(gòu)的對比
于是我們從模型中加載前饋網(wǎng)絡(luò)的權(quá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])別忘了前饋層之后還有一次相加。
layer_0_embedding = embedding_after_edit+output_after_feedforward
layer_0_embedding.shapetorch.Size([17, 4096])以上就是一個完整Transformer層的實現(xiàn)，最終輸出的向量維度為[17x4096]，相當(dāng)于為句子中每個token重新計算了一個長度為4096的embedding向量。
預(yù)測下一個輸出
之后的每一個Transformer層都會編碼出越來越復(fù)雜的查詢，直到最后一層的輸出的embedding可以預(yù)測句子下一個token。
因此需要再嵌套一個外層循環(huán)，將Transformer層的流程重復(fù)32次。
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最后一個Transformer層的輸出維度與第一層相同，依舊是[17x4096]。
final_embedding = rms_norm(final_embedding, model["norm.weight"])
final_embedding.shapetorch.Size([17, 4096])此時需要利用輸出解碼器，將最后一層輸出的embedding先進(jìn)行歸一化處理，再轉(zhuǎn)換為token。
final_embedding = rms_norm(final_embedding, model["norm.weight"])
final_embedding.shape
# torch.Size([17, 4096])model["output.weight"].shape
# torch.Size([128256, 4096])logits = torch.matmul(final_embedding[-1], model["output.weight"].T)
logits.shapetorch.Size([128256])輸出的向量維度與分詞器中詞匯數(shù)量相同，每個值代表了下一個token的預(yù)測概率。
模型預(yù)測下一個詞是42?
和《銀河系漫游指南》的夢幻聯(lián)動（不知道是不是作者故意設(shè)置成這樣的）
next_token = torch.argmax(logits, dim=-1)
next_tokentensor(2983)tokenizer.decode([next_token.item()])'42'至此，我們就完成了Llama 3對輸入句子進(jìn)行下一個token預(yù)測的全過程。
