1、前言
上一篇《从0构造Transformer文本生成模型》简单介绍了Transfomer文本生成模型的基本方法，本文将做进一步深入讲解。
github上有一个项目《llama3-from-scratch》，以Jupter笔记的方式深入剖析了llama3的工作原理，但是其中很多概念都没有讲述清楚，一般小白很难看懂。本文将在此基础上，通过提取llama3的模型参数，从0构造类llama3的Transfomer文本生成模型。
2、总体架构
3、基础概念
3.1 TikToken
使用 TikToken 的原因主要涉及其在自然语言处理（NLP）和大规模语言模型（如 GPT 系列模型）中的关键功能和优势。以下是使用 TikToken 的几个主要原因：
- 高效的标记化：
TikToken 提供了高效的文本标记化方法，可以快速将文本分割成标记（tokens）。这种高效的标记化对于处理大量文本或在实时应用中至关重要。 - 模型兼容性：
TikToken 专门设计用于与 OpenAI 的 GPT 模型和其他类似的 NLP 模型兼容。这意味着它可以直接处理并生成与这些模型匹配的标记格式，确保模型能够正确理解和处理输入的文本。 - 节省计算资源：
通过优化的标记化过程，TikToken 可以减少计算资源的使用。这对于需要处理大量文本数据或在资源受限的环境中运行模型的应用尤为重要。 - 支持复杂的语言特性：
TikToken 能够处理复杂的语言特性，包括多语言支持、特殊字符处理和子词标记化。这样可以提高模型对不同语言和文本格式的理解和处理能力。 - 简化开发过程：
TikToken 提供了易于使用的接口和工具，使开发者能够更简便地进行文本预处理和标记化。这减少了开发和调试的工作量，加速了模型开发和部署的过程。 - 提高模型性能：
通过优化标记化和预处理步骤，TikToken 可以提高模型的性能和生成结果的质量。更好的标记化方法可以使模型生成更加连贯和准确的文本输出。 - 社区和支持：
作为一个与 OpenAI 模型兼容的工具，TikToken 享有广泛的社区支持和丰富的文档资源。这使得开发者可以更容易地找到解决方案和最佳实践。
综合来看，TikToken 通过提供高效、优化和兼容的标记化方法，为 NLP 和大规模语言模型的开发和使用提供了强有力的支持，从而显著提升了文本处理的效率和模型性能。
3.2 token转embedding后归一化处理
在自然语言处理（NLP）和机器学习的上下文中，输入参数经过 TikToken 编码成 token 并转换成嵌入（embedding）后，进行归一化处理有几个重要的原因：
- 稳定训练过程
归一化处理可以使数据在相同的尺度范围内，这有助于稳定和加速模型的训练过程。未经归一化的输入可能具有不同的尺度，这会导致模型在训练过程中需要更长时间来找到最佳参数，甚至可能导致训练不稳定或发散。
- 提高数值稳定性
在神经网络中，尤其是深度网络，数值稳定性是一个关键问题。归一化可以防止数值溢出或下溢，确保在训练过程中参数的更新不会过大或过小，从而使模型能够更好地学习数据的特征。
- 优化器效果更好
大多数优化器（如 SGD、Adam 等）对数据的尺度非常敏感。如果输入数据的尺度不统一，优化器可能需要花费更多的时间来调整权重。归一化可以使优化器在训练过程中更加高效，缩短收敛时间。
- 防止过拟合
归一化有助于防止过拟合。在一些情况下，输入数据的不同尺度可能会导致模型对某些特征过度拟合，而忽略其他特征。归一化可以使模型更均衡地学习不同特征，从而提高泛化能力。
- 提高相似性计算效果
在许多 NLP 任务中，嵌入向量的相似性计算是一个重要步骤，如计算余弦相似度等。未经归一化的嵌入向量可能会因为尺度差异影响相似性计算的结果。归一化可以确保嵌入向量在相同的尺度上，从而提高相似性计算的准确性和可靠性。
- 防止梯度消失和爆炸
深度神经网络中的梯度消失和梯度爆炸问题是常见的问题。归一化可以帮助缓解这些问题，使梯度在反向传播过程中保持在适当的范围内，从而有助于网络的训练。
归一化的常用方法
- L2 归一化：将嵌入向量除以其 L2 范数，使得所有向量的 L2 范数都为 1。
- Min-Max 归一化：将数据缩放到 [0, 1] 的范围内。
- 标准化（Z-score 归一化）：将数据减去均值再除以标准差，使数据具有零均值和单位方差。
综上所述，归一化处理在嵌入向量的使用中起着关键作用，能够提高训练的效率和模型的性能。
3.3 归一化处理
在 Transformer 架构中，RMS 归一化和 softmax 归一化函数用于不同的目的，并在模型的不同部分中起作用。理解这两者的区别和它们在模型中的角色，可以帮助我们理解为什么需要同时使用它们。
3.3.1 RMS 归一化 (Root Mean Square Normalization)
RMS 归一化是一种归一化方法，用于对输入的张量进行标准化，以稳定训练过程并提升模型性能。RMS 归一化通过对张量进行归一化处理，使得张量的值分布更加均匀，防止数值不稳定。这种归一化方法在模型的每一层中应用，以确保模型的各层输出在同一尺度上。
在 Transformer 中的作用：
- 稳定训练过程：通过将张量的值归一化，RMS 归一化可以减少梯度爆炸或梯度消失的问题。
- 提高模型性能：归一化后的张量可以让模型更容易学习到有效的特征表示。
3.3.2 Softmax 归一化
Softmax函数是一种用于将一个向量中的每个元素转换为概率的数学函数，特别适用于多分类问题。它的输出是一个概率分布，所有输出的和为1。Softmax函数的数学表达式如下：
其中：
是自然对数的底，大约等于2.71828。
是对输入向量中所有元素的指数函数值的和。
Softmax 归一化是一种将向量转换为概率分布的函数，通常用于多分类任务中。它将输入的实数向量转换为一个概率向量，其中每个元素表示该类别的概率，并且所有元素的和为1。
在 Transformer 中的作用：
- 计算注意力权重：在自注意力机制中，softmax 用于将注意力分数（通过查询和键的点积计算得到的分数）转换为注意力权重，这些权重用于加权和输入的值。
- 输出概率分布：在生成模型的输出层中，softmax 用于将最后一层的 logits 转换为各个 token 的概率分布，以便于选择下一个 token。
PS: 后面我将单独出一篇文章详解Softmax函数。
3.3.3 结合使用
在 Transformer 模型中，同时使用 RMS 归一化和 softmax 归一化是为了在不同的阶段解决不同的问题：
- RMS 归一化用于层间输出：确保每一层的输出在进入下一层时具有稳定的分布，防止梯度爆炸或消失。
- softmax 归一化用于注意力机制和输出层：将注意力分数转换为注意力权重，以实现加权和；将最终的 logits 转换为概率分布，以生成模型的输出。
RMS 归一化用于确保模型各层输出的稳定性，而 softmax 归一化用于将注意力分数和输出 logits 转换为概率分布，两者在 Transformer 模型中协同工作，共同提升模型的性能和稳定性。
3.4 Q、K、V矩阵
在Transformer架构中，Q（Query）、K（Key）、V（Value）矩阵是自注意力机制（Self-Attention Mechanism）的核心组成部分。自注意力机制使得Transformer能够灵活地建模输入序列中不同位置的依赖关系。以下是对Q、K、V矩阵的详细解释及其在Transformer中的作用：
3.4.1 定义与生成
- Q矩阵（Query）：查询矩阵，表示查询向量的集合。
- K矩阵（Key）：键矩阵，表示键向量的集合。
- V矩阵（Value）：值矩阵，表示值向量的集合。
在Transformer中，输入序列首先通过嵌入层转换为一系列向量，然后这些向量通过三个不同的线性变换层生成Q、K、V矩阵。这些变换层的权重是可以学习的参数。具体来说，对于每个输入向量 ( x_i )：
其中，( W_Q )、( W_K )、( W_V ) 是相应的权重矩阵。
3.4.2 作用与计算
自注意力机制的计算过程可以概述如下：
- 计算注意力分数（Attention Scores）：首先，计算每个查询向量与所有键向量之间的点积，以衡量查询和键的相似性。这些相似度称为注意力分数。
- 归一化注意力分数：然后，将这些分数通过Softmax函数归一化，以便得到每个查询向量对所有键向量的注意力权重。
- 加权求和：最后，用这些注意力权重对值向量进行加权求和，得到每个查询的最终表示。
3.4.3 在Transformer中的作用
- 捕捉依赖关系：
自注意力机制使每个位置的表示能够根据输入序列中的所有其他位置来计算。这种全局依赖捕捉能力使得Transformer能够处理长距离依赖关系，而不是像RNN那样只能处理局部依赖关系。 - 并行计算：
Transformer的自注意力机制允许并行计算，因为每个位置的注意力计算与其他位置的计算是独立的。这使得Transformer比RNN更高效，特别是在处理长序列时。 - 多头注意力机制（Multi-Head Attention）：
为了增强模型的能力，Transformer使用多头注意力机制，即将Q、K、V矩阵分成多个头，每个头独立地进行注意力计算，然后将这些头的结果连接起来。这样做可以使模型从不同的子空间中捕捉信息，增强表示能力。
Q、K、V矩阵在Transformer架构中的自注意力机制中起到了核心作用。通过这些矩阵，Transformer能够计算每个输入位置与其他位置的相似性，并根据这些相似性灵活地聚合信息，从而实现对序列数据的高效处理和长距离依赖的捕捉。自注意力机制的引入，使得Transformer在许多NLP任务中取得了显著的性能提升。
3.5 复数和实数转换的应用
在 PyTorch 中，view_as_complex 和 view_as_real 是用于在实数和复数表示之间转换的函数。它们非常有用，特别是在处理涉及复数运算的信号处理、傅里叶变换等领域。以下是这两个函数的详细功能和使用方法：
3.5.1 view_as_complex
torch.view_as_complex(input) 将形状为 (*, 2) 的实数张量视为复数张量。这里的 * 表示任意维度，最后一个维度的大小必须为 2，它代表复数的实部和虚部。
使用方法：
假设有一个形状为 (m, n, 2) 的实数张量 x，其中 x[..., 0] 是复数的实部，x[..., 1] 是复数的虚部。你可以使用 view_as_complex 将其转换为复数张量。
import torch
# 创建一个形状为 (m, n, 2) 的实数张量
x = torch.randn(3, 4, 2)
# 将实数张量视为复数张量
x_complex = torch.view_as_complex(x)
print(x_complex)3.5.2 view_as_real
torch.view_as_real(input) 将复数张量视为实数张量。复数张量将被表示为形状为 (*, 2) 的实数张量，其中最后一个维度包含复数的实部和虚部。
使用方法：
假设有一个复数张量 y，你可以使用 view_as_real 将其转换为形状为 (*, 2) 的实数张量。
import torch
# 创建一个复数张量
y = torch.randn(3, 4, dtype=torch.complex64)
# 将复数张量视为实数张量
y_real = torch.view_as_real(y)
print(y_real)3.5.3 一般使用场景
- 信号处理：
在进行傅里叶变换等操作时，通常需要在实数和复数表示之间进行转换。 - 数据准备和转换：
当数据以复数形式存储或传输时，需要将其转换为实数张量进行进一步处理。 - 模型处理：
一些模型可能需要处理复数数据，此时需要在实数表示和复数表示之间进行转换，以便利用 PyTorch 的复数支持。
3.5.4 示例
以下是一个完整的示例，展示了如何在实数和复数张量之间进行转换：
import torch
# 创建一个形状为 (2, 2, 2) 的实数张量，其中最后一个维度表示实部和虚部
x = torch.tensor([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]])
# 将实数张量视为复数张量
x_complex = torch.view_as_complex(x)
# 打印复数张量
print("复数张量:")
print(x_complex)
# 将复数张量视为实数张量
x_real = torch.view_as_real(x_complex)
# 打印实数张量
print("实数张量:")
print(x_real)输出：
复数张量:
tensor([[1.0 + 2.0j, 3.0 + 4.0j],
[5.0 + 6.0j, 7.0 + 8.0j]], dtype=torch.complex64)
实数张量:
tensor([[[1.0, 2.0], [3.0, 4.0]],
[[5.0, 6.0], [7.0, 8.0]]])通过这种方式，可以方便地在实数和复数表示之间进行转换，以适应不同的计算需求和操作。
3.6 前馈网络
在Transformer架构中，前馈网络（Feed-Forward Network, FFN）是每个Transformer层中的一个重要组成部分。前馈网络的作用是对每个位置上的特征进行非线性变换，从而增强模型的表达能力。在这种上下文中，像SwiGLU（Switching Gated Linear Unit）这样的前馈网络变体引入了一些新的激活函数和结构，旨在进一步提升模型的性能和效率。
3.6.1 前馈网络的作用
- 非线性变换：
前馈网络通过线性变换和非线性激活函数，对输入特征进行复杂的非线性变换。这使得模型能够捕捉更丰富和复杂的特征，而不仅仅是线性关系。 - 增强表示能力：
每个Transformer层中的前馈网络对每个位置的特征独立地进行处理，使得模型能够更好地表示和处理输入数据的局部信息。 - 引入非线性：
非线性激活函数（如ReLU、GELU等）帮助模型在不同的层之间引入非线性，增加模型的表达能力，使其能够拟合更复杂的函数。
3.6.2 SwiGLU等前馈网络变体
SwiGLU（Switching Gated Linear Unit）是一个新的前馈网络变体，它的设计旨在提高模型的性能。其主要思想是使用门控机制来控制信息流动，从而增强网络的表达能力。
3.6.3 SwiGLU的结构
SwiGLU由以下几个部分组成：
- 线性变换：
对输入进行两次线性变换，产生两个不同的表示。 - 门控机制：
使用一个门控单元来控制这两个表示的结合。 - 非线性激活：
对组合后的表示应用非线性激活函数。
具体来说，SwiGLU可以表示为：
其中：
- X 是输入。
- W1 和 W2 是线性变换的权重矩阵。
- 圈点 表示逐元素相乘。
- GELU（Gaussian Error Linear Unit）是一种非线性激活函数。
3.6.4 SwiGLU的优势
- 更有效的信息流动：
门控机制允许模型动态地选择和控制信息流动，使得模型能够更灵活地调整特征的组合，从而提高模型的表达能力。 - 更强的非线性表达：
SwiGLU通过引入更多的非线性变换，能够更好地捕捉复杂的特征关系。 - 性能提升：
在一些实验中，SwiGLU被证明能够在保持计算复杂度的情况下，提升模型的性能和效率。
3.6.5 为什么要使用前馈网络
前馈网络是Transformer架构中不可或缺的一部分，主要原因有以下几点：
- 增强模型能力：
前馈网络能够显著增强模型的表达能力，使其能够处理更复杂和多样的输入数据。 - 独立处理每个位置：
前馈网络在每个位置上独立操作，这意味着它能够并行处理输入序列中的所有位置，提高计算效率。 - 引入非线性变换：
通过非线性激活函数，前馈网络能够引入非线性变换，使得模型能够学习和拟合复杂的关系。
前馈网络在Transformer架构中起着关键作用，通过对每个位置上的特征进行非线性变换来增强模型的表达能力。像SwiGLU这样的前馈网络变体通过引入门控机制和新的激活函数，进一步提升了模型的性能和效率。前馈网络的引入，使得Transformer能够更好地处理复杂的序列数据，并在许多NLP任务中取得了显著的成功。
4、核心代码
from pathlib import Path
import tiktoken
from tiktoken.load import load_tiktoken_bpe
import torch
import json
import matplotlib.pyplot as plt
# 定义 tokenizer 模型路径
tokenizer_path = "Meta-Llama-3-8B/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)]
# 加载 BPE 合并规则
mergeable_ranks = load_tiktoken_bpe(tokenizer_path)
# 初始化 tokenizer
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)},
)
# 测试编码和解码
tokenizer.decode(tokenizer.encode("hello world!"))
# 加载预训练模型
model = torch.load("Meta-Llama-3-8B/consolidated.00.pth")
print(json.dumps(list(model.keys())[:20], indent=4))
# 加载模型配置参数
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"])
# 定义输入提示词并进行编码
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)
# 初始化嵌入层并复制模型的权重
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):
# rms = (tensor.pow(2).mean(-1, keepdim=True) + norm_eps)**0.5
# return tensor * (norm_weights / rms)
def rms_norm(tensor, norm_weights):
return (tensor * torch.rsqrt(tensor.pow(2).mean(-1, keepdim=True) + norm_eps)) * norm_weights
# 开始进行 Transformer 层的计算
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"])
# 获取 QKV 权重并进行重塑
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"]
# 计算每个注意力头的 QKV
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、K、V 向量
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)
# 处理旋转嵌入（RoPE）
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 =
# 应用 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"])
final_embedding.shape
model["output.weight"].shape
# 计算输出层的 logits
logits = torch.matmul(final_embedding[-1], model["output.weight"].T)
logits.shape
# 获取下一个 token
next_token = torch.argmax(logits, dim=-1)
next_token
# 解码下一个 token
tokenizer.decode([next_token.item()])解释
- 初始化和导入：导入所需的库，并设置 tokenizer 和模型路径。
- 特殊令牌和 BPE 规则：定义特殊令牌并加载 BPE 规则，初始化 tokenizer。
- 加载模型和配置：加载预训练的模型权重和配置文件，提取必要的超参数。
- 处理提示词：对输入提示词进行编码，并准备嵌入层。
- RMS 归一化函数：定义 RMS 归一化函数，用于层间归一化。
- Transformer 层计算：循环计算每一层的注意力和前馈网络，更新最终的嵌入。
- 计算输出和解码：计算输出层的 logits，获取并解码下一个 token。
