Skip to content

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

Description

@LiuCharming

▎ 分支地址: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

▎ 欢迎大家试用和提建议!

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions