ID 平衡：通过基于 PID 的负载控制对极其稀疏的 MoE 进行稳定训练

通过专家混合 (MoE) 扩展大语言模型 (LLM) 可以实现大规模参数增长，并且每个词元的计算几乎恒定。然而，进一步扩展参数数量需要越来越稀疏的路由，其中专家负载不平衡变得更加严重。这种不平衡会降低参数利用率和训练效率，并可能破坏训练稳定性，成为可靠扩展的瓶颈。在这项工作中，我们将两种代表性的辅助无损失方法统一为不完全比例积分微分（PID）控制器：DeepSeek的无损失方法充当固定步长积分控制器，而Kimi K3的分位数平衡充当广义比例控制器。基于这种控制视角，我们提出了 ID Balancing，一种积分微分控制器。它会根据负载误差缩放积分项，并仅在不平衡恶化时才激活其导数项，从而能够对较大或恶化的误差进行更强的修正，并在接近平衡时进行较小的更新。经过对超过 $768$ 的 Top-$10$、Top-$5$ 和 Top-$3$ 路由专家进行评估，ID Balancing 相对于 Top-$3$ 设置中的最佳基线，分别将最坏情况主干 MaxVio 和训练平均主干 MinVio 减少了超过 $50\%$ 和 $12\%$。当总参数计数从 $18.9$B 增加到 $69.9$B（Top-$10$-of-$768$）时，ID Balancing 的最坏情况主干 MaxVio 几乎保持不变，并且比辅助损失基线低约 $89.6\%$。 ID Balancing 还保持有竞争力的语言建模和下游性能。 ID Balancing 的优势随着稀疏性的增加而增强，使其成为扩展更大、更稀疏的 MoE 模型的有前途的解决方案。