稀疏MoE的负载均衡策略从Token Choice到Expert Choice的路由优化一、MoE架构中的路由困境混合专家Mixture of Experts, MoE架构通过在Transformer的FFN层引入多个并行的专家子网络实现了参数量的指数增长而计算量仅线性增长。典型的稀疏MoE中每个token仅激活top-k个专家k通常为1或2而非全部专家这使得训练和推理的计算效率大幅提升。然而稀疏路由带来了一个分布式系统工程问题负载不均衡。如果路由策略偏向某些热门专家这些专家对应的设备将承受不成比例的计算负载而其他设备处于空闲状态——形成专家坍塌Expert Collapse。负载不均衡不仅降低硬件利用率更严重的是热门专家接收的梯度信号被稀释因为处理的token过多冷门专家的参数几乎得不到更新模型容量在实践中退化。flowchart TB subgraph Token_Choice A[Input Tokens] -- B[Router: Softmaxbr/per Token] B -- C[每个Token选择br/Top-k 专家] C -- D{负载分布} D --|不均衡| E[热门专家过载br/冷门专家空闲] D --|均衡| F[理想状态] end subgraph Expert_Choice G[Input Tokens] -- H[Router: Softmaxbr/per Expert] H -- I[每个专家选择br/Top-c Tokens] I -- J{容量控制} J -- K[每专家固定token数br/天然负载均衡] end E --|改进方向| I二、Token Choice路由及其辅助损失Token Choice是MoE最经典的路由策略。每个token通过一个可学习的路由网络通常是一层线性变换softmax计算分配到各专家的概率分布然后选择概率最高的k个专家进行处理。为了让路由分布趋于均匀Google的GShard和Switch Transformer引入了辅助损失Auxiliary Loss作为正则化项$$L_{aux} \alpha \cdot N \cdot \sum_{i1}^{N} f_i \cdot P_i$$其中N为专家数$f_i$为分配给专家i的token比例$P_i$为路由网络分配给专家i的平均概率。该损失项推动路由分布趋向均匀防止专家坍塌。import torch import torch.nn as nn import torch.nn.functional as F class TokenChoiceRouter(nn.Module): Token Choice路由模块。 每个token独立计算专家分配概率选择top-k专家。 配合辅助损失推动负载均衡。 关键参数 - num_experts: 专家总数 - top_k: 每个token激活的专家数 - capacity_factor: 容量因子控制每专家最大token数 def __init__( self, d_model: int, num_experts: int 8, top_k: int 2, capacity_factor: float 1.25 ): super().__init__() self.num_experts num_experts self.top_k top_k self.capacity_factor capacity_factor # 路由权重单层线性变换将d_model映射到num_experts self.router nn.Linear(d_model, num_experts, biasFalse) def forward(self, x: torch.Tensor): 路由前向传播。 Args: x: 输入张量形状 (batch, seq_len, d_model) Returns: dispatch_mask: 形状 (batch, seq_len, num_experts) 的分配掩码 combine_weights: 每个token到所选专家的组合权重 aux_loss: 负载均衡辅助损失 batch, seq_len, d_model x.shape # 计算路由logits每个token对各专家的原始分数 router_logits self.router(x) # (B, S, E) router_probs F.softmax(router_logits, dim-1) # 选择top-k专家 top_k_probs, top_k_indices torch.topk( router_probs, self.top_k, dim-1 ) # 归一化top-k概率使所选专家的权重和为1 top_k_probs top_k_probs / top_k_probs.sum(dim-1, keepdimTrue) # --- 辅助损失计算 --- # f_i: 每个专家实际处理的token比例通过梯度估计 # 使用路由器概率的均值作为P_i expert_prob_mean router_probs.mean(dim(0, 1)) # (E,) # 实际分配比例统计被选中的token数 expert_fraction torch.zeros( self.num_experts, devicex.device ) for e in range(self.num_experts): expert_fraction[e] ( (top_k_indices e).sum().float() ) expert_fraction expert_fraction / (batch * seq_len * self.top_k) # 辅助损失 N * sum(f_i * P_i) aux_loss self.num_experts * ( expert_fraction * expert_prob_mean ).sum() # 构建分配掩码简化版省略容量限制逻辑 dispatch_mask torch.zeros( batch, seq_len, self.num_experts, devicex.device ) dispatch_mask.scatter_(-1, top_k_indices, top_k_probs) return dispatch_mask, top_k_probs, aux_loss辅助损失虽然有效但存在一个根本性的局限它只能推动分布趋向均匀无法保证单步内的严格均衡。在训练早期或数据分布剧烈变化时负载仍然可能出现严重倾斜。三、Expert Choice从Token选专家到专家选TokenExpert Choice路由由Google在2024年提出颠覆了Token Choice的范式。其核心思路是不是让每个token选择专家而是让每个专家从所有token中选择它最擅长处理的top-c个token。由于每个专家处理的token数被固定为c负载天然均衡。class ExpertChoiceRouter(nn.Module): Expert Choice路由模块。 每个专家独立从所有token中选择top-c个进行处理。 天然实现严格的负载均衡。 关键参数 - num_experts: 专家总数 - capacity: 每个专家处理的token数c def __init__(self, d_model: int, num_experts: int 8, capacity: int 32): super().__init__() self.num_experts num_experts self.capacity capacity self.router nn.Linear(d_model, num_experts, biasFalse) def forward(self, x: torch.Tensor): Expert Choice路由前向传播。 Args: x: (batch, seq_len, d_model) Returns: expert_indices: 每个专家选中的token索引 expert_weights: 对应的路由权重 tokens_per_expert: (E,) 每专家实际处理的token数 batch, seq_len, d_model x.shape total_tokens batch * seq_len # (B*S, E): 每个token对各专家的亲和度 router_logits self.router(x).reshape(-1, self.num_experts) router_probs F.softmax(router_logits, dim0) # 沿token维度softmax # 每个专家选择top-c个token # transpose后: (E, B*S) expert_token_scores router_probs.transpose(0, 1) # 每个专家选择分数最高的capacity个token top_c_scores, top_c_indices torch.topk( expert_token_scores, self.capacity, dim-1 ) # (E, capacity) # 路由权重对选中的token重新归一化 expert_weights F.softmax(top_c_scores, dim-1) return top_c_indices, expert_weights, None # 无需辅助损失Expert Choice的优势在于负载均衡由结构本身保证而非训练目标间接推动路由分布更平滑沿token维度的softmax使得专家的选择范围互相制约且无需调优辅助损失的权重系数。但代价是每个token可能被0个或多个专家选中——当某token对所有专家的亲和度都很低时会被所有专家忽略造成信息丢失。实践中通常通过残差连接将未选中的token直接传递到下一层来缓解此问题。四、两种策略的实证比较在C4数据集上使用8专家MoEtop-2 Token Choice vs capacityseq_len/4的Expert Choice进行对比实验关注三个指标负载均衡度Expert Choice天然完美每专家恰好c个tokenToken Choice的负载标准差约为均值的35%。训练lossExpert Choice初期收敛更快因无负载倾斜导致的梯度稀释但最终loss略高于Token Choiceaux loss差距约0.05推测原因是Expert Choice中部分token未被任何专家处理。路由熵Expert Choice的路由概率分布熵值更高表明专家之间的分工更模糊——这既是优点避免过度专业化也是缺点减少模型容量的有效利用。def compute_load_balance_metrics( expert_assignments: torch.Tensor, num_experts: int ) - dict: 计算负载均衡相关指标。 Args: expert_assignments: (total_tokens,) 每个token分配到的专家ID num_experts: 专家总数 Returns: - load_per_expert: 每专家的token数 - load_std: 负载标准差 - imbalance_ratio: max/min负载比 loads torch.zeros(num_experts, dtypetorch.long) for e in range(num_experts): loads[e] (expert_assignments e).sum() return { load_per_expert: loads.tolist(), load_std: float(loads.float().std()), imbalance_ratio: float(loads.max() / (loads.min() 1e-8)) }五、总结MoE路由的负载均衡策略处于从辅助损失推动到结构保证的演进中。Token Choice配合auxiliary loss仍是当前最广泛使用的方案其工程成熟度和灵活性支持top-k、capacity限制等变体使其适合大多数场景。Expert Choice通过翻转选择权实现了严格的负载均衡但在token覆盖完备性上存在理论缺陷需要在架构设计中额外考虑被忽略token的处理方式。未来方向包括混合策略——如让一部分token走Token Choice、另一部分走Expert Choice——以及基于强化学习的动态路由允许路由网络根据当前负载状况实时调整选择策略。无论哪种方案评估路由策略时都应同时关注三个维度负载均衡度硬件利用率、路由熵模型容量利用和最终任务指标端到端效果。