【Happy-LLM】以llama2为例动手搭建一个大模型结构
1 Transformer结构
1.1 Attention机制
学习参考：《动手学习深度学习第10章》
1.2 Encoder-Decoder架构
学习参考：
1.3 transformer架构
2 Llama2 大模型架构与动手搭建
2.1 Llama2模型架构
huggingFace Llama2 文档
llama1
llama2
昇思MindSpore技术公开课 大咖深度解析LLaMA2 模型架构
与Transformer的对比：Transformer主要分为编码器（encoder）和解码器（decoder）两部分。相较之下，LLaMA仅使用了Transformer的解码器部分，采取了一个仅解码器（decoder-only）的结构。这种结构现在被大多数生成型的语言模型所采用。
- 在结构上，与Transformer模型相比，LLaMA2的主要变化是将其中的层标准化（LayerNorm）替换为了均方根标准化（RMSNorm）；传统的transformer将norm层放在attention之后，而llama将其前置。
- 多头注意力（Multi-Head Attention）换成了分组查询注意力（GQA，在LLaMA中则是多查询注意力MQA）
- 将位置编码（Positional Encoding）替换为了旋转嵌入（Rotary Position Embedding，RoPE）
- MLP 层更新 - 增加了 UP 和 Donw 的线性层并使用 SiLU 激活
2.2 动手实现llama大模型
（1）超参数和模型参数类：
自定义一个ModelConfig类，来存储和记录我们的超参数，这里我们继承了PretrainedConfig类，这是transformers库中的参数类，我们可以通过继承这个类来方便的使用transformers库中的一些功能，也方便在后续导出Hugging Face模型
from transformers import PretrainedConfig
class ModelConfig(PretrainedConfig):
model_type = "Tiny-K"
def __init__(
self,
dim: int = 768, # 模型维度
n_layers: int = 12, # Transformer的层数
n_heads: int = 16, # 注意力机制的头数
n_kv_heads: int = 8, # 键值头的数量
vocab_size: int = 6144, # 词汇表大小
hidden_dim: int = None, # 隐藏层维度
multiple_of: int = 64,
norm_eps: float = 1e-5, # 归一化层的eps
max_seq_len: int = 512, # 最大序列长度
dropout: float = 0.0, # dropout概率
flash_attn: bool = True, # 是否使用Flash Attention
**kwargs,
):
self.dim = dim
self.n_layers = n_layers
self.n_heads = n_heads
self.n_kv_heads = n_kv_heads
self.vocab_size = vocab_size
self.hidden_dim = hidden_dim
self.multiple_of = multiple_of
self.norm_eps = norm_eps
self.max_seq_len = max_seq_len
self.dropout = dropout
self.flash_attn = flash_attn
super().__init__(**kwargs)
（2）RMSNorm层的实现：
RMSNorm可以用如下的数学公式表示：
\[\text{RMSNorm}(x) = \frac{x}{\sqrt{\frac{1}{n}\sum_{i=1}^{n}x_i^2 + \epsilon}} \cdot \gamma
\]
其中：
- \(x_i\) 是输入向量的第 \(i\) 个元素
- \(\gamma\) 是可学习的缩放参数（对应代码中的
self.weight） - \(n\) 是输入向量的维度数量
- \(\epsilon\) 是一个小常数，用于数值稳定性（以避免除以零的情况）
class RMSNorm(torch.nn.Module):
def __init__(self, dim: int, eps: float = 1e-6):
"""
Initialize the RMSNorm normalization layer.
Args:
dim (int): The dimension of the input tensor.
eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
Attributes:
eps (float): A small value added to the denominator for numerical stability. 防止除以0
weight (nn.Parameter): Learnable scaling parameter.可学习化的参数，初始化为1
"""
super().__init__()
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim))
def _norm(self, x):
# 计算RMSNorm的核心部分
# x.pow(2).mean(-1, keepdim=True)计算了输入x的平方的均值
# torch.rsqrt是平方根的倒数，这样就得到了RMSNorm的分母部分，再加上eps防止分母为0
# 最后乘以x，得到RMSNorm的结果
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
output = self._norm(x.float()).type_as(x)
return output * self.weight
（3）构建attention机制（分组查询注意力机制（Grouped-Query Attention，GQA））
1、repeat_kv
将键和值的维度扩展到和查询的维度一样，这样才能进行注意力计算
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
"""torch.repeat_interleave(x, dim=2, repeats=n_rep)"""
bs, slen, n_kv_heads, head_dim = x.shape
if n_rep == 1:
return x
return (
x[:, :, :, None, :]
.expand(bs, slen, n_kv_heads, n_rep, head_dim)
.reshape(bs, slen, n_kv_heads * n_rep, head_dim)
)
在上述代码中：
-
首先，获取输入张量的形状：首先，代码通过 x.shape 获取输入张量的形状，包括批量大小（bs）、序列长度（slen）、键/值对头的数量（n_kv_heads）以及每个头的维度大小（head_dim）。
-
然后，检查重复次数：接着，代码检查重复次数 n_rep 是否为1。如果是1，则说明不需要对键和值进行重复，直接返回原始张量 x。
-
最后，扩展和重塑张量：
- 在第三个维度（即键/值对头的维度）之后添加一个新的维度，形成
x[:, :, :, None, :]。 - 使用
expand方法将新添加的维度扩展到n_rep大小，实现键/值对的重复效果。 - 最后，通过 reshape 方法重新塑形，将扩展后的维度合并回键/值对头的数量中，即
x.reshape(bs, slen, n_kv_heads * n_rep, head_dim)，这样最终的张量形状就达到了与查询维度一致的效果。
- 在第三个维度（即键/值对头的维度）之后添加一个新的维度，形成
2、Rotate Embedding旋转嵌入
（1）实虚部分离：
def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
# torch.arange(0, dim, 2)[: (dim // 2)].float()生成了一个从0开始，步长为2的序列，长度为dim的一半
# 然后每个元素除以dim，再取theta的倒数，得到频率freqs
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
# 生成时间序列长度，生成一个从0到end的序列，长度为end。end通常是序列的最大长度。
t = torch.arange(end, device=freqs.device) # type: ignore
freqs = torch.outer(t, freqs).float() # 计算外积，得到一个二维矩阵，每一行是t的元素乘以freqs的元素
freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64
return freqs_cis
（2）张量形状调整
调整 freqs_cis 的形状，使其在进行广播操作时与 x 的维度对齐，从而能够进行正确的张量运算：
def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
ndim = x.ndim
assert 0 <= 1 < ndim
assert freqs_cis.shape == (x.shape[1], x.shape[-1])
shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
return freqs_cis.view(*shape)
（3）RoPE旋转编码：
def apply_rotary_emb(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor):
xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
"""
.reshape(*xq.shape[:-1], -1, 2):
*xq.shape[:-1]：保持除最后一维外的所有维度
-1, 2：将最后一维 D 重塑为 (D/2, 2)
最后转化为复数张量，取最后一维的两个值作为复数的实部和虚部
原始方法：xq_real, xq_img = xq.float().reshape(x_q.shape[:-1] + (-1, 2)).unbind(-1)
"""
xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
return xq_out.type_as(xq), xk_out.type_as(xk)
【stack操作学习】
# 示例代码，准备实虚部数据
# 实部数据 (xq_out_r)
xq_out_r = torch.tensor([[
[[1.0, 3.0, 5.0], # head 0, position 0
[2.0, 4.0, 6.0]], # head 1, position 0
[[7.0, 9.0, 11.0], # head 0, position 1
[8.0, 10.0, 12.0]] # head 1, position 1
]])
# 虚部数据 (xq_out_i)
xq_out_i = torch.tensor([[
[[1.1, 3.1, 5.1], # head 0, position 0
[2.1, 4.1, 6.1]], # head 1, position 0
[[7.1, 9.1, 11.1], # head 0, position 1
[8.1, 10.1, 12.1]] # head 1, position 1
]])
# 在最后一个维度（dim=-1）上堆叠
stacked = torch.stack([xq_out_r, xq_out_i], dim=-1)
print("Stacked shape:", stacked.shape)
print("Stacked tensor:")
print(stacked)
输出结果
Stacked shape: torch.Size([1, 2, 2, 3, 2])
解释：[batch=1, seq_len=2, num_heads=2, head_dim/2=3, 实虚部=2]
Stacked tensor:
tensor([[[[[ 1.0000, 1.1000],
[ 3.0000, 3.1000],
[ 5.0000, 5.1000]],
[[ 2.0000, 2.1000],
[ 4.0000, 4.1000],
[ 6.0000, 6.1000]]],
[[[ 7.0000, 7.1000],
[ 9.0000, 9.1000],
[11.0000, 11.1000]],
[[ 8.0000, 8.1000],
[10.0000, 10.1000],
[12.0000, 12.1000]]]]])
3、attention模块
Llama2采用不同条件下的注意力头机制：n_heads: 查询头数量（Query Heads）、n_kv_heads: 键值头数量（Key-Value Heads）
- 当 n_kv_heads < n_heads 时，实现 分组查询注意力（GQA）
- 当 n_kv_heads = 1 时，实现 多查询注意力（MQA）
- 当 n_kv_heads = n_heads 时，是标准的 多头注意力（MHA）
```python
class Attention(nn.Module):
def __init__(self, args: ModelConfig):
super().__init__()
# 根据是否指定n_kv_heads，确定用于键（key）和值（value）的头的数量。
self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
# 确保总头数可以被键值头数整除。
assert args.n_heads % self.n_kv_heads == 0
# 模型并行处理大小，默认为1。
model_parallel_size = 1
# 本地计算头数，等于总头数除以模型并行处理大小。
self.n_local_heads = args.n_heads // model_parallel_size
# 本地键值头数，等于键值头数除以模型并行处理大小。
self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
# 重复次数，用于扩展键和值的尺寸。
self.n_rep = self.n_local_heads // self.n_local_kv_heads
# 每个头的维度，等于模型维度除以头的总数。
self.head_dim = args.dim // args.n_heads
# 定义权重矩阵。
self.wq = nn.Linear(args.dim, args.n_heads * self.head_dim, bias=False)
self.wk = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
self.wv = nn.Linear(args.dim, self.n_kv_heads * self.head_dim, bias=False)
# 输出权重矩阵。
self.wo = nn.Linear(args.n_heads * self.head_dim, args.dim, bias=False)
# 定义dropout。
self.attn_dropout = nn.Dropout(args.dropout)
self.resid_dropout = nn.Dropout(args.dropout)
# 保存dropout概率。
self.dropout = args.dropout
# 检查是否使用Flash Attention（需要PyTorch >= 2.0）。
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')
if not self.flash:
# 若不支持Flash Attention，则使用手动实现的注意力机制，并设置mask。
print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
# 创建一个上三角矩阵，用于遮蔽未来信息。
mask = torch.full((1, 1, args.max_seq_len, args.max_seq_len), float("-inf"))
mask = torch.triu(mask, diagonal=1)
# 注册为模型的缓冲区
self.register_buffer("mask", mask)
def forward(self, x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor):
# 获取批次大小和序列长度，[batch_size, seq_len, dim]
bsz, seqlen, _ = x.shape
# 计算查询（Q）、键（K）、值（V）。
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
# 调整形状以适应头的维度。
xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
# 应用旋转位置嵌入（RoPE）。
xq, xk = apply_rotary_emb(xq, xk, freqs_cos, freqs_sin)
# 对键和值进行扩展以适应重复次数。
xk = repeat_kv(xk, self.n_rep)
xv = repeat_kv(xv, self.n_rep)
# 将头作为批次维度处理。
xq = xq.transpose(1, 2)
xk = xk.transpose(1, 2)
xv = xv.transpose(1, 2)
# 根据是否支持Flash Attention，选择实现方式。
if self.flash:
# 使用Flash Attention。
output = torch.nn.functional.scaled_dot_product_attention(xq, xk, xv, attn_mask=None, dropout_p=self.dropout if self.training else 0.0, is_causal=True)
else:
# 使用手动实现的注意力机制。
scores = torch.matmul(xq, xk.transpose(2, 3)) / math.sqrt(self.head_dim)
assert hasattr(self, 'mask')
scores = scores + self.mask[:, :, :seqlen, :seqlen]
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
scores = self.attn_dropout(scores)
output = torch.matmul(scores, xv)
# 恢复时间维度并合并头。
output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
# 最终投影回残差流。
output = self.wo(output)
output = self.resid_dropout(output)
return output
llama源代码，采用kv chace高采样：
class Attention(nn.Module):
"""Multi-head attention module."""
def __init__(self, args: ModelArgs):
super().__init__()
self.n_kv_heads = args.n_heads if args.n_kv_heads is None else args.n_kv_heads
model_parallel_size = fs_init.get_model_parallel_world_size()
self.n_local_heads = args.n_heads // model_parallel_size
self.n_local_kv_heads = self.n_kv_heads // model_parallel_size
self.n_rep = self.n_local_heads // self.n_local_kv_heads
self.head_dim = args.dim // args.n_heads
# 输入在列维度切分，适合并行计算
self.wq = ColumnParallelLinear(
args.dim,
args.n_heads * self.head_dim,
bias=False,
gather_output=False,
init_method=lambda x: x,
)
self.wk = ColumnParallelLinear(
args.dim,
self.n_kv_heads * self.head_dim,
bias=False,
gather_output=False,
init_method=lambda x: x,
)
self.wv = ColumnParallelLinear(
args.dim,
self.n_kv_heads * self.head_dim,
bias=False,
gather_output=False,
init_method=lambda x: x,
)
self.wo = RowParallelLinear(
args.n_heads * self.head_dim,
args.dim,
bias=False,
input_is_parallel=True,
init_method=lambda x: x,
)
# 缓存过去的键，避免重复计算，KV cache加速attention参数训练
# 形状: [最大批次, 最大序列长度, 本地KV头数, 头维度]
self.cache_k = torch.zeros(
(
args.max_batch_size,
args.max_seq_len,
self.n_local_kv_heads,
self.head_dim,
)
).cuda()
self.cache_v = torch.zeros(
(
args.max_batch_size,
args.max_seq_len,
self.n_local_kv_heads,
self.head_dim,
)
).cuda()
def forward(
self,
x: torch.Tensor,
start_pos: int,
freqs_cis: torch.Tensor,
mask: Optional[torch.Tensor],
):
bsz, seqlen, _ = x.shape
xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
# 输入: [bsz, seqlen, dim]
# 输出:
# xq: [bsz, seqlen, n_heads * head_dim]
# xk/xv: [bsz, seqlen, n_kv_heads * head_dim]
# 然后对q，k，v重塑为: [批次, 序列长度, 头数, 头维度]
xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
# 应用RoPE旋转位置编码
xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
# 存储k ,v键值到kv cache当中，并对缓存进行时序更新
self.cache_k = self.cache_k.to(xq)
self.cache_v = self.cache_v.to(xq)
self.cache_k[:bsz, start_pos : start_pos + seqlen] = xk
self.cache_v[:bsz, start_pos : start_pos + seqlen] = xv
keys = self.cache_k[:bsz, : start_pos + seqlen]
values = self.cache_v[:bsz, : start_pos + seqlen]
# repeat k/v heads if n_kv_heads < n_heads 应用多头注意力机制
keys = repeat_kv(keys, self.n_rep) # (bs, cache_len + seqlen, n_local_heads, head_dim)
values = repeat_kv(values, self.n_rep) # (bs, cache_len + seqlen, n_local_heads, head_dim)
# 调整维度顺序
xq = xq.transpose(1, 2) # (bs, n_local_heads, seqlen, head_dim)
keys = keys.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim)
values = values.transpose(1, 2) # (bs, n_local_heads, cache_len + seqlen, head_dim)
# 注意力分数
scores = torch.matmul(xq, keys.transpose(2, 3)) / math.sqrt(self.head_dim)
# 添加注意力掩码（用于防止看到未来信息）
if mask is not None:
scores = scores + mask # (bs, n_local_heads, seqlen, cache_len + seqlen)
# Softmax归一化
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
# 注意力加权求和
output = torch.matmul(scores, values) # (bs, n_local_heads, seqlen, head_dim)
output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
return self.wo(output)
4、MLP模块
forward函数的实现，首先，输入 x 通过第一层线性变换 self.w1 和 SILU 激活函数，然后，结果乘以输入 x 通过第三层线性变换 self.w3 的结果，最后，通过第二层线性变换 self.w2 和 dropout 层，得到最终输出。
class MLP(nn.Module):
def __init__(self, dim: int, hidden_dim: int, multiple_of: int, dropout: float):
super().__init__()
# 如果没有指定隐藏层的维度，我们将其设置为输入维度的4倍
# 然后将其减少到2/3，最后确保它是multiple_of的倍数
if hidden_dim is None:
hidden_dim = 4 * dim
hidden_dim = int(2 * hidden_dim / 3)
hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
# 定义第一层线性变换，从输入维度到隐藏维度
self.w1 = nn.Linear(dim, hidden_dim, bias=False)
# 定义第二层线性变换，从隐藏维度到输入维度
self.w2 = nn.Linear(hidden_dim, dim, bias=False)
# 定义第三层线性变换，从输入维度到隐藏维度
self.w3 = nn.Linear(dim, hidden_dim, bias=False)
# 定义dropout层，用于防止过拟合
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# 前向传播函数
# 首先，输入x通过第一层线性变换和SILU激活函数
# 然后，结果乘以输入x通过第三层线性变换的结果
# 最后，通过第二层线性变换和dropout层
return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
对比llama源代码实现：这里考虑张量并行的方法进行前向传播，用于加速预训练
ColumnParallelLinear和RowParallelLinear
Dive into Tensor Parallelism: Building ColumnParallelLinear and RowParallelLinear from Scratch
class FeedForward(nn.Module):
def __init__(
self,
dim: int,
hidden_dim: int,
multiple_of: int,
ffn_dim_multiplier: Optional[float],
):
"""
Initialize the FeedForward module.
Args:
dim (int): Input dimension.
hidden_dim (int): Hidden dimension of the feedforward layer.
multiple_of (int): Value to ensure hidden dimension is a multiple of this value.
ffn_dim_multiplier (float, optional): Custom multiplier for hidden dimension. Defaults to None.
Attributes:
w1 (ColumnParallelLinear): Linear transformation for the first layer.
w2 (RowParallelLinear): Linear transformation for the second layer.
w3 (ColumnParallelLinear): Linear transformation for the third layer.
"""
super().__init__()
hidden_dim = int(2 * hidden_dim / 3)
# custom dim factor multiplier
if ffn_dim_multiplier is not None:
hidden_dim = int(ffn_dim_multiplier * hidden_dim)
hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
self.w1 = ColumnParallelLinear(
dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
)
self.w2 = RowParallelLinear(
hidden_dim, dim, bias=False, input_is_parallel=True, init_method=lambda x: x
)
self.w3 = ColumnParallelLinear(
dim, hidden_dim, bias=False, gather_output=False, init_method=lambda x: x
)
def forward(self, x):
return self.w2(F.silu(self.w1(x)) * self.w3(x))
5、Decoder模块
由搭建好的attention和mlp层完成Decoder层的搭建：
class DecoderLayer(nn.Module):
def __init__(self, layer_id: int, args: ModelConfig):
super().__init__()
# 定义多头注意力的头数
self.n_heads = args.n_heads
# 定义输入维度
self.dim = args.dim
# 定义每个头的维度，等于输入维度除以头数
self.head_dim = args.dim // args.n_heads
# 定义LLaMA2Attention对象，用于进行多头注意力计算
self.attention = Attention(args)
# 定义LLaMAMLP对象，用于进行前馈神经网络计算
self.feed_forward = MLP(
dim=args.dim,
hidden_dim=args.hidden_dim,
multiple_of=args.multiple_of,
dropout=args.dropout,
)
# 定义层的ID
self.layer_id = layer_id
# 定义注意力计算的归一化层
self.attention_norm = RMSNorm(args.dim, eps=args.norm_eps)
# 定义前馈神经网络计算的归一化层
self.ffn_norm = RMSNorm(args.dim, eps=args.norm_eps)
def forward(self, x, freqs_cos, freqs_sin):
# 前向传播函数
# 首先，输入x经过注意力归一化层，然后进行注意力计算，结果与输入x相加得到h
# 然后，h经过前馈神经网络归一化层，然后进行前馈神经网络计算，结果与h相加得到输出
h = x + self.attention.forward(self.attention_norm(x), freqs_cos, freqs_sin)
out = h + self.feed_forward.forward(self.ffn_norm(h))
return out
6、Transformer模块搭建
class Transformer(PreTrainedModel):
config_class = ModelConfig # 配置类
last_loss: Optional[torch.Tensor] # 记录最后一次计算的损失
def __init__(self, args: ModelConfig = None):
super().__init__(args)
# 初始化模型参数
self.args = args
# 词汇表大小
self.vocab_size = args.vocab_size
# 层数
self.n_layers = args.n_layers
# 词嵌入层
self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim)
# Dropout层
self.dropout = nn.Dropout(args.dropout)
# Decoder层
self.layers = torch.nn.ModuleList()
for layer_id in range(args.n_layers):
self.layers.append(DecoderLayer(layer_id, args))
# 归一化层
self.norm = RMSNorm(args.dim, eps=args.norm_eps)
# 输出层
self.output = nn.Linear(args.dim, args.vocab_size, bias=False)
# 将词嵌入层的权重与输出层的权重共享
self.tok_embeddings.weight = self.output.weight
# 预计算相对位置嵌入的频率
freqs_cos, freqs_sin = precompute_freqs_cis(self.args.dim // self.args.n_heads, self.args.max_seq_len)
self.register_buffer("freqs_cos", freqs_cos, persistent=False)
self.register_buffer("freqs_sin", freqs_sin, persistent=False)
# 初始化所有权重
self.apply(self._init_weights)
# 对残差投影进行特殊的缩放初始化
for pn, p in self.named_parameters():
if pn.endswith('w3.weight') or pn.endswith('wo.weight'):
torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * args.n_layers))
# 初始化最后一次前向传播的损失属性
self.last_loss = None
self.OUT = CausalLMOutputWithPast() # 输出容器
self._no_split_modules = [name for name, _ in self.named_modules()] # 不分割的模块列表
def _init_weights(self, module):
# 初始化权重的函数
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
def forward(self, tokens: torch.Tensor, targets: Optional[torch.Tensor] = None, **kwargs) -> torch.Tensor:
"""
- tokens: Optional[torch.Tensor], 输入 token 张量。
- targets: Optional[torch.Tensor], 目标 token 张量。
- kv_cache: bool, 是否使用键值缓存。
- kwargs: 其他关键字参数。
- self.OUT: CausalLMOutputWithPast, 包含 logits 和损失。
"""
if 'input_ids' in kwargs:
tokens = kwargs['input_ids']
if 'attention_mask' in kwargs:
targets = kwargs['attention_mask']
# 前向传播函数
_bsz, seqlen = tokens.shape
# 通过词嵌入层和Dropout层
h = self.tok_embeddings(tokens)
h = self.dropout(h)
# 获取相对位置嵌入的频率
freqs_cos = self.freqs_cos[:seqlen]
freqs_sin = self.freqs_sin[:seqlen]
# 通过Decoder层
for layer in self.layers:
h = layer(h, freqs_cos, freqs_sin)
# 通过归一化层
h = self.norm(h)
if targets is not None:
# 如果给定了目标，计算损失
logits = self.output(h)
self.last_loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=0, reduction='none')
else:
# 推理时的小优化：只对最后一个位置的输出进行前向传播
logits = self.output(h[:, [-1], :])
self.last_loss = None
# 设置输出
self.OUT.__setitem__('logits', logits)
self.OUT.__setitem__('last_loss', self.last_loss)
return self.OUT
@torch.inference_mode()
def generate(self, idx, stop_id=None, max_new_tokens=256, temperature=1.0, top_k=None):
"""
给定输入序列 idx（形状为 (bz,seq_len) 的长整型张量），通过多次生成新 token 来完成序列。
在 model.eval() 模式下运行。效率较低的采样版本，没有使用键k/v cache。
"""
index = idx.shape[1]
for _ in range(max_new_tokens):
# 如果序列上下文过长，截断它到最大长度
idx_cond = idx if idx.size(1) <= self.args.max_seq_len else idx[:, -self.args.max_seq_len:]
# 前向传播获取序列中最后一个位置的 logits
logits = self(idx_cond).logits
logits = logits[:, -1, :] # 只保留最后一个时间步的输出
if temperature == 0.0:
# 选择最有可能的索引
_, idx_next = torch.topk(logits, k=1, dim=-1)
else:
# 缩放 logits 并应用 softmax
logits = logits / temperature
if top_k is not None:
v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
logits[logits < v[:, [-1]]] = -float('Inf')
probs = F.softmax(logits, dim=-1)
idx_next = torch.multinomial(probs, num_samples=1)
if idx_next == stop_id:
break
# 将采样的索引添加到序列中并继续
idx = torch.cat((idx, idx_next), dim=1)
return idx[:, index:] # 只返回生成的token
（4）tokenizer搭建
在自然语言处理 (NLP) 中，Tokenizer 是一种将文本分解为较小单位（称为 token）的工具。这些 token 可以是词、子词、字符，甚至是特定的符号。Tokenization 是 NLP 中的第一步，直接影响后续处理和分析的效果。不同类型的 tokenizer 适用于不同的应用场景，以下是几种常见的 tokenizer 及其特点。
- Word-based Tokenizer
- Word-based Tokenizer 是最简单和直观的一种分词方法。它将文本按空格和标点符号分割成单词。这种方法的优点在于其简单和直接，易于实现，且与人类对语言的直觉相符。
- 缺点：
- 无法处理未登录词（OOV，out-of-vocabulary）和罕见词，对复合词（如“New York”）或缩略词（如“don't”）的处理也不够精细
- Word-based Tokenizer 在处理不同语言时也会遇到挑战，因为一些语言（如中文、日文）没有显式的单词分隔符。
Input: "Hello, world! There is Datawhale."
Output: ["Hello", ",", "world", "!", "There", "is", "Datawhale", "."]
- Character-based Tokenizer
- Character-based Tokenizer 将文本中的每个字符视为一个独立的 token。这种方法能非常精细地处理文本，适用于处理拼写错误、未登录词或新词。由于每个字符都是一个独立的 token，因此这种方法可以捕捉到非常细微的语言特征。
- 这种方法也会导致 token 序列变得非常长，增加了模型的计算复杂度和训练时间
- 字符级的分割可能会丢失一些词级别的语义信息，使得模型难以理解上下文。
Input: "Hello"
Output: ["H", "e", "l", "l", "o"]
- Subword Tokenizer
- Subword Tokenizer 介于词和字符之间，能够更好地平衡分词的细粒度和处理未登录词的能力。Subword Tokenizer 的关键思想是将文本分割成比单词更小的单位，但又比字符更大，这样既能处理未知词，又能保持一定的语义信息。
- Byte Pair Encoding (BPE)
- BPE 是一种基于统计方法，通过反复合并频率最高的字符或字符序列对来生成子词词典。这种方法的优点在于其简单和高效，能够有效地处理未知词和罕见词，同时保持较低的词典大小。
Input: "lower" Output: ["low", "er"]
- WordPiece:
- WordPiece 是另一种基于子词的分词方法，最初用于谷歌的 BERT 模型。与 BPE 类似，WordPiece 通过最大化子词序列的似然函数来生成词典，但在合并子词时更注重语言模型的优化。
- WordPiece 会优先选择能够最大化整体句子概率的子词，使得分词结果在语言模型中具有更高的概率。
- Unigram:
- Unigram 分词方法基于概率模型，通过选择具有最高概率的子词来分割文本。
- Unigram 词典是通过训练语言模型生成的，可以处理多种语言和不同类型的文本。Unigram 模型会为每个子词分配一个概率，然后根据这些概率进行最优分割。
- Byte Pair Encoding (BPE)
如何训练一个tokenizer：datawhale happy-llm 5,2,4
这里我们选择使用 BPE 算法来训练一个 Subword Tokenizer。BPE 是一种简单而有效的分词方法，能够处理未登录词和罕见词，同时保持较小的词典大小。我们将使用 Hugging Face 的 tokenizers 库来训练一个 BPE Tokenizer。
Step 1: 安装和导入依赖库
首先，我们需要安装 tokenizers 库，除此之外还需要安装 datasets 和 transformers 库，用于加载训练数据和加载训练完成后的 Tokenizer。
pip install tokenizers datasets transformers
然后，导入所需的库。
import random
import json
import os
from transformers import AutoTokenizer, PreTrainedTokenizerFast
from tokenizers import (
decoders,
models,
pre_tokenizers,
trainers,
Tokenizer,
)
from tokenizers.normalizers import NFKC
from typing import Generator
Step 2: 加载训练数据
这里我们使用与预训练相同的数据集（出门问问序列猴子开源数据集）训练tokenizer，可使用code/download_dataset.sh 和 code/deal_dataset.py 下载和预处理数据集。
注：由于数据集过大，可能会导致在训练过程中内存不足。因为本项目为学习目的，建议学习者手动分割小部分数据集用于训练验证，笔者也在 Github 仓库中存放了训练好的 tokenizer，可以直接使用。
def read_texts_from_jsonl(file_path: str) -> Generator[str, None, None]:
"""读取JSONL文件并安全提取文本数据"""
with open(file_path, 'r', encoding='utf-8') as f:
for line_num, line in enumerate(f, 1):
try:
data = json.loads(line)
if 'text' not in data:
raise KeyError(f"Missing 'text' field in line {line_num}")
yield data['text']
except json.JSONDecodeError:
print(f"Error decoding JSON in line {line_num}")
continue
except KeyError as e:
print(e)
continue
Step 3: 创建配置文件
在训练 BPE Tokenizer 之前，我们需要创建一个完整的 Tokenizer 配置文件，包括 tokenizer_config.json 和 special_tokens_map.json。这些配置文件定义了 Tokenizer 的参数和特殊标记，用于训练和加载 Tokenizer。此处的chat_template我们与Qwen2.5模型保持一致。
def create_tokenizer_config(save_dir: str) -> None:
"""创建完整的tokenizer配置文件"""
config = {
"add_bos_token": False,
"add_eos_token": False,
"add_prefix_space": False,
"bos_token": "<|im_start|>",
"eos_token": "<|im_end|>",
"pad_token": "<|im_end|>",
"unk_token": "<unk>",
"model_max_length": 1000000000000000019884624838656,
"clean_up_tokenization_spaces": False,
"tokenizer_class": "PreTrainedTokenizerFast",
"chat_template": (
"{% for message in messages %}"
"{% if message['role'] == 'system' %}"
"<|im_start|>system\n{{ message['content'] }}<|im_end|>\n"
"{% elif message['role'] == 'user' %}"
"<|im_start|>user\n{{ message['content'] }}<|im_end|>\n"
"{% elif message['role'] == 'assistant' %}"
"<|im_start|>assistant\n{{ message['content'] }}<|im_end|>\n"
"{% endif %}"
"{% endfor %}"
"{% if add_generation_prompt %}"
"{{ '<|im_start|>assistant\n' }}"
"{% endif %}"
)
}
# 保存主配置文件
with open(os.path.join(save_dir, "tokenizer_config.json"), "w", encoding="utf-8") as f:
json.dump(config, f, ensure_ascii=False, indent=4)
# 创建special_tokens_map.json
special_tokens_map = {
"bos_token": "<|im_start|>",
"eos_token": "<|im_end|>",
"unk_token": "<unk>",
"pad_token": "<|im_end|>",
"additional_special_tokens": ["<s>", "</s>"]
}
with open(os.path.join(save_dir, "special_tokens_map.json"), "w", encoding="utf-8") as f:
json.dump(special_tokens_map, f, ensure_ascii=False, indent=4)
Step 4: 训练 BPE Tokenizer
在训练 BPE Tokenizer 之前，我们需要定义一个训练函数，用于训练 Tokenizer 并保存训练好的 Tokenizer 文件。这里我们使用 tokenizers 库中的 Tokenizer 类来训练 BPE Tokenizer。
可以看到我们在训练 Tokenizer 时，配置了一些特殊的 token，如 <unk>、<s>、</s>、<|im_start|> 和 <|im_end|>。这些 token 用于标记未知词、句子的开始和结束，以及对话的开始和结束。这些特殊 token 可以帮助模型更好地理解文本数据，提高模型的泛化能力和效果。
def train_tokenizer(data_path: str, save_dir: str, vocab_size: int = 8192) -> None:
"""训练并保存自定义tokenizer"""
os.makedirs(save_dir, exist_ok=True)
# 初始化tokenizer
tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
tokenizer.normalizer = NFKC() # 添加文本规范化
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
tokenizer.decoder = decoders.ByteLevel()
# 配置特殊token
special_tokens = [
"<unk>",
"<s>",
"</s>",
"<|im_start|>",
"<|im_end|>"
]
# 配置训练器
trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
special_tokens=special_tokens,
min_frequency=2, # 提高低频词过滤
show_progress=True,
initial_alphabet=pre_tokenizers.ByteLevel.alphabet()
)
# 训练tokenizer
print(f"Training tokenizer with data from {data_path}")
texts = read_texts_from_jsonl(data_path)
tokenizer.train_from_iterator(texts, trainer=trainer, length=os.path.getsize(data_path))
# 验证特殊token映射
try:
assert tokenizer.token_to_id("<unk>") == 0
assert tokenizer.token_to_id("<s>") == 1
assert tokenizer.token_to_id("</s>") == 2
assert tokenizer.token_to_id("<|im_start|>") == 3
assert tokenizer.token_to_id("<|im_end|>") == 4
except AssertionError as e:
print("Special tokens mapping error:", e)
raise
# 保存tokenizer文件
tokenizer.save(os.path.join(save_dir, "tokenizer.json"))
# 创建配置文件
create_tokenizer_config(save_dir)
print(f"Tokenizer saved to {save_dir}")
Step 5: 使用训练好的 Tokenizer
我们可以使用训练好的 Tokenizer 来处理文本数据，如编码、解码、生成对话等。下面是一个简单的示例，展示了如何使用训练好的 Tokenizer 来处理文本数据。
def eval_tokenizer(tokenizer_path: str) -> None:
"""评估tokenizer功能"""
try:
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
except Exception as e:
print(f"Error loading tokenizer: {e}")
return
# 测试基本属性
print("\n=== Tokenizer基本信息 ===")
print(f"Vocab size: {len(tokenizer)}")
print(f"Special tokens: {tokenizer.all_special_tokens}")
print(f"Special token IDs: {tokenizer.all_special_ids}")
# 测试聊天模板
messages = [
{"role": "system", "content": "你是一个AI助手。"},
{"role": "user", "content": "How are you?"},
{"role": "assistant", "content": "I'm fine, thank you. and you?"},
{"role": "user", "content": "I'm good too."},
{"role": "assistant", "content": "That's great to hear!"},
]
print("\n=== 聊天模板测试 ===")
prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
# add_generation_prompt=True
)
print("Generated prompt:\n", prompt, sep="")
# 测试编码解码
print("\n=== 编码解码测试 ===")
encoded = tokenizer(prompt, truncation=True, max_length=256)
decoded = tokenizer.decode(encoded["input_ids"], skip_special_tokens=False)
print("Decoded text matches original:", decoded == prompt)
# 测试特殊token处理
print("\n=== 特殊token处理 ===")
test_text = "<|im_start|>user\nHello<|im_end|>"
encoded = tokenizer(test_text).input_ids
decoded = tokenizer.decode(encoded)
print(f"Original: {test_text}")
print(f"Decoded: {decoded}")
print("Special tokens preserved:", decoded == test_text)
eval_tokenizer('your tokenizer path')
OUT:
=== Tokenizer基本信息 ===
Vocab size: 6144
Special tokens: ['<|im_start|>', '<|im_end|>', '<unk>', '<s>', '</s>']
Special token IDs: [3, 4, 0, 1, 2]
=== 聊天模板测试 ===
Generated prompt:
<|im_start|>system
你是一个AI助手。<|im_end|>
<|im_start|>user
How are you?<|im_end|>
<|im_start|>assistant
I'm fine, thank you. and you?<|im_end|>
<|im_start|>user
I'm good too.<|im_end|>
<|im_start|>assistant
That's great to hear!<|im_end|>
=== 编码解码测试 ===
Decoded text matches original: False
=== 特殊token处理 ===
Original: <|im_start|>user
Hello<|im_end|>
Decoded: <|im_start|> user
Hello<|im_end|>
Special tokens preserved: False
浙公网安备 33010602011771号