我要投稿

从零复现 Llama3

发布日期：2024-06-08 06:29:07 浏览次数： 2524 来源：NLP轻松谈

前摩根士丹利 LLM 研究员 naklecha基于开源的 Llama3 模型参数权重和karpathy的开源分词器 minbpe 详细实现了 Llama3。github地址：https://github.com/naklecha/llama3-from-scratch。

通过动画式图解讲解多头注意力机制、位置编码和介于两者之间的其他层的每个矩阵乘法。

Tokenizer

Andrej Karpathy 有一个非常好的实现，地址：https://github.com/karpathy/minbpe

from pathlib import Pathimport tiktokenfrom tiktoken.load import load_tiktoken_bpeimport torchimport jsonimport 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!"))

Reading the model file

在正式读取之前，需要直接从 Meta 为 llama3 提供的模型文件加载张量，您需要在运行此文件之前下载权重。以下是下载github的链接：https://llama.meta.com/llama-downloads/,我们将一次读取一个张量文件.

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}

该模型有 32 个变压器层
每个多头注意力块有 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"])

将文本转换为令牌，我们使用 Tiktoken为分词器

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

将令牌转换为其嵌入

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

然后，使用 RMS 规范化对嵌入进行规范化

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

构建transformers的第一层

Normalization

从模型字典访问 layer.0，在归一化后，我们的形状仍然 [17x4096] 与嵌入相同。

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

从零开始实施

让我们加载变压器第一层的注意头，当我们从模型加载查询、键、值和输出向量时，我们注意到形状为 [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)

我们将分解来自多个注意力头的查询，生成的形状为 [32x128x4096]，32 是 llama3 中的注意力头数，128 是查询向量的大小，4096 是令牌嵌入的大小

q_layer0 = model["layers.0.attention.wq.weight"]head_dim = q_layer0.shape[0] // n_headsq_layer0 = q_layer0.view(n_heads, head_dim, dim)q_layer0.shape

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

q_layer0_head0 = q_layer0[0]q_layer0_head0.shape

将查询权重与令牌嵌入相乘，以接收对令牌的查询,可以看到生成的形状是 [17x128]，这是因为我们有 17 个标记，每个标记都有一个 128 长度的查询。

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

Positioning encoding

使用 RoPE（旋转位置嵌入）,可以观看https://www.youtube.com/watch?v=o29P0Kpobz0&t=530s

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

我们将查询向量拆分为对，我们对每对应用旋转角度偏移，我们现在有一个大小为 [17x64x2] 的向量，这是提示中每个标记的 128 个长度查询，分为 64 对！这 64 对中的每一对都将由 m*（theta）旋转，其中 m 是我们旋转查询的令牌的位置！

使用复数的点积旋转向量

zero_to_one_split_into_64_parts = torch.tensor(range(64))/64freqs = 1.0 / (rope_theta ** zero_to_one_split_into_64_parts)

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_cisvalue = 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()

每个token的查询元素都有一个复数，我们可以将查询转换为复数，然后点积根据位置旋转查询

q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)q_per_token_as_complex_numbers_rotated = q_per_token_as_complex_numbers * freqs_cis

获得旋转向量后，我们可以通过再次将复数视为实数来将查询作为对返回

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

旋转的对现在被合并，我们现在有一个新的查询向量（旋转的查询向量），其形状为 [17x128]，其中 17 是标记数，128 是查询向量。

q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)q_per_token_rotated.shape

键值和query基本上是一致的，但是键的权重数只有查询权重的 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

在此阶段，现在每个令牌都具有查询和键的旋转值;每个查询和键现在都是 [17x128] 的形状。

注意力得分矩阵（qk_per_token）的形状为 [17x17]，其中 17 是提示中的标记数

qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T)/(head_dim)**0.5qk_per_token.shape

在 llama3 的训练过程中，未来的token的QK 分数被屏蔽。在训练过程中，我们只学习使用过去的令牌来预测令牌。因此，在推理过程中，我们将未来标记设置为零。

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)display_qk_heatmap(qk_per_token_after_masking_after_softmax)

values就像键一样，值权重也是每 4 个注意力头共享一次的（以节省计算）;因此，下面的值权重矩阵的形状为 [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

values值得计算，其大小为 [17x128]，其中 17 是提示中的标记数，128 是每个标记的值向量

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

attention的计算

每个标记的值相乘后的结果注意力向量形状为 [17*128]

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

Multi head attention

现在有了第一层和第一头的注意力值；现在将运行一个循环并执行与上面单元格完全相同的数学运算

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 headsv_layer0_head = v_layer0[head//4] # value weights are shared across 4 headsq_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.5mask = 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 + maskqk_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 个头部提供了一个qkv_attention矩阵，接下来把所有注意力分数合并到一个大小为 [17x4096] 的大矩阵中

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

最后一步，进行权重的合并

对于第 0 层的注意力，要做的最后一件事是做一个简单的线性变换

w_layer0 = model["layers.0.attention.wo.weight"]embedding_delta = torch.matmul(stacked_qkv_attention, w_layer0.T)embedding_delta.shap

在注意后对嵌入值进行了更改，这应该添加到原始令牌嵌入中

embedding_after_edit = token_embeddings_unnormalized + embedding_deltaembedding_after_edit.shap

对嵌入增量进行归一化，然后运行前馈神经网络

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

加载 FF 权重并实现前馈网络,在 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

在第一层之后，我们终于为每个标记提供了新的编辑嵌入

layer_0_embedding = embedding_after_edit+output_after_feedforwardlayer_0_embedding.shape

完整的32层查询向量，键值，values值得计算

final_embedding = token_embeddings_unnormalizedfor 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.5mask = 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 + maskqk_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_deltaembedding_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做出的最佳猜测

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

最后，让我们解码嵌入到token值中；我们将使用输出解码器将最终嵌入转换为token

model["output.weight"].shape

使用最后一个标记的嵌入来预测下一个值

logits = torch.matmul(final_embedding[-1], model["output.weight"].T)next_token = torch.argmax(logits, dim=-1)tokenizer.decode([next_token.item()])

53AI，企业落地应用大模型首选服务商

产品：大模型应用平台+智能体定制开发+落地咨询服务

承诺：先做场景POC验证，看到效果再签署服务协议。零风险落地应用大模型，已交付160+中大型企业

160+中大型企业正在使用53AI

立即咨询预约演示

百度智能云邀53AI：共创AI新纪元，启航智能新时代

2024-05-27

钉钉恒星计划：53AI与百余位企业家及钉钉生态伙伴，共议“AI 浪潮下的新机遇”

2024-05-22

热点资讯

本地部署大模型？看这篇就够了，Ollama 部署和实战

2024-08-13

一文带你了解大模型——智能体（Agent）

2024-05-28

全面对比dify、coze、streamlit、chainlit

2024-04-26

Cursor 一个真正让程序员产生危机感的 AI 编程工具

2024-08-21

16个具有“联网搜索”功能的AI：总有一个适合你！

2024-06-13

50+个AI大模型在不同领域的应用案例

2024-08-04

【深度】AI搜索产品深度分析-搜索原理和商业模式分析

2024-07-09

OpenAI o1与GPT4o的对比分析

2024-09-23

Claude中国官网正式上线！1:1还原官网，非套壳，国内直连

2024-07-18

“大数据+”医疗

2024-04-11

大家都在问

2025 AI Agent迷局：谁在玩真的，谁在演戏？

2025-01-10

杨芳贤｜AI 2.0时代，如何拥抱与驾驭大模型？

2025-01-09

深度长文｜AI的“巴别塔”：多Agent协同为何如此之难？