从零到一训练一个 0.6B 的 MoE 大语言模型
大家好。今天,我想和大家分享一个我个人从零开始训练的 MoE(Mixture-of-Experts)大语言模型项目——Cortex。这个项目不仅包含了一个 0.6B 参数量的模型,更涵盖了从数据准备、预训练、指令微调、思维链训练,到 DPO 和 GRPO 偏好对齐的全套流程和代码。
项目的初衷是完整地走一遍现代 LLM 的训练全流程,并把这个过程中的经验和成果分享出来。希望这个项目能为对 LLM 训练感兴趣的同学提供一些参考和便利。
项目亮点速览:
在线体验:https://s.c1ns.cn/cortex (服务可能不稳定,推荐本地部署)
GitHub 源码:https://github.com/qibin0506/Cortex
模型下载 (ModelScope):https://www.modelscope.cn/models/qibin0506/Cortex
Cortex 是什么?
Cortex 是一个拥有 0.6B(6亿)总参数量,但在推理时仅激活 0.2B(2亿)参数的 稀疏混合专家(MoE)模型。MoE 架构允许模型在保持较低计算成本的同时,扩展其总参数量,从而获得更强大的能力。
除了模型本身,Cortex 项目最大的特点是其训练过程的完整性和透明性。我将整个训练过程分为了五个核心阶段,并为每个阶段提供了独立的训练脚本和数据支持:
- 预训练 (Pretrain)
- 指令微调 (SFT)
- 思维链训练 (Reasoning)
- 直接偏好优化 (DPO)
- 分组奖励策略优化 (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