MOE模型的过去、现状和未来是怎样的？

zenRRan

于 2024-09-27 22:27:11 发布

阅读量2

点赞数

原文链接：https://mp.weixin.qq.com/s?__biz=MzI3ODgwODA2MA==&mid=2247531522&idx=2&sn=b20f0543964a18d4f2a922bbe1790a9b&chksm=eaa9e4d071157a7e1bed0f9814dbe4c9bf1e8831da4f58c48916af4a0290988d4db7f24e57df&scene=126&sessionid=0

版权

知乎：朱小霖
链接：https://www.zhihu.com/question/664040671/answer/3655141787

过去就先不说了，这里我 survey 得不够全面，主要来聊一下我认为的现在和未来。

首先说现状。

MoE 目前的架构基本集中在于将原先 GPT 每层的 FFN 复制多份作为 n 个 expert，并增加一个 router，用来计算每个 token 对应到哪个 FFN（一般采用每个 token 固定指派 n 个 expert 的方案），也就是类似 Mixtral 7x8B 的结构。之后 deepspeed 和 qwen 都陆续采用了更细的 granularity，也就是在不改变参数数量的情况下，将单个 FFN 变窄，FFN 数量变多，以及采用了 shared expert+，也就是所有 token 都会共享一部分 FFN 的方案。这方面推荐阅读 deepspeed 的这篇论文：《DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models》

https://arxiv.org/abs/2401.06066

那么基于这样的一个方案，我认为从算法层面和系统层面上逐渐在收敛至这样的工作：

算法层面

目前已有的开源模型普遍有这样的一个特点，就是当使用相同的预训练数据从零训练时，一个参数量为 N 的 dense 模型与一个参数量在 2N，激活参数量在 0.4N 的 MoE 模型能力基本相仿。这里的能力主要指在常规 benchmark，如 MMLU、C-Eval 上的分数。这种对比在 Qwen-1.5-MoE-A2.7B 和 Qwen2-57B-A14B 中最为明显。在我个人的小规模经验中，也基本看到了这样的结论。

https://huggingface.co/Qwen/Qwen1.5-MoE-A2.7B
https://huggingface.co/Qwen/Qwen2-57B-A14B

那么我们可以认为当前的 MoE 架构就是一个用显存换训练时长/推理延迟的架构。而对于非端侧模型，这样的 trade off 无疑是很值得的。所以我认为在算法层面上，对现阶段有指导价值的工作是围绕这个 trade off 的 Pareto 曲线，类似于如果我们可以对标一个 7B dense 模型，那么激活参数数多少的 MoE 需要多少参数，vice versa。这样的指标对于 LLM 的实际应用会很有帮助。

工程层面

主要是需要实现更高效的 MoE 训练/推理基建。我对推理不太熟，主要说下训练方向，我认为主要是分 2 个方向。一块是如何优化高稀疏度、高 granularity 的 grouped matmul kernel，让 MoE 训练的端到端速度逐步追赶同激活参数数的 dense 模型。这方面可能还是要看 cutlass 官方的一些进展，如这里：hopper_grouped_gemm。如果想自己上手的话，可以考虑类似 together.ai 的这篇博文：Supercharging NVIDIA H200 and H100 GPU Cluster Performance With Together Kernel Collection，或者看 flash attention 3 的流水的方式来整体优化一下 FFN。这里完全对齐 dense 应该是很难的，毕竟存多了，但从我个人的经验来看，如果稀疏度不是太离谱，估计做到 dense 端到端的 80%-90% 还是很常规的。

https://github.com/NVIDIA/cutlass/tree/main/examples/57_hopper_grouped_gemm
https://www.together.ai/blog/nvidia-h200-and-h100-gpu-cluster-performance-together-kernel-collection
https://tridao.me/blog/2024/flash3/

另一块是带 expert parallel 的训练通信优化，以及配合着的 load balancing。这里我没有训练足够大的模型的经验，所以不太熟悉。我的看法是，对于绝大多数 MoE 模型，全部机内做 tp 可能就是足够平衡，也足够快的方案了。对此方向感兴趣的朋友还是推荐 deepspeed 的论文：《DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model》。

https://arxiv.org/abs/2405.04434

再来聊聊未来。

对未来的预测我肯定是做不到的，所以主要说 2 个我比较期待有发展的点。

一个是 MoE 的激活参数变化在哪些方面有负面影响。因为上面的 trade-off，即明显看得出换激活参数，实在是太值了，相当于训练也快了，推理也省了。从一个不相信免费的午餐的角度，我认为是常规的 benchmark 没有暴露出激活参数少的缺陷。这问题可能和一些 llm 这样特别的重要问题相关，比如说，我们是否可以认为，神经网络很多参数并不记忆，只需要很少的参数来做 reasoning 呢？我比较期待去对这个方面理解得深入！这样的工作方向研究前景深远：如果 0-1 mini 真的会是一个很小的模型，它替我们掌握了 reasoning 的好的 scaling law，可能这方面会有一些新的突破，对于业务应用来说，Pareto 曲线可能又可以多一条维度。

另一个是目前 MoE 架构的稳定性问题。因为上面提到的，MoE 看起来太诱人了，而 llama3 没有采用，甚至最先推出 Mixtral 7x8B 的 mistral 的 Mistral-Large-Instruct-2407 也是个 dense 模型，可能只有都市传说中的 GPT-4 是个 1.8T MoE，非常可疑。这里坊间传言是大的 MoE 不好训，很容易崩，如果是这样，那么我认为有可能是因为目前的 router 设计不够稳定导致的，毕竟就这个是新东西，而且一个 nn.Linear(hidden_state, 1) 加上 topk，看起来就好瘦弱……不过具体的原因是不像夏大提到的 "Dio-晶：民科 2：不使用 MoE，非不为也，实不能也"，我对硬件 failure 没有数量级的概念，所以就不贴笑大方了。但是我还是比较期待这里能有一些新的发展的。