Skip to content

MoE 推理优化账:路由、负载均衡与专家并行的显存通信 ​

混合专家(MoE)用"稀疏激活"换来"大参数 + 低算力"——但推理侧的账远不止算力:所有专家都常驻显存,路由可能失衡,跨卡取专家要付 All-to-All 通信。本文把 MoE 推理的三笔账算清,衔接 DeepSeek MoE/MLA。

1. 显存账:稀疏激活,稠密驻留 ​

MoE 每 token 只激活少数专家(如 top-2),但所有专家权重都在显存里:

MemMoE=Wdense⏟attention/FFN 共享层+E×Wexpert⏟全部 E 个专家,活跃算力仅 k/EMem_{MoE} = \underbrace{W_{dense}}_{\text{attention/FFN 共享层}} + \underbrace{E \times W_{expert}}_{\text{全部 E 个专家}}, \quad \text{活跃算力仅 } k/E

模型总参数量活跃参数/token显存
稠密 70B70B70B140 GB+
MoE(8 专家 top-2)~50B~5B~100 GB(全部专家)

判断:MoE 省的是算力(激活少),不是显存(驻留全)——显存账用总参数量算,算力账用活跃参数量算,两本账别混。

2. 路由与负载均衡:失衡的代价 ​

路由(router)把 token 分给专家,若分配失衡:

  • 热点专家:个别专家排队,成为延迟瓶颈;
  • 冷门专家:闲置,浪费显存驻留;
  • 训练时的辅助损失(load balancing loss)缓解失衡,但推理负载变化仍会漂移。

工程对策:上线前用真实请求分布测路由分布,热点超阈值(如单专家 >2× 均值)时调整路由策略或重训路由。

3. 通信账:专家并行与 All-to-All ​

专家分布在不同卡上(专家并行 EP),token 路由到目标卡需要 All-to-All:

Comm∝跨卡 token 比例×激活大小,Talltoall≈(N−1)N×MsgBWComm \propto \text{跨卡 token 比例} \times \text{激活大小}, \quad T_{alltoall} \approx \frac{(N-1)}{N} \times \frac{Msg}{BW}

关键权衡:

  • EP 越大:单卡专家越少、显存压力越小,但跨卡 token 越多、通信越重;
  • EP 越小:通信轻,但每卡要驻留更多专家;
  • 结论:EP 规模 = 显存 vs 通信的折中,用 集合通信账 的 All-to-All 公式先算。

4. 优化手段:prefetch、卸载与量化 ​

手段原理收益
Expert prefetch提前加载下一层可能用到的专家隐藏取专家延迟
专家卸载(offload)冷门专家放 CPU/远端,用到再取显存大幅降、延迟略升
专家量化低精度存专家权重显存减半、质量略降
路由缓存相似 prompt 复用路由结果省路由计算与通信

判断:显存吃紧时"冷专家卸载 + 热专家常驻"是性价比最高的组合;prefetch 对低并发场景收益明显。

5. 可复算 Python:EP 通信与显存 ​

python
def moe_ep_ledger(n_experts, expert_gb, world, cross_ratio, msg_mb, bw_gbps):
    """专家并行:每卡显存(GB)与 All-to-All 耗时(ms)"""
    mem = n_experts * expert_gb / world
    t_ms = (world - 1) / world * msg_mb * cross_ratio / (bw_gbps / 8) / 1e3
    return round(mem, 1), round(t_ms, 2)

moe_ep_ledger(n_experts=8, expert_gb=12, world=8, cross_ratio=0.4, msg_mb=8, bw_gbps=400)
# (12 GB/卡, 0.56 ms/token) → EP8 显存可接受、通信可控

6. 交付清单 ​

  1. MoE 显存账按总参数量算,算力账按活跃参数量算,两本账分开;
  2. 上线前实测路由分布,热点专家超阈值(>2× 均值)先处理;
  3. EP 规模用 All-to-All 公式先算通信占比(>30% 缩 EP 或改拓扑);
  4. 显存吃紧:冷专家卸载 + 热专家常驻,再叠加专家量化;
  5. 每次改动跑 验收基线,记录延迟与显存。

本文为本站整理思考后撰写,用于知识库学习交流;公式与算例为本站自行推导整理,具体版本与实测数据以现场为准。

AI-HPC 智算超算专家 · Contact: openaihpc@gmail.com