从零开始实现 LLaMa3 模型：核心架构与代码解析

LLaMA3 模型通常包含 32 个 Transformer 层，每个多头注意力块有 32 个头，输入维度为 4096。

分词器 (Tokenizer) 实现

作者未自行实现分词器，而是借用了 Andrej Karpathy 的 minbpe 实现方式，基于 TikToken 库处理 BPE 分词。

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|>",
    # ... 其他特殊 token
] + [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)},
)

# 测试分词
result = tokenizer.decode(tokenizer.encode("hello world!"))
print(result)

文本嵌入与归一化

将文本转换为 Token ID 后，需要通过嵌入层映射到向量空间，并进行 RMSNorm 归一化。

1. 嵌入层

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)

2. RMSNorm 归一化

LLaMA3 使用均方根归一化（RMSNorm）替代了传统的 LayerNorm，计算更高效。

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

注意力机制实现

这是 Transformer 的核心部分，包括查询（Query）、键（Key）、值（Value）的计算以及旋转位置编码（RoPE）。

1. 查询与键的投影

对于第一层，我们提取对应的权重矩阵并重塑形状以匹配多头结构。

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)

2. 旋转位置编码 (RoPE)

RoPE 通过复数点积来旋转向量，从而引入位置信息。

# 生成频率向量
freqs = 1.0 / (rope_theta ** torch.arange(0, head_dim, 2).float() / head_dim)
freqs_cis = torch.polar(torch.ones_like(freqs), freqs)

# 应用旋转
q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)
q_per_token_rotated = torch.view_as_real(q_per_token_as_complex_numbers * freqs_cis)

3. 注意力分数计算与掩码

计算 Q 和 K 的点积，并应用因果掩码（Causal Mask），确保预测时只能看到过去的 token。

qk_per_token = torch.matmul(q_per_token_rotated, k_per_token_rotated.T) / (head_dim)**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
attention_scores = torch.nn.functional.softmax(qk_per_token_after_masking, dim=1).to(torch.bfloat16)

4. 多头注意力聚合

遍历所有注意力头，计算加权后的值向量，并拼接回原始维度。

qkv_attention_store = []
for head in range(n_heads):
    # 计算单个头的注意力
    # ... (省略中间步骤)
    qkv_attention_store.append(qkv_attention)

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

最后乘以输出权重矩阵 wo.weight 得到注意力输出。

前馈神经网络 (FFN)

LLaMA3 采用 SwiGLU 激活函数的前馈网络，相比传统 ReLU 具有更好的非线性表达能力。

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
)

完整推理循环

将上述步骤封装在循环中，依次处理所有 Transformer 层。

final_embedding = token_embeddings_unnormalized
for layer in range(n_layers):
    # 1. 归一化
    layer_embedding_norm = rms_norm(final_embedding, model[f"layers.{layer}.attention_norm.weight"])
    
    # 2. 注意力机制
    # ... (复用注意力逻辑)
    
    # 3. 残差连接
    embedding_after_edit = final_embedding + embedding_delta
    
    # 4. FFN 与前向传播
    # ... (复用 FFN 逻辑)
    
    # 5. 更新最终嵌入
    final_embedding = embedding_after_edit + output_after_feedforward

输出解码

经过所有层处理后，对最后一个 token 的嵌入进行线性变换，得到 logits，并通过 Softmax 或 Argmax 预测下一个 token。

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

在示例中，输入提示词为 "the answer to the ultimate question of life..."，模型成功预测出答案 "42"，验证了实现的正确性。

总结

本文详细拆解了 LLaMA3 模型的从零实现过程，涵盖了从数据加载、分词、嵌入、注意力机制到前馈网络的完整链路。通过手动实现这些组件，开发者可以更深入地理解大语言模型的运作原理，为后续的优化和微调打下坚实基础。