2 Aug

基于树莓派Zero2W搭建一个随身旁路由

前段时间搞了个很迷你的开发板树莓派Zero2W(下面简称“Pi”),还搭配了个USB Key转接板,这几天折腾了一下,用于实现一个随身的旁路由。本文记录一下关键技术点,供有同样需求的读者参考。

树莓派Zero2W

树莓派Zero2W

点击阅读全文...

12 Jul

QK-Clip:让Muon在Scaleup之路上更进一步

四个月前,我们发布了Moonlight,在16B的MoE模型上验证了Muon优化器的有效性。在Moonlight中,我们确认了给Muon添加Weight Decay的必要性,同时提出了通过Update RMS对齐来迁移Adam超参的技巧,这使得Muon可以快速应用于LLM的训练。然而,当我们尝试将Muon进一步拓展到千亿参数以上的模型时,遇到了新的“拦路虎”——MaxLogit爆炸。

为了解决这个问题,我们提出了一种简单但极其有效的新方法,我们称之为“QK-Clip”。该方法从一个非常本质的角度去看待和解决MaxLogit现象,并且无损模型效果,这成为我们最新发布的万亿参数模型“Kimi K2”的关键训练技术之一。

问题描述

我们先来简单介绍一下MaxLogit爆炸现象。回顾Attention的定义
\begin{equation}\boldsymbol{O} = softmax(\boldsymbol{Q}\boldsymbol{K}^{\top})\boldsymbol{V}\end{equation}

点击阅读全文...

10 Jul

Transformer升级之路:21、MLA好在哪里?(下)

在文章《Transformer升级之路:20、MLA好在哪里?(上)》中,我们对MLA相比常见MHA、GQA、MQA的一些变化分别做了消融实验,其中的变化包括“增大head_dims”、“Partial RoPE”和“KV共享”,实验的初步结果是这三个变化很可能都是MLA效果优异的原因。

本文我们将从一个更加偏理论的角度出发,来理解MLA的成功之处。

部分旋转

首先,我们把最终的断言放在前面:

在相同训练成本和推理成本下,MLA可能是效果最好的Full Attention变体。

点击阅读全文...

20 Jun

线性注意力简史:从模仿、创新到反哺

在中文圈,本站应该算是比较早关注线性Attention的了,在2020年写首篇相关博客《线性Attention的探索:Attention必须有个Softmax吗?》时,大家主要讨论的还是BERT相关的Softmax Attention。事后来看,在BERT时代考虑线性Attention并不是太明智,因为当时训练长度比较短,且模型主要还是Encoder,用线性Attention来做基本没有优势。对此,笔者也曾撰文《线性Transformer应该不是你要等的那个模型》表达这一观点。

直到ChatGPT的出世,倒逼大家都去做Decoder-only的生成式模型,这跟线性Attention的RNN形式高度契合。同时,追求更长的训练长度也使得Softmax Attention的二次复杂度瓶颈愈发明显。在这样的新背景下,线性Attention越来越体现出竞争力,甚至出现了“反哺”Softmax Attention的迹象。

点击阅读全文...

26 May

众所周知,生成速度慢是扩散模型一直以来的痛点,而为了解决这个问题,大家可谓“八仙过海,各显神通”,提出了各式各样的解决方案,然而长久以来并没一项工作能够脱颖而出,成为标配。什么样的工作能够达到这个标准呢?在笔者看来,它至少满足几个条件:

1、数学原理清晰,能够揭示出快速生成的本质所在;

2、能够单目标从零训练,不需要对抗、蒸馏等额外手段;

3、单步生成接近SOTA,可以通过增加步数提升效果。

根据笔者的阅读经历,几乎没有一项工作能同时满足这三个标准。然而,就在几天前,arXiv出了一篇《Mean Flows for One-step Generative Modeling》(简称“MeanFlow”),看上去非常有潜力。接下来,我们将以此为契机,讨论一下相关思路和进展。

点击阅读全文...

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的额外自由度,恰好可以用来简单地实现这个目标。

点击阅读全文...