📌 本文亮点:深入解析 MiniMind 预训练的每一步细节------模型架构、数据准备、训练脚本、超参配置、Loss 曲线与训练指标,帮你真正理解 LLM 预训练的全过程!
前言
大语言模型(LLM)的预训练,是整个训练链路中最核心、最耗时的阶段。然而动辄数十亿参数的模型规模,让绝大多数人只能停留在"用 LoRA 微调"的层面,从未真正理解预训练到底在做什么、Loss 是怎么收敛的、训练指标意味着什么。
MiniMind 项目提供了一个绝佳的入口:仅需单卡 3090 + 约 1.21 小时 ,即可从零完成一个 64M 参数语言模型的预训练。更关键的是,所有核心代码均基于 PyTorch 原生实现,没有 transformers 的高层封装,每一步都可以被理解、被追踪。
本文将带你完整走一遍 MiniMind 预训练的全流程,从模型架构到数据格式,从训练脚本到 Loss 指标,真正做到"知其然,更知其所以然"。
一、预训练的核心目标
LLM 预训练做的事情,本质上就是让模型先埋头读大量文本,从中学习事实知识、语言模式以及上下文之间的统计关系。这个阶段是无监督的:人类不需要逐条告诉模型哪里对、哪里错,而是让它自己从海量文本里总结规律。
更直白地说,模型在这一阶段的核心目标就是学会高质量地词语接龙。例如输入"秦始皇",它要能够继续生成"是中国历史上的第一位皇帝"这类符合语义与常识的后续内容。
┌──────────────────────────────────────────────────┐
│ 预训练 (Pretrain) │
│ │
│ 输入: "天空之所以看起来是蓝色的,主要是因为" │
│ 目标: 预测下一个最可能的 token │
│ 输出: "太阳" / "光" / "散射" ... │
│ │
│ 训练信号: 自身的文本(无需人工标注) │
│ 损失函数: Cross-Entropy (Next Token Prediction) │
└──────────────────────────────────────────────────┘
二、模型架构
2.1 架构总览
minimind-3 Dense 使用 Transformer Decoder-Only 结构,整体配置已向 Qwen3 生态对齐:
- 预标准化(Pre-Norm) + RMSNorm
- SwiGLU 激活函数
- RoPE 旋转位置编码,支持 YaRN 外推
- GQA (Grouped Query Attention):
q_heads=8、kv_heads=4 max_position_embeddings=32768,rope_theta=1e6
2.2 模型参数配置
| 配置项 | minimind-3 | minimind-3-moe | 说明 |
|---|---|---|---|
| 参数量 | 64M | 198M-A64M | MoE 激活参数仅 64M |
vocab_size |
6,400 | 6,400 | 精简词表,压缩 embedding 占比 |
hidden_size |
768 | 768 | 隐藏层维度 |
num_hidden_layers |
8 | 8 | Transformer 层数 |
num_attention_heads |
8 | 8 | Query 头数 |
num_key_value_heads |
4 | 4 | KV 头数(GQA) |
head_dim |
96 | 96 | 每个注意力头的维度 |
intermediate_size |
3,776 | 3,776 | FFN 中间层维度(SwiGLU) |
max_position_embeddings |
32,768 | 32,768 | 最大位置编码长度 |
rope_theta |
1e6 | 1e6 | RoPE 基频 |
| MoE experts | - | 4 | 专家数量 |
| top-k routing | - | 1 | 每个 token 激活的专家数 |
2.3 架构设计理念
💡 关于
d_model与n_layers的取舍:当前minimind-3选择dim=768, n_layers=8,本质上是一种工程取舍------更浅的网络训练更快,同时dim不至于过小而导致模式崩溃,在训练效率、稳定性与最终效果之间取得相对均衡。
核心结构代码如下(简化版):
python
class MiniMindBlock(nn.Module):
def __init__(self, layer_id, config):
self.self_attn = Attention(config) # GQA + RoPE
self.input_layernorm = RMSNorm(config.hidden_size) # Pre-Norm
self.post_attention_layernorm = RMSNorm(config.hidden_size)
self.mlp = FeedForward(config) # SwiGLU
def forward(self, hidden_states, position_embeddings, ...):
residual = hidden_states
# 1. 注意力层(Pre-Norm)
hidden_states, present_kv = self.self_attn(
self.input_layernorm(hidden_states), position_embeddings, ...
)
hidden_states += residual # 残差连接
# 2. 前馈层(Pre-Norm)
hidden_states = hidden_states + self.mlp(
self.post_attention_layernorm(hidden_states)
)
return hidden_states, present_kv
三、数据准备
3.1 预训练数据集
MiniMind-3 提供两种规模的预训练数据:
| 数据集 | 大小 | 适用场景 | 推荐 max_seq_len |
|---|---|---|---|
pretrain_t2t_mini.jsonl |
1.2GB ✨ | 快速复现 | ≈768 |
pretrain_t2t.jsonl |
10GB | 完整训练 MiniMind-3 | ≈380 |
✨ 为推荐必选项,适合首次复现
3.2 数据格式
预训练数据采用统一的 text -> next token prediction 格式,每行一条纯文本:
jsonl
{"text": "如何才能摆脱拖延症?治愈拖延症并不容易,但以下建议可能有所帮助。"}
{"text": "清晨的阳光透过窗帘洒进房间,桌上的书页被风轻轻翻动。"}
{"text": "Transformer 通过自注意力机制建模上下文关系,是现代大语言模型的重要基础结构。"}
3.3 数据来源
数据来源包括但不限于:
- 通用文本语料
- 对话整理语料
- 蒸馏补充语料
- 各类宽松开源协议可用的数据集
主要公开数据源:匠数大模型数据集、Magpie-Align
3.4 Tokenizer
MiniMind 使用自定义的 BPE + ByteLevel 分词器,词表大小仅 6,400:
| Tokenizer | 词表大小 | 来源 |
|---|---|---|
| MiniMind | 6,400 | 自定义 |
| Yi | 64,000 | 01万物 |
| Mistral | 32,000 | Mistral AI |
| Llama 3 | 128,000 | Meta |
| Qwen2 | 151,643 | 阿里云 |
⚠️ 词表精简是为了显著压缩
embedding层和输出层的参数占比,更适合小模型的体积约束。中文约1.5~1.7 字符/token,纯英文约4~5 字符/token。
3.5 数据加载流程
预训练数据的加载逻辑定义在 PretrainDataset 类中:
python
class PretrainDataset(Dataset):
def __init__(self, data_path, tokenizer, max_length=512):
self.tokenizer = tokenizer
self.max_length = max_length
# 直接从 jsonl 加载,无需预处理
self.samples = load_dataset('json', data_files=data_path, split='train')
def __getitem__(self, index):
sample = self.samples[index]
# 1. 分词 + 截断(保留 max_length - 2 给 BOS/EOS)
tokens = self.tokenizer(
str(sample['text']), add_special_tokens=False,
max_length=self.max_length - 2, truncation=True
).input_ids
# 2. 添加 BOS/EOS token
tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id]
# 3. Padding 到 max_length
input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens))
# 4. 构造 labels:pad 位置标记为 -100(不参与 loss 计算)
labels = input_ids.clone()
labels[input_ids == self.tokenizer.pad_token_id] = -100
return input_ids, labels
关键设计:
- BOS/EOS 包裹 :每条文本前后添加
<|im_start|>和<|im_end|>标记 - 右填充 :短于
max_seq_len的样本用pad_token_id填充 - Label 掩码 :
pad位置的label设为-100,不参与交叉熵计算 - 直接加载 :无需预处理成
.bin,直接从jsonl在线读取
四、训练脚本详解
4.1 训练命令
bash
# 单卡训练
cd trainer && python train_pretrain.py
# 多卡训练(N 为 GPU 数量)
cd trainer && torchrun --nproc_per_node N train_pretrain.py
# 断点续训
cd trainer && python train_pretrain.py --from_resume 1
# 使用 wandb/swanlab 可视化
cd trainer && python train_pretrain.py --use_wandb
4.2 默认超参数配置
| 超参数 | 默认值 | 说明 |
|---|---|---|
--epochs |
2 | 训练轮数 |
--batch_size |
32 | 每个 GPU 的 batch size |
--learning_rate |
5e-4 | 初始学习率 |
--accumulation_steps |
8 | 梯度累积步数 |
--max_seq_len |
340 | 最大序列长度(tokens) |
--hidden_size |
768 | 隐藏层维度 |
--num_hidden_layers |
8 | Transformer 层数 |
--dtype |
bfloat16 | 混合精度类型 |
--grad_clip |
1.0 | 梯度裁剪阈值 |
--log_interval |
100 | 日志打印间隔(步) |
--save_interval |
1000 | 模型保存间隔(步) |
--data_path |
pretrain_t2t_mini.jsonl |
预训练数据路径 |
💡 等效 batch size =
batch_size × accumulation_steps × GPU数量=32 × 8 × 1= 256(单卡)
4.3 学习率调度
MiniMind 使用**余弦退火(Cosine Annealing)**学习率调度:
python
def get_lr(current_step, total_steps, lr):
return lr * (0.1 + 0.45 * (1 + math.cos(math.pi * current_step / total_steps)))
调度特点:
- 初始学习率 :
lr × 0.55≈2.75e-4(warm-up 阶段内置) - 最终学习率 :
lr × 0.1=5e-5(衰减到 10%) - 过渡方式:平滑余弦曲线,无骤降
- 无需单独 warmup:公式天然包含了从较高值起步的过渡
4.4 训练循环核心逻辑
python
def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
for step, (input_ids, labels) in enumerate(loader, start=start_step + 1):
# 1. 动态学习率
lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
for param_group in optimizer.param_groups:
param_group['lr'] = lr
# 2. 前向传播(混合精度)
with autocast_ctx:
res = model(input_ids, labels=labels)
loss = res.loss + res.aux_loss # 主 loss + MoE 辅助 loss
loss = loss / args.accumulation_steps # 梯度累积归一化
# 3. 反向传播
scaler.scale(loss).backward()
# 4. 梯度累积 + 梯度裁剪 + 参数更新
if step % args.accumulation_steps == 0:
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad(set_to_none=True)
# 5. 日志记录
if step % args.log_interval == 0:
current_loss = loss.item() * args.accumulation_steps
Logger(f'Epoch:[{epoch+1}/{args.epochs}]({step}/{iters}), '
f'loss: {current_loss:.4f}, lr: {current_lr:.8f}')
# 6. 模型保存
if step % args.save_interval == 0:
torch.save(model.state_dict(), ckp_path)
4.5 损失函数
预训练的损失函数为交叉熵(Cross-Entropy),用于 Next Token Prediction:
python
# 在 MiniMindForCausalLM.forward() 中
if labels is not None:
x = logits[..., :-1, :].contiguous() # 去掉最后一个 token 的预测
y = labels[..., 1:].contiguous() # 去掉第一个 token 的标签
loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
对于 MoE 架构,还有额外的**辅助损失(aux_loss)**用于负载均衡:
python
# 总 loss = 主 loss + MoE 辅助 loss
loss = res.loss + res.aux_loss
- Dense 模型 :
aux_loss = 0,仅有交叉熵损失 - MoE 模型 :
aux_loss = router_aux_loss_coef × load_balance_loss,默认router_aux_loss_coef = 5e-4
五、训练指标详解
5.1 核心 Loss 指标
预训练过程中,MiniMind 记录以下核心指标:
| 指标 | 含义 | 正常范围 |
|---|---|---|
loss |
总损失 = logits_loss + aux_loss |
逐步下降 |
logits_loss |
交叉熵损失(Next Token Prediction) | 逐步下降 |
aux_loss |
MoE 负载均衡辅助损失 | 接近 0(Dense 模型为 0) |
learning_rate |
当前学习率 | 余弦衰减 |
5.2 Loss 曲线

从曲线中可以观察到:
- 初始 Loss :约 8~9 (接近
-ln(1/vocab_size) = -ln(1/6400) ≈ 8.76,即随机猜测水平) - 中期 Loss :快速下降至 3~4 区间
- 最终 Loss :收敛至 2.5~3.0 左右
- 收敛速度:前几个 epoch 下降最快,后续趋于平缓
5.3 Loss 数值的含义
💡 如何理解交叉熵 Loss 的数值?
Loss ≈ ln(vocab_size):模型等价于随机猜测,还没学到任何东西Loss ≈ 3.0:平均每个 token 的预测,模型在约e^3 ≈ 20个候选中就能猜对Loss ≈ 1.0:模型对下一个 token 的预测已经相当精准Loss < 0.5:通常意味着过拟合或数据泄露
对于 MiniMind(vocab_size=6400),随机猜测的 Loss 为 ln(6400) ≈ 8.76。预训练后 Loss 收敛到 2.5~3.0,意味着模型从"在 6400 个词中盲猜"进化到了"在约 12~20 个候选中就能选对"。
5.4 MoE 辅助损失
对于 MoE 模型,aux_loss 用于确保各专家的负载均衡,避免所有 token 都被路由到同一个专家:
python
# 负载均衡损失计算
load = F.one_hot(topk_idx, num_experts).float().mean(0) # 各专家被选中的频率
aux_loss = (load * scores.mean(0)).sum() * num_experts * router_aux_loss_coef
router_aux_loss_coef默认为 5e-4,是一个非常小的权重- 正常训练中
aux_loss应保持接近 0,如果突然增大说明路由出现严重不均衡
5.5 训练日志示例
训练过程中,每 log_interval(默认 100)步打印一次日志:
text
Epoch:[1/2](100/5483), loss: 6.2341, logits_loss: 6.2341, aux_loss: 0.0000, lr: 0.00027498, epoch_time: 12.3min
Epoch:[1/2](200/5483), loss: 5.1027, logits_loss: 5.1027, aux_loss: 0.0000, lr: 0.00027481, epoch_time: 11.8min
Epoch:[1/2](300/5483), loss: 4.5123, logits_loss: 4.5123, aux_loss: 0.0000, lr: 0.00027447, epoch_time: 11.5min
...
Epoch:[1/2](5000/5483), loss: 2.8945, logits_loss: 2.8945, aux_loss: 0.0000, lr: 0.00005123, epoch_time: 3.2min
Epoch:[2/2](5483/5483), loss: 2.6712, logits_loss: 2.6712, aux_loss: 0.0000, lr: 0.00005000, epoch_time: 0.1min
六、训练开销参考
6.1 硬件与时间
训练开销参考(单卡 3090):
| 模型 | 数据集 | 预训练时间 | 预训练费用 |
|---|---|---|---|
| minimind-3 (64M) | pretrain_t2t_mini |
≈1.21 小时 | ≈1.57 元 |
| minimind-3-moe (198M-A64M) | pretrain_t2t_mini |
≈1.69 小时 | ≈2.20 元 |
💡 3090 租卡单价约
1.3¥/h(7¥ ≈ 1 美元)
6.2 从零到 Zero 模型的完整开销
| 阶段 | 数据集 | 时间(minimind-3) | 费用 |
|---|---|---|---|
| 预训练 | pretrain_t2t_mini |
≈1.21h | ≈1.57¥ |
| SFT | sft_t2t_mini |
≈1.10h | ≈1.43¥ |
| 合计 | mini 组合 | ≈2.31h | ≈3.0¥ |
七、预训练结果验证
7.1 推理测试
预训练完成后,可以使用 eval_llm.py 对预训练结果做简单测试:
bash
# 测试预训练权重
python eval_llm.py --weight pretrain
预期输出示例:
text
💬: 为什么天空是蓝色的
🧠: 天空之所以看起来是蓝色的,主要是因为太阳光进入大气层后,短波长的蓝光更容易被空气分子散射,因此人眼从各个方向接收到的蓝光会更多。
💬: 解释什么是机器学习
🧠: 机器学习是人工智能的一个重要分支,它通过数据训练模型,使系统能够自动学习规律,并在分类、预测、推荐、自然语言处理等任务中持续改进效果。
⚠️ 预训练模型只具备"续写"能力,还不具备对话/问答能力。如需对话,需继续进行 SFT 微调。
7.2 权重文件
训练后的模型权重保存在 ./out/ 目录:
./out/
├── pretrain_768.pth # Dense 模型预训练权重
├── pretrain_768_moe.pth # MoE 模型预训练权重(如启用 MoE)
权重命名规则:{save_weight}_{hidden_size}[_moe].pth
pretrain:预训练权重768:隐藏层维度_moe:MoE 架构后缀(仅当--use_moe 1时)
八、断点续训
所有训练脚本均支持检查点保存与恢复,添加 --from_resume 1 参数即可自动检测并恢复训练进度:
bash
# 启用断点续训
python train_pretrain.py --from_resume 1
续训机制说明:
- 训练过程自动在
./checkpoints/目录保存完整检查点(模型、优化器、训练进度等) - 检查点文件命名:
{weight}_{dim}_resume.pth(如pretrain_768_resume.pth) - 支持跨 GPU 数量恢复:自动调整 step 计数
- 支持 wandb 训练记录连续性:自动恢复同一个 run
九、完整训练流程图
┌─────────────────────────────────────────────────────────────────┐
│ MiniMind 预训练完整流程 │
├─────────────────────────────────────────────────────────────────┤
│ │
│ 1. 数据准备 │
│ ┌──────────────────┐ │
│ │ pretrain_t2t_ │ ┌────────────────┐ │
│ │ mini.jsonl (1.2G)│────▶│ PretrainDataset│ │
│ │ (text格式数据) │ │ - BOS/EOS添加 │ │
│ └──────────────────┘ │ - Padding │ │
│ │ - Label掩码 │ │
│ └───────┬────────┘ │
│ │ │
│ 2. 模型初始化 ▼ │
│ ┌──────────────────┐ ┌────────────────┐ │
│ │ MiniMindConfig │────▶│MiniMindForCausal│ │
│ │ dim=768, layers=8│ │ LM (64M) │ │
│ └──────────────────┘ └───────┬────────┘ │
│ │ │
│ 3. 训练配置 ▼ │
│ ┌──────────────────┐ ┌────────────────┐ │
│ │ lr=5e-4 │ │ AdamW 优化器 │ │
│ │ cosine schedule │────▶│ 梯度累积 ×8 │ │
│ │ bfloat16 混合精度 │ │ 梯度裁剪 1.0 │ │
│ └──────────────────┘ └───────┬────────┘ │
│ │ │
│ 4. 训练循环 ▼ │
│ ┌──────────────────────────────────────────┐ │
│ │ for epoch in range(2): │ │
│ │ for batch in loader: │ │
│ │ loss = CE(logits, labels) + aux_loss │ │
│ │ loss.backward() │ │
│ │ clip_grad_norm_(1.0) │ │
│ │ optimizer.step() │ │
│ └─────────────┬────────────────────────────┘ │
│ │ │
│ ▼ │
│ 5. 输出 │
│ ┌──────────────────┐ │
│ │ pretrain_768.pth │ ← Loss 从 ~8.76 降至 ~2.7 │
│ │ (预训练权重) │ ← 约 1.21 小时 / 单卡 3090 │
│ └──────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────────┘
十、关键代码文件索引
| 文件 | 作用 |
|---|---|
trainer/train_pretrain.py |
预训练主脚本 |
model/model_minimind.py |
模型定义(MiniMindConfig + MiniMindForCausalLM) |
dataset/lm_dataset.py |
数据加载(PretrainDataset) |
trainer/trainer_utils.py |
训练工具函数(学习率调度、检查点、分布式) |
model/tokenizer.json |
BPE 分词器词表 |
eval_llm.py |
推理测试脚本 |
总结
本文核心要点:
- 预训练目标:通过 Next Token Prediction 让模型学会"高质量词语接龙",从随机猜测(Loss ≈ 8.76)收敛到有意义的预测(Loss ≈ 2.7)
- 极低门槛:单卡 3090 + 1.21 小时 ≈ 1.57 元,即可完成 MiniMind-3 的预训练
- 全程透明:所有核心代码基于 PyTorch 原生实现,数据加载、损失计算、学习率调度均可追踪
- 关键超参:等效 batch size=256,学习率 5e-4 余弦退火,梯度累积 8 步,bfloat16 混合精度
- 预训练只是起点:预训练模型只具备续写能力,后续还需 SFT 微调才能成为对话模型
📌 项目地址 :https://github.com/jingyaogong/minimind
📂 数据集下载 :ModelScope | HuggingFace