从零到一训练一个 0.6B 的 MoE 大语言模型

大家好。今天，我想和大家分享一个我个人从零开始训练的 MoE（Mixture-of-Experts）大语言模型项目——Cortex。这个项目不仅包含了一个 0.6B 参数量的模型，更涵盖了从数据准备、预训练、指令微调、思维链训练，到 DPO 和 GRPO 偏好对齐的全套流程和代码。

项目的初衷是完整地走一遍现代 LLM 的训练全流程，并把这个过程中的经验和成果分享出来。希望这个项目能为对 LLM 训练感兴趣的同学提供一些参考和便利。

项目亮点速览：

GitHub 源码：https://github.com/qibin0506/Cortex
在线体验：https://s.c1ns.cn/cortex (服务可能不稳定，推荐本地部署)
模型下载 (ModelScope)：https://www.modelscope.cn/models/qibin0506/Cortex

---

Cortex 是什么？

Cortex 是一个拥有 0.6B（6亿）总参数量，但在推理时仅激活 0.2B（2亿）参数的 稀疏混合专家（MoE）模型。MoE 架构允许模型在保持较低计算成本的同时，扩展其总参数量，从而获得更强大的能力。

除了模型本身，Cortex 项目最大的特点是其训练过程的完整性和透明性。我将整个训练过程分为了五个核心阶段，并为每个阶段提供了独立的训练脚本和数据支持：

1. 预训练 (Pretrain)
2. 指令微调 (SFT)
3. 思维链训练 (Reasoning)
4. 直接偏好优化 (DPO)
5. 分组奖励策略优化 (GRPO)
   无论你是想直接体验最终模型，还是想深入其中某个训练环节，都能在项目中找到你需要的资源。

技术核心：MoE 架构与训练细节

让我们深入了解一下 Cortex 的技术实现。

1、模型架构
模型的核心配置定义在 utils.py 文件中。

# From utils.py
def get_model_config():return ModelConfig(vocab_size=TrainerTools().tokenizer.vocab_size,hidden_size=768,intermediate_size=2048,moe_intermediate_size=1024,moe_n_dense_layer=1,num_hidden_layers=24,num_attention_heads=12,num_key_value_heads=4,max_position_embeddings=max_seq_len,attention_implementation='auto',rope_config=RoPEConfig(rope_theta=1e6),moe_config=MoEConfig(num_experts_per_tok=2,n_routed_experts=8,n_shared_experts=1,aux_loss_alpha=0.1,seq_aux=True,norm_topk_prob=True))

从配置中可以看到几个关键点：

MoE 配置：模型包含 8 个专家（n_routed_experts=8），每个 Token 的计算会由一个路由器（Router）选择最合适的 2 个专家（num_experts_per_tok=2）来处理。这种稀疏激活的机制正是 MoE 模型高效的原因。

• 注意力机制：采用了 Grouped-Query Attention (GQA) 的变体（num_attention_heads=12, num_key_value_heads=4），在保证性能的同时有效降低了 KV 缓存的占用。
• 位置编码：使用了 RoPE（旋转位置编码），并将 rope_theta 设置为 1e6，以增强模型在长文本上的表现。

2、完整的训练流程
Cortex 的训练是一场精心设计的“马拉松”，每个阶段都有其独特的使命。

第一阶段：预训练 (Pretrain)
这是模型学习语言基础的阶段。我使用了大量的中英文语料，通过 train_pretrain.py 脚本 对模型进行基础训练，使其掌握语言的统计规律。

第二阶段：指令微调 (SFT)
为了让模型听懂并遵循指令，我使用了大量的指令数据对其进行微调。train_sft.py 脚本负责这个过程，让模型从一个语言模型转变为一个对话助手。

第三阶段：思维链训练 (Reasoning)
为了提升模型的逻辑推理能力，我引入了思维链（Chain-of-Thought）的训练。通过 train_reasoning.py，模型被引导学习一种“思考->回答”的模式。这种模式使用了 和 等特殊 Token 来显式地分离思考过程和最终答案。

第四、五阶段：偏好对齐 (DPO & GRPO)
为了让模型的回答更符合人类的偏好（更有用、更无害），我先后使用了 DPO 和 GRPO 两种对齐技术。

DPO (train_dpo.py)：通过“更好”和“更差”的回答对比，直接优化模型，使其学会“择优录取”。
GRPO (train_grpo.py)：这是一个更有趣的阶段。我编写了一个自定义的奖励函数 reward_func，它会根据模型输出的格式规范性、思考过程的丰富度和答案的准确性来给出一个综合评分，从而更精细地引导模型的行为。

# Snippet from train_grpo.py
def get_reward(completion_text: str, correct_answer: str)-> float:# ...reasoning_score = min(2.0, len(reasoning_text) / 75.0)# ...if response_last_number == correct_answer:answer_score = 8.0# ...if answer_score > 0:reward = answer_score + reasoning_scoreelse:reward = reasoning_score * 0.5return reward

交互式 Web UI：洞察 AI 的“思考”

除了核心模型和训练代码，项目还提供了一个基于 Bottle 和 Tailwind CSS 构建的现代化 Web 界面，让你可以直观地与 Cortex 互动。

界面最大的亮点是**“思考模式”和“思考预算”**。

思考模式：开启后，Cortex 会在给出最终答案前，先展示一段它的“内心独白”（即标签中的内容）。这让你能清晰地看到它是如何一步步分析问题，并最终得出结论的，极大地增强了模型的可解释性。
思考预算：这个功能允许你为模型的“思考过程”设定一个计算量上限（单位：token），有助于在保证推理质量和控制响应时间之间找到平衡。

立刻上手体验

想亲手试试 Cortex 吗？非常简单！

1、本机部署（推荐）
README.md 中 提供了详细的部署步骤。总的来说，只需要四步：
安装依赖:

# 克隆项目
git clone https://github.com/qibin0506/Cortex.git
cd Cortex# 安装基础依赖
pip3 install -r requirements.txt# 安装作者封装的模型和训练库
pip3 install project_llm_model
pip3 install project_llm_trainer

下载模型: 从 ModelScope 下载 last_checkpoint.bin 文件，并将其放置在项目根目录。
运行后端: python3 app.py
访问: 打开浏览器，访问 http://0.0.0.0:8080/ 即可开始聊天。

2、自己动手训练
如果你想在我的基础上继续训练，或者体验某个特定的训练阶段，也非常方便。

项目中的 file_dataset.py 实现了一个非常智能的数据集管理类，它会自动从 ModelScope 下载所需的训练文件，并在训练过程中管理本地磁盘空间，你无需手动下载和整理数据。

训练命令非常直观（在安装好依赖后）：

# 1. 预训练
smart_train train_pretrain.py# 2. SFT
smart_train train_sft.py# 3. 推理能力训练
smart_train train_reasoning.py# 4. DPO
smart_train train_dpo.py# 5. GRPO
smart_train train_grpo.py

你可以在 utils.py 中调整超参数，以适应你自己的硬件环境。

结语

Cortex 项目是我在 LLM 领域学习和探索的一个阶段性总结。从零开始构建和训练一个模型，过程虽然充满挑战，但也收获巨大。我将它完全开源，希望能降低大家进行 LLM 研究和实践的门槛。
GitHub 源码：https://github.com/qibin0506/Cortex