本文详细介绍了如何使用 Python 和 PyTorch 从零开始实现 LLaMA3 大模型的核心架构。内容涵盖模型权重加载、BPE 分词器集成、RMSNorm 归一化、旋转位置编码 (RoPE)、多头注意力机制及 SwiGLU 前馈网络的代码实现。通过逐步解析注意力矩阵计算、因果掩码应用及多层 Transformer 堆叠逻辑，展示了模型推理的全过程。最终通过示例验证了模型能够正确预测特定上下文下的下一个 token，为理解大模型底层技术提供了实践参考。
Meta 近期发布了开源大模型 LLaMA3 系列，在多个关键基准测试中表现优异，尤其在代码生成任务上展现了强大的能力。为了深入理解其内部机制，开发者们开始尝试本地部署和从零实现。本文将详细介绍如何从零开始使用 Python 和 PyTorch 实现 LLaMA3 的核心架构，包括注意力机制、位置编码和前馈网络等关键组件。
该项目由开发者 Nishant Aklecha 发布，旨在通过模块嵌套和相互调用，清晰地展示 LLaMA3 的底层逻辑。项目地址为 https://github.com/nishant-aklecha/llama3-from-scratch（注：此处为通用示例链接，实际请以官方为准）。该实现涵盖了跨多头注意力矩阵乘法、位置编码以及每一层的详细解释。
首先，我们需要从 Meta 提供的官方渠道下载 LLaMA3 的模型文件。由于是从头开始实现，代码将逐个读取张量文件。
import torch
import json
from pathlib import Path
# 加载模型文件
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 等。同时，我们需要读取配置文件 params.json 来获取模型超参数。
with open("Meta-Llama-3-8B/params.json", "r") as f:
config = json.load(f)
config
配置信息通常包含维度、层数、头数等关键参数：
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"])
LLaMA3 模型通常包含 32 个 Transformer 层，每个多头注意力块有 32 个头，输入维度为 4096。
作者未自行实现分词器，而是借用了 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 归一化。
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)
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）。
对于第一层，我们提取对应的权重矩阵并重塑形状以匹配多头结构。
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)
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)
计算 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)
遍历所有注意力头，计算加权后的值向量，并拼接回原始维度。
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 得到注意力输出。
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 模型的从零实现过程，涵盖了从数据加载、分词、嵌入、注意力机制到前馈网络的完整链路。通过手动实现这些组件，开发者可以更深入地理解大语言模型的运作原理，为后续的优化和微调打下坚实基础。
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
生成新的随机RSA私钥和公钥pem证书。 在线工具，RSA密钥对生成器在线工具，online
基于 Mermaid.js 实时预览流程图、时序图等图表，支持源码编辑与即时渲染。 在线工具，Mermaid 预览与可视化编辑在线工具，online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具，curl 转代码在线工具，online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online