From f74117e5d9631f69c6a3efb6055febe7d6b6a638 Mon Sep 17 00:00:00 2001 From: 0x0059 Date: Mon, 5 May 2025 21:59:08 +0800 Subject: [PATCH] =?UTF-8?q?docs(model):=20=E4=B8=BAMiniMind=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B=E6=B7=BB=E5=8A=A0=E8=AF=A6=E7=BB=86=E4=BB=A3=E7=A0=81?= =?UTF-8?q?=E6=B3=A8=E9=87=8A?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 为MiniMind模型的各个模块和函数添加详细的中文注释，包括类、方法、参数和核心逻辑的解释，以提高代码的可读性和可维护性。 --- model/model_minimind.py | 208 +++++++++++++++++++++++++++++----------- 1 file changed, 154 insertions(+), 54 deletions(-) diff --git a/model/model_minimind.py b/model/model_minimind.py index 4b43610..dc2cceb 100644 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -2,28 +2,33 @@ # MiniMind Config # 📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘📘 +# MiniMind模型配置类，继承自transformers的PretrainedConfig +# 包含了模型的所有超参数设置，包括基础参数和MoE（混合专家）相关参数 + from transformers import PretrainedConfig class MiniMindConfig(PretrainedConfig): + # 模型类型标识符 model_type = "minimind" def __init__( self, - dropout: float = 0.0, - bos_token_id: int = 1, - eos_token_id: int = 2, - hidden_act: str = 'silu', - hidden_size: int = 512, - intermediate_size: int = None, - max_position_embeddings: int = 32768, - num_attention_heads: int = 8, - num_hidden_layers: int = 8, - num_key_value_heads: int = 2, - vocab_size: int = 6400, - rms_norm_eps: float = 1e-05, - rope_theta: int = 1000000.0, - flash_attn: bool = True, + # 基础模型参数 + dropout: float = 0.0, # dropout比率，用于防止过拟合 + bos_token_id: int = 1, # 序列开始标记的ID + eos_token_id: int = 2, # 序列结束标记的ID + hidden_act: str = 'silu', # 激活函数类型 + hidden_size: int = 512, # 隐藏层维度 + intermediate_size: int = None, # 前馈网络中间层维度 + max_position_embeddings: int = 32768, # 最大位置编码长度 + num_attention_heads: int = 8, # 注意力头数量 + num_hidden_layers: int = 8, # Transformer层数 + num_key_value_heads: int = 2, # KV注意力头数量（用于分组查询注意力） + vocab_size: int = 6400, # 词表大小 + rms_norm_eps: float = 1e-05, # RMSNorm层的epsilon值 + rope_theta: int = 1000000.0, # 旋转位置编码的theta参数 + flash_attn: bool = True, # 是否使用Flash Attention优化 #################################################### # Here are the specific configurations of MOE # When use_moe is false, the following is invalid @@ -82,41 +87,74 @@ from transformers.modeling_outputs import CausalLMOutputWithPast class RMSNorm(torch.nn.Module): + """RMSNorm层实现 + 相比LayerNorm，RMSNorm去掉了均值归一化，只保留方差归一化，计算更简单且效果相当 + """ def __init__(self, dim: int, eps: float = 1e-5): super().__init__() - self.eps = eps - self.weight = nn.Parameter(torch.ones(dim)) + self.eps = eps # 数值稳定性参数 + self.weight = nn.Parameter(torch.ones(dim)) # 可学习的缩放参数 def _norm(self, x): + # 计算RMS（均方根）归一化 return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) def forward(self, 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): + """预计算旋转位置编码（RoPE）的三角函数值 + + Args: + dim: 隐藏层维度 + end: 最大序列长度 + theta: 位置编码的缩放因子 + + Returns: + freqs_cos: 余弦位置编码 + freqs_sin: 正弦位置编码 + """ + # 计算频率基础值 freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) t = torch.arange(end, device=freqs.device) + # 计算外积得到位置编码矩阵 freqs = torch.outer(t, freqs).float() + # 生成最终的三角函数值 freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) return freqs_cos, freqs_sin def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """应用旋转位置编码（RoPE）到查询和键向量 + + 通过将位置信息编码到注意力的查询(Q)和键(K)向量中，使模型能够感知token的相对位置关系 + """ def rotate_half(x): + # 将向量的后半部分取负并与前半部分交换位置 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)) k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim)) return q_embed, k_embed 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 if n_rep == 1: return x + # 扩展并重塑张量以实现重复 return ( x[:, :, :, None, :] .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): + """多头注意力机制实现 + + 支持分组查询注意力（GQA）和Flash Attention优化 + 分组查询注意力通过减少键值头的数量来降低计算和内存开销 + """ def __init__(self, args: MiniMindConfig): 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 assert args.num_attention_heads % self.num_key_value_heads == 0 - self.n_local_heads = args.num_attention_heads - self.n_local_kv_heads = self.num_key_value_heads - self.n_rep = self.n_local_heads // self.n_local_kv_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.o_proj = nn.Linear(args.num_attention_heads * self.head_dim, args.hidden_size, bias=False) - self.attn_dropout = nn.Dropout(args.dropout) - self.resid_dropout = nn.Dropout(args.dropout) + self.n_local_heads = args.num_attention_heads # 查询头总数 + self.n_local_kv_heads = self.num_key_value_heads # 键值头数量 + self.n_rep = self.n_local_heads // self.n_local_kv_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.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 + + # 是否使用Flash Attention优化 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, x: torch.Tensor, - position_embeddings: Tuple[torch.Tensor, torch.Tensor], # 修改为接收cos和sin - past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, - use_cache=False, - attention_mask: Optional[torch.Tensor] = None): + position_embeddings: Tuple[torch.Tensor, torch.Tensor], # 位置编码的cos和sin值 + past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # KV缓存 + use_cache=False, # 是否使用KV缓存 + attention_mask: Optional[torch.Tensor] = None): # 注意力掩码 + """前向传播函数 + + 实现多头注意力的核心计算逻辑，包括： + 1. 线性变换生成查询、键、值向量 + 2. 应用旋转位置编码 + 3. 处理KV缓存（用于加速自回归生成） + 4. 执行注意力计算（支持Flash Attention优化） + """ bsz, seq_len, _ = x.shape + # 生成查询、键、值向量 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) 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) + # 应用旋转位置编码 cos, sin = position_embeddings xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len]) - # kv_cache实现 + # 处理KV缓存，用于加速自回归生成 if past_key_value is not None: - xk = torch.cat([past_key_value[0], xk], dim=1) - xv = torch.cat([past_key_value[1], xv], dim=1) + xk = torch.cat([past_key_value[0], xk], dim=1) # 拼接历史键向量 + xv = torch.cat([past_key_value[1], xv], dim=1) # 拼接历史值向量 past_kv = (xk, xv) if use_cache else None + # 调整维度顺序并重复键值向量以匹配查询头数量 xq, xk, xv = ( - xq.transpose(1, 2), - repeat_kv(xk, self.n_rep).transpose(1, 2), - repeat_kv(xv, self.n_rep).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(xv, self.n_rep).transpose(1, 2) # 重复值向量 ) if self.flash and seq_len != 1: @@ -200,36 +261,53 @@ class Attention(nn.Module): class FeedForward(nn.Module): + """前馈神经网络实现 + + 使用SwiGLU激活函数的变体，包含门控机制，可以更好地控制信息流动 + 结构：输入 -> (门控路径 + 上投影路径) -> 下投影 -> 输出 + """ def __init__(self, config: MiniMindConfig): super().__init__() + # 如果未指定中间层维度，则根据隐藏层维度计算 if config.intermediate_size is None: - intermediate_size = int(config.hidden_size * 8 / 3) - config.intermediate_size = 64 * ((intermediate_size + 64 - 1) // 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.dropout = nn.Dropout(config.dropout) - self.act_fn = ACT2FN[config.hidden_act] + intermediate_size = int(config.hidden_size * 8 / 3) # 设置为隐藏层维度的8/3倍 + 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.dropout = nn.Dropout(config.dropout) # dropout层 + self.act_fn = ACT2FN[config.hidden_act] # 激活函数（默认为SiLU） def forward(self, x): + # 实现SwiGLU的变体：激活(门控路径) * 上投影 -> 下投影 -> dropout return self.dropout(self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))) class MoEGate(nn.Module): + """混合专家（MoE）的门控机制实现 + + 负责为每个token选择最合适的专家组合，包括： + 1. 计算每个token与专家的相关性分数 + 2. 选择top-k个专家处理每个token + 3. 计算辅助损失以平衡专家的使用 + """ def __init__(self, config: MiniMindConfig): super().__init__() self.config = config - self.top_k = config.num_experts_per_tok - self.n_routed_experts = config.n_routed_experts + self.top_k = config.num_experts_per_tok # 每个token选择的专家数量 + self.n_routed_experts = config.n_routed_experts # 可路由的专家总数 - self.scoring_func = config.scoring_func - self.alpha = config.aux_loss_alpha - self.seq_aux = config.seq_aux + self.scoring_func = config.scoring_func # 专家选择的评分函数 + self.alpha = config.aux_loss_alpha # 辅助损失的权重 + self.seq_aux = config.seq_aux # 是否在序列级别计算辅助损失 - self.norm_topk_prob = config.norm_topk_prob - self.gating_dim = config.hidden_size - self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim))) - self.reset_parameters() + self.norm_topk_prob = config.norm_topk_prob # 是否对top-k概率进行归一化 + self.gating_dim = config.hidden_size # 用于门控的特征维度 + self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim))) # 门控权重矩阵 + self.reset_parameters() # 初始化参数 def reset_parameters(self) -> None: import torch.nn.init as init @@ -273,14 +351,22 @@ class MoEGate(nn.Module): class MOEFeedForward(nn.Module): + """混合专家（MoE）前馈网络实现 + + 包含多个专家网络和一个门控机制，每个token的处理会被动态路由到最合适的专家组合 + 可选地包含共享专家，它们会处理所有token + """ def __init__(self, config: MiniMindConfig): super().__init__() self.config = config + # 初始化可路由的专家网络 self.experts = nn.ModuleList([ FeedForward(config) for _ in range(config.n_routed_experts) ]) + # 初始化门控机制 self.gate = MoEGate(config) + # 可选：初始化共享专家网络 if config.n_shared_experts > 0: self.shared_experts = nn.ModuleList([ FeedForward(config) @@ -359,17 +445,31 @@ class MiniMindBlock(nn.Module): class MiniMindModel(nn.Module): + """MiniMind模型的主体实现 + + 一个基于Transformer的语言模型，支持以下特性： + 1. 分组查询注意力（GQA）降低计算开销 + 2. 旋转位置编码（RoPE）处理位置信息 + 3. 混合专家机制（MoE）提升模型容量 + 4. RMSNorm归一化提升训练稳定性 + """ def __init__(self, config: MiniMindConfig): super().__init__() self.config = config + # 基础配置 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.dropout = nn.Dropout(config.dropout) + # 初始化Transformer层 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) + # 预计算旋转位置编码的三角函数值 freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.hidden_size // config.num_attention_heads, end=config.max_position_embeddings, theta=config.rope_theta) + # 注册为非持久化缓冲区（不会被保存到模型文件中） self.register_buffer("freqs_cos", freqs_cos, persistent=False) self.register_buffer("freqs_sin", freqs_sin, persistent=False)