#801·minimind

[Feature] 独立可插拔 MoE V2 模块 + moe_type 切换机制

Author: LiuCharmingCreated Jun 23, 2026Updated Aug 26, 2026

▎ 分支地址:https://github.com/LiuCharming/minimind/tree/feature/moe-v2 ▎ ▎ 改动概述 ▎ ▎ 从 language_model_moe.py 中提取了一个独立、可插拔的 MoE 模块 (model/moe.py,约670行),与原有 MoE 通过 moe_type 配置项切换,不影响现有功能。 ▎ ▎ 核心变更 ▎ ▎ 1. 新增 model/moe.py — 独立 MoE 模块 ▎ - 零项目依赖,仅需 torch,可单独复用到其他项目 ▎ - 包含:MoEConfig、SwiGLU FFN 专家、Constant/Copy/Zero 轻量专家、负载均衡器、张量化路由、MoEBlock 等 ▎ - 专家结构:[FFN × (N-2-C)] + [Constant × C] + [Copy × 1] + [Zero × 1] ▎ - 训练时按专家预排序 token → 批量处理 → index_add_ 聚合,高效且显存友好 ▎ - 支持 Mixtral 式 gating (top-k logits softmax) 和 Switch Transformer 式负载均衡 ▎ ▎ 2. 修改 model/model_minimind.py ▎ - MiniMindConfig 新增 moe_type 字段("v1"=原始 / "v2"=新 MoEBlock) ▎ - MiniMindBlock 自动按 use_moe + moe_type 选择 FFN/MoE-v1/MoE-v2 ▎ - 新增 MOEFeedForwardV2 包装类作为 drop-in replacement ▎ ▎ 3. 所有训练脚本增加 CLI 参数 ▎ - train_pretrain.py、train_full_sft.py、train_lora.py、train_dpo.py、train_ppo.py、train_grpo.py、train_agent.py、train_distillation.py 均已添加 --moe_type、--num_experts、--num_experts_per_tok ▎ ▎ 4. eval_llm.py 增加对应参数,确保推理时正确加载 V2 权重 ▎ ▎ 5. trainer_utils.py 增加路径解析,修复 ../model 相对路径在不同工作目录下的加载问题 ▎ ▎ 使用方式 ▎ ▎ # V2 MoE 预训练 ▎ python trainer/train_pretrain.py --use_moe 1 --moe_type v2 --num_experts 6 --num_experts_per_tok 1 ▎ ▎ # V2 MoE 推理 ▎ python eval_llm.py --use_moe 1 --moe_type v2 --num_experts 6 --num_experts_per_tok 1 --weight pretrain ▎ ▎ 欢迎大家试用和提建议!