【升级版】从零到一训练一个 0.6B 的 MoE 大语言模型

前文：从零到一训练一个 0.6B 的 MoE 大语言模型，本次升级完全重新从零开始重新训练。主要升级如下：

1. 替换预训练数据集，使用序列猴子通用文本数据集进行预训练。
2. 使用更先进的训练方法。
3. 新增思考模式控制，可通过添加/think和/no think控制是否思考。
4. 新增思考预算功能，可控制思考token长度。

效果展示

快速开始

1. 确保本机已安装python3
2. clone或下载项目：https://github.com/qibin0506/Cortex
3. 安装依赖 pip3 install -r requirements.txt
4. 下载checkpoint https://www.modelscope.cn/models/qibin0506/Cortex-V2/resolve/master/dpo.bin，并放置到项目根目录
5. 执行 python3 app.py运行项目，访问链接http://0.0.0.0:8080/ 即可体验

训练细节

Cortex V2采用更加先进的训练方式进行训练，具体情况如下；

预训练

预训练过程采用两阶段训练模式

阶段1	阶段2
train_pretrain_stage0.py	train_pretrain_stage1.py
上下文长度为512，在较短训练文本上进行训练	采用YaRN技术将上下文扩展至2048，并在长文本序列上继续训练

后训练

后训练过程采用四阶段训练模式

阶段1	阶段2	阶段3	阶段4
train_cot.py	train_grpo.py	train_mix.py	train_dpo.py
在纯COT数据集上进行SFT，让模型原生支持思考模式	采用GSPO技术，提升模式的逻辑思考能力	使用COT和非COT混合数据集上进行SFT，让模式支持思考控制和思考预算能力	使用DPO进行对齐训练

继续训练

本项目提供各个阶段训练完成后的checkpoint, 可根据自己需求选择checkpoint继续训练。
checkpoint下载：https://www.modelscope.cn/models/qibin0506/Cortex-V2/files
训练方式：

1. 确定继续训练的阶段，修改file_dataset.py中对应阶段的FileDataset中的文件，然后使用smart_train进行训练，例如重新进行dpo，则执行smart_train train_dpo.py
2. 本项目GSPO阶段是在4x5090进行训练，其他阶段都是在4x4090进行训练，同时utils.py中的配置数据也是按照对应硬件配置确定，如有不同的训练设备可自行修改utils.py进行适配。
3. file_dataset.py文件用来管理数据集文件，可按需修改，数据集文件会自动下载，使用完成后会自动删除，无需人工管理。

详细介绍参考项目Github: https://github.com/qibin0506/Cortex