github
/
minimind
mirror of https://github.com/jingyaogong/minimind.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

docs(model): 为MiniMind模型添加详细代码注释 

为MiniMind模型的各个模块和函数添加详细的中文注释，包括类、方法、参数和核心逻辑的解释，以提高代码的可读性和可维护性。
pull/394/head
0x0059 2025-05-05 21:59:08 +08:00
parent 561979c7e3
commit f74117e5d9
1 changed files with 154 additions and 54 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

208
model/model_minimind.py
View File

@ -2,28 +2,33 @@
# MiniMind Config # MiniMind Config
# 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘 # 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘
# MiniMind模型配置类，继承自transformers的PretrainedConfig
# 包含了模型的所有超参数设置，包括基础参数和MoE（混合专家）相关参数
from transformers import PretrainedConfig from transformers import PretrainedConfig
class MiniMindConfig(PretrainedConfig): class MiniMindConfig(PretrainedConfig):
# 模型类型标识符
model_type = "minimind" model_type = "minimind"
def __init__( def __init__(
self, self,
dropout: float = 0.0, # 基础模型参数
bos_token_id: int = 1, dropout: float = 0.0, # dropout比率，用于防止过拟合
eos_token_id: int = 2, bos_token_id: int = 1, # 序列开始标记的ID
hidden_act: str = 'silu', eos_token_id: int = 2, # 序列结束标记的ID
hidden_size: int = 512, hidden_act: str = 'silu', # 激活函数类型
intermediate_size: int = None, hidden_size: int = 512, # 隐藏层维度
max_position_embeddings: int = 32768, intermediate_size: int = None, # 前馈网络中间层维度
num_attention_heads: int = 8, max_position_embeddings: int = 32768, # 最大位置编码长度
num_hidden_layers: int = 8, num_attention_heads: int = 8, # 注意力头数量
num_key_value_heads: int = 2, num_hidden_layers: int = 8, # Transformer层数
vocab_size: int = 6400, num_key_value_heads: int = 2, # KV注意力头数量（用于分组查询注意力）
rms_norm_eps: float = 1e-05, vocab_size: int = 6400, # 词表大小
rope_theta: int = 1000000.0, rms_norm_eps: float = 1e-05, # RMSNorm层的epsilon值
flash_attn: bool = True, rope_theta: int = 1000000.0, # 旋转位置编码的theta参数
flash_attn: bool = True, # 是否使用Flash Attention优化
#################################################### ####################################################
# Here are the specific configurations of MOE # Here are the specific configurations of MOE
# When use_moe is false, the following is invalid # When use_moe is false, the following is invalid
@ -82,41 +87,74 @@ from transformers.modeling_outputs import CausalLMOutputWithPast
class RMSNorm(torch.nn.Module): class RMSNorm(torch.nn.Module):
"""RMSNorm层实现
相比LayerNorm，RMSNorm去掉了均值归一化，只保留方差归一化，计算更简单且效果相当
"""
def __init__(self, dim: int, eps: float = 1e-5): def __init__(self, dim: int, eps: float = 1e-5):
super().__init__() super().__init__()
self.eps = eps self.eps = eps # 数值稳定性参数
self.weight = nn.Parameter(torch.ones(dim)) self.weight = nn.Parameter(torch.ones(dim)) # 可学习的缩放参数
def _norm(self, x): def _norm(self, x):
# 计算RMS（均方根）归一化
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x): def forward(self, x):
# 应用缩放并保持输入数据类型
return self.weight * self._norm(x.float()).type_as(x) return self.weight * self._norm(x.float()).type_as(x)
def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6): def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), theta: float = 1e6):
"""预计算旋转位置编码（RoPE）的三角函数值
Args:
dim: 隐藏层维度
end: 最大序列长度
theta: 位置编码的缩放因子
Returns:
freqs_cos: 余弦位置编码
freqs_sin: 正弦位置编码
"""
# 计算频率基础值
freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
t = torch.arange(end, device=freqs.device) t = torch.arange(end, device=freqs.device)
# 计算外积得到位置编码矩阵
freqs = torch.outer(t, freqs).float() freqs = torch.outer(t, freqs).float()
# 生成最终的三角函数值
freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1)
freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1)
return freqs_cos, freqs_sin return freqs_cos, freqs_sin
def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
"""应用旋转位置编码（RoPE）到查询和键向量
通过将位置信息编码到注意力的查询(Q)和键(K)向量中，使模型能够感知token的相对位置关系
"""
def rotate_half(x): def rotate_half(x):
# 将向量的后半部分取负并与前半部分交换位置
return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1) return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
# 应用旋转变换到查询和键向量
q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim)) q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim)) k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
return q_embed, k_embed return q_embed, k_embed
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
"""torch.repeat_interleave(x, dim=2, repeats=n_rep)""" """重复键值向量以匹配注意力头的数量
在分组查询注意力中，键值向量的头数少于查询向量，需要将键值向量重复多次
Args:
x: 输入张量，形状为[batch_size, seq_len, num_kv_heads, head_dim]
n_rep: 重复次数
"""
bs, slen, num_key_value_heads, head_dim = x.shape bs, slen, num_key_value_heads, head_dim = x.shape
if n_rep == 1: if n_rep == 1:
return x return x
# 扩展并重塑张量以实现重复
return ( return (
x[:, :, :, None, :] x[:, :, :, None, :]
.expand(bs, slen, num_key_value_heads, n_rep, head_dim) .expand(bs, slen, num_key_value_heads, n_rep, head_dim)
@ -125,49 +163,72 @@ def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
class Attention(nn.Module): class Attention(nn.Module):
"""多头注意力机制实现
支持分组查询注意力（GQA）和Flash Attention优化
分组查询注意力通过减少键值头的数量来降低计算和内存开销
"""
def __init__(self, args: MiniMindConfig): def __init__(self, args: MiniMindConfig):
super().__init__() super().__init__()
# 设置键值头数量，支持分组查询注意力
self.num_key_value_heads = args.num_attention_heads if args.num_key_value_heads is None else args.num_key_value_heads self.num_key_value_heads = args.num_attention_heads if args.num_key_value_heads is None else args.num_key_value_heads
assert args.num_attention_heads % self.num_key_value_heads == 0 assert args.num_attention_heads % self.num_key_value_heads == 0
self.n_local_heads = args.num_attention_heads self.n_local_heads = args.num_attention_heads # 查询头总数
self.n_local_kv_heads = self.num_key_value_heads self.n_local_kv_heads = self.num_key_value_heads # 键值头数量
self.n_rep = self.n_local_heads // self.n_local_kv_heads self.n_rep = self.n_local_heads // self.n_local_kv_heads # 每个键值头需要重复的次数
self.head_dim = args.hidden_size // args.num_attention_heads self.head_dim = args.hidden_size // args.num_attention_heads # 每个注意力头的维度
self.q_proj = nn.Linear(args.hidden_size, args.num_attention_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) # 定义线性变换层
self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) self.q_proj = nn.Linear(args.hidden_size, args.num_attention_heads * self.head_dim, bias=False) # 查询变换
self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False) self.k_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) # 键变换
self.attn_dropout = nn.Dropout(args.dropout) self.v_proj = nn.Linear(args.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) # 值变换
self.resid_dropout = nn.Dropout(args.dropout) self.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False) # 输出变换
# Dropout层
self.attn_dropout = nn.Dropout(args.dropout) # 注意力权重的dropout
self.resid_dropout = nn.Dropout(args.dropout) # 残差连接的dropout
self.dropout = args.dropout self.dropout = args.dropout
# 是否使用Flash Attention优化
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and args.flash_attn
# print("WARNING: using slow attention. Flash Attention requires PyTorch >= 2.0")
def forward(self, def forward(self,
x: torch.Tensor, x: torch.Tensor,
position_embeddings: Tuple[torch.Tensor, torch.Tensor], # 修改为接收cos和sin position_embeddings: Tuple[torch.Tensor, torch.Tensor], # 位置编码的cos和sin值
past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # KV缓存
use_cache=False, use_cache=False, # 是否使用KV缓存
attention_mask: Optional[torch.Tensor] = None): attention_mask: Optional[torch.Tensor] = None): # 注意力掩码
"""前向传播函数
实现多头注意力的核心计算逻辑，包括：
1. 线性变换生成查询、键、值向量
2. 应用旋转位置编码
3. 处理KV缓存（用于加速自回归生成）
4. 执行注意力计算（支持Flash Attention优化）
"""
bsz, seq_len, _ = x.shape bsz, seq_len, _ = x.shape
# 生成查询、键、值向量
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x) xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
# 重塑张量维度以适应多头注意力格式
xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim) xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
# 应用旋转位置编码
cos, sin = position_embeddings cos, sin = position_embeddings
xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len]) xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len])
# kv_cache实现 # 处理KV缓存，用于加速自回归生成
if past_key_value is not None: if past_key_value is not None:
xk = torch.cat([past_key_value[0], xk], dim=1) xk = torch.cat([past_key_value[0], xk], dim=1) # 拼接历史键向量
xv = torch.cat([past_key_value[1], xv], dim=1) xv = torch.cat([past_key_value[1], xv], dim=1) # 拼接历史值向量
past_kv = (xk, xv) if use_cache else None past_kv = (xk, xv) if use_cache else None
# 调整维度顺序并重复键值向量以匹配查询头数量
xq, xk, xv = ( xq, xk, xv = (
xq.transpose(1, 2), xq.transpose(1, 2), # [batch, num_heads, seq_len, head_dim]
repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), # 重复键向量
repeat_kv(xv, self.n_rep).transpose(1, 2) repeat_kv(xv, self.n_rep).transpose(1, 2) # 重复值向量
) )
if self.flash and seq_len != 1: if self.flash and seq_len != 1:
@ -200,36 +261,53 @@ class Attention(nn.Module):
class FeedForward(nn.Module): class FeedForward(nn.Module):
"""前馈神经网络实现
使用SwiGLU激活函数的变体，包含门控机制，可以更好地控制信息流动
结构：输入 -> (门控路径 + 上投影路径) -> 下投影 -> 输出
"""
def __init__(self, config: MiniMindConfig): def __init__(self, config: MiniMindConfig):
super().__init__() super().__init__()
# 如果未指定中间层维度，则根据隐藏层维度计算
if config.intermediate_size is None: if config.intermediate_size is None:
intermediate_size = int(config.hidden_size * 8 / 3) intermediate_size = int(config.hidden_size * 8 / 3) # 设置为隐藏层维度的8/3倍
config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64) config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 64) # 调整为64的倍数
self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) # 定义三个线性变换层
self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) # 门控路径
self.dropout = nn.Dropout(config.dropout) self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) # 降维投影
self.act_fn = ACT2FN[config.hidden_act] self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) # 升维投影
self.dropout = nn.Dropout(config.dropout) # dropout层
self.act_fn = ACT2FN[config.hidden_act] # 激活函数（默认为SiLU）
def forward(self, x): def forward(self, x):
# 实现SwiGLU的变体：激活(门控路径) * 上投影 -> 下投影 -> dropout
return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))) return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)))
class MoEGate(nn.Module): class MoEGate(nn.Module):
"""混合专家（MoE）的门控机制实现
负责为每个token选择最合适的专家组合，包括：
1. 计算每个token与专家的相关性分数
2. 选择top-k个专家处理每个token
3. 计算辅助损失以平衡专家的使用
"""
def __init__(self, config: MiniMindConfig): def __init__(self, config: MiniMindConfig):
super().__init__() super().__init__()
self.config = config self.config = config
self.top_k = config.num_experts_per_tok self.top_k = config.num_experts_per_tok # 每个token选择的专家数量
self.n_routed_experts = config.n_routed_experts self.n_routed_experts = config.n_routed_experts # 可路由的专家总数
self.scoring_func = config.scoring_func self.scoring_func = config.scoring_func # 专家选择的评分函数
self.alpha = config.aux_loss_alpha self.alpha = config.aux_loss_alpha # 辅助损失的权重
self.seq_aux = config.seq_aux self.seq_aux = config.seq_aux # 是否在序列级别计算辅助损失
self.norm_topk_prob = config.norm_topk_prob self.norm_topk_prob = config.norm_topk_prob # 是否对top-k概率进行归一化
self.gating_dim = config.hidden_size self.gating_dim = config.hidden_size # 用于门控的特征维度
self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim))) self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim))) # 门控权重矩阵
self.reset_parameters() self.reset_parameters() # 初始化参数
def reset_parameters(self) -> None: def reset_parameters(self) -> None:
import torch.nn.init as init import torch.nn.init as init
@ -273,14 +351,22 @@ class MoEGate(nn.Module):
class MOEFeedForward(nn.Module): class MOEFeedForward(nn.Module):
"""混合专家（MoE）前馈网络实现
包含多个专家网络和一个门控机制，每个token的处理会被动态路由到最合适的专家组合
可选地包含共享专家，它们会处理所有token
"""
def __init__(self, config: MiniMindConfig): def __init__(self, config: MiniMindConfig):
super().__init__() super().__init__()
self.config = config self.config = config
# 初始化可路由的专家网络
self.experts = nn.ModuleList([ self.experts = nn.ModuleList([
FeedForward(config) FeedForward(config)
for _ in range(config.n_routed_experts) for _ in range(config.n_routed_experts)
]) ])
# 初始化门控机制
self.gate = MoEGate(config) self.gate = MoEGate(config)
# 可选：初始化共享专家网络
if config.n_shared_experts > 0: if config.n_shared_experts > 0:
self.shared_experts = nn.ModuleList([ self.shared_experts = nn.ModuleList([
FeedForward(config) FeedForward(config)
@ -359,17 +445,31 @@ class MiniMindBlock(nn.Module):
class MiniMindModel(nn.Module): class MiniMindModel(nn.Module):
"""MiniMind模型的主体实现
一个基于Transformer的语言模型，支持以下特性：
1. 分组查询注意力（GQA）降低计算开销
2. 旋转位置编码（RoPE）处理位置信息
3. 混合专家机制（MoE）提升模型容量
4. RMSNorm归一化提升训练稳定性
"""
def __init__(self, config: MiniMindConfig): def __init__(self, config: MiniMindConfig):
super().__init__() super().__init__()
self.config = config self.config = config
# 基础配置
self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
# 词嵌入层
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
self.dropout = nn.Dropout(config.dropout) self.dropout = nn.Dropout(config.dropout)
# 初始化Transformer层
self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)]) self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
# 最终的归一化层
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
# 预计算旋转位置编码的三角函数值
freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.hidden_size // config.num_attention_heads, freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.hidden_size // config.num_attention_heads,
end=config.max_position_embeddings, theta=config.rope_theta) end=config.max_position_embeddings, theta=config.rope_theta)
# 注册为非持久化缓冲区（不会被保存到模型文件中）
self.register_buffer("freqs_cos", freqs_cos, persistent=False) self.register_buffer("freqs_cos", freqs_cos, persistent=False)
self.register_buffer("freqs_sin", freqs_sin, persistent=False) self.register_buffer("freqs_sin", freqs_sin, persistent=False)