跳转至

MoE 架构完整指南:路由、负载均衡、EP、训练稳定性、Serving

更新日期:2026-04-14


一、Expert Parallelism(EP)深入

1.1 为什么 MoE 需要专门的并行维度

MoE 的通信模式与 Dense 模型完全不同。Dense 模型的 TP 是 AllReduce(所有 GPU 处理相同的 token),而 MoE 的 EP 是 All-to-All(每个 GPU 处理不同的 token 子集)。

1.2 EP 的 All-to-All 通信详解

def moe_with_expert_parallelism(x, router, experts, ep_group):
    # 假设 EP=4: 256 个专家分布在 4 个 GPU 上, 每个 GPU 有 64 个专家
    # x: [local_batch * seq_len, d_model] — 每个 GPU 上的本地 token

    # Step 1: Router 在本地计算(每个 GPU 独立)
    weights, expert_ids = router(x)  # 每个 token 选 top-k 个专家

    # Step 2: Dispatch — All-to-All 通信
    # 问题: token 在 GPU0 上, 但它需要的专家可能在 GPU2 上
    # 需要把 token 发到正确的 GPU

    # 构造发送缓冲区: 按目标 GPU 分组
    send_buffers = [[] for _ in range(ep_size)]
    for token_idx, (w, eid) in enumerate(zip(weights, expert_ids)):
        for k in range(top_k):
            target_gpu = eid[k] // experts_per_gpu
            send_buffers[target_gpu].append((token_idx, eid[k], w[k], x[token_idx]))

    # All-to-All: 每个 GPU 同时发送和接收
    recv_buffers = all_to_all(send_buffers, ep_group)
    # GPU0 收到: 所有需要 expert 0-63 处理的 token (来自所有 GPU)

    # Step 3: 本地专家计算
    local_outputs = {}
    for (src_token_idx, expert_id, weight, token_data) in recv_buffers:
        local_expert = experts[expert_id % experts_per_gpu]
        local_outputs[src_token_idx] = weight * local_expert(token_data)

    # Step 4: Combine — 再次 All-to-All
    # 把结果发回原来的 GPU
    results = all_to_all(local_outputs, ep_group)  # 反向发送

    # Step 5: 加权聚合
    output = aggregate(results)
    return output

1.3 EP 通信量分析

EP 的 All-to-All 通信量随 batch_size 和 seq_len 线性增长,且无法像 TP 的 AllReduce 那样通过 Ring 优化。这是 MoE 训练的主要通信瓶颈。

1.4 EP 与 TP 的配合

# DeepSeek-V3 的并行策略
# 2048 GPU, 每节点 8 GPU (NVLink)
config = {
    'TP': 1,   # MoE 层不用 TP! (因为 EP 已经切分了)
    'EP': 64,  # 256 个专家 / 64 = 每 GPU 4 个专家
    'PP': 16,  # 16 个流水线阶段 
    'DP': 2,   # 2048 / (1 × 64 × 16) = 2
}

# 关键设计:
# 1. Attention 层: 用 TP=8 (节点内 NVLink)
# 2. MoE 层: 用 EP (跨节点, 需要 IB)
# 3. 两种并行在同一模型中交替使用!
# 
# 这意味着在 Attention → MoE 的边界需要通信重组:
# Attention 结束: 每个 TP group 有完整 token 的部分维度
# MoE 开始: 每个 EP rank 需要完整的 token
# → 需要 AllGather 操作

# DeepSeek 的优化: 将 AllGather 和 MoE Dispatch 重叠

1.5 EP 的负载不均衡问题

正常情况: 每个 GPU 处理 BStop_k / EP 个 token 实际情况: 路由不均匀 → 某些 GPU 收到远多于平均的 token

GPU0: 收到 1.5x 平均量 → 计算时间长 → 其他 GPU 等待 GPU1: 收到 0.5x 平均量 → 计算完了等 GPU0

解决方案: 1. Capacity Factor: 每个 GPU 最多处理 C × 平均量, 多余的 drop 2. Aux-Loss-Free: 动态偏置让路由更均匀 3. 异步 All-to-All: 不等所有 GPU 完成, 用流水线方式


二、视频训练对 MoE 的不稳定性

2.1 问题描述

视频/多模态训练中 MoE 的不稳定性是一个已知的工业级痛点,但公开资料极少。

2.2 解决方案

class ModalityAwareMoE:
    def __init__(self, d_model, n_experts, top_k):
        self.text_router = Router(d_model, n_experts)
        self.vision_router = Router(d_model, n_experts)
        # 或者: 共享 router 但不同 bias
        self.vision_bias = zeros(n_experts)  # 视觉专用偏置

    def forward(self, x, modality_mask):
        text_tokens = x[modality_mask == 'text']
        vision_tokens = x[modality_mask == 'vision']

        text_weights, text_ids = self.text_router(text_tokens)

        # 视觉 token: 更高 temperature → 更均匀分布
        vision_logits = self.vision_router.gate(vision_tokens) / temperature
        vision_logits = vision_logits + self.vision_bias
        vision_weights, vision_ids = topk_softmax(vision_logits, self.top_k)

        # 分别计算, 然后合并
        ...

2.3 Cosmos/Megatron 的视频 MoE 训练

注意 NVIDIA 的 Cosmos 视频生成管线用 DiT (Diffusion Transformer),而不是标准 LLM。DiT 中的 MoE 有额外挑战:


三、MoE vs Dense "智力"对比(深入版)

3.1 严格控制变量的对比

3.2 MoE "降智"的场景

MoE 并非总是更好。以下场景 MoE 可能不如同激活参数的 Dense:

3.3 什么时候选 MoE vs Dense

选 MoE 当:训练预算有限但需要大知识容量(如知识密集型问答)、推理吞吐量要求高(同等质量下 MoE 推理更快)、有足够 GPU 部署。 选 Dense 当:部署 GPU 有限(Dense 模型更小)、需要极致推理深度(如数学竞赛)、工程简单性优先。


四、路由崩塌诊断与修复(工程实战版)

4.1 实时监控脚本

class MoEMonitor:
    def __init__(self, model, log_interval=100):
        self.log_interval = log_interval
        self.step = 0
        self.history = defaultdict(list)

    def log(self, model):
        self.step += 1
        if self.step % self.log_interval != 0:
            return

        for layer_idx, layer in enumerate(model.moe_layers):
            if not hasattr(layer, '_last_expert_counts'):
                continue
            counts = layer._last_expert_counts  # 需要在 forward 中记录
            total = counts.sum()
            load = counts / total

            # 指标 1: 变异系数 (CV)
            cv = load.std() / load.mean()

            # 指标 2: 死专家比例
            dead_ratio = (load < 0.0001).float().mean()

            # 指标 3: 最大负载 / 平均负载
            max_over_avg = load.max() / load.mean()

            # 指标 4: 有效专家数 (perplexity of load distribution)
            entropy = -(load * log(load + 1e-10)).sum()
            effective_experts = exp(entropy)

            self.history[f'layer{layer_idx}/cv'].append(cv)
            self.history[f'layer{layer_idx}/dead_ratio'].append(dead_ratio)
            self.history[f'layer{layer_idx}/effective_experts'].append(effective_experts)

            # 告警
            if cv > 0.5:
                print(f"WARNING: Layer {layer_idx} load CV={cv:.2f} > 0.5")
            if dead_ratio > 0.1:
                print(f"CRITICAL: Layer {layer_idx} {dead_ratio:.0%} experts are dead!")
            if max_over_avg > 5:
                print(f"WARNING: Layer {layer_idx} max/avg load = {max_over_avg:.1f}")

4.2 专家重置(最后手段)

def reset_dead_experts(model, threshold=0.001):
    for layer in model.moe_layers:
        load = get_expert_load(layer)
        dead_mask = load < threshold
        alive_mask = ~dead_mask

        if dead_mask.sum() == 0:
            continue

        # 用最活跃的专家的权重重置死专家
        alive_indices = alive_mask.nonzero().flatten()
        dead_indices = dead_mask.nonzero().flatten()

        for dead_idx in dead_indices:
            # 随机选一个活跃专家
            donor_idx = alive_indices[randint(0, len(alive_indices))]
            # 复制权重 + 加噪声 (避免完全相同)
            for p_dead, p_donor in zip(
                layer.experts[dead_idx].parameters(),
                layer.experts[donor_idx].parameters()
            ):
                p_dead.data = p_donor.data.clone()
                p_dead.data += 0.01 * randn_like(p_dead.data)  # 小扰动

            # 重置 router bias
            layer.router.bias[dead_idx] = layer.router.bias[donor_idx] + 0.1

        print(f"Reset {dead_mask.sum()} dead experts in layer")

五、MoE Serving 不降智(深入版)

5.1 量化对 MoE 的特殊挑战

MoE 量化不能用统一的 calibration 策略,因为不同专家看到的数据分布完全不同。

def per_expert_quantization(model, calibration_data):
    for layer in model.moe_layers:
        for expert_id, expert in enumerate(layer.experts):
            # 只用路由到这个专家的 token 做 calibration
            expert_tokens = []
            for batch in calibration_data:
                _, indices = layer.router(batch)
                mask = (indices == expert_id).any(dim=-1)
                expert_tokens.append(batch[mask])

            expert_calibration = concat(expert_tokens)

            # 关键: 冷门专家可能 calibration 数据很少!
            if len(expert_calibration) < 100:
                # 冷门专家: 用全局 calibration 数据 + 更保守的量化
                expert_calibration = global_calibration_subset
                quantize(expert, expert_calibration, bits=8)  # 保守: 用 8 bit
            else:
                quantize(expert, expert_calibration, bits=4)   # 热门: 可以 4 bit

5.2 推理时的路由一致性

2025 年的重要发现:训练时用的 router bias 在推理时去掉会导致质量下降。

# 错误做法:
serving_model = load_model('deepseek-v3')
# router.bias 被忽略了 → 路由分布偏移 → 降智

# 正确做法:
serving_model = load_model('deepseek-v3')
for layer in serving_model.moe_layers:
    # 保留训练时的 bias!
    layer.router.bias = load_training_bias(layer)

# 更好的做法: 用 serving calibration 数据重新调整 bias
for layer in serving_model.moe_layers:
    layer.router.bias = calibrate_bias(layer, serving_data, target_load=uniform)

5.3 MoE 推理的 Expert Offloading

# 问题: 671B 模型需要 1.3TB (FP16) 内存, 即使 8×H100 也放不下
# 方案: 热门专家常驻 GPU, 冷门专家在 CPU/NVMe 上

class ExpertOffloadingMoE:
    def __init__(self, experts, top_k_hot=32, gpu_capacity=64):
        # 统计历史负载, 找出 top_k_hot 个最热门的专家
        hot_expert_ids = get_hottest_experts(experts, top_k_hot)

        self.gpu_experts = {eid: experts[eid].to('cuda') for eid in hot_expert_ids}
        self.cpu_experts = {eid: experts[eid] for eid in range(len(experts)) if eid not in hot_expert_ids}
        self.cache = LRUCache(capacity=gpu_capacity - top_k_hot)  # GPU 上的 LRU 缓存

    def forward(self, x, expert_ids):
        for eid in expert_ids:
            if eid in self.gpu_experts:
                # 热门专家: 直接在 GPU 上计算
                output = self.gpu_experts[eid](x)
            elif eid in self.cache:
                # 缓存中: GPU 上计算
                output = self.cache[eid](x)
            else:
                # 冷门专家: 从 CPU 加载到 GPU
                expert = self.cpu_experts[eid].to('cuda')
                self.cache.put(eid, expert)
                output = expert(x)
        return output

    # 效果: 常见查询用 GPU 上的热门专家 → 快
    #        罕见查询需要加载冷门专家 → 慢, 但不频繁
    # 类似于操作系统的缓存层级: L1 (GPU) → L2 (CPU) → L3 (NVMe)

六、MoE 前沿方向

方向 核心思想 状态
MoE++ (异构专家) 不同专家不同大小: 重要 token 用大专家 ICLR 2025
Expert Merging 训练完后合并相似专家 → 减少部署成本 研究中
Dynamic top-k 不同 token 不同 k: 简单 token k=2, 复杂 token k=8 研究中
MoE in Attention 不只 FFN, Attention 也用 MoE 前沿
Orthogonality Loss 鼓励专家处理不同类型 token, 增加多样性 2025
Continual Learning 新增专家处理新知识, 旧专家冻结 研究中

参考链接


上级 · A. 基础理论