机器学习 | 从0开发大模型—译llama3-from-scratch

最近在看一篇Github上大佬的文章，从0开始训练llama3，觉得对于《从0开发大模型》有点帮助，于是翻译一下，发现其中很多内容当前系列文章的知识点相似。原文：https://github.com/naklecha/llama3-from-scratch其中meta-llama/Meta-Llama-3-8B文件地址：https://huggingface.co/meta-llama/Meta-Llama-3-8B/tree/main/original

1、Tokenizer

原始代码没有实现tokenizer，而是使用llama3的 tokenizer.model，实现代码如下：

# 执行：pip install blobfile
# 执行：pip install tiktoken

from pathlib import Path
import tiktoken
from tiktoken.load import load_tiktoken_bpe

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)},
)

print(tokenizer.decode(tokenizer.encode("hello world!")))

## 输出
hello world!

这里用了字节对编码（BPE），和我们训练的tokenzier使用的方式一样。

2、读取模型文件

将模型文件下载到 Meta-Llama-3-8B 文件夹中，然后读取模型文件，代码如下：

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

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

## 输出
[
    "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"
]

{
    '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
}

其中输出的配置看：

• n_layers=32：表示该模型有32个Transformer层
• n_heads=32：表示每个Transformer层有32个注意力头
• vobac_size=128256：表示词汇表大小为128256

3、文本转换为token

使用 tiktoken（openai的库）作为 tokenizer，实现如下：

llama3-scratch

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', ' ']

其中，128000是 <|begin_of_text|> 的token，还包括如下特殊token：

"<|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|>"

4、将token转换为embedding

将上面的 token 通过 embedding 层，[17X1] 转换为 [17X4096]，即 17 个 embeding（每个token一个），长度为 4096。

llama3-scratch

代码如下：

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])

5、使用RMS对embedding进行归一化

当前步骤不会改变形状，只是做标准化处理，需要用到norm_eps，代码如下：

def rms_norm(tensor, norm_weights):
    return (tensor * torch.rsqrt(tensor.pow(2).mean(-1, keepdim=True) + norm_eps)) * norm_weights

llama3-scratch

6、构建Transformer第一层

6.1、归一化

获取模型的第一层权重（layer.0），进行归一化处理：

llama3-scratch

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

## 输出
torch.Size([17, 4096])

6.2、注意力机制

先加载Transformer的第一层注意力头：

llama3-scratch

• 从模型中加载wq，wk，wv，wo的权重，分别是[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
)

## 输出
torch.Size([4096, 4096]) torch.Size([1024, 4096]) torch.Size([1024, 4096]) torch.Size([4096, 4096])

6.3、展开query

将多个注意力头展开query，可以得到的形状为 [32X128X4096]，其中32是llama3的头数，128是查询向量的大小，4096是token emebdding的大小，代码如下：

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])

6.4、实现第一层的第一个head

访问第一层的query权重矩阵第一个head，权重矩阵的大小是[128x4096]，打印一下：

q_layer0_head0 = q_layer0[0]
q_layer0_head0.shape

## 输出
torch.Size([128, 4096])

6.5、query权重和token embedding相乘

token embeding的形状是[17x128]，这是因为我们有17个标记，每个标记都有一个128长度的查询，将查询的权重和token embeding层相乘，代码如下：

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

## 输出
torch.Size([17, 128])

llama3-scratch

6.6、位置编码

当前阶段我们的token都有一个query向量，但是仔细想想——单独的query向量根本不知道提示词的位置，比如：

query: "the answer to the ultimate question of life, the universe, and everything is "

在 prompt 中，使⽤了三次 "the" ，需要根据它们在prompt中的位置为每个 "the" token ⽣成不同的query向量（每个⻓度为 128），可以使⽤ RoPE （旋转位置编码）来实现这⼀点，具体RoPE实现可以参考之前的文章，也可以看看这个视频：https://www.youtube.com/watch?v=o29P0Kpobz0&t=530s。

llama3-scratch

q_per_token_split_into_pairs = q_per_token.float().view(q_per_token.shape[0], -1, 2)
q_per_token_split_into_pairs.shape

## 输出
torch.Size([17, 64, 2])

从上面的步骤中，将query向量分成几对，每对使用旋转角度移位，现在将一个大小为[17X64X2]的向量，这个是prompt中每个标记分成64对的128个长度的query向量， 64对中的每一对都将旋转 m*theta，其中 m 是我们要旋转查询的 token 位置。

llama3-scratch

6.7、复数的点积来计算旋转向量

llama3-scratch

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

## 输出
torch.Size([17, 64])

6.8、现在在每个token的query元素都有⼀个复数（⻆度变化向量）

我们可以将 query（分成两对的查询）转换为复数，然后使用点积根据位置旋转查询，代码如下：

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])

6.9、获得旋转后的向量

可以通过再次将复数看作实数来返回成对的 query，代码如下：

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])

旋转对现在已合并，现在有了⼀个新的query向量（旋转query向量），其shape为[17x128]，其中17是token的数量，128是query向量的维度。

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])

6.10、keys（操作query一样）

llama3-scratch

我太懒了，所以我不打算对key进行数学运算，唯一需要记住的是：

• key生成的key向量是128维
• key的权重数量只有查询的1/4，因为key的权重一次在4个头部之间共享，以减少所需的计算次数
• key也会选择添加位置信息，和查询一样

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])

完成当前步骤，每个token的query和key都是[17X128]维度。

6.11、将query和key相乘

这样做会给我们一个分数，将每个token相互映射，这个分数描述了每个token的查询与每个token的键的关联程度，这是自我注意力。注意力得分矩阵 (qk_per_token) 的形状是 [17x17]，其中17是提示中的token数量。

llama3-scratch

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])

7、屏蔽QK分数

7.1、屏蔽预测的token

在llama3的训练过程中，未来的token qk分数被屏蔽，为什么？因为在训练过程中，只学习使⽤过去的token来预测token，因此，在推理过程中，将未来的token设置为零。

llama3-scratch

执行如下代码输出：

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)

llama3-scratch

mask = torch.full((len(tokens), len(tokens)), float("-inf"), device=tokens.device)
mask = torch.triu(mask, diagnotallow=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)

7.2、values（注意力机制最后部分）

图片

llama3-scratch

注意力分数用于确定每个token的使用多少矩阵，就像key一样，value的权重也是每4个注意力头之间共享，因此value矩阵形状是[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.shape

## 输出
torch.Size([8, 128, 4096])

第一层，第一个头值权重矩阵如下：

v_layer0_head0 = v_layer0[0]
v_layer0_head0.shape

## 输出
torch.Size([128, 4096])

7.3、value向量

llama3-scratch

我们现在使用value权重来获取每个token的注意力值，其大小为[17x128]，其中17是提示中的token数量，128是每个token的值向量的大小。

v_per_token = torch.matmul(token_embeddings, v_layer0_head0.T)
v_per_token.shape

## 输出
torch.Size([17, 128])

7.4、最终的注意力输出

llama3-scratch

与每个token的值相乘后得到的注意向量形状为[17*128]，输出如下：

qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)
qkv_attention.shape

## 输出
torch.Size([17, 128])

8、多头注意力机制

llama3-scratch

现在获取了第一层和第一个头的注意力输出，只需要按照n_heads重复这个步骤，对第一层的每个头部执行与上面的单元格完全相同的数学运算。

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

llama3-scratch

现在又有了第一层的32个head的qkv_attention矩阵，接下来将所有的注意力分数合并为一个大小为[17x4096]的大矩阵，如下所示：

stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)
stacked_qkv_attention.shape

## 输出
torch.Size([17, 4096])

9、最后一步：计算权重矩阵

对于第0层注意力，最后要做的事情之一是将权重矩阵相乘：

w_layer0 = model["layers.0.attention.wo.weight"]
w_layer0.shape

## 输出
torch.Size([4096, 4096])

9.1、简单的线性层，所以我们只需要matmul

llama3-scratch

embedding_delta = torch.matmul(stacked_qkv_attention, w_layer0.T)
embedding_delta.shape

## 输出
torch.Size([17, 4096])

现在embeding值变化了，添加到原始的token中嵌入：

embedding_after_edit = token_embeddings_unnormalized + embedding_delta
embedding_after_edit.shape

## 输出
torch.Size([17, 4096])

9.2、归一化，然后通过embedding增量运⾏⼀个前馈神经⽹络

llama3-scratch

embedding_after_edit_normalized = rms_norm(embedding_after_edit, model["layers.0.ffn_norm.weight"])
embedding_after_edit_normalized.shape

## 输出
torch.Size([17, 4096])

9.3、加载FFN权重并实现前馈⽹络

llama3-scratch

在llama3中，使用了SwiGLU前馈网络，这种网络架构非常适合在模型需要时添加非线性，如今，在llms中使用这种前馈网络架构相当常见。

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])

10、在第⼀层之后，终于为每个token生成新的embeding

在我们完成之前只需要再经过31层（一个 for 循环），可以想象这个生成过的embeding具有关于第一层上提出的所有查询的信息。现在，对所有提出的问题每⼀层都会对query进⾏越来越复杂的编码，直到得到⼀个embedding，其中包含了需要的下⼀个token的所有信息。

layer_0_embedding = embedding_after_edit+output_after_feedforward
layer_0_embedding.shape

## 输出
torch.Size([17, 4096])

for循环对每一层都执行相同的逻辑，最终获得final_embedding，如下所示：

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

11、有了final_embedding，这是模型对下一个标记做出的最佳猜测

embeding的形状和token embeding的形状相同，都是[17X4096]，其中17是token的数量，4096是嵌入的维度。

final_embedding = rms_norm(final_embedding, model["norm.weight"])
final_embedding.shape

## 输出
torch.Size([17, 4096])

12、最后，让我们将embeding解码为token

llama3-scratch

我们将使用输出解码器将最终的嵌入转换为token，如下所示：

model["output.weight"].shape

## 输出
torch.Size([128256, 4096])

最终我们希望这个问题the answer to the ultimate question of life, the universe, and everything is ，答案是42，现在执行如下代码：

logits = torch.matmul(final_embedding[-1], model["output.weight"].T)
logits.shape

## 输出
torch.Size([128256])

预测下一个token的代码：

next_token = torch.argmax(logits, dim=-1)
next_token

## 输出
tensor(2983)

llama3-scratch

最终输出结果：

tokenizer.decode([next_token.item()])

## 输出
42

最终我整理了一下代码：https://github.com/linkxzhou/mylib/blob/master/llm/0-llama3-from-scratch.py

参考

（1）https://github.com/naklecha/llama3-from-scratch