Skip to content

[DEV] Add MoE inference and expert parallel support #443

Description

@qinyiqun

目标版本
main

功能描述
为 InfiniLM 增加通用 MoE 推理模块,并接入 Qwen3-MoE 的 TP/EP 推理路径。

主要目标:

  1. 抽象通用 MoE layer。

    • Router
    • Experts
    • Runner
    • Dispatcher
    • EP backend
    • Legacy MoE compatibility
  2. 支持 Qwen3-MoE 推理。

    • 复用通用 SparseMoeBlock
    • 支持 TopKRouter
    • 支持 packed expert weights
    • 支持 fused MoE runner
  3. 支持基础 Expert Parallel。

    • local_allreduce
    • allgather_reducescatter
    • deepep placeholder
    • expert weight sharding
    • expert id remap
  4. 优化 MoE 大模型加载。

    • 使用 model.safetensors.index.json 做加载计划
    • rank-local 参数过滤
    • EP expert 权重只加载到对应 rank
  5. 保持兼容。

    • 非 MoE 模型默认禁用 EP
    • DeepSeek-V2 继续使用 legacy MoE
    • Qwen3/Qwen3Next 支持 num_key_value_heads < tp_size 的 GQA TP 场景

预期结果:

  • Qwen3-30B-A3B 可以在 TP=2/EP=2 下推理。
  • Qwen3-235B-A22B 可以在 TP=8/EP=8 下加载和推理。
  • MoE EP 路径可以后续替换为 DeepEP 或更高性能 fused MoE kernel。

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions