从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([email protected])+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_feedforwardDecoder层
对剩下的层都做一样的操作
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)