MiniMind 预训练详解:从零训练一个 64M 参数的语言模型

📌 本文亮点:深入解析 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=8kv_heads=4
  • max_position_embeddings=32768rope_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_modeln_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.552.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¥/h7¥ ≈ 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 推理测试脚本

总结

本文核心要点:

  1. 预训练目标:通过 Next Token Prediction 让模型学会"高质量词语接龙",从随机猜测(Loss ≈ 8.76)收敛到有意义的预测(Loss ≈ 2.7)
  2. 极低门槛:单卡 3090 + 1.21 小时 ≈ 1.57 元,即可完成 MiniMind-3 的预训练
  3. 全程透明:所有核心代码基于 PyTorch 原生实现,数据加载、损失计算、学习率调度均可追踪
  4. 关键超参:等效 batch size=256,学习率 5e-4 余弦退火,梯度累积 8 步,bfloat16 混合精度
  5. 预训练只是起点:预训练模型只具备续写能力,后续还需 SFT 微调才能成为对话模型

📌 项目地址https://github.com/jingyaogong/minimind

📂 数据集下载ModelScope | HuggingFace


相关推荐
下班走回家1 小时前
Qwen2.5 模型架构解读:国产大模型的进化
人工智能·架构
跟风舞烟学编程1 小时前
Hermes Agent 从入门到企业实战-01:Hermes-Agent核心架构
人工智能·ai agent·hermes agent·自进化 agent
深圳市晶科鑫实业有限公司1 小时前
国产TCXO温补晶振是否可以完美替代欧美日系主流型号
人工智能·stm32·单片机·物联网·51单片机·信息与通信
cyyt1 小时前
深度学习周报(6.8~6.14)
人工智能·深度学习
带娃的IT创业者1 小时前
深度解析:当 MLX 遇上视觉语言模型,Mac 本地推理的新范式
人工智能·macos·语言模型·视觉语言模型·apple silicon·mlx·mac本地推理
沪漂阿龙1 小时前
LangChain 系列之Tools:让大模型真正连接业务系统
人工智能·python·langchain
竹叶青lvye1 小时前
ROS2自定义接口消息、参数服务案例
人工智能·ros2·具身智能·接口消息·参数服务
AI科技星1 小时前
数术工坊·第八卷 大道归一录・番外・下篇 零界封神・万法归元终章
网络·人工智能·算法·几何学·拓扑学