MiniMind 预训练学习指南

这份文档专门介绍 MiniMind 项目的预训练部分,帮助你真正理解:

  • train_pretrain.py 在训练什么
  • pretrain_hq.jsonl 这种数据为什么能拿来训练
  • tokenizer、input_ids、embedding、hidden_size 分别是什么
  • “下一 token 预测”到底怎么做
  • 训练日志里的 loss / logits_loss / aux_loss 应该怎么理解

本文默认对应当前仓库中的:

  • 训练脚本:trainer/train_pretrain.py
  • 数据读取:dataset/lm_dataset.py
  • 模型定义:model/model_minimind.py
  • tokenizer 配置:model/tokenizer_config.json
  • tokenizer 本体:model/tokenizer.json

MiniMind代码仓库:https://github.com/jingyaogong/minimind.git


1. 预训练在做什么

预训练本质上是在做一件事:

给模型一段文本,让它根据前文去预测“下一个 token 是什么”。

这叫做:

  • Causal Language Modeling
  • 自回归语言模型训练
  • 下一 token 预测

例如给模型一句话:

1
今天天气很好

模型不是一次性理解整句话后输出答案,而是反复学习:

  • 看到“今天”后,后面大概率会出现什么
  • 看到“今天天气”后,后面大概率会出现什么
  • 看到“今天天气很”后,后面大概率会出现什么

训练久了,模型就会逐渐学会语言的局部规律、常见搭配、语法结构、问答模式和一些常识性分布。

2. 先看完整训练流程

在理解 tokenizer、embedding、loss 这些局部概念之前,先看一眼一条样本在预训练里会经历什么。

flowchart LR
    A[jsonl中的一条text] --> B[tokenizer编码]
    B --> C[input_ids]
    C --> D[补bos/eos]
    D --> E[padding到固定长度]
    E --> F[得到labels并屏蔽pad位置]
    F --> G[Embedding层]
    G --> H[Transformer多层计算]
    H --> I[lm_head输出logits]
    I --> J[shift后做下一token预测]
    J --> K[交叉熵loss]
    K --> L[反向传播更新参数]

如果把 batch 维度也加进去,大致会看到这样的张量流动:

flowchart TD
    A["原始文本 x batch_size"] --> B["tokenizer后: input_ids [B, T]"]
    B --> C["Embedding后: hidden_states [B, T, H]"]
    C --> D["Transformer后: hidden_states [B, T, H]"]
    D --> E["输出logits: [B, T, V]"]
    E --> F["shift_logits / shift_labels"]
    F --> G["cross_entropy(ignore_index=-100)"]

其中:

  • B = batch size
  • T = max sequence length
  • H = hidden size
  • V = vocab size

先把这个主线记住,后面的每个概念都只是这个流程中的一个环节。

3. train_pretrain.py 用的是什么数据

train_pretrain.py 默认读取:

1
parser.add_argument("--data_path", type=str, default="../dataset/pretrain_hq.jsonl", help="预训练数据路径")

所以默认训练集是:pretrain_hq.jsonl

其格式非常简单,每一行都是一条 JSON:

1
{"text": "如何才能摆脱拖延症? 治愈拖延症并不容易,但以下建议可能有所帮助..."}

也就是说,预训练阶段最核心的数据字段只有一个:

  • text
    • 一条原始文本样本
    • 可能是一段连续文本
    • 也可能是多个短片段拼接成的一条长文本

4. 为什么 pretrain_hq.jsonl 里会有 <|im_end|>

你会在数据里看到类似这样的内容:

1
{"text": "你好<|im_end|>请帮我查天气<|im_end|>北京今天天气晴...<|im_end|>"}

这里的 <|im_end|> 不是普通字符串,而是 tokenizer 里预定义好的特殊 token。

在当前项目中:

  • bos_token = <|im_start|>
  • eos_token = <|im_end|>
  • pad_token = <|endoftext|>

对应配置见:model/tokenizer_config.json

这些特殊 token 的作用可以先直观理解成“给文本结构打标记”:

  • bos_token = <|im_start|>
    • bos 是 beginning of sequence
    • 表示“一条输入序列从这里开始”
    • 在当前项目里,预训练数据读取时会手动加在每条样本最前面
  • eos_token = <|im_end|>
    • eos 是 end of sequence
    • 表示“一条片段或一条样本到这里结束”
    • 在样本内部它常常表示片段边界,在样本最后它表示整条样本结束
  • pad_token = <|endoftext|>
    • pad 是 padding token
    • 用于把不同长度的样本补到同一个长度,方便组成 batch
    • 它不是“真实文本内容”,只是一个占位符

可以把 <|im_end|> 理解成“结束边界”:

  • 在样本内部,它可以表示一个片段结束、一个问答轮次结束
  • 在样本最后,它也可以表示整条样本结束

所以它本质上的语义是一致的:都是“到这里结束”。

5. 为什么不把每个短片段都拆成独立样本

因为预训练通常会设定一个固定的 max_seq_lentrainer/train_pretrain.py中:

1
parser.add_argument('--max_seq_len', default=340, type=int, help="训练的最大截断长度")

如果每条样本都特别短,例如只有 10 个 token,但 max_seq_len=340,那其余 330 个位置都会变成 padding,算力浪费很大。

这里的 padding 可以理解成“补空位”。

例如一条样本真实只有:

1
[101, 205, 88, 32, 9]

但训练脚本要求所有样本都整理成长度 10,那么就会在后面补上 pad_token_id

1
[101, 205, 88, 32, 9, 0, 0, 0, 0, 0]

这些后补的 0 就是 padding。
它们的作用只是把样本长度对齐,好让很多样本可以一起送进模型。

因此很多预训练数据会把多个短片段打包成一条更长的文本:

1
片段A<|im_end|>片段B<|im_end|>片段C<|im_end|>

这样做的好处:

  • 提高每条样本的有效 token 占比
  • 降低 padding 浪费
  • 让模型学到片段边界
  • 提升训练吞吐

看到这里,一个很自然的问题会出现:

训练脚本拿到的是原始字符串,但神经网络只能处理数字,那中间这一步是谁来负责把文本变成数字的?

答案就是 tokenizer。

所以 tokenizer 不是一个需要孤立记忆的名词,它是在训练流程里“被逼出来”的:

1
2
3
4
5
原始文本
-> 模型不能直接吃字符串
-> 必须先切成token
-> 再把token映射成整数id
-> 这一步就由tokenizer负责

换句话说,tokenizer 是“原始文本进入模型之前的第一道转换器”。

7. tokenizer 是什么

神经网络不能直接吃字符串,只能处理数字。

tokenizer 的作用就是:

  • 把原始文本切成 token
  • 再把 token 映射成整数 id

当前项目里的 tokenizer 是一个:

  • PreTrainedTokenizerFast
  • ByteLevel BPE tokenizer

可以从 model/tokenizer.json 里看到:

  • pre_tokenizer.type = ByteLevel
  • model.type = BPE

它不是“按汉字一个一个切”的简单规则,也不是“按空格 split”。
它会根据自己训练好的词表和 merge 规则,把文本自动切成 token。

8. tokenizer 和 embedding 不是一回事

这两个概念特别容易混。

tokenizer 负责:

  • 文本 -> token -> id

例如:

1
你好<|im_end|>世界

可能被编码成:

1
[5134, 2, 2219]

而 embedding 负责:

  • id -> 向量

例如:

1
2
3
5134 -> 一个 512 维向量
2 -> 一个 512 维向量
2219 -> 一个 512 维向量

所以:

  • tokenizer 是“离散化工具”
  • embedding 是“向量化层”

9. 先从训练时的形状变化理解参数

假设当前训练配置是:

  • batch_size = 32
  • max_seq_len = 340
  • hidden_size = 512

那么一批数据在训练过程中,会大致经历这样的形状变化:

1
2
3
4
5
6
7
原始文本
-> tokenizer 编码
-> input_ids: [32, 340]
-> embedding lookup
-> hidden_states: [32, 340, 512]
-> lm_head 投影
-> logits: [32, 340, vocab_size]

这时再回头理解这些参数就会更自然。

batch_size 是什么

表示每次同时喂给模型多少条样本。

例如:

  • batch_size = 32

那就表示一次迭代同时训练 32 条样本。

max_seq_len 是什么

表示一条样本最多保留多少个 token。

例如:

  • max_seq_len = 340

那么每条样本最终都会被整理成长度固定为 340 的 input_idslabels

hidden_size 是什么

表示每个 token 在模型内部会被表示成多少维的向量。

例如:

  • hidden_size = 512

那每个 token 在 embedding 之后,就会变成一个 512 维向量。

注意:

  • hidden_size 不是文本长度
  • hidden_size 也不是“整条文本被压成 512 维”
  • 它是“每个 token 的向量维度”

vocab_size 是什么

表示 tokenizer 词表里一共有多少个 token。

例如在当前 MiniMind 配置中:

  • vocab_size = 6400

这意味着模型最终不是在“无限多种文本”里直接选答案,而是在:

  • 6400 个候选 token

里预测“下一个 token 是哪一个”。

所以当模型输出 logits 时,最后一个维度就是 vocab_size

1
[B, T, V]

如果:

  • B = 32
  • T = 340
  • V = 6400

那么 logits 的形状大致就是:

1
[32, 340, 6400]

可以把它理解成:

  • 32 条样本
  • 每条样本 340 个位置
  • 每个位置都要对 6400 个 token 分别打分

vocab_size 还会直接影响模型大小,因为 embedding 表的形状就是:

1
[vocab_size, hidden_size]

例如:

1
[6400, 512]

所以:

  • 词表越大
  • embedding 参数通常也越多
  • 输出层的参数量也会随之变大

num_hidden_layers 是什么

表示 Transformer block 的层数。

例如:

  • num_hidden_layers = 8

就表示每个 token 的表示会被连续加工 8 层。

这里的“层”不是简单的一层线性层,而是一整个 Transformer block。

在 MiniMind 里,一层 block 的定义可以参考:

  • model/model_minimind.py

它主要包含这几个部分:

  • input_layernorm
  • self_attn
  • 第一条残差连接
  • post_attention_layernorm
  • mlpmoe-mlp
  • 第二条残差连接

可以画成:

flowchart TD
    A[hidden_states输入] --> B[RMSNorm]
    B --> C[Self-Attention]
    C --> D[残差相加]
    D --> E[RMSNorm]
    E --> F[MLP 或 MoE-MLP]
    F --> G[残差相加]
    G --> H[hidden_states输出]

如果用更贴近代码的方式理解:

1
2
3
4
5
6
7
8
输入 hidden_states
-> 先做一次归一化
-> 做 self-attention,让每个 token 看见前文
-> 和原输入做残差相加
-> 再做一次归一化
-> 送进前馈网络 MLP
-> 再和上一阶段结果做一次残差相加
-> 得到这一层 block 的输出

一层里的每个部分分别做什么

1. input_layernorm

作用是把进入 attention 前的表示做归一化,让数值更稳定。

MiniMind 这里用的是 RMSNorm,不是传统的 LayerNorm

你可以先把它理解成:

  • 在做 attention 之前,先把向量尺度整理得更稳定一些

2. self_attn

这是最核心的部分,它让每个 token 能根据上下文重新理解自己。

如果你想把“self-attention”具体落到 MiniMind 源码里看,最关键的代码就在:

  • model/model_minimind.py

其中 Q / K / V 对应关系非常直接:

  • Q 对应 q_proj
  • K 对应 k_proj
  • V 对应 v_proj

它们在前向传播里真正生成 Q / K / V 的位置是:

1
xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)

这里:

  • x 是当前层输入的 hidden_states
  • xq 就是 Query
  • xk 就是 Key
  • xv 就是 Value

接着,MiniMind 会把它们 reshape 成多头 attention 需要的形状:

1
2
3
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)

然后对 QK 应用 RoPE 位置编码:

1
2
cos, sin = position_embeddings
xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)

再之后,attention 的计算主干分两种情况:

  • 如果环境满足条件,直接走 PyTorch 的 scaled_dot_product_attention
  • 否则走手写的显式计算流程

手写流程最关键的几行是:

1
2
3
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
scores = F.softmax(scores.float(), dim=-1).type_as(xq)
output = scores @ xv

这三行就对应了最经典的 attention 公式:

1
Attention(Q, K, V) = softmax(QK^T / sqrt(d)) V

最后再通过输出投影把多头结果合并回隐藏维度:

1
output = self.resid_dropout(self.o_proj(output))

直觉上:

  • token 不再只看自己的 embedding
  • 而是会参考序列中前面的 token
  • 从而形成“带上下文的表示”

例如句子:

1
苹果很好吃

和:

1
苹果发布新手机

虽然“苹果”这个 token 一样,但经过 self-attention 后,它会结合上下文变成不同的 hidden state。

在预训练里,这一步特别重要,因为模型需要靠上下文来判断“下一个 token 最可能是什么”。

你也可以把 MiniMind 里的 Q / K / V 流程简化记成:

flowchart LR
    A[hidden_states] --> B[q_proj]
    A --> C[k_proj]
    A --> D[v_proj]
    B --> E[Q]
    C --> F[K]
    D --> G[V]
    E --> H[QK^T / sqrt(d)]
    F --> H
    H --> I[softmax]
    I --> J[attention weights]
    J --> K[加权求和 V]
    G --> K
    K --> L[o_proj]
    L --> M[新的hidden_states]

在训练时,这整套 Q / K / V 计算都属于前向传播的一部分。
loss 算出来之后,反向传播会更新:

  • q_proj
  • k_proj
  • v_proj
  • o_proj

这些线性层的参数。

一个单头 attention 的手算例子

如果你已经知道 Q / K / V 的名字,但对:

1
QK^T -> softmax -> 加权 V

这条链路还是感觉抽象,可以先看一个极简例子。

我们假设:

  • 只有 2 个 token
  • 每个向量维度是 2
  • 只看单头 attention

设这两个 token 的 Q / K / V 分别是:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
Q =
[
[1, 0], # token1 的 query
[0, 1] # token2 的 query
]

K =
[
[1, 0], # token1 的 key
[1, 1] # token2 的 key
]

V =
[
[10, 0], # token1 的 value
[0, 20] # token2 的 value
]

这里的目标是:

  • 看 token1 最终会从两个 token 身上拿到什么信息
  • 看 token2 最终会从两个 token 身上拿到什么信息

第一步:计算 QK^T

先把 K 转置:

1
2
3
4
5
K^T =
[
[1, 1],
[0, 1]
]

然后做矩阵乘法:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
QK^T =
[
[1, 0],
[0, 1]
]
×
[
[1, 1],
[0, 1]
]
=
[
[1, 1],
[0, 1]
]

这个结果的含义是:

  • 第一行:token1 对 token1、token2 的相关性分数
  • 第二行:token2 对 token1、token2 的相关性分数

也就是说:

1
2
3
4
5
scores =
[
[1, 1], # token1 看自己和 token2
[0, 1] # token2 看 token1 和自己
]

第二步:除以 sqrt(d)

这里向量维度 d = 2,所以:

1
sqrt(d) = sqrt(2) ≈ 1.414

所以缩放后的分数约为:

1
2
3
4
5
6
7
8
9
[
[1/1.414, 1/1.414],
[0/1.414, 1/1.414]
]
=
[
[0.707, 0.707],
[0.000, 0.707]
]

第三步:对每一行做 softmax

第一行:

1
softmax([0.707, 0.707]) = [0.5, 0.5]

第二行:

1
softmax([0.000, 0.707])

先取指数:

1
[e^0, e^0.707] ≈ [1.0, 2.028]

归一化后:

1
[1.0 / 3.028, 2.028 / 3.028] ≈ [0.33, 0.67]

所以注意力权重大约是:

1
2
3
4
5
attention_weights =
[
[0.50, 0.50],
[0.33, 0.67]
]

含义是:

  • token1 平均关注自己和 token2
  • token2 更关注 token2 自己

第四步:用注意力权重去加权 V

现在计算:

1
output = attention_weights × V

也就是:

1
2
3
4
5
6
7
8
9
[
[0.50, 0.50],
[0.33, 0.67]
]
×
[
[10, 0],
[0, 20]
]

先看第一行:

1
2
3
0.50 * [10, 0] + 0.50 * [0, 20]
= [5, 0] + [0, 10]
= [5, 10]

再看第二行:

1
2
3
0.33 * [10, 0] + 0.67 * [0, 20]
= [3.3, 0] + [0, 13.4]
= [3.3, 13.4]

所以最终输出约为:

1
2
3
4
5
output =
[
[5.0, 10.0],
[3.3, 13.4]
]

这个结果怎么理解

原来两个 token 的 V 是:

1
2
token1: [10, 0]
token2: [0, 20]

经过 attention 之后,得到的新表示是:

1
2
token1 -> [5.0, 10.0]
token2 -> [3.3, 13.4]

也就是说,每个 token 的新表示不再只是“自己的值”,而是:

  • 根据自己对别人的关注程度
  • 从所有 token 身上收集信息
  • 形成新的上下文表示

这就是 attention 最核心的直觉:

  • 不是孤立地处理每个 token
  • 而是让每个 token 根据相关性,从上下文中动态取信息

3. 第一条残差连接

代码里是:

1
hidden_states += residual

残差连接的作用可以先简单理解成:

  • 不让新计算结果完全覆盖旧信息
  • 保留输入信息
  • 让深层网络更容易训练

它是 Transformer 能堆很多层仍然稳定训练的重要原因之一。

4. post_attention_layernorm

这是 attention 之后、MLP 之前再做一次归一化。

作用还是:

  • 稳定数值
  • 让后面的 MLP 更容易处理输入

5. mlp

这部分是前馈网络,默认情况下是普通 FeedForward

可以把它理解成:

  • 对每个 token 的表示单独做一轮更强的非线性变换
  • 补充 attention 不擅长的特征变换能力

如果把 attention 理解成“让 token 彼此交流信息”,那么 MLP 更像是:

  • token 自己在拿到上下文之后,重新加工自己

6. 第二条残差连接

代码里是:

1
hidden_states = hidden_states + self.mlp(...)

作用和前面的残差连接类似:

  • 保留前一阶段结果
  • 让 MLP 的更新是“增量式”的
  • 避免深层训练时信息退化

为什么要堆很多层

一层 block 能做的事情有限。

堆多层的意义在于:

  • 第 1 层学到比较浅层的局部模式
  • 后面的层逐渐整合更长范围、更抽象的上下文信息

可以把它粗略理解成:

  • 前几层更像“看字词和短语”
  • 中间层更像“看句子结构”
  • 更后面的层更像“看更复杂的语义和上下文关系”

这不是绝对划分,但作为学习阶段的直觉是有帮助的。

所以:

  • num_hidden_layers 越大
  • 模型可加工上下文的深度通常越强
  • 参数量和计算量也会随之上升

use_moe 是什么

表示是否启用 MoE(Mixture of Experts)结构。

  • use_moe = 0
    • 普通 Dense 模型
  • use_moe = 1
    • 使用专家路由结构

当前默认是 0,所以你看到的训练日志里 aux_loss 通常是 0.0000

从学习顺序上,也更建议先把默认 Dense 模型理解透,再去看 MoE。

10. 文本是怎么变成 input_ids

预训练数据读取逻辑在:

核心过程:

1
2
3
4
5
tokens = self.tokenizer(str(sample['text']), add_special_tokens=False, max_length=self.max_length - 2, truncation=True).input_ids
tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]
input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens))
labels = input_ids.clone()
labels[input_ids == self.tokenizer.pad_token_id] = -100

可以拆成这几步:

  1. 用 tokenizer 把文本编码成 id
  2. 手动在开头加 bos_token_id
  3. 手动在结尾加 eos_token_id
  4. padding 到固定长度
  5. 复制一份得到 labels
  6. 把 padding 对应的 label 改成 -100

11. 一个极简样本例子

假设原始文本是:

1
你好<|im_end|>世界

假设 tokenizer 编码后得到:

1
[5134, 2, 2219]

其中:

  • 5134 对应某个 token
  • 2 就是 <|im_end|>
  • 2219 对应某个 token

再假设:

  • bos_token_id = 1
  • eos_token_id = 2
  • pad_token_id = 0
  • max_length = 10

那么处理过程就是:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
原始 text:
你好<|im_end|>世界

tokenizer(..., add_special_tokens=False) 后:
[5134, 2, 2219]

手动加 bos/eos:
[1, 5134, 2, 2219, 2]

再 pad 到长度 10:
input_ids = [1, 5134, 2, 2219, 2, 0, 0, 0, 0, 0]

labels 初始复制 input_ids:
labels = [1, 5134, 2, 2219, 2, 0, 0, 0, 0, 0]

把 pad 位置改成 -100:
labels = [1, 5134, 2, 2219, 2, -100, -100, -100, -100, -100]

12. 为什么 labels 先复制 input_ids

因为语言模型训练的目标是:

给定当前位置之前的内容,预测下一个 token。

目标序列和输入序列其实就是同一串文本,只是训练时会错开一位。

模型里的关键代码在:

1
2
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()

也就是说:

  • 第 0 个位置的输出去预测第 1 个 token
  • 第 1 个位置的输出去预测第 2 个 token
  • 依此类推

所以 labels = input_ids.clone() 是很自然的写法。

13. 为什么 pad 位置要改成 -100

因为 padding 只是为了对齐 batch,并不是真实文本。

如果让模型也去学习这些 padding 位置,就会让训练目标被污染。

当前项目在计算交叉熵时使用了:

1
ignore_index=-100

这表示:

  • labels 中所有等于 -100 的位置
  • 直接跳过,不参与 loss 计算

所以把 pad 位置改成 -100 的作用是:

  • 保留固定长度 batch
  • 但不让 padding 影响训练

这里可以把 pad_token-100 区分开:

  • pad_token
    • 还在 input_ids
    • 模型仍然会看到它这个位置的输入
  • -100
    • 出现在 labels
    • 表示这个位置不参与 loss

所以:

  • pad_token 是输入对齐方案
  • -100 是 loss 屏蔽方案

14. token 是怎么变成向量的

这一步不是 tokenizer 做的,而是模型里的 embedding 层做的。

代码在:

1
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)

如果:

  • vocab_size = 6400
  • hidden_size = 512

那它本质上就是一个形状约为:

1
[6400, 512]

的可训练矩阵。

你可以把它理解成一张大表:

  • 每个 token id 对应一行
  • 每一行就是一个 512 维向量

例如:

1
2
3
4
id 0    -> [ ... 512个数 ... ]
id 1 -> [ ... 512个数 ... ]
id 5134 -> [ ... 512个数 ... ]
id 2219 -> [ ... 512个数 ... ]

当输入:

1
input_ids = [1, 5134, 2, 2219, 2]

embedding lookup 之后就得到:

1
2
3
4
5
6
7
[
vec(1),
vec(5134),
vec(2),
vec(2219),
vec(2)
]

形状为:

1
[5, 512]

15. train_pretrain.py 默认训练的是多大的模型

到这里再看“模型大小”和配置参数,会更顺手一些,因为你已经知道这些参数会在训练流程的哪个环节起作用。

默认参数在:

1
2
3
hidden_size=512
num_hidden_layers=8
use_moe=0

这对应 README 里的:

大约是:

  • 26M 参数

README 中的常见规模:

  • MiniMind2-Small: 26M
  • MiniMind2: 104M
  • MiniMind2-MoE: 145M

可以先这样粗略理解模型大小是怎么“涨起来”的:

  • hidden_size 越大
    • 每层线性变换矩阵越大
    • embedding 表也更大
  • num_hidden_layers 越多
    • Transformer block 叠得越多
    • 参数量近似线性增长
  • use_moe=1
    • FFN 部分会引入更多专家参数
    • 总参数量会进一步增大

如果只做最粗糙的直觉估计,可以先记住:

1
参数量 ~ 每层参数量 × 层数 + embedding/lm_head参数

在 MiniMind 默认配置里:

  • 词表大小约 6400
  • hidden_size = 512
  • num_hidden_layers = 8

所以 embedding 表本身大约就有:

1
6400 × 512 ≈ 3.28M

再加上 8 层 Transformer 里的 attention、MLP、norm 等参数,总体就来到约 26M。

这里不用一开始就精确手推到每一项,只要先建立这个直觉:

  • 模型大不大,本质上取决于“词表有多大、每个 token 用多宽的向量表示、层数有多少”

16. 为什么反向传播会学出“语义相近”

模型一开始并不知道什么叫语义相近。

它只是反复做:

  • 根据前文预测下一个 token
  • 预测错了就根据 loss 调整参数

如果两个 token 经常出现在相似上下文里,它们就会收到相似方向的梯度更新。

例如:

  • “小猫在睡觉”
  • “小狗在睡觉”

为了把这两句话都预测好,模型会逐渐把“猫”和“狗”的表示学得更接近。
所以模型学到的不是词典定义,而是“上下文中的行为相似性”。

17. hidden_states -> logits 的迷你例子

为了帮助理解,我们构造一个极简模型:

  • 词表大小 vocab_size = 3
  • 向量维度 hidden_size = 4

3 个 token:

  • id 0 = 猫
  • id 1 = 狗
  • id 2 = 车

设 embedding 矩阵为:

1
2
3
4
5
6
E =
[
[1.0, 0.0, 1.0, 0.0], # 猫
[0.9, 0.1, 1.0, 0.0], # 狗
[0.0, 1.0, 0.0, 1.0] # 车
]

输入:

1
input_ids = [0, 2]

embedding lookup 后:

1
2
3
4
5
hidden_states =
[
[1.0, 0.0, 1.0, 0.0], # 猫
[0.0, 1.0, 0.0, 1.0] # 车
]

假设输出层使用同权重投影,则:

1
logits = hidden_states × E^T

对第一个位置:

1
[2.0, 1.9, 0.0]

对第二个位置:

1
[0.0, 0.1, 2.0]

于是整条序列的 logits 为:

1
2
3
4
[
[2.0, 1.9, 0.0],
[0.0, 0.1, 2.0]
]

这说明:

  • “猫”这个位置和“狗”的相似度也很高
  • “车”和“猫/狗”差异更大

这就是向量空间表示的直觉来源。

18. 交叉熵损失是什么

在每个位置上,模型都会输出一个长度为 vocab_size 的分数向量,也就是 logits。

例如某个位置输出:

1
[2.0, 1.9, 0.0]

如果词表只有 3 个 token,这表示模型认为:

  • token0 的可能性较高
  • token1 的可能性也较高
  • token2 的可能性很低

把 logits 做 softmax 后,就会变成概率分布。

假设真实答案是 token1,softmax 后它的概率是 0.44,那该位置的交叉熵就是:

1
-log(0.44)

真实答案概率越高,loss 越低。

在当前项目中,交叉熵定义在:

1
2
3
4
5
loss = F.cross_entropy(
shift_logits.view(-1, shift_logits.size(-1)),
shift_labels.view(-1),
ignore_index=-100
)

19. loss / logits_loss / aux_loss 各是什么意思

训练日志中会打印:

  • loss
  • logits_loss
  • aux_loss

定义位置见:

  • trainer/train_pretrain.py
  • trainer/train_pretrain.py

logits_loss

主任务损失,也就是语言模型的交叉熵损失。
这是最重要的指标。

aux_loss

辅助损失,主要在开启 MoE 时才有意义。
它用于约束专家路由更均衡,避免个别专家负载过重。

如果你当前是:

1
use_moe = 0

那么通常会看到:

1
aux_loss = 0.0000

loss

总损失:

1
loss = logits_loss + aux_loss

如果不用 MoE,那通常就是:

1
loss = logits_loss

20. 预训练脚本的完整主线

把整个流程连起来,就是:

  1. pretrain_hq.jsonl 读取 text
  2. tokenizer 把文本变成 token ids
  3. 手动加 bos/eos
  4. pad 到固定长度,得到 input_ids
  5. 复制生成 labels,把 pad 位置改成 -100
  6. input_ids 进入 embedding 层,变成 token 向量
  7. 经过多层 Transformer 得到 hidden_states
  8. lm_head 投影到词表大小,得到 logits
  9. shift_logitsshift_labels 做下一 token 预测
  10. 计算交叉熵损失
  11. 反向传播更新参数

21. 预训练阶段最容易混淆的几个点

误区 1:hidden_size 是输入长度

不是。
输入长度主要由 max_seq_len 决定,hidden_size 是每个 token 的向量维度。

误区 2:tokenizer 就是 embedding

不是。
tokenizer 负责文本转 id,embedding 负责 id 转向量。

误区 3:labelsinput_ids 一样就没有意义

不是。
真正训练时会做 shift,对齐到“预测下一个 token”。

误区 4:<|im_end|> 只是普通字符串

不是。
它在当前 tokenizer 里是特殊 token,也是 eos_token

误区 5:先背参数定义,再理解训练过程

不太推荐。
更自然的顺序通常是:

  1. 先看一条样本在训练里如何流动
  2. 再看 input_ids / labels / logits / loss
  3. 最后再回头理解 hidden_size / max_seq_len / batch_size / num_hidden_layers

这样参数就不容易变成抽象名词。

22. 从预训练到 SFT 的自然衔接

预训练学到的是:

  • 基础语言规律
  • 基础续写能力
  • 一些常见知识分布

但它还不等于“好用的聊天模型”。

因为预训练阶段只是让模型学会:

  • 文本如何继续写

而 SFT 阶段会进一步让模型学会:

  • 如何按“用户-助手”的对话格式回答
  • 哪些内容应该由 assistant 来生成
  • 如何更像一个有指令跟随能力的助手

所以:

  • train_pretrain.py 是打基础
  • train_full_sft.py 是把基础模型进一步调成“会对话”的模型

23. 一句话总结

MiniMind 的预训练阶段,本质上就是:

pretrain_hq.jsonl 中的原始文本,借助 tokenizer 变成 token id 序列,再通过 embedding 和 Transformer 预测“下一个 token”,并用交叉熵不断纠正预测错误。

如果你已经理解了本文中的:

  • tokenizer
  • input_ids / labels
  • embedding
  • hidden_size / max_seq_len / batch_size
  • shift_logits / shift_labels
  • 交叉熵损失

那么你已经掌握了 MiniMind 预训练阶段最核心的脉络。