MiniMind - 微型大语言模型

基于 PyTorch 从零实现的微型大语言模型，完整支持预训练、SFT 微调和推理部署。

特性

• 完整架构实现: RoPE、RMSNorm、SwiGLU、GQA
• KV-Cache 优化: 推理吞吐量提升 ~7x
• MoE 支持: 混合专家模型
• 全链路训练: 预训练 → SFT → RLHF
• 轻量高效: 64M 参数，单卡可训练

模型规模

模型	参数量	Q头	KV头	层数
minimind_64m	52M	8	4	8
minimind_26m	26M	8	2	8
minimind_198m_moe	198M	8	4	8

安装

pip install torch numpy tqdm

快速开始

测试模型

python demo.py

预训练

cd trainer
python pretrain.py --hidden_size 768 --batch_size 8

SFT 微调

cd trainer
python sft.py --pretrain_weights ../out/pretrain_768.pth

对话推理

python inference/chat.py

项目结构

mini_llm_project/
├── model/
│   ├── config.py      # 模型配置
│   └── model.py       # 核心架构
├── trainer/
│   ├── pretrain.py   # 预训练
│   └── sft.py        # 监督微调
├── inference/
│   └── chat.py       # 对话推理
└── demo.py           # 测试演示

核心组件

RoPE (旋转位置编码)

from model import precompute_freqs_cis, apply_rotary_pos_emb
cos, sin = precompute_freqs_cis(dim, max_len)
q, k = apply_rotary_pos_emb(q, k, cos, sin)

GQA (分组查询注意力)

from model import Attention, MiniMindConfig
config = MiniMindConfig(num_attention_heads=8, num_key_value_heads=4)
attention = Attention(config)

KV-Cache

logits, past_kv = model(input_ids, past_key_values=past_kv, use_cache=True)

MoE

from model import MOEFeedForward, MiniMindConfig
config = MiniMindConfig(use_moe=True, num_experts=4)
moe = MOEFeedForward(config)

性能

• KV-Cache 加速: ~7x
• 单卡训练: 3090 约 2-3 小时
• 推理显存: ~2GB

参考

基于 minimind 设计理念实现。