16 May

MoE环游记:5、均匀分布的反思

如果说Meta的LLAMA系列为Dense模型确立了标准架构,那么DeepSeek或许就是MoE标准架构的奠基者。当然,这并非指DeepSeek首创了MoE,也不是说它的MoE不可超越,而是指DeepSeek对MoE所提的一些改进,很可能都是效果增益比较显著的方向,从而逐渐成为MoE的标配。这其中,包括我们在《MoE环游记:3、换个思路来分配》介绍的Loss-Free负载均衡方案,还有本文将要介绍的Shared Expert、Fine-Grained Expert策略。

说到负载均衡,它无疑是MoE一个极为重要的目标,本系列的第2~4篇,可以说都在围绕着它展开。然而,已有读者逐渐意识到,这里边有个尚未回答的本质问题:抛开效率上的需求不谈,均匀分布就一定是效果最好的方向吗?本文就带着这个疑问,去理解Shared Expert、Fine-Grained Expert。

共享专家

让我们再次回顾MoE的基本形式
\begin{equation}\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i\end{equation}

点击阅读全文...

4 May

Transformer升级之路:20、MLA究竟好在哪里?

自从DeepSeek爆火后,它所提的Attention变体MLA(Multi-head Latent Attention)也愈发受到关注。MLA通过巧妙的设计实现了MHA与MQA的自由切换,使得模型可以根据训练和推理的不同特性(Compute-Bound or Memory-Bound)选择最佳的形式,尽可能地达到效率最大化。

诚然,MLA很有效,但也有观点认为它不够优雅,所以寻找MLA替代品的努力一直存在,包括我们也有在尝试。然而,经过一段时间的实验,我们发现很多KV Cache相同甚至更大的Attention变体,最终效果都不如MLA。这不得不让我们开始反思:MLA的出色表现背后的关键原因究竟是什么?

接下来,本文将详细介绍笔者围绕这一问题的思考过程以及相关实验结果。

观察

MLA提出自DeepSeek-V2,本文假设读者已经熟悉MLA,至少了解之前的博客《缓存与效果的极限拉扯:从MHA、MQA、GQA到MLA》所介绍的内容,因此MLA自身的细节将不会过多展开。

点击阅读全文...

28 Mar

MoE环游记:4、难处应当多投入

前两篇文章我们都在讨论负载均衡,其中在《MoE环游记:3、换个思路来分配》介绍Loss-Free方案时,笔者留了一个悬念:它引入的Bias项有一个冗余的自由度,这个自由度可以用来做另外有趣的事情。这篇文章我们就来讨论这件事。

我们知道,MoE是为每个Token只选择最匹配的$k$个Expert来进行计算,从而在增大参数量的同时还节省了计算量。然而,当我们仔细思考就会发现,这个策略实际上有明显的可改进之处:直观来看,每个Token的难度并不一样,所以更合理的方案应该是难的Token分配更多的计算资源,简单的token分配更少的资源,这样或许能在同样有限的资源下将效果最大化。

而刚才提到的Bias的额外自由度,恰好可以用来简单地实现这个目标。

点击阅读全文...

5 Mar

MoE环游记:3、换个思路来分配

这篇文章我们继续探讨MoE的负载均衡问题。在上一篇文章《MoE环游记:2、不患寡而患不均》中,我们主要讨论了通过Aux Loss来促进负载均衡的思路。Aux Loss固然简单直观,但它也有一个明显的缺点——权重不好调——调低了无法促进均衡,调高了容易损害LM Loss,所以业界一直有寻找替代方案的尝试。

本文要分享的是名为“Loss-Free”的方案,由DeepSeek在《Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts》提出。和DeepSeek众多耀眼的开源作品相比,这篇论文也许不算起眼,但在笔者看来,它潜在的学术影响力可能远超其他工作,因为所提方法不仅简单有效,而且极具普适性,堪称经典。

方法大意

面对负载不均衡,Aux Loss的应对思路是通过额外的损失引导Router给出均衡的打分,而Loss-Free的想法则是换个新的分配思路,即不改变Router现有打分结果,而是改变$\mathop{\text{argtop}}_k \boldsymbol{\rho}$这个分配方式。

点击阅读全文...

27 Feb

Muon续集:为什么我们选择尝试Muon?

本文解读一下我们最新的技术报告《Muon is Scalable for LLM Training》,里边分享了我们之前在《Muon优化器赏析:从向量到矩阵的本质跨越》介绍过的Muon优化器的一次较大规模的实践,并开源了相应的模型(我们称之为“Moonlight”,目前是一个3B/16B的MoE模型)。我们发现了一个比较惊人的结论:在我们的实验设置下,Muon相比Adam能够达到将近2倍的训练效率。

Muon的Scaling Law及Moonlight的MMLU表现

Muon的Scaling Law及Moonlight的MMLU表现

优化器的工作说多不多,但说少也不少,为什么我们会选择Muon来作为新的尝试方向呢?已经调好超参的Adam优化器,怎么快速切换到Muon上进行尝试呢?模型Scale上去之后,Muon与Adam的性能效果差异如何?接下来将分享我们的思考过程。

点击阅读全文...

21 Feb

MoE环游记:2、不患寡而患不均

在上一篇文章《MoE环游记:1、从几何意义出发》中,我们介绍了MoE的一个几何诠释,旨在通过Dense模型的最佳逼近出发来推导和理解MoE。同时在文末我们也说了,给出MoE的计算公式仅仅是开始,训练一个实际有效的MoE模型还有很多细节补,比如本文要讨论的负载均衡(Load Balance)问题。

负载均衡,即“不患寡而患不均”,说白了就是让每个Expert都在干活,并且都在干尽可能一样多的活,避免某些Expert浪费算力。负载均衡既是充分利用训练算力的需求,也是尽可能发挥MoE大参数量潜力的需求。

需求分析

我们知道,MoE的基本形式是
\begin{equation}\boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i\end{equation}

点击阅读全文...

14 Feb

生成扩散模型漫谈(二十九):用DDPM来离散编码

笔者前两天在arXiv刷到了一篇新论文《Compressed Image Generation with Denoising Diffusion Codebook Models》,实在为作者的天马行空所叹服,忍不住来跟大家分享一番。

如本文标题所述,作者提出了一个叫DDCM(Denoising Diffusion Codebook Models)的脑洞,它把DDPM的噪声采样限制在一个有限的集合上,然后就可以实现一些很奇妙的效果,比如像VQVAE一样将样本编码为离散的ID序列并重构回来。注意这些操作都是在预训练好的DDPM上进行的,无需额外的训练。

有限集合

由于DDCM只需要用到一个预训练好的DDPM模型来执行采样,所以这里我们就不重复介绍DDPM的模型细节了,对DDPM还不大了解的读者可以回顾我们《生成扩散模型漫谈》系列的(一)(二)(三)篇。

点击阅读全文...

8 Feb

MoE环游记:1、从几何意义出发

前两年福至心灵之下,开了一个“Transformer升级之路”系列,陆续分享了主流Transformer架构的一些改进工作和个人思考,得到了部份读者的认可。这篇文章开始,我们沿着同样的风格,介绍当前另一个主流架构MoE(Mixture of Experts)。

MoE的流行自不必多说,近来火出圈的DeepSeek-V3便是MoE架构,传言GPT-4也是MoE架构,国内最近出的一些模型也有不少用上了MoE。然而,虽然MoE的研究由来已久,但其应用长时间内都不愠不火,大致上是从去年初的《Mixtral of Experts》开始,MoE才逐渐吸引大家的注意力,其显著优点是参数量大,但训练和推理成本都显著低。

但同时MoE也有一些难题,如训练不稳定、负载不均衡、效果不够好等,这也是它早年没有流行起来的主要原因。不过随着这两年关注度的提升,这些问题在很大程度上已经得到解决,我们在接下来的介绍中会逐一谈到这些内容。

点击阅读全文...