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
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions