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

最近在看一篇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!1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.

这里用了字节对编码（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
}1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.

其中输出的配置看：

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', ' ']1.
2.
3.
4.
5.
6.
7.
8.
9.
10.

其中，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|>"1.
2.
3.
4.
5.
6.
7.
8.
9.
10.

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])1.
2.
3.
4.
5.
6.
7.

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_weights1.
2.

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])1.
2.
3.
4.
5.

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])1.
2.
3.
4.
5.
6.
7.
8.
9.

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])1.
2.
3.
4.
5.
6.
7.

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

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

q_layer0_head0 = q_layer0[0]
q_layer0_head0.shape

## 输出
torch.Size([128, 4096])1.
2.
3.
4.
5.

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])1.
2.
3.
4.
5.

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])1.
2.
3.
4.
5.

从上面的步骤中，将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])1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.

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])1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.

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])1.
2.
3.
4.
5.

旋转对现在已合并，现在有了⼀个新的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])1.
2.
3.
4.
5.

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])1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.

完成当前步骤，每个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])1.
2.
3.
4.
5.

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)1.
2.
3.
4.
5.
6.
7.
8.
9.
10.

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)1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.

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)1.
2.

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])1.
2.
3.
4.
5.
6.

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

v_layer0_head0 = v_layer0[0]
v_layer0_head0.shape

## 输出
torch.Size([128, 4096])1.
2.
3.
4.
5.

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])1.
2.
3.
4.
5.

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])1.
2.
3.
4.
5.

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)

## 输出
321.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.

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])1.
2.
3.
4.
5.

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

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

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

## 输出
torch.Size([4096, 4096])1.
2.
3.
4.
5.

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

llama3-scratch

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

## 输出
torch.Size([17, 4096])1.
2.
3.
4.
5.

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

embedding_after_edit = token_embeddings_unnormalized + embedding_delta
embedding_after_edit.shape

## 输出
torch.Size([17, 4096])1.
2.
3.
4.
5.

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])1.
2.
3.
4.
5.

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])1.
2.
3.
4.
5.
6.
7.
8.

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])1.
2.
3.
4.
5.

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_feedforward1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.

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])1.
2.
3.
4.
5.

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

llama3-scratch

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

model["output.weight"].shape

## 输出
torch.Size([128256, 4096])1.
2.
3.
4.

最终我们希望这个问题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])1.
2.
3.
4.
5.

预测下一个token的代码：

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

## 输出
tensor(2983)1.
2.
3.
4.
5.

llama3-scratch

最终输出结果：

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

## 输出
421.
2.
3.
4.

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

参考

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