月之暗面MoBA覈心作者自述：一個“新晋大模型訓練師”的三入思過崖

文章來源：晚點LatePost

Image source: Generated by 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