专家并行 (EP)

专家并行将混合专家模型（MoE）的专家分布到多个 GPU 上，每个 rank 只持有部分专家。这降低了单卡显存占用，使大规模 MoE 模型的训练成为可能。

概览

概念	说明
ExpertParallelConfig	控制 EP 行为的配置数据类
apply_expert_parallel()	入口函数，负责分片专家并替换前向传播
shard_experts()	将专家均匀分配到各 EP rank
patch_forward()	将 MoE block 的 forward 替换为带 all-to-all 通信的 EP 版本

配置

from twinkle.model.transformers.moe.expert_parallel import ExpertParallelConfig

config = ExpertParallelConfig(
    enabled=True,              # 启用专家并行
    router_dtype='fp32',       # 路由计算精度：'fp32', 'bf16', 'fp16'
    keep_router_logits=True,   # 在输出中保留路由 logits
    ignore_shared_experts=False,# 跳过共享专家计算（如 DeepSeek）
    ep_size=None,              # EP 并行度（由 TransformersModel 使用）
)

配合 DeviceMesh 使用

在 DeviceMesh.from_sizes() 中设置 ep_size 即可激活 EP。框架会在模型初始化时自动调用 apply_expert_parallel()。

from twinkle.utils import DeviceMesh

# 8 卡：2 路 EP × 4 路数据并行
device_mesh = DeviceMesh.from_sizes(
    world_size=8,
    dp_size=4,
    ep_size=2,
)

EP + FSDP 组合分片：

# 8 卡：2 路 EP，每个 EP 组内 2 路 FSDP
device_mesh = DeviceMesh.from_sizes(
    world_size=8,
    dp_size=2,
    ep_size=2,
    ep_fsdp_size=2,
)

通信模式

EP 前向传播遵循 4 阶段流水线：

预处理 — 计算每个专家的 token 数量和分割大小
Token Pre-All2All — 按专家分配排列 token，然后在 EP rank 间执行 all-to-all 交换
专家计算 — 每个 rank 在接收到的 token 上运行本地专家
Token Post-All2All — all-to-all 交换结果，反排列并应用路由权重

输入 token → 路由器 → [预处理] → [pre_all2all] → [本地专家] → [post_all2all] → 输出

要求

num_experts 必须能被 ep_size 整除
torch.distributed 必须已初始化
MoE block 必须定义 gate/router 模块和 experts（支持 nn.ModuleList 或张量形式的 gate_up_proj/down_proj）
共享专家（如 DeepSeek MoE）会自动处理，除非设置 ignore_shared_experts=True