▎ 分支地址: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
▎
▎ 欢迎大家试用和提建议!
▎ 分支地址: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
▎
▎ 欢迎大家试用和提建议!