从零实现Llama3：深入解析Transformer架构与实现细节-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00824/article/details/148377626

从零实现Llama3：深入解析Transformer架构与实现细节

llama3-from-scratch llama3 一次实现一个矩阵乘法。项目地址: https://gitcode.com/gh_mirrors/ll/llama3-from-scratch

引言

本文将深入探讨如何从零开始实现Llama3语言模型。我们将从最基本的张量操作开始，逐步构建完整的Transformer架构。通过这个过程，读者不仅能理解Llama3的工作原理，还能掌握现代大型语言模型的核心实现技术。

模型架构概览

Llama3采用了标准的Transformer架构，包含以下核心组件：

词嵌入层(Embedding Layer)
多层Transformer块(包含自注意力机制和前馈网络)
RMSNorm归一化层
RoPE位置编码

模型配置参数如下：

隐藏层维度(dim): 4096
层数(n_layers): 32
注意力头数(n_heads): 32
KV头数(n_kv_heads): 8
词表大小(vocab_size): 128256
归一化epsilon(norm_eps): 1e-05
RoPE基数(rope_theta): 500000.0

分词器实现

虽然本文不从头实现分词器，但我们使用基于BPE(Byte Pair Encoding)的分词方案。分词器将输入文本转换为一系列token ID，例如：

prompt = "the answer to the ultimate question of life, the universe, and everything is "
tokens = [128000] + tokenizer.encode(prompt)
# 输出: [128000, 1820, 4320, 311, 279, 17139, 3488, 315, 2324, 11, 279, 15861, 11, 323, 4395, 374, 220]

特殊token如<|begin_of_text|>(ID 128000)用于标记文本开始。

词嵌入层

词嵌入层将token ID转换为高维向量表示：

embedding_layer = torch.nn.Embedding(vocab_size, dim)
token_embeddings = embedding_layer(tokens)  # 形状: [17, 4096]

这里17是token数量，4096是嵌入维度。

RMSNorm归一化

Llama3使用RMSNorm而非传统的LayerNorm，计算公式为：

output = (input * torch.rsqrt(input.pow(2).mean(-1, keepdim=True) + eps)) * weight

实现代码：

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

自注意力机制

Llama3的自注意力机制包含查询(Query)、键(Key)、值(Value)三个投影矩阵。这些矩阵被"打包"存储以提高并行效率：

q_layer0 = model["layers.0.attention.wq.weight"]  # 形状: [4096, 4096]
k_layer0 = model["layers.0.attention.wk.weight"]  # 形状: [1024, 4096]
v_layer0 = model["layers.0.attention.wv.weight"]  # 形状: [1024, 4096]

多头注意力拆分

对于32个注意力头，我们将查询矩阵拆分为：

q_layer0 = q_layer0.view(n_heads, head_dim, dim)  # 形状: [32, 128, 4096]

单个头的查询计算：

q_per_token = torch.matmul(token_embeddings, q_layer0[0].T)  # 形状: [17, 128]

RoPE位置编码

RoPE(Rotary Positional Embedding)通过旋转矩阵将位置信息编码到查询和键向量中。实现步骤：

将128维查询向量拆分为64个2维对：

q_per_token_split_into_pairs = q_per_token.view(q_per_token.shape[0], -1, 2)  # 形状: [17, 64, 2]

计算旋转角度：

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

应用旋转：

# 构造旋转矩阵
cos = torch.cos(m * freqs)  # m为位置索引
sin = torch.sin(m * freqs)
# 应用旋转
rotated_q = q_per_token_split_into_pairs * torch.stack([cos, sin], dim=-1)

完整层实现

一个完整的Transformer层包含：

注意力归一化
多头自注意力
前馈网络归一化
前馈网络(包含门控机制)

前馈网络实现：

w1 = model["layers.0.feed_forward.w1.weight"]  # 形状: [14336, 4096]
w3 = model["layers.0.feed_forward.w3.weight"]  # 形状: [14336, 4096]
w2 = model["layers.0.feed_forward.w2.weight"]  # 形状: [4096, 14336]

hidden = torch.matmul(token_embeddings, w1.T) * torch.matmul(token_embeddings, w3.T)
output = torch.matmul(hidden, w2.T)