您现在的位置是：快讯快讯

月之暗面 MoBA 核心作者自述：一个 “新晋大模型训练师” 的三入思过崖

梦游仙境 2025-02-21 18:22:41 快讯已有人查阅

导读“从开源论文、开源代码出发，现在已经进化到开源思维链了嘛！”...

欧意最新版本

欧意最新版本app是一款安全、稳定、可靠的数字货币交易平台。

APP下载官网地址

“从开源论文、开源代码出发，现在已经进化到开源思维链了嘛！”

作者：Andrew Lu，晚点团队

图片来源：由无界AI生成

2 月 18 日，Kimi 和 DeepSeek 同一天发布新进展，分别是 MoBA 和 NSA，二者都是对 “注意力机制”（Attention Mechanism）的改进。

今天，MoBA 的一位主要研发同学 Andrew Lu 在知乎发帖，自述研发过程的三次踩坑，他称为 “三入思过崖”。他在知乎的签名是“新晋 LLM 训练师”。

这条回答下的一个评论是：“从开源论文、开源代码出发，现在已经进化到开源思维链了嘛。”

注意力机制之所以重要，是因为它是当前大语言模型（LLM）的核心机制。回到 2017 年 6 月那篇开启 LLM 革命的 Transformer 八子论文，标题就是：Attention Is All You Need（注意力就是你所需要的一切），该论文被引用次数至今已达 15.3 万。

注意力机制能让 AI 模型像人类一样，知道在处理信息时该 “重点关注” 什么、“忽略” 什么，抓住信息中最关键的部分。

在大模型的训练阶段和使用（推理）阶段，注意力机制都会发挥作用。它的大致工作原理是，当输入一段数据，如 “我喜欢吃苹果”，大模型会计算句子中每一个词（Token）与其他词的关系，从而理解语义等信息。

而当大模型需要处理的上下文越来越长，标准 Transformer 最初采用的 Full Attention（全注意力机制）对计算资源的占用变得不可忍受，因为原初流程是，需要全部计算所有输入词的重要性分数，再加权计算后得到哪些是最重要的词，它的计算复杂度会随文本变长而平方级（非线性）增长。如 MoBA 论文的 “摘要” 部分所写：

“传统注意力机制中固有的计算复杂性的平方增加，带来令人望而却步的计算开销。”

同时，研究者们又追求大模型处理的上下文能足够长——多轮对话、复杂推理、记忆能力……这些设想中 AGI 应该具备的特性都需要更长长长长的上下文能力。

如何找到一个既没那么占用计算资源和内存，又不损失模型性能的注意力机制优化方法，于是成为大模型研究的重要课题。

这是数家公司将注意力交汇到 “注意力” 上的技术背景。

在 DeepSeek NSA 和 Kimi MoBA 之外，今年 1 月中旬，另一家中国大模型创业公司 MiniMax 也在其首个开源模型 MiniMax-01 中大规模实现了一种新的注意力机制。MiniMax 创始人闫俊杰当时告诉我们，这是 MiniMax-01 最主要的创新点之一。

面壁智能联合创始人、清华大学计算机系副教授刘知远的团队也在 2024 年发表过 InfLLM，其中也涉及一种稀疏注意力改进，该论文被 NSA 的论文引用。

这几个成果中，NSA、MoBA、InfLLm 里的注意力机制都属于 “稀疏注意力机制”（Sparse Attention)；而 MiniMax-01 的尝试则主要是另一个方向：“线性注意力机制”（Linear Attention)。

SeerAttention 作者之一，微软亚研院高级研究员曹士杰告诉我们：总体来说，线性注意力机制对标准注意力机制的改动更多、更激进，想直接解决随文本变长，计算度平方爆炸（所以是非线性的）的问题，可能的一个代价是，会损失对长上下文的复杂依赖关系的捕捉；稀疏注意力机制则是利用注意力固有的稀疏性，尝试寻找一种较为稳健的优化方式。

同时在此推荐曹士杰老师在知乎上关于注意力机制的高赞回答：https://www.zhihu.com/people/cao-shi-jie-67/answers

（他回答了 “梁文锋参与发表的 DeepSeek 新论文 NSA 注意力机制，有哪些信息值得关注？会带来哪些影响？” 这个问题。）

MoA（Mixture of Sparse Attention）的共同一作，清华大学 NICS-EFC 实验室博士傅天予说，在稀疏注意力机制的大方向下：“NSA 和 MoBA 都引入了动态注意力方法，即能动态选择需要计算细粒度注意力的 KV Cache 块，相比一些使用静态方法的稀疏注意力机制，能提升模型性能。这两个方法也都是在模型训练中就引入了稀疏注意力，而非仅在推理时引入，这也进一步提升了模型性能。”

（注：KV Cache 块是存储之前计算的 Key 标签和 Value 值的缓存；其中 Key 标签指注意力机制相关计算中，用于标识数据特征或数据位置等信息的标识标签，以便在计算注意力权重时，能与其他数据匹配和关联，Value 值则与 Key 标签对应，通常包含了实际要处理的数据内容，如单词或短语的语义向量等信息。）

同时，这次月之暗面除了发布详细的 MoBA 技术论文，也在 GitHub 该项目网站上发布了 MoBA 工程代码，该套代码已在月之暗面自己的产品 Kimi 中线上使用了一年多。

* 以下为 Andrew Lu 在知乎上的自述，已获得作者授权。原文中有多处 AI 术语，（）中灰色文字解释部分均为编者注。原贴链接：https://www.zhihu.com/people/deer-andrew

Andrew Lu 的研发自述

应章老师（清华大学助理教授章明星）邀请，来答一下之前做 MoBA 起起伏伏的心路历程，我戏称为 “三入思过崖”。（Andrew Lu 回答的问题是：“如何评价 Kimi 开源的稀疏注意力框架 MoBA？与 DeepSeek 的 NSA 相比，二者各有哪些亮点？”）
MoBA 的开始

MoBA 这个项目开始得非常早，在 2023 年 5 月底月之暗面刚建立不久的时候，入职报到当天就被 Tim （月之暗面联合创始人周昕宇）拉到了小房间里，和裘老师（浙江大学 / 之江实验室裘捷中，MoBA idea 的提出者）以及 Dylan（月之暗面研究员）一起开始搞 Long Context Training（长上下文训练）。这里首先要感谢一下 Tim 的耐心和教导，对一个 LLM 新手给予厚望并乐意培养，研发各种上线模型和模型相关技术的诸位大佬里面，很多人和我一样基本是从零开始接触 LLM。

当时业界普遍水平也不是很高，大家都在 4K 预训练（模型能处理的输入输出长度约在 4000 个 Token，几千个汉字），项目一开始叫 16K on 16B，意思是在 16B （模型参数 160 亿）上做 16K 长度的 Pre-train（预训练）即可，当然后来很快的这个需求在 8 月变成了需要支持 128K 下 Pre-train 。这也是 MoBA 设计时的第一个要求，能 From Scratch（从头开始）快速的训练出一个能支持 128K 长度下的模型，此时还不需要 Continue Training（持续训练，在已训模型基础上继续训练）。

这里也引申出一个有趣的问题，23 年 5/6 月的时候，业界普遍认为训长得长，端到端训练长文本（直接用长文本训练模型）的效果好于训练一个较短的模型再想办法搞长它。这种认知到 23 年下半年 long Llama （Meta 开发的、支持长文本处理的大模型）出现的时候才发生了转变。我们自己也进行了严格的验证，实际上短文本训练长度激活具有更好的 token efficiency（每个 token 贡献的有效信息量提升，意味着模型能用更少的 token 完成更高质量的任务）。自然 MoBA 设计中的第一个功能就成了时代的眼泪。

在这个时期，MoBA 的结构设计也是更为 “激进” 的，相较于现在的 “极简化” 结果，初提出的 MoBA 是一个带 cross attention （处理两段不同文本数据之间关系的注意力机制）的两层注意力机制串行方案，gate （控制输入数据如何在各个专家网络之间分配权重）本身是一个无参的结构（没有参数，不需要数据训练），但是为了更好的学习历史的 token，我们在每个 Transformer 层加了一个机器间的 cross attention 以及对应的参数（可以更好地记住历史信息）。此时的 MoBA 设计已经结合了后面为大家熟知的 Context Parallel 思想（完整的上下文序列被存放到不同节点上，在需要计算的时候才集中在一起），我们将整个上下文序列平铺在数据并行节点之间，将每个数据并行节点内的 context 看成是一个 MoE （Mixture of Experts，专家混合系统）中的 expert（专家），将需要 attention 的 token 发送到对应 expert 上进行 cross attention 再把结果通信回来。我们将 fastmoe （一种早期 MoE 训练的框架）的工作整合进了 Megatron-LM（来自 Nvidia 的现在通用大模型训练框架）来支持 expert 间通讯能力。

这个思路我们称之为 MoBA v0.5。

（编者注：MoBA 的灵感来自当前主流的大模型 MoE 结构。MoE 指大模型工作时，每次只激活部分专家的参数，而不是全部，从而节省算力；MoBA 的核心思路是 “每次只看最相关上下文，而不是全部上下文，从而节省计算和成本”。）

随着时间推进到 23 年 8 月初，主模型 Pre-Train 已经训练了大量 token，再来一次成本不低。显著改变了结构并增加了额外参数的 MoBA 至此第一次进入思过崖。

一个非常简单的 MoBA v0.5 的示意图

编者注：

History Tokens（历史标记）——在自然语言处理等场景中，代表之前处理过的文本单元集合。

Gate（门）——在神经网络中，用于控制信息流动的结构

Input（输入）——模型接收的数据或信息

V（Value）——在注意力机制中，包含实际要处理或关注的数据内容，比如语义向量等

K（Key 标签）——在注意力机制中，用于标识数据特征或位置等信息的标识标签，以便与其他数据匹配关联

Q（Querry 查询）——在注意力机制中，用于从键 - 值对中检索相关信息的向量

Cross Attention（交叉注意力）——一种注意力机制，关注不同来源的输入，如将输入与历史信息关联

Self Attention（自注意力）——一种注意力机制，模型对自身输入进行关注，捕捉输入内部的依赖关系

一入思过崖

进入思过崖当然是一种戏称，是停下来寻找改进方案的时间，也是深入理解新结构的时间。第一次进思过崖悟道，进去的快，出来的也快。Tim 作为月之暗面点子王掏出了新的改进思路，将 MoBA 从串行两层注意力方案改并行的单层注意力方案。MoBA 不再增加额外的模型参数，而是利用现有注意力机制参数，同步学习一个序列里面的所有信息，这样就可以尽可能不变动当前结构进行 Continue Training。

这个思路我们称之为 MoBA v1。

MoBA v1 实际上是 Sparse Attention（稀疏注意力）Context Parallel 的产物，在当时 Context Parallel 并不是大行其道的时候，MoBA v1 体现了极高的端到端加速能力。我们在 3B，7B 上都验证了它有效之后，在更大模型 scale 水平上撞墙了，训练过程中出现了非常大的 loss spike（模型训练时出现的异常现象）。我们初版合并 block attention output（注意力模块处理完数据后输出结果）的方式过于粗浅，只是简单累加，导致完全无法和 Full Attention 对分进行 debug，没有 ground truth（标准答案，这里指 Full Attention 的结果）的 debug 是极其困难的，我们用尽了各种当时的稳定性手段都不能解决。由于在较大的模型上训练出了问题，MoBA 至此二入思过崖。

一个非常简单的 MoBA v1 的示意图

编者注：

Self Attention to History（对历史的自注意力）——一种注意力机制，模型关注历史标记，捕捉当前输入与历史信息间的依赖关系

Share weights（共享权重）——神经网络中不同部分使用相同的权重参数，以减少参数数量和提高模型泛化能力

FFN（Feed - Forward Neural Network，前馈神经网络）——一种基本的神经网络结构，数据沿单一方向从输入层经过隐藏层流向输出层

Weighted Sum（加权和）——将多个值按照各自的权重进行求和的操作

二入思过崖

第二次留在思过崖时间比较长，从 23 年 9 月开始，至出思过崖的时候已经到了 24 年初。但是在思过崖里并不意味着被放弃了，我得以体会到在月之暗面工作的第二大特色，饱和救援。

除了一直就在强力输出的 Tim 和裘老师，苏神（苏剑林，月之暗面研究员）、远哥（Jingyuan Liu，月之暗面研究员）以及各路大佬都参与进行激烈的讨论，开始拆解并修正 MoBA，首先被修正的就是那个简单的 Weighted Sum（加权和）叠加。我们这里尝试过各种和 Gate Matrix 相乘相加的叠加搞法之后，Tim 从故纸堆里掏出了 Online Softmax（不是看到所有数据才能计算，而是来一个数据处理一个），说这个应该能行。其中最大的一个好处是使用 Online Softmax 之后我们可以通过将稀疏度降低至 0（全选所有的分块），来和一个数学等价的 Full Attention 进行严格对照 debug，这解决了大部分实现中遇到的难题。但是上下文拆分在数据并行节点之间的这个设计依然会导致不均衡问题，一个数据 sample 在数据并行间平铺之后，第一个数据并行 rank 上的头部几个 token 会被后续茫茫多的 Q 发送进行 attend（注意力计算过程），带来极差的平衡性，进而拖慢加速效率。这个现象也有一个更广为人知的名字——Attention Sink（注意力汇聚点）。

此时章老师到访，听完我们的想法后提出了新的思路，将 Context Parallel 能力和 MoBA 切分开。Context Parallel 是 Context Parallel，MoBA 是 MoBA ，MoBA 回归一个 Sparse Attention 本身而不是分布式 Sparse Attention 训练框架。只要显存放得下，完全可以在单机上处理全部上下文，用 MoBA 来进行计算加速，通过 Context Parallel 的方式来组织和传递机器间的上下文。因此我们重新实现了 MoBA v2，基本已经是也是当前大家见到的 MoBA 的样子。

现在的 MoBA 设计

编者注：

MoBA Gating（MoBA 门控）——MoBA 中的特定门控机制

RoPE（Rotary Position Embedding 旋转位置嵌入）——一种为序列添加位置信息的技术

Partition to blocks（划分为块）——将数据分割成不同的块

Mean Pooling（平均池化）——一种在深度学习中对数据进行降采样的操作，计算区域内数据的平均值

MatMul（Matrix - Multiply 矩阵乘法）——一种数学运算，用于计算两个矩阵的乘积

TopK Gating（Top - K 门控）——一种门控机制，选择前 K 个重要的元素等操作

Selected Block Index（选定的块索引）——表示被选中的块的编号

Index Select（索引选择）——根据索引从数据中选取相应元素

Varlen Flash-Attention（可变长度快速注意力）——一种注意力机制，适用于可变长度序列且计算效率较高

Attention Output（注意力输出）——注意力机制计算后的输出结果

MoBA v2 稳定可训，短文本和 Full Attention 可以完全对齐，Scaling Law 看起来非常可靠，而且比较丝滑的支持推广到线上模型上。我们因此加了更多的资源进来，在经过了一系列的 debug 并消耗了 n 把 infra 组同学头发后，我们可以将经过 MoBA 激活后的 Pretrain 模型做到捞针测试全绿（大模型处理长文本能力测试中达标），在这一步我们已经觉得非常不错了，开始进行上线。

但是最不意外的，只有意外。SFT （监督微调，在预训练模型的基础上，针对特定任务对模型进行进一步训练，以提高模型在该任务上的性能）阶段部分数据带着非常稀疏的 loss mask（使得只有 1% 甚至更少的 token 有训练用的梯度）（loss mask 指挑选哪些部分参与衡量模型预测结果与标准答案计算的技术），这导致 MoBA 在大部分 SFT 任务上表现良好，但是越是长文总结类型任务，loss mask 越是稀疏，反应在出来的学习效率越是低下。MoBA 在准上线流程中被按下暂停键，第三次进入思过崖。

三入思过崖

第三次进入思过崖的时候其实最紧张，此时整个项目已经有了巨大的沉没成本，公司付出了大量的计算资源和人力资源，如果端到端最后长文应用场景出现问题，那么前期研究接近于打水漂。幸运的是，由于 MoBA 本身优秀的数学性质，在新一轮饱和救援的实验 ablation（消融实验，通过去除模型的某些部分或改变某些设置来研究其对模型性能的影响）中，我们发现去掉 loss mask 表现就非常良好，带上 loss mask 表现就不尽如人意，进而意识到是带有 gradient （梯度，在机器学习中用于更新模型参数的方向和步长的一个值）的 token 在 SFT 阶段过于稀疏，从而带来的学习效率低下的问题，因此通过将最后几层修改为 Full Attention，来提高反向传播时带 gradient token 的密度，改善特定任务的学习效率。后续其它实验证明，这种切换并不会显著影响切换回来的 Sparse Attention 效果，在 1M （100 万）长度上和同结构 Full attention 各项指标持平。MoBA 再次从思过崖回归，并成功上线服务用户。

最后的最后，感谢各路大神的拔刀相助，感谢公司的大力支持以及巨量显卡。现在我们开放的出来的就是我们在线上用的代码，是一个已经经过了长期验证，因为实际需求砍掉了各种额外设计，保持极简结构但同时具备足够效果的 Sparse Attention 结构。希望 MoBA 以及它诞生的 CoT （Chain of Thought，思维链）能给大家带来一些帮助和价值。

FAQ

顺带借地方回答一些这两天频繁被人问到的问题，这两天基本麻烦章老师和苏神当客服回答问题了，实在过意不去，这里提取了几个常见问题一并回答一下。

1.MoBA 对 Decoding （模型推理阶段的文本生成过程）无效吗？

MoBA 对 Decoding 是有效的，对 MHA（Multi-Head Attention，多头注意力）很有效，对 GQA （Grouped Query Attention，分组查询注意力）效果降低，对 MQA （Multi-Query Attention，多查询注意力）效果最差。原理其实很简单，MHA 的情况下，每个 Q 有一个自己对应的 KV cache，那么 MoBA 的 gate 在理想情况下是可以通过摊余计算在 prefill （首次处理输入时的计算阶段）算好并存储每个 block（数据块）的代表 token，这个 token 在后续都不会变动，因此所有的 IO （输入输出操作）基本可以做到只来自 index select（通过索引选择数据的操作）之后的 KV cache，这种情况下 MoBA 的稀疏程度就决定了 IO 减少的程度。

但是对于 GQA 和 MQA，由于一组 Q Head 实际上在共享同样的一份 KV cache，那么在每个 Q Head 能自由选择感兴趣 Block 的情况下，很有可能填满了由稀疏性带来的 IO 优化。比如我们思考这么一个场景：16 个 Q Head 的 MQA，MoBA 刚好切分整个序列到 16 份，这意味着当最坏情况每个 Q head 感兴趣的分别是序号从 1 到 16 的每个上下文块，节省 IO 的优势就会被磨平。能自由选择 KV Block 的 Q Head 越多，效果越差。

由于 “自由选择 KV Block 的 Q Head” 这么个现象存在，天然的改进想法就是合并，假设大家都选一样的 Block，那不就净赚 IO 优化了么。对，但是在我们实际测试下，尤其是对于已经支付了大量成本的预训练模型，每个 Q head 都有自己独特的 “品味”，强行合并不如从头重新训练。

2.MoBA 默认是必选 self attention（自注意力机制），那么 self 的邻居会必选吗？

不会必选，这个是已知会产生些许疑惑的地方，我们最后选择相信 SGD（Stochastic Gradient Descent，随机梯度下降）。现在的 MoBA gate 实现非常直接，有感兴趣的同学可以简单改造 gate 使其必选上一个 chunk（数据块），但我们自己测试这个改动带来的收益比较 margin（收益不大）。

3.MoBA 有 Triton （编写高性能 GPU 代码的框架，由 OpenAI 开发）实现么？

我们实现过一个版本，端到端性能提升 10% ，但是 Triton 的实现想要持续维护跟上主线的成本比较高昂，我们因此多次迭代后暂缓进一步优化它。

*文章开篇提及的几个成果的项目地址（GitHub 页面中均包含技术 paper 链接，DeepSeek 暂未上线 NSA 的 GitHub 页面）：

MoBA GitHub 页面：https://github.com/MoonshotAI/MoBA

NSA 技术论文：https://arxiv.org/abs/2502.11089

MiniMax-01 GitHub 页面：https://github.com/MiniMax-AI/MiniMax-01

InfLLM GitHub 页面：https://github.com/thunlp/InfLLM?tab=readme-ov-file

SeerAttention GitHub 页面：https://github.com/microsoft/SeerAttention