从零实现llama3

在这个文件中，我一步一步地从零实现了llama3,每一步都只涉及张量和矩阵乘法。

此外，我将直接从Meta提供的llama3模型文件中加载张量，你需要在运行此文件之前下载权重。 以下是官方下载权重的链接:https://llama.meta.com/llama-downloads/

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/19cb348d-c1f6-4f5c-a4d6-05fcfd668b3f.png"/> </div>

分词器

我不会实现一个bpe分词器(但是Andrej Karpathy有一个非常简洁的实现) <br> 他的实现链接:https://github.com/karpathy/minbpe

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/9414c5c3-6e6c-4bf6-8773-153b2d17852d.png" width="600"/> </div>

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!'

读取模型文件

通常,这取决于模型类的编写方式和其中的变量名。 <br> 但由于我们是从零实现llama3,我们将逐个张量地读取文件。

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/753cd9d9-0c48-48d9-acf2-d5c95f37969e.png" width="600"/> </div>

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}

我们使用这个配置来推断模型的详细信息,比如

1. 模型有32个transformer层
2. 每个多头注意力块有32个头
3. 词汇表大小等等

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

将文本转换为tokens

这里我们使用tiktoken(我认为是一个OpenAI库)作为分词器

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/3cb6c97a-f20e-4cb4-bbe0-743f305910cb.png" width="600"/> </div>

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

将tokens转换为它们的嵌入

抱歉,这是代码库中唯一使用内置神经网络模块的部分 <br> 总之,我们的[17x1] tokens现在变成了[17x4096],即17个嵌入(每个token一个),长度为4096 <br> <br> 注意:记住形状,这会让理解一切变得更容易

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/20b4201f-7dbd-400a-a2e7-1e4df7f3e7ca.png" width="600"/> </div>

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

然后我们使用rms归一化来归一化嵌入

请注意,在这一步之后形状不会改变,只是值被归一化了 <br> 需要记住的是,我们需要一个norm_eps(来自配置),因为我们不想意外地将rms设置为0并除以0 <br> 这里是公式:

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/41b6a6f6-b26e-4935-955d-a8f1d42770a8.png" width="600"/> </div>

# 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

构建transformer的第一层

归一化

你会看到我从模型字典中访问layer.0(这是第一层) <br> 总之,归一化后我们的形状仍然是[17x4096],与嵌入相同,但已经归一化

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/756c431e-7ea1-4835-8be8-010f630adbda.png" width="600"/> </div>

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

torch.Size([17, 4096])

从零实现注意力机制

让我们加载transformer第一层的注意力头

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/13c92c12-735b-4e9c-9238-a050d7e34edb.png" width="600"/> </div> <br>

当我们从模型中加载查询、键、值和输出向量时,我们注意到它们的形状分别为[4096x4096]、[1024x4096]、[1024x4096]和[4096x4096]

<br> 乍一看这很奇怪,因为理想情况下我们希望每个头都有各自的q、k、v和o

<br> > 代码作者将它们打包在一起是因为这样做很容易,有助于并行化注意力头乘法。 <br> > 我要把所有东西都展开...

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] <br><br> 这里，32是llama3中注意力头的数量，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.shape

torch.Size([32, 128, 4096])

我将实现第一层的第一个头

这里我访问第一层第一个头的查询权重矩阵，这个查询权重矩阵的大小是[128x4096]

q_layer0_head0 = q_layer0[0]
q_layer0_head0.shape

torch.Size([128, 4096])

现在我们将查询权重与token嵌入相乘，得到token的查询

这里你可以看到结果形状是[17x128]，这是因为我们有17个token，每个token都有一个长度为128的查询。

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/981ee546-eaca-40b2-bbcb-c16c0cd56918.png" width="600"/> </div>

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

torch.Size([17, 128])

位置编码

我们现在处于一个阶段，我们为提示中的每个token都有一个查询向量，但如果你仔细思考 -- 单个查询向量并不知道它在提示中的位置。 <br><br> 查询："the answer to the ultimate question of life, the universe, and everything is " <br><br> 在我们的提示中，我们使用了三次"the"，我们需要所有3个"the" token的查询向量（每个大小为[1x128]）根据它们在查询中的位置有不同的查询向量。我们使用RoPE（旋转位置嵌入）执行这些旋转。 <br><br>

RoPE

观看这个视频（这是我看的）来理解数学原理。 https://www.youtube.com/watch?v=o29P0Kpobz0&t=530s

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/eef1b838-8d8c-41df-ab9b-3aa905192572.png" width="600"/> </div>

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

在上面的步骤中，我们将查询向量分成对，我们对每一对应用旋转角度偏移！ <br><br> 现在我们有一个大小为[17x64x2]的向量，这是128长度的查询分成64对，对应提示中的每个token！这64对中的每一对都将被旋转m*(theta)，其中m是我们正在旋转查询的token的位置！

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/cd3eeaa3-501d-418e-ad29-815764f7a61b.png" width="600"/> </div>

使用复数的点积来旋转向量

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/984587b1-280f-44a8-a856-7d968f0aa9ab.png" width="600"/> </div>

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

# 查看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('实部')
plt.ylabel('虚部')
plt.title('freqs_cis的一行的图')
plt.show()

现在我们为每个token的查询元素都有了一个复数（角度变化向量）

我们可以将我们的查询（我们分成对的那个）转换为复数，然后进行点积，根据位置旋转查询 <br> 说实话，想到这一点真是美妙 :)

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

现在旋转后的对已经合并，我们得到了一个新的查询向量（旋转后的查询向量），其形状为[17x128]，其中17是标记数量，128是查询向量的维度。

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

键（几乎与查询相同）

我很懒，所以不会详细介绍键的数学计算，你只需要记住以下几点：

键也生成128维的键向量 键的权重数量只有查询的1/4，这是因为键的权重在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.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])

在这个阶段，我们现在有了每个标记的查询和键的旋转值。

每个查询和键的形状都是[17x128]。

下一步我们将乘以查询和键矩阵

这样做会得到一个分数映射，描述每个标记之间的关系 这个分数描述了每个标记的查询与每个标记的键的相关程度。 这就是自注意力 :) 注意力分数矩阵（qk_per_token）的形状是[17x17]，其中17是提示中的标记数量

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

现在我们需要掩蔽查询键分数

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

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

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)

值（注意力的最后阶段）

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/cc2ec252-f7c2-4a5f-98fb-f2950da51676.png" width="600px"/> </div> 这些分数（0-1）用于确定每个token使用多少价值矩阵 <br> > 就像键一样，价值权重也在每4个注意力头之间共享（以节省计算） <br> > 因此，下面价值权重矩阵的形状是 [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])

价值向量

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/9953b5cc-9509-4f79-93a8-b2409dae8bab.png" width="600px"/> </div> 现在我们使用价值权重来获取每个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])

注意力

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/d00b7d44-620c-4689-aeac-b555ebce9ed4.png" width="600px"/> </div> 与每个token的价值相乘后得到的注意力向量形状为 [17*128]

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

torch.Size([17, 128])

多头注意力

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/12144cf1-d7c3-4b87-ab58-eb091925f204.png" width="600px"/> </div> 我们现在有了第一层和第一个头的注意力值 <br> 现在我将运行一个循环，对第一层的每个头执行与上面单元格完全相同的数学运算

qkv_attention_store = []

for head in range(n_heads):
    q_layer0_head = q_layer0[head]
    k_layer0_head = k_layer0[head//4] # 键权重在4个头之间共享
    v_layer0_head = v_layer0[head//4] # 值权重在4个头之间共享
    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)

32

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/5d3ad8c2-2087-4159-8f7b-124d4fcc97f9.png" width="600px"/> </div> 我们现在有了第一层所有32个头的qkv_attention矩阵，接下来我将把所有注意力分数合并成一个大小为 [17x4096] 的大矩阵 <br> 我们快要结束了 :)

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

torch.Size([17, 4096])

权重矩阵，最后几步之一

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/f3b17a23-2cb7-4dad-82d3-305c8fce9b93.png" width="600px"/> </div> 第0层注意力的最后一件事是乘以权重矩阵

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.shape

torch.Size([17, 4096])

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/d5d166b9-2fcb-493c-8ae0-b40250eb6175.png" width="600px"/> </div> 我们现在有了注意力后嵌入值的变化，应该添加到原始token嵌入中

embedding_after_edit = token_embeddings_unnormalized + embedding_delta
embedding_after_edit.shape

torch.Size([17, 4096])

我们对嵌入delta进行归一化，然后通过前馈神经网络运行

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/7b47f583-6379-4b68-b99d-cd5bc7a4259d.png" width="600px"/> </div>

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

torch.Size([17, 4096])

加载前馈权重并实现前馈网络

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/0a04d43c-0b75-4022-8bb7-f45d3d5c6b09.png" width="600px"/> </div> 在llama3中，他们使用了SwiGLU前馈网络，这种网络架构在模型需要时非常擅长添加非线性。 <br> 现在在大型语言模型中使用这种前馈网络架构已经很普遍了

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

我们终于得到了第一层后每个token的新编辑嵌入

还剩31层就完成了（只差一个for循环） <br> 你可以想象这个编辑后的嵌入包含了第一层上所有查询的信息 <br> 现在每一层都会对所问的问题进行更复杂的编码，直到我们得到一个包含我们需要的下一个token的所有信息的嵌入。

layer_0_embedding = embedding_after_edit+output_after_feedforward
layer_0_embedding.shape

torch.Size([17, 4096])

天哪，一切都在同时发生

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/6ce6a90c-b859-4e91-8b43-5c4cd1f156f2.png" width="600px"/> </div> 没错，就是这样。我们之前做的一切，一次性全部完成，对每一层都是如此。 <br>

祝你阅读愉快 :)

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

我们现在得到了最终的嵌入向量，这是模型对下一个token的最佳猜测 嵌入向量的形状与常规token嵌入相同[17x4096]，其中17是token数量，4096是嵌入维度

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/1a755bb0-00d6-4cf7-aa92-49ef06bc4056.png" width="600px"/> </div>

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

torch.Size([17, 4096])

最后，让我们将嵌入向量解码为token值

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/2840fe95-59d4-4107-9492-0e9219ccce22.png" width="600px"/> </div> 我们将使用输出解码器将最终嵌入转换为token

model["output.weight"].shape

torch.Size([128256, 4096])

我们使用最后一个token的嵌入来预测下一个值 希望在我们的例子中是42 :) 注意：42是"生命、宇宙和一切的终极问题的答案"，根据《银河系漫游指南》一书，大多数现代语言模型在这里会回答42，这应该可以验证我们的整个代码！祝我好运 :)

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

torch.Size([128256])

模型预测下一个token编号为2983，这是否是42的token编号？

我正在为你加油，这是最后一个代码单元，希望你玩得开心 :)

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

tensor(2983)

太棒了

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/831c8389-604a-4df6-812f-bdce8f110ea3.png" width="600px"/> </div>

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

'42'

谢谢，我爱你 :)

这就是结尾了。希望你喜欢阅读这篇文章！

如果你想支持我的工作

1. 在Twitter上关注我 https://twitter.com/naklecha
2. 或者，给我买杯咖啡 https://www.buymeacoffee.com/naklecha

说实话，如果你能看到这里，你已经让我很开心了 :)

是什么激励着我？

我和我的朋友们正在进行一项使命 - 让研究更容易获取！ 我们创建了一个名为A10的研究实验室 - AAAAAAAAAA.org

A10 Twitter - https://twitter.com/aaaaaaaaaaorg

我们的论点：

<div> <img src="https://yellow-cdn.veclightyear.com/835a84d5/43080beb-a743-4f4b-9c5c-aa67b0e30a09.png" width="600px"/> </div>