零基础解读MiniMind代码：基座模型全解析
·
MiniMind-O 逐行代码解读 · 第一章：基座模型 model_minimind.py
🎯 目标：让零基础读者也能看懂每一行代码在做什么 📁 文件：
model/model_minimind.py（287行） 🔖 这是整个项目的根基，所有多模态扩展都建立在这个基座模型之上
📚 本章导读
这个文件实现了一个完整的小型语言模型（LLM），叫 MiniMind。它包含了：
- 配置类 — 定义模型的所有超参数
- RMSNorm — 一种比 LayerNorm 更快的归一化方法
- RoPE 位置编码 — 让模型知道"每个词在第几个位置"
- Attention — 注意力机制，模型理解上下文的核心
- FeedForward — 前馈网络，对信息做非线性变换
- MoE — 混合专家系统（可选），让不同 token 走不同"专家"
- TransformerBlock — 把注意力+前馈组合成一个完整层
- MiniMindModel — 把多层 Block 堆叠成完整模型
- MiniMindForCausalLM — 加上语言模型头，能做文本生成
1️⃣ 导入库（第1-5行）
import math, torch, torch.nn.functional as F
from torch import nn
from transformers.activations import ACT2FN
from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
from transformers.modeling_outputs import MoeCausalLMOutputWithPast
| 导入 | 干什么用的 |
|---|---|
math |
数学函数，比如算 π、开方 |
torch |
PyTorch 核心，张量运算 |
torch.nn.functional as F |
函数式API，如 softmax、cross_entropy |
torch.nn as nn |
神经网络模块基类，所有层都继承它 |
ACT2FN |
HuggingFace 的激活函数字典，如 ACT2FN['silu'] 得到 SiLU 函数 |
PreTrainedModel |
HuggingFace 预训练模型基类，提供 save/load 接口 |
GenerationMixin |
HuggingFace 生成混入类，提供 .generate() 方法 |
PretrainedConfig |
HuggingFace 配置基类，提供 save/load 配置接口 |
MoeCausalLMOutputWithPast |
MoE 因果语言模型的输出数据类，包含 loss/logits/past_key_values/aux_loss |
2️⃣ 配置类 MiniMindConfig（第10-45行）
class MiniMindConfig(PretrainedConfig):
model_type = "minimind"
什么是配置类？
想象你在拼乐高——配置类就是"说明书"，告诉程序：
- 模型有多宽（hidden_size）
- 有几层（num_hidden_layers）
- 有多少个注意力头（num_attention_heads）
- ……等等
逐行解读
class MiniMindConfig(PretrainedConfig): # 继承 HuggingFace 的配置基类
model_type = "minimind" # 模型类型标识符，HuggingFace 用它来识别模型
def __init__(self, hidden_size=768, num_hidden_layers=8, use_moe=False, **kwargs):
参数解释：
hidden_size=768：模型的"宽度"。每个 token 会被表示为一个 768 维的向量。数字越大，模型越"聪明"但也越慢。768 是 BERT-base 的标准尺寸。num_hidden_layers=8：Transformer 的层数。层数越多，模型越"深"，能学到越复杂的模式。GPT-3 有96层，这里只有8层，所以是"迷你"模型。use_moe=False：是否使用混合专家（Mixture of Experts）。开启后，每个 token 会被"路由"到不同的"专家"处理，增加模型容量但不等比例增加计算量。**kwargs：其他参数，传给父类PretrainedConfig。
super().__init__(**kwargs) # 调用父类初始化，处理通用配置
接下来是所有配置参数的定义：
self.hidden_size = hidden_size # 隐藏层维度 = 768
self.num_hidden_layers = num_hidden_layers # 层数 = 8
self.use_moe = use_moe # 是否使用 MoE = False
self.dropout = kwargs.get("dropout", 0.0)
- dropout：随机丢弃一些神经元输出的概率。0.0 = 不丢弃。训练时防止过拟合。
self.vocab_size = kwargs.get("vocab_size", 6400)
- vocab_size：词表大小。6400 意味着模型只能识别 6400 个不同的 token。GPT-4 有 100000+，6400 非常小，但够用于中文基础对话。
self.bos_token_id = kwargs.get("bos_token_id", 1) # 句子开头标记的 ID
self.eos_token_id = kwargs.get("eos_token_id", 2) # 句子结束标记的 ID
- bos = Begin Of Sentence，eos = End Of Sentence。模型用这些特殊标记知道一句话从哪开始、到哪结束。
self.flash_attn = kwargs.get("flash_attn", True)
- flash_attn：是否使用 Flash Attention。Flash Attention 是一种优化算法，能大幅加速注意力计算并节省显存。默认开启。
self.num_attention_heads = kwargs.get("num_attention_heads", 8)
- 注意力头数：注意力机制被分成8个"头"，每个头关注不同的信息模式。就像8个人同时看一篇文章，每个人关注不同方面。
self.num_key_value_heads = kwargs.get("num_key_value_heads", 4)
- KV 头数：Key 和 Value 只有4个头，而不是8个。这就是 GQA（Grouped Query Attention）——8个查询头共享4组 KV 头（每2个查询头共享1组）。好处：减少 KV cache 的显存占用。
self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)
- 每个头的维度：768 / 8 = 96。每个注意力头的向量长度。
self.hidden_act = kwargs.get("hidden_act", 'silu')
- 激活函数：SiLU（Sigmoid Linear Unit），公式：
x * sigmoid(x)。比 ReLU 更平滑，训练更稳定。
self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * math.pi / 64) * 64)
- FFN 中间层维度：
⌈768 × π / 64⌉ × 64 = ⌈37.7⌉ × 64 = 38 × 64 = 2432...768 × 3.14159 / 64 = 37.70→ceil(37.70) = 38→38 × 64 = 2432- 实际代码默认不是这个，是因为
**kwargs里没有intermediate_size时用默认值 - 这个公式来自 Llama 的做法：π 倍缩放，再对齐到 64 的倍数，提升 GPU 计算效率
self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768)
- 最大位置编码长度：模型最多能处理 32768 个 token 的序列。超出则需要长度外推。
self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)
- RMSNorm 的 epsilon：防止除零的小常数。1e-6 是常用值。
self.rope_theta = kwargs.get("rope_theta", 1e6)
- RoPE 的基础频率：1e6（一百万）。这个值越大，低频成分的波长越长，模型对远距离位置的区分能力越强。
self.tie_word_embeddings = kwargs.get("tie_word_embeddings", True)
- 权重绑定：输入嵌入层和输出层的权重共享。节省参数，且在小型模型上效果不错。
self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)
- 推理时的 RoPE 缩放：是否启用 YaRN 长度外推算法。
self.rope_scaling = {
"beta_fast": 32, # 高频分量的边界参数
"beta_slow": 1, # 低频分量的边界参数
"factor": 16, # 长度缩放因子（16倍外推）
"original_max_position_embeddings": 2048, # 原始训练长度
"attention_factor": 1.0, # 注意力分数缩放
"type": "yarn" # 使用 YaRN 算法
} if self.inference_rope_scaling else None
- YaRN（Yet another RoPE extensioN）算法配置：训练时只看了 2048 长度，但推理时通过缩放 RoPE 频率，可以处理 2048×16=32768 长度的序列。
self.num_experts = kwargs.get("num_experts", 4) # MoE 专家数量
self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1) # 每个 token 选几个专家
self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size) # 专家 FFN 中间层维度
self.norm_topk_prob = kwargs.get("norm_topk_prob", True) # 是否归一化 top-k 概率
self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 5e-4) # 路由辅助损失系数
- 这些都是 MoE 专用参数，
use_moe=False时不会生效。
3️⃣ RMSNorm 归一化（第50-60行）
什么是归一化？
想象考试分数：有人考了 100 分，有人考了 50 分。归一化就是把所有分数"拉"到同一尺度，让模型更容易学习。
RMSNorm 比 LayerNorm 更快，因为它不需要计算均值，只算均方根。
class RMSNorm(torch.nn.Module):
def __init__(self, dim: int, eps: float = 1e-5):
super().__init__()
self.eps = eps # 防除零的小常数
self.weight = nn.Parameter(torch.ones(dim)) # 可学习的缩放参数，初始值全1
def norm(self, x):
# x.pow(2) → 每个元素平方
# .mean(-1, keepdim=True) → 在最后一维求均值
# torch.rsqrt(...) → 取平方根的倒数（1/√x）
# + self.eps → 防止除零
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
# 1. 先在 float32 下计算归一化（数值更稳定）
# 2. 乘以可学习的 weight 参数
# 3. 转回输入的原始数据类型
return (self.weight * self.norm(x.float())).type_as(x)
为什么用 float32 计算？ 因为 bfloat16 精度低，做归一化时容易溢出。先转 float32 算完再转回去。
对比 LayerNorm：
- LayerNorm:
(x - mean) / sqrt(var + eps) * weight + bias - RMSNorm:
x / sqrt(mean(x²) + eps) * weight - RMSNorm 省掉了
mean计算和bias参数，速度快 ~10%
4️⃣ RoPE 位置编码（第62-84行）
为什么需要位置编码？
Transformer 的注意力机制是"无序"的——它不知道"我爱你"和"你爱我"的区别。位置编码就是给每个位置加上一个"标签"，让模型知道每个词在第几个位置。
RoPE（Rotary Position Embedding）的做法很巧妙：用旋转矩阵给 Q 和 K 加上位置信息。
def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: dict = None):
参数：
| 参数 | 说明 |
|---|---|
dim |
旋转的维度（等于 head_dim） |
end |
预计算的最大位置数（32768） |
rope_base |
RoPE 基础频率（1e6） |
rope_scaling |
YaRN 缩放配置（None = 不缩放） |
返回值：(freqs_cos, freqs_sin)，形状都是 (end, dim)，预计算的余弦和正弦表。
freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0
这行做了很多事，拆解：
torch.arange(0, dim, 2)→ 生成[0, 2, 4, ..., dim-2]，共dim//2个频率索引.float() / dim→ 归一化到 [0, 1)rope_base ** (...)→ 计算θ_i = 1000000^(2i/d)1.0 / (...)→ 取倒数，得到频率f_i = 1/θ_i
低维度对应低频（变化慢，适合编码远距离），高维度对应高频（变化快，适合编码近距离）。
if rope_scaling is not None: # YaRN 长度外推
如果启用了 YaRN：
- 计算高频和低频的边界索引
low和high - 生成一个从 0 到 1 的线性坡道
ramp - 对低频部分（ramp=1）做
1/factor缩放，高频部分（ramp=0）不缩放 - 中间部分线性插值过渡
t = torch.arange(end, device=freqs.device) # 位置索引 [0, 1, 2, ..., end-1]
freqs = torch.outer(t, freqs).float() # 外积：位置 × 频率 → (end, dim//2)
- 这就是
θ_i * pos，把位置信息编码到频率空间。
freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor # (end, dim)
freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor # (end, dim)
- 在最后一个维度拼接两份相同的余弦/正弦值，使其维度从
dim//2扩展到dim - 这是为了配合
apply_rotary_pos_emb中的旋转操作
apply_rotary_pos_emb（第80-84行）
def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
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))).to(q.dtype)
k_embed = ((k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))).to(k.dtype)
return q_embed, k_embed
旋转操作的本质：二维旋转矩阵 [[cos θ, -sin θ], [sin θ, cos θ]] 作用于向量对 (x₁, x₂)：
x₁' = x₁ * cos θ - x₂ * sin θ
x₂' = x₁ * sin θ + x₂ * cos θ
RoPE 把这个旋转推广到高维——相邻的两个维度组成一对，分别旋转。这样 Q·K 的点积自然包含位置差信息。
5️⃣ repeat_kv 函数（第86-89行）
def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
bs, slen, num_key_value_heads, head_dim = x.shape
if n_rep == 1: return x # 不需要重复，直接返回
# 在第4维扩展 n_rep 次，再 reshape 合并
return (x[:, :, :, None, :].expand(bs, slen, num_key_value_heads, n_rep, head_dim)
.reshape(bs, slen, num_key_value_heads * n_rep, head_dim))
为什么需要这个？ GQA 中 KV 头数少于 Q 头数。比如 8 个 Q 头，4 个 KV 头，每个 KV 头要"复制"2 次，才能和 8 个 Q 头一一配对做注意力计算。
KV heads: [K0, K1, K2, K3]
repeat 2x: [K0, K0, K1, K1, K2, K2, K3, K3] ← 现在有8个，能和8个Q头配对
6️⃣ Attention 注意力机制（第91-134行）
什么是注意力？
简单说：对于当前词，找出上下文中哪些词最相关，然后把相关信息汇总过来。
比如句子"猫坐在垫子上，它很舒服"，当模型处理"它"这个词时，注意力机制会让"它"重点关注"猫"，因为"它"指代"猫"。
class Attention(nn.Module):
def __init__(self, config: MiniMindConfig):
super().__init__()
初始化——定义所有需要的层
self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
- 如果没指定 KV 头数，就用 Q 头数（退化为标准 MHA）
- 指定了就用指定值（GQA 模式）
self.n_local_heads = config.num_attention_heads # Q 头数 = 8
self.n_local_kv_heads = self.num_key_value_heads # KV 头数 = 4
self.n_rep = self.n_local_heads // self.n_local_kv_heads # 重复次数 = 8/4 = 2
self.head_dim = config.head_dim # 每头维度 = 96
self.is_causal = True # 因果注意力（不能看未来）
四个线性投影层：
self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
| 层 | 输入维度 | 输出维度 | 说明 |
|---|---|---|---|
| q_proj | 768 | 8×96=768 | Query 投影，每个头96维 |
| k_proj | 768 | 4×96=384 | Key 投影，KV头只有4个 |
| v_proj | 768 | 4×96=384 | Value 投影，同上 |
| o_proj | 768 | 768 | 输出投影，合并多头结果 |
QK-Norm（稳定训练的关键技巧）：
self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) # Q 归一化
self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) # K 归一化
- 在计算注意力分数之前，对 Q 和 K 做归一化
- 防止 Q·K^T 的值过大导致 softmax 溢出
- 在大型模型训练中非常重要
self.attn_dropout = nn.Dropout(config.dropout) # 注意力 dropout
self.resid_dropout = nn.Dropout(config.dropout) # 残差 dropout
self.dropout = config.dropout # dropout 概率
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn
self.flash：检测 PyTorch 是否支持 Flash Attention，且配置开启
forward——注意力计算的核心
def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
参数：
| 参数 | 说明 |
|---|---|
x |
输入 hidden_states，形状 (batch, seq_len, 768) |
position_embeddings |
RoPE 的 (cos, sin) |
past_key_value |
上一轮的 KV cache（生成时用） |
use_cache |
是否缓存当前 KV（生成时用） |
attention_mask |
注意力遮罩（padding 时用） |
bsz, seq_len, _ = x.shape # batch_size, 序列长度
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) # (B, T, 8, 96)
xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) # (B, T, 4, 96)
xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim) # (B, T, 4, 96)
- reshape 成多头格式
xq, xk = self.q_norm(xq), self.k_norm(xk) # QK-Norm
cos, sin = position_embeddings
xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin) # 应用 RoPE 位置编码
KV Cache 处理：
if past_key_value is not None:
xk = torch.cat([past_key_value[0], xk], dim=1) # 拼接历史 K
xv = torch.cat([past_key_value[1], xv], dim=1) # 拼接历史 V
past_kv = (xk, xv) if use_cache else None # 缓存当前 KV
- KV Cache 是自回归生成的核心优化：生成第 N 个 token 时，不需要重新计算前 N-1 个 token 的 K 和 V，直接复用缓存。
xq, xk, xv = (
xq.transpose(1, 2), # (B, 8, T, 96)
repeat_kv(xk, self.n_rep).transpose(1, 2), # (B, 8, T, 96) ← 从4头扩展到8头
repeat_kv(xv, self.n_rep).transpose(1, 2) # (B, 8, T, 96)
)
注意力计算——两条路径：
if self.flash and (seq_len > 1) and ... :
# Flash Attention 路径（快）
output = F.scaled_dot_product_attention(
xq, xk, xv,
dropout_p=self.dropout if self.training else 0.0,
is_causal=self.is_causal
)
else:
# 手动计算路径（慢但灵活）
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim) # Q·K^T / √d
if self.is_causal:
# 因果遮罩：下三角矩阵，上三角为 -inf
scores[:, :, :, -seq_len:] += torch.full(
(seq_len, seq_len), float("-inf"), device=scores.device
).triu(1)
if attention_mask is not None:
scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
因果遮罩的可视化：
pos1 pos2 pos3 pos4
pos1 [ 0 -inf -inf -inf ] ← pos1 只能看自己
pos2 [ 0 0 -inf -inf ] ← pos2 能看 pos1 和自己
pos3 [ 0 0 0 -inf ] ← pos3 能看 pos1, pos2, 自己
pos4 [ 0 0 0 0 ] ← pos4 能看所有
output = output.transpose(1, 2).reshape(bsz, seq_len, -1) # 合并多头
output = self.resid_dropout(self.o_proj(output)) # 输出投影 + dropout
return output, past_kv
7️⃣ FeedForward 前馈网络（第136-146行）
class FeedForward(nn.Module):
def __init__(self, config, intermediate_size=None):
super().__init__()
intermediate_size = intermediate_size or config.intermediate_size
self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False) # 768 → 2432
self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False) # 2432 → 768
self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False) # 768 → 2432
self.act_fn = ACT2FN[config.hidden_act] # SiLU 激活函数
def forward(self, x):
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
SwiGLU 结构：
x → gate_proj → SiLU → * ← up_proj ← x
↓
down_proj → output
gate_proj和up_proj是两个独立投影- SiLU 作用于 gate，然后和 up 做逐元素乘法
- 最后 down_proj 投影回原始维度
这是目前最流行的 FFN 结构，比简单的 ReLU(W2 · ReLU(W1 · x)) 效果更好。
8️⃣ MOEFeedForward 混合专家（第148-176行）
什么是 MoE？
想象一个医院：以前所有病人都看同一个医生，现在有4个专科医生，分诊台（Router）根据病情把病人分给最合适的1个医生。
class MOEFeedForward(nn.Module):
def __init__(self, config):
super().__init__()
self.config = config
self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) # 路由器
self.experts = nn.ModuleList([FeedForward(config) for _ in range(config.num_experts)]) # 4个专家
self.act_fn = ACT2FN[config.hidden_act]
forward 流程：
def forward(self, x):
batch_size, seq_len, hidden_dim = x.shape
x_flat = x.view(-1, hidden_dim) # 展平为 (B*T, D)
# 1. 路由器打分
scores = F.softmax(self.gate(x_flat), dim=-1) # (B*T, num_experts)
# 2. 选 top-1 专家
topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
if self.config.norm_topk_prob:
topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20) # 归一化
# 3. 计算每个专家的输出，加权求和
y = torch.zeros_like(x_flat)
for i, expert in enumerate(self.experts):
mask = (topk_idx == i) # 哪些 token 选择了专家 i
if mask.any():
token_idx = mask.any(dim=-1).nonzero().flatten() # 这些 token 的索引
weight = topk_weight[mask].view(-1, 1) # 路由权重
y.index_add_(0, token_idx, (expert(x_flat[token_idx]) * weight).to(y.dtype))
elif self.training:
y[0, 0] += 0 * sum(p.sum() for p in expert.parameters()) # dummy梯度，防止未选中的专家参数不更新
# 4. 计算辅助损失（防止所有token都选同一个专家）
if self.training and self.config.router_aux_loss_coef > 0:
load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0) # 每个专家被选的频率
self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
else:
self.aux_loss = scores.new_zeros(1).squeeze()
return y.view(batch_size, seq_len, hidden_dim)
aux_loss 的意义：如果所有 token 都选专家0，其他专家就"失业"了。aux_loss 惩罚这种不均衡，让每个专家都有机会被选中。
9️⃣ MiniMindBlock Transformer 块（第178-194行）
class MiniMindBlock(nn.Module):
def __init__(self, layer_id: int, config: MiniMindConfig):
super().__init__()
self.self_attn = Attention(config) # 注意力层
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # 注意力前的归一化
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # FFN前的归一化
self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config) # FFN或MoE
def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
residual = hidden_states # 保存输入用于残差连接
hidden_states, present_key_value = self.self_attn( # 注意力计算
self.input_layernorm(hidden_states), position_embeddings, # Pre-Norm: 先归一化再计算
past_key_value, use_cache, attention_mask
)
hidden_states += residual # 残差连接
hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states)) # FFN + 残差
return hidden_states, present_key_value
Pre-Norm Transformer Block 的数据流：
输入 x
│
├→ RMSNorm → Attention → + x (残差)
│ │
│ RMSNorm → FFN → + (残差)
│ │
└──────────────────────────────→ 输出
🔟 MiniMindModel 模型主体（第196-232行）
class MiniMindModel(nn.Module):
def __init__(self, config: MiniMindConfig):
super().__init__()
self.config = config
self.vocab_size = config.vocab_size # 6400
self.num_hidden_layers = config.num_hidden_layers # 8
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) # 词嵌入层
self.dropout = nn.Dropout(config.dropout)
self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)]) # 8层
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) # 最终归一化
# 预计算 RoPE 的 cos/sin 表
freqs_cos, freqs_sin = precompute_freqs_cis(...)
self.register_buffer("freqs_cos", freqs_cos, persistent=False) # 不保存到 state_dict
self.register_buffer("freqs_sin", freqs_sin, persistent=False)
register_buffer：把 tensor 注册为模型的一部分（会随模型 .to(device) 移动），但不作为可训练参数。persistent=False 表示不保存到 checkpoint。
forward
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
batch_size, seq_length = input_ids.shape
# 处理 HuggingFace DynamicCache 兼容性
if hasattr(past_key_values, 'layers'): past_key_values = None
past_key_values = past_key_values or [None] * len(self.layers)
# 计算当前起始位置（从 KV cache 推断）
start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
# 词嵌入 + dropout
hidden_states = self.dropout(self.embed_tokens(input_ids))
# 修复 meta-device 初始化时丢失的 RoPE 缓冲区
if self.freqs_cos[0, 0] == 0:
freqs_cos, freqs_sin = precompute_freqs_cis(...)
self.freqs_cos, self.freqs_sin = freqs_cos.to(hidden_states.device), freqs_sin.to(hidden_states.device)
# 取当前序列对应的 RoPE 位置编码
position_embeddings = (
self.freqs_cos[start_pos:start_pos + seq_length],
self.freqs_sin[start_pos:start_pos + seq_length]
)
# 逐层通过 Transformer
presents = []
for layer, past_key_value in zip(self.layers, past_key_values):
hidden_states, present = layer(
hidden_states, position_embeddings,
past_key_value=past_key_value,
use_cache=use_cache,
attention_mask=attention_mask
)
presents.append(present)
# 最终归一化
hidden_states = self.norm(hidden_states)
# 汇总 MoE 辅助损失
aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], ...)
return hidden_states, presents, aux_loss
1️⃣1️⃣ MiniMindForCausalLM 语言模型（第234-288行）
class MiniMindForCausalLM(PreTrainedModel, GenerationMixin):
config_class = MiniMindConfig
_tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
- 继承
PreTrainedModel：获得 HuggingFace 的 save/load 能力 - 继承
GenerationMixin：获得.generate()方法（但这里重写了） _tied_weights_keys：声明权重绑定关系
def __init__(self, config=None):
self.config = config or MiniMindConfig()
super().__init__(self.config)
self.model = MiniMindModel(self.config) # 模型主体
self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False) # 语言模型头
if self.config.tie_word_embeddings:
self.model.embed_tokens.weight = self.lm_head.weight # 权重绑定！
self.post_init() # HuggingFace 后初始化
权重绑定：输入嵌入层和输出分类层共享权重。直觉：如果一个词的嵌入向量好，那它反过来预测这个词的概率也应该高。
forward——训练时调用
def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False,
logits_to_keep=0, labels=None, **kwargs):
hidden_states, past_key_values, aux_loss = self.model(input_ids, ...)
# logits_to_keep: 只计算最后 N 个位置的 logits（节省计算）
slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(hidden_states[:, slice_indices, :])
# 如果提供了 labels，计算交叉熵损失
loss = None
if labels is not None:
x, y = logits[..., :-1, :].contiguous(), labels[..., 1:].contiguous()
loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
return MoeCausalLMOutputWithPast(loss=loss, aux_loss=aux_loss, logits=logits, ...)
labels[..., 1:] vs logits[..., :-1, :]：经典的"下一个 token 预测"——用位置 t 的输出预测位置 t+1 的 token。
generate——推理时调用
@torch.inference_mode()
def generate(self, input_ids, eos_token_id=2, max_new_tokens=8192, temperature=0.85,
top_p=0.85, top_k=50, stream=False, use_cache=True, ...):
参数：
| 参数 | 默认值 | 说明 |
|---|---|---|
temperature |
0.85 | 采样温度。1.0=正常，<1=更确定，>1=更随机 |
top_p |
0.85 | 核采样阈值。只从累积概率前 85% 的 token 中选 |
top_k |
50 | 只从概率最高的 50 个 token 中选 |
repetition_penalty |
1.0 | 重复惩罚。>1 时降低已出现 token 的概率 |
自回归生成循环：
for _ in range(max_new_tokens):
# 1. 前向传播
outputs = self.forward(input_ids[:, past_len:], ...)
logits = outputs.logits[:, -1, :] / temperature # 取最后一个位置的 logits
# 2. 重复惩罚
if repetition_penalty != 1.0:
# 已出现的 token：正数除以 rp，负数乘以 rp
logits[i, seen] = torch.where(score > 0, score / rp, score * rp)
# 3. Top-K 过滤
if top_k > 0:
logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
# 4. Top-P (Nucleus) 过滤
if top_p < 1.0:
# 排序 → 累积概率 → 超过 top_p 的设为 -inf
# 5. 采样
next_token = torch.multinomial(softmax(logits), 1)
# 6. 拼接到输入
input_ids = torch.cat([input_ids, next_token], dim=-1)
# 7. 检查是否全部结束
if (next_token == eos_token_id).all(): break
📋 本章变量/函数速查表
全局函数
| 函数 | 参数 | 返回值 | 作用 |
|---|---|---|---|
precompute_freqs_cis |
dim, end, rope_base, rope_scaling | (cos, sin) | 预计算 RoPE 位置编码 |
apply_rotary_pos_emb |
q, k, cos, sin | (q_embed, k_embed) | 给 Q/K 应用旋转位置编码 |
repeat_kv |
x, n_rep | x_repeated | 复制 KV 头以匹配 Q 头数 |
类一览
| 类 | 行数 | 参数量 | 作用 |
|---|---|---|---|
MiniMindConfig |
35行 | 0 | 模型配置 |
RMSNorm |
11行 | dim | 归一化 |
Attention |
44行 | 2.4M | 多头注意力 |
FeedForward |
11行 | 5.9M | SwiGLU 前馈网络 |
MOEFeedForward |
29行 | 23.6M | 4个专家的 MoE |
MiniMindBlock |
17行 | — | 一个 Transformer 块 |
MiniMindModel |
37行 | — | 多层 Transformer 主体 |
MiniMindForCausalLM |
55行 | 113M | 完整的语言模型 |
✅ 第一章完成 | 下一篇：
model_omni.py多模态扩展
脑启社区是一个专注类脑智能领域的开发者社区。欢迎加入社区，共建类脑智能生态。社区为开发者提供了丰富的开源类脑工具软件、类脑算法模型及数据集、类脑知识库、类脑技术培训课程以及类脑应用案例等资源。
更多推荐
28
0- 0
已为社区贡献1条内容
所有评论(0)