从0实现llama3

分享一下从0实现llama的过程

流程如下：

word --> embedding layer --> n * decoder layer --> final linear layer --> output

分词器

在embedding之前，需要进行分词，将句子分成单词。llama3采用了基于BPE算法的分词器。

这个链接实现了一个非常简洁的BPE分词器简易分词器实现

BPE分词器（选看）

1) 训练 tokenizer 词汇表并合并给定文本，2) 将文本编码为 token，3) 将 token 解码为文本。项目结构如下：

minbpe/base.py：实现Tokenizer类，即基类。它包含train、encode和decode存根、保存/加载功能，还有一些常用的实用函数。
minbpe/basic.py：实现BasicTokenizer直接在文本上运行的 BPE 算法的最简单实现。
minbpe/regex.py：实现RegexTokenizer通过正则表达式模式进一步拆分输入文本，这是一个预处理阶段，在标记化之前按类别（例如：字母、数字、标点符号）拆分输入文本。这确保不会发生跨类别边界的合并。这是在 GPT-2 论文中引入的，并且从 GPT-4 开始仍在使用。此类还可以处理特殊标记。
minbpe/gpt4.py：实现GPT4Tokenizer。此类是RegexTokenizer（上面的 2）的轻量级包装器，可精确重现tiktoken库中 GPT-4 的标记化。

Base.py文件

共实现了一个tokenizer基类和一些辅助函数

get_stats():这个函数计算一个整数列表中连续元素对的出现次数，并返回一个字典 
#input: [1, 2, 3, 1, 2]
#output: ((1, 2): 2, (2, 3): 1, (3, 1): 1)。

merge():在列表中替换所有连续出现的 pair 元组为新整数 idx。
#input: ids=[1, 2, 3, 1, 2]，pair=(1, 2)，idx=4
#output: [4, 3, 4]。

replace_control_characters():去除字符串中的控制字符（例如换行符等），并用 Unicode 转义表示

render_token():将字节序列 t 解码为字符串，并使用 replace_control_characters 转义控制字符。

Class Tokenizer():
    def init():
        """
        初始化
	    merges：一个字典，表示合并规则（例如：两个整数对合并成一个新的整数）。
	    pattern：分词器的模式字符串，暂时为空字符串。
	    special_tokens：一个字典，保存特殊 token 对应的整数 ID。
	    vocab：通过 _build_vocab 方法生成的词汇表，初始包含 256 个基本字节。
        """

    def train、encode、decode均为抽象方法。

    def build_vocab():构建词汇表。根据 merges 中的合并规则和 special_tokens 中的特殊 token，创建词汇表。初始化时，词汇表包含 256 个字节（0-255）。
    
    def save():保存分词器模型。
    
    def load():加载分词器模型，恢复 merges、special_tokens 和 vocab。

Basic.py文件

通过BPE算法来合并最常见的字节对，从而构建词汇表并对文本进行编码和解码。

def Train():
"""
流程：
1）训练Tokenizer，从给定的text中构建词汇表，vocab_size
2）文本预处理：将输入文本 text 编码为 UTF-8 字节流 text_bytes，然后将每个字节转换为一个整数（范围 0–255），存储在列表 ids 中。
3）合并最常见的字节对：初始化一个空的 merges 字典和初始的词汇表 vocab（包含 0-255 的字节），然后进行迭代，通过合并最常见的字节对逐步扩展词汇表，直到达到目标词汇表大小。
    a）通过统计 ids 中连续字节对出现的次数，选择出现频率最高的字节对，将其合并为一个新的 token，并更新 merges 和 vocab。
"""

def Decode():
"""
根据 ids 查找 vocab 中对应的字节序列。
将所有字节序列连接成一个字节流 text_bytes，并尝试使用 UTF-8 解码为字符串。
若解码过程中出现问题（例如无法解码的字节），则使用 errors="replace" 进行替换，保证不会出现解码错误。
"""

def encode():
"""
1)将输入文本 text 编码为字节流 text_bytes，并转换为整数列表 ids。
2)合并操作：在每次循环中，统计连续字节对的频率，并选择 merges 字典中具有最小合并索引的字节对进行合并。该过程会持续到没有更多可以合并的字节对。
"""

text = "i like pigs"
text_bytes = text.encode("utf-8")
list(text_bytes) #[105, 32, 108, 105, 107, 101, 32, 112, 105, 103, 115]

加载BPE分词器

Llama3定义了一些特殊token，如begin/end_of_text用来标记文本开始结束，

reserved_special_token_n：预留标记用来完成特定任务。

load_tiktoken_bpe()函数：用来加载BPE合成规则

import tiktoken
tokenizer = tiktoken.load.load_tiktoken_bpe("./mytest")#{b'mytoken': 1, b'foo': 2, b'bar': 3}

tokenizer_path = "/model/Llama3-8B/org/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)},
)

读取模型

直接用torch的load方法来加载模型参数

model = torch.load("/model/Llama3-8B/org/consolidated.00.pth")
print(json.dumps(list(model.keys())[:20]))

with open("/model/Llama3-8B/org/original_params.json", "r") as f:
    config = json.load(f)
print（config）

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

{'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}

使用配置文件来设置模型

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"])#rope_theta 控制位置编码中的缩放因子

word->tokens

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

调用了nn.embedding，并且加载了llama中的embedding参数。

然后对tokens进行了embedding得到token_embeddings_unnormalized

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

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

#(torch.Size([17, 4096]), torch.Size([4096]))

注意力层

实现Query

将model中的wq层权重作为query层。实现多头注意力。llama的注意力头数量为32，注意力头的形状为4096/32

q_layer0 = model["layers.0.attention.wq.weight"]#torch.Size([4096, 4096])
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])
q_layer0[0].shape #torch.Size([128, 4096])

将embedding后的向量与Query矩阵相乘，得到X的查询矩阵Q

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

ROPE

原理：十分钟读懂旋转编码（RoPE）

首先将查询矩阵Q两两分成对，并且对每一对都进行角度（复数表示）偏移

q_per_token_split_into_pairs = q_per_token.float().view(q_per_token.shape[0], -1, 2)#torch.Size([17, 64, 2])
#因为有64个对，因此需要64个角度信息
#计算旋转角度
zero_to_one_split_into_64_parts = torch.tensor(range(64))/64
#
freqs = 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)

将每个Q转换为复数，然后进行ROPE旋转，然后再转变为成对的Q的形式，然后再转变为正常的Q

q_per_token_as_complex_numbers = torch.view_as_complex(q_per_token_split_into_pairs)#torch.Size([17, 64])

q_per_token_as_complex_numbers_rotated = q_per_token_as_complex_numbers * freqs_cis#torch.Size([17, 64])

q_per_token_split_into_pairs_rotated = torch.view_as_real(q_per_token_as_complex_numbers_rotated)#torch.Size([17, 64])


q_per_token_rotated = q_per_token_split_into_pairs_rotated.view(q_per_token.shape)#torch.Size([17, 64])

Q的处理就完成了！！！！！

实现Key（与Query类似）

Key所生成的向量也是128维的，但是它的权重只有Query的1/4，因为Llama3底层用到了GQA，所以Key的权重是4个头共享的，减少了计算量

#加载Key层
k_layer0 = model["layers.0.attention.wk.weight"]
print(k_layer0.shape)#torch.Size([1024, 4096])
#多头注意力机制
k_layer0 = k_layer0.view(n_kv_heads, k_layer0.shape[0] // n_kv_heads, dim)#torch.Size([8, 128, 4096])
k_layer0_head0 = k_layer0[0]#torch.Size([128, 4096])

#将embedding转换为k
k_per_token = torch.matmul(token_embeddings, k_layer0_head0.T)#torch.Size([17,128])

#转换为对，对每个对应用偏移
k_per_token_split_into_pairs = k_per_token.float().view(k_per_token.shape[0], -1, 2)#torch.Size([17,64，2])
#将K转变为复数
k_per_token_as_complex_numbers = torch.view_as_complex(k_per_token_split_into_pairs)##torch.Size([17,64])

#进行Rope旋转后转换为实数
k_per_token_split_into_pairs_rotated = torch.view_as_real(k_per_token_as_complex_numbers * freqs_cis)#torch.Size([17,64，2])
#合并
k_per_token_rotated = k_per_token_split_into_pairs_rotated.view(k_per_token.shape)#torch.Size([17,128])

实现Softmax(Q@K.T)+mask

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

#构建上三角的mask
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
display_qk_heatmap(qk_per_token_after_masking)
#进行softmax操作
qk_per_token_after_masking_after_softmax = torch.nn.functional.softmax(qk_per_token_after_masking, dim=1).to(torch.bfloat16)

实现Value（与key类似）

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_head0 = v_layer0[0]

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

#score*value
qkv_attention = torch.matmul(qk_per_token_after_masking_after_softmax, v_per_token)

多头注意力机制

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 heads
    v_layer0_head = v_layer0[head//4] # value weights are shared across 4 heads
    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_store.append(qkv_attention)

#合并多头矩阵
stacked_qkv_attention = torch.cat(qkv_attention_store, dim=-1)

对线性层和注意力层进行相乘

embedding_delta = torch.matmul(stacked_qkv_attention, w_layer0.T)
#残差连接
embedding_after_edit = token_embeddings_unnormalized + embedding_delta


#进行归一化
embedding_after_edit_normalized = rms_norm(embedding_after_edit, model["layers.0.ffn_norm.weight"])
embedding_after_edit_normalized.shape

前馈神经网络FFN层

Llama3使用了SwiGLU前馈神经网络

SwiGLU与Vanilla的区别就是

FFN层添加了门控单元，使用SwiGLU作为激活函数后得到W2然后与升维矩阵W1进行了点乘然后与降维矩阵W3处理后得到最终结果

w1 = model["layers.0.feed_forward.w1.weight"]
w2 = model["layers.0.feed_forward.w2.weight"]
w3 = model["layers.0.feed_forward.w3.weight"]

#经过FFN后结果
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)

#得到最后嵌入层，它有全部的信息
layer_0_embedding = embedding_after_edit+output_after_feedforward

Decoder层

对剩下的层都做一样的操作

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

#对最终的嵌入层进行归一化
final_embedding = rms_norm(final_embedding, model["norm.weight"])

即可

#插入一个线性层进行分类
logits = torch.matmul(final_embedding[-1], model["output.weight"].T)