- [2026-06-13] 新增架构实验，提供一个可交互的环境用于快速测试不同架构设计的预训练 loss
- [2026-05-21] 实现
mini_deepseekv4模型 - [2026-04-27] 增加 triton 实现的 Flash Attention 前向
- [2026-04-19] 增加 YaRN、DPO、GRPO
- [2026-02-01] 实现
mini_qwen3_next模型；优化多轮对话数据构造；优化mini_models结构。 - [2025-12-29] 完成兼容 transformers 的项目重构；实现
mini_llama3、mini_deepseekv3模型；实现 Pretrain、SFT。
本项目旨在基于较小的算力，复现当前主流开源模型的架构，实现一个 100-200M 参数量版本的迷你模型。项目将数据集、训练流程等基础设施尽可能固定下来，以便在学习新的模型架构时能够快速复现，从而将主要精力聚焦在模型架构的学习和复现上。
主要目标：
- 学习并复现当前主流开源模型架构
- 从零实现常用的训练和推理流程
为了实现这一目标，在先前版本的 Mini-LLM 中，我们完全自定义实现了 model 包，其中包括 BaseModel 和 BaseModelArgs 等基类。后来发现，这样的构建思路与 transformers 库的 PreTrainedModel 和 PretrainedConfig 类似。基于这种相似性，为了更好地与 HuggingFace 生态兼容，我们直接重构了项目结构。当前版本实现的模型完全兼容 transformers 库，可以直接使用 from_pretrained、generate 等方法进行模型加载和推理。同时，为了深入理解训练和推理原理，项目仍然提供了一套独立的训练代码和生成代码实现。早期版本的 Mini-LLM 已移动到 legacy 分支。
本项目建立时，是基于 transformers 4.x 版本，当前 transformers 已处于 5.x 版本，transformers 中有一些函数签名变化了，尚不确定会不会存在兼容性问题，后续会逐渐过渡到 transformers 5.x 接口，目前固定到 4.56.1 版本保证兼容性。
- 克隆项目
git clone https://github.com/WKQ9411/Mini-LLM.git
cd Mini-LLM- 初始化环境
执行以下脚本，将自动检测环境并安装依赖。模型训练需要 CUDA 环境，如果仅尝试使用训练好的权重进行推理，可以使用 CPU 环境。
# Linux
bash ./scripts/setup.sh
# Windows
.\scripts\setup.ps1目前使用的数据集包括：
主要使用该数据集中得分为4-5分的高质量语料部分，共9745个 parquet 文件（编号是从0-9766，中间似乎有缺失编号，但从 ModelScope 上确实下载的是9745个文件），总共约70GB。主要用途：
- 使用该数据集的5%采样作为 tokenizer 训练数据
- 使用该数据集的20%采样作为预训练数据
- 如果算力足够，也可采用全部数据进行预训练
- 使用该数据集的0.1%采样作为YaRN微调数据
大约16GB，主要用途：
- 使用该数据集的全部作为预训练数据
- 采样50000条作为 SFT 数据
使用其中的 ultrafeedback-chinese-binarized-lowest 子集作为 DPO 数据集，可根据 DPO 评分范围筛选训练子集。
- 自行合成的 GRPO 数据集。
数据集由 Gemini 3.1 Pro 合成，任务是修复或修改 JSON 格式，分为 prompt、thinking、response 三个部分。prompt 用于给出一个有错误或需修改的简单 JSON，thinking 是合成的思维链，response 是不包含额外解释的 JSON 输出。
数据集相关脚本位于scripts文件夹中，可以通过以下命令下载需要的数据集：
# Linux
bash ./scripts/download_data.sh
# Windows
.\scripts\download_data.ps1界面如下图所示，可选择需要下载的数据集序号进行下载：
其中：
- [1] 【Tokenizer】下载按照5%的比例对OpenCSG Fineweb-Edu-Chinese-V2.1数据集进行采样的.parquet格式数据子集，用于训练tokenizer（你也可以直接使用训练好的tokenizer，文件位于项目的
mini_tokenizer文件夹内） - [2] 【Pretrain】下载原始OpenCSG Fineweb-Edu-Chinese-V2.1数据集中得分4-5的全部.parquet文件，用于预训练
- [3] 【Pretrain】下载按照20%的比例对OpenCSG Fineweb-Edu-Chinese-V2.1数据集进行采样的.parquet格式数据子集，用于预训练，采样按照类别等比例采样，分布与原始数据集一致：
- [4] 【Pretrain】下载按照20%的比例对OpenCSG Fineweb-Edu-Chinese-V2.1数据集进行采样的.bin格式数据子集，用于预训练（已通过
mini_tokenizer处理为token ids） - [5] 【Pretrain】下载OpenCSG Fineweb-Edu-Chinese-V2.1数据集中得分4-5的全部.bin格式数据文件，用于预训练（已通过
mini_tokenizer处理为token ids） - [6] 【Pretrain】下载匠数科技大模型数据集的全部.bin格式数据文件，用于预训练（已通过
mini_tokenizer处理为token ids） - [7] 【YaRN】下载按照0.1%的比例对OpenCSG Fineweb-Edu-Chinese-V2.1数据集进行采样的.bin格式数据子集，用于YaRN微调（已通过
mini_tokenizer处理为token ids） - [8] 【SFT】下载原始匠数科技大模型数据集的全部.jsonl格式数据文件，用于SFT
- [9] 【SFT】下载经过处理的匠数科技大模型数据集的.parquet格式数据文件，用于SFT（已通过
mini_tokenizer处理为token ids，其中包括：(a) 转化为parquet格式的全部符合条件的SFT数据；(b) 采样的50000条采样数据和200条自我认知数据；(c) 将 (b) 进行packing处理后的数据；推荐使用 (c) 进行SFT），采样后的数据长度分布如下：
- [10] 【DPO】下载经过处理的DPO数据集
- [11] 【GRPO】下载合成的GRPO数据集
- [12] 【Architecture Lab】架构实验室数专用数据集，从 OpenCSG Fineweb-Edu-Chinese-V2.1数据集中采样 1% 并预处理。
你可以选择直接下载处理好的数据进行训练（推荐），也可以选择下载原始数据自行处理，处理数据的代码位于：
./scripts/prepare_tokenizer_data.py
./scripts/prepare_pretrain_data.py
./scripts/prepare_sft_data.py
./scripts/prepare_dpo_data.py
./scripts/prepare_grpo_data.py
./scripts/prepare_architecture_lab_data.py本项目当前使用：[1] 训练 tokenizer，[4]+[6] 进行预训练（通过prepare_pretrain_data.py中的merge_pretrain_data函数合并.bin文件），[9]中的(c)进行 SFT。
新版本的 mini_tokenizer 与 Qwen 保持一致，使用的特殊 token 包括：<|endoftext|>, <|im_start|>, <|im_end|>, <think>, </think>。
基础词表大小为32000（包括<|endoftext|>），<|im_start|>, <|im_end|>, <think>, </think>作为 added tokens 添加到词表中，因此词表大小为32004。tokenizer 使用方法可以参考example/tokenizer_example.ipynb。聊天模板位于data/tokenizer_data/chat_template.jinja2。
可直接使用训练好的mini_tokenizer，或重新训练。重新训练执行：
python ./train/train_tokenizer.py如果需要重新训练 tokenizer，建议确保 CPU 具有较大的 RAM。如果5%采样数据对于所训练的 tokenizer 仍然较大，可以在scripts/prepare_tokenizer_data.py中使用更小的sample_ratio采样更小一点的 tokenizer 数据集。
模型结构参考论文、官方仓库源码、transformers 实现等，其中的hidden_states形状统一为: (B, H, L, D)，其中B为 batch size，H为头数，L为序列长度，D为每个头的维度。
模型结构参考本人 GitHub 博客：
博客中的
mini_llama3和mini_deepseekv3的代码部分是基于早期版本的 Mini-LLM 实现的，与当前版本不完全一致，但核心思想是相同的。
mini_llama3，Dense Model:mini_deepseekv3，MoE Model:mini_qwen3_next，Linear Model:mini_deepseekv4，MoE Model:
使用单卡进行训练：
python ./train/pretrain.py --model_name=mini_deepseekv3 --max_batch_size=32使用 DDP 进行训练：
CUDA_VISIBLE_DEVICES=0,1 torchrun --nproc_per_node=2 ./train/pretrain.py --model_name=mini_deepseekv3 --max_batch_size=32更多训练参数说明，请参考train/pretrain.py中的parse_args()函数。
训练开始后，在新的终端内开启 TensorBoard 用于监控训练进度：
tensorboard --logdir=output/如果使用的是云服务器，根据不同平台使用文档，配置 TensorBoard 端口等参数并设置公开访问等，从而能够在本地监控训练进度，例如：
tensorboard --logdir=output/ --port=8080 --bind_all训练会记录通用的learning_rate、loss、ppl等指标，此外，以mini_deepseekv3模型为例，还额外记录了包括专家负载情况、序列级辅助损失、MTP 损失等指标，如下图所示：
其中，专家负载曲线通过记录每层的所有专家最大/最小激活次数的比值，趋近于1表示负载均衡，越大表示负载不均衡。
由于模型参数量基本在100-200M，SFT 训练数据又相对较少，仅使用单卡训练即可：
python ./train/sft.py --model_name=mini_deepseekv3 --max_batch_size=32更多训练参数说明，请参考train/sft.py中的parse_args()函数。
SFT 数据集可选择是否使用 packing 数据集，开启 packing 后，能够有效利用算力，同时让每个 batch 的实际有效 token 长度尽可能一致，从而避免梯度稀释问题。使用 packing 后，每个 batch 需要构造相应的attention_mask，可视化如下（pack了两条数据）：
packing 后，SFT 曲线相对更加平滑。
- packing 的曲线：
- 未 packing 的曲线：
由于 Linear Model 的特点，当前暂不支持 mini_qwen3_next 的 packing SFT，见Issues #3。 另外，由于 DeepSeek-V4 Compressor 的特性，packing 数据需要使用非常复杂的逻辑来处理不同 packed sample 之间的 compressing 情况，因此 mini_deepseekv4 目前同样不支持 packing SFT。
关于 YaRN 的理论部分，可以参考博客：YaRN 论文解读
在模型配置中传入 rope_scaling 参数，例如：
rope_scaling = {
"rope_type": "yarn",
"factor": 4.0,
"attention_factor": None, # 默认为 None，内部自动计算
"beta_fast": 32,
"beta_slow": 1,
}可选对少量长文本进行微调：
python ./train/yarn.py --model_name mini_llama3 --max_seq_len 2048更多训练参数说明，请参考 train/yarn.py 中的 parse_args() 函数。
对加入YaRN之前、加入YaRN但不微调、加入YaRN并微调三种情况计算长文本的PPL：
python ./eval/eval_yarn.py --base_model_path output/pretrained_mini_llama3 --yarn_finetuned_model_path output/yarn_mini_llama3结果如下：
可见加入YaRN后，相比原始模型，长文本的PPL有明显下降。此外，加入YaRN并微调后，长文本的PPL进一步下降，但在微调长度范围之外，PPL会逐渐上升，这可能是因为微调后模型学习到了新的位置编码语义，其外推能力反而弱于仅插入 YaRN 而不微调的模型。因此，若目标是固定长度范围内的长文本建模，可以选择对少量长文本继续微调；若更关注超过训练长度后的外推表现，则可以只在推理或评估时加入 YaRN，而不进行额外微调。
关于 DPO 的理论部分，可以参考博客：DPO 论文解读
运行如下代码：
python ./train/dpo.py --model_name mini_llama3更多训练参数说明，请参考 train/dpo.py 中的 parse_args() 函数。由于小模型DPO很容易训崩，因此通常使用较小的学习率和参数
其中，DPOP 是 DPO 的一个变体，在 DPO Loss 后加入了一个针对 chosen response 的正样本约束项：当 policy model 对 chosen response 的 log probability 低于 reference model 时，额外加入 lambda * max(ref_chosen_logp - policy_chosen_logp, 0) 作为惩罚，从而避免模型为了拉大 chosen 和 rejected 的偏好差距而降低 chosen 本身的概率。在本项目的小模型设置下，DPOP 通常比直接 DPO 更稳定，可以通过 --loss_type dpop --dpop_lambda 1.0 开启。
DPO 训练结果如下：
具体结果对比请参考 eval/eval_dpo.ipynb。
关于 GRPO 的理论部分，可以参考博客：从策略梯度到PPO再到GRPO
本项目中的 GRPO 示例任务是 JSON 修复或修改：模型根据 prompt 中给出的错误 JSON 或修改要求，先在 <think>...</think> 中给出简短思考过程，再在 JSON 代码块中输出最终 JSON。冷启动时，为了使模型仍然保留一定的指令遵循能力，避免模型输出迅速坍缩到 JSON 任务的狭窄模式里，加入了一部分通用数据。最终，冷启动数据包含800条通用数据和1200条 JSON 任务数据。
当前奖励函数主要由四部分组成：输出格式奖励、思考长度奖励、JSON 可解析奖励和正确性奖励。其中正确性奖励权重最高，用于鼓励模型输出与标准答案一致的 JSON；其余奖励用于约束输出结构，避免小模型在 RL 阶段产生不可解析或格式混乱的结果。
运行如下代码：
python ./train/grpo.py --model_name mini_llama3 --max_batch_size 4 --cold_start_sft --sft_batch_size 16 --sft_epochs 3 --grpo_epochs 2更多训练参数说明，请参考 train/grpo.py 中的 parse_args() 函数。默认情况下，GRPO 会从 output/sft_{model_name} 加载 SFT 后的模型作为初始 policy；开启 --cold_start_sft 后，会先使用 data/grpo_data/cold_start.jsonl 做一轮任务格式冷启动 SFT，再进入 GRPO 训练。若希望在 JSON 任务语料上继续预训练，也可以额外开启 --mid_training（不过可能影响不大）。
训练结果如下：
从结果看，GRPO 后模型在格式遵循和 JSON 可解析性上进一步提升，同时正确率相比仅进行 cold-start SFT 有明显提高。具体结果对比请参考 eval/eval_grpo.ipynb。
相关推理 demo 的代码位于example文件夹中，可以使用项目自定义的Generator类进行推理，也可以使用 transformers 原生的generate方法进行推理。
在终端中运行：
python ./example/test_terminal.py --model_name=mini_deepseekv3更多推理参数说明，请参考example/test_terminal.py中的parse_args()函数。
也可以通过 API 方式进行推理，提供给流行的前端进行对话（通过wrap-openai封装 OpenAI 兼容 API，可参考我的另一个仓库wrap-openai），通过如下命令启动后端：
python ./example/test_api.py --model_name=mini_deepseekv3以CherryStudio为例，配置好OpenAI兼容API后，对话效果如下：
此外，本项目的模型参数已上传至 HuggingFace，可直接下载使用，调用方法见example/use_example.ipynb。
由于模型参数量较小，虽然可能一定程度上较好的预测下一个token，但是并不等同于它具备了良好的泛化能力、知识储备或推理能力。小模型更容易“记住”训练数据中的表面模式（比如特定短语、句子结构、格式），而不是真正“理解”其含义。这导致它们在面对需要知识、推理或稍微偏离训练模式的prompt时，容易产生幻觉和不连贯的输出。
添加 --enable_flash_attention 参数后，可以启用 Triton 实现的 Flash Attention 前向推理，Flash Attention 原理可参考博客：Flash Attention。在 eval/eval_flash_attention.ipynb 中与 naive pytorch 实现的 attention 实现进行了对比，效果如下：
由于我们在 example 中的推理示例仅进行单 batch 的短序列推理，因此体感差异不大。
需确保安装 Node.js，建议 20 及以上版本，并执行 scripts 中的 setup，然后使用如下命令启动：
# Linux
bash ./architecture_lab/start.sh
# Windows
.\architecture_lab\start.ps1为快速进行架构实验，相比 mini_tokenizer，专门设置了 vocab_size=10003 的更小的 tokenizer，以减小 embedding 和 lm_head 的参数量。如果需要调整可以在 scripts/prepare_architecture_lab_data.py 中调整这一设置（需重新训练 tokenizer）。另外需确保使用 scripts 中的 download_data 脚本下载架构实验室专用数据集，脚本会自动下载数据至默认路径。
在架构实验室中，可以自由组合当前已实现的主要模块，构建自己的模型架构并进行预训练 loss 测试：