重访 MoE 与稠密模型的速度-精度比较：LLM 训练视角
机器学习
2024-07-02 v2 人工智能
摘要
混合专家(MoE) 通过保持计算成本不变而增加模型容量，从而实现性能提升。在与稠密模型进行比较时，先前的工作通常采用以下设置：1) 使用 FLOPs 或激活参数作为模型复杂度的度量标准；2) 训练所有模型到相同的 token 数。我们认为该设置偏好 MoE，因为 FLOPs 和激活参数并不准确地衡量稀疏层中的通信开销，导致 MoE 拥有更大的实际训练预算。本文通过采用步长时间作为模型复杂度的更准确度量标准，并在 Chinchilla 计算最优设置下确定总计算预算来重新审视该设置。为了高效运行 MoE，我们采用 3D 分片方法，将稠密到 MoE 的步长时间增加控制在健康范围内。我们在九个 0-shot 和两个 1-shot 英文任务上评估了 MoE 和稠密 LLM，以及 MMLU 5-shot 和 GSM8K 8-shot，覆盖三个模型规模：6.4B、12.6B 和 29.6B。实验结果显示，在这些设置下，MoE 在速度-精度权衡曲线上始终优于稠密 LLM，且存在显著差距。我们的完整模型实现和分片策略已发布于~\url{https://github.com/apple/axlearn}
引用
@article{arxiv.2405.15052,
title = {Revisiting MoE and Dense Speed-Accuracy Comparisons for LLM Training},
author = {Xianzhi Du and Tom Gunter and Xiang Kong and Mark Lee and Zirui Wang and Aonan Zhang and Nan Du and Ruoming Pang},
journal= {arXiv preprint arXiv:2405.15052},
year = {2024}
}
备注
8 pages