跳到內容

Technology

新聞|微軟公開 MAI-Base-2 預訓練細節:2.8 兆參數 MoE、source scaling 與 activation 梯度裁剪

微軟 AI 10/15 的技術文公開 MAI-Base-2 架構與訓練細節:2.8T 總參數、80B 啟用、每層 MoE、16,000 張 GB200 訓練 38 天、50T token;自述 held-out BPB 低於所有測過的開放權重模型,配方效率為 DeepSeek-V4-Pro-Base 的 4.2–12.3 倍。

Published: at 09:00 AM
編輯本頁下載圖卡

重點

規格

技術文列出 MAI-Base-1 與 MAI-Base-2 的對照:

項目MAI-Base-1MAI-Base-2
啟用/總參數35B/1T80B/2.8T
層數78102
hidden 維度66568704
MoE 配置dense FFN 與 MoE 交替每層都是 MoE
專家路由0 個共享 + 從 512 個中選 8 個1 個共享 + 從 512 個中選 7 個
attention head(KV/Q)global 8/80(每 6 層一次)、local 8/80global 2/96(每 6 層一次)、local 16/96
embedding綁定不綁定
訓練 token30T50T
optimizerAdamWMuon(非 embedding 的 2D 參數)+ AdamW
FP8tensor-wise、delayed scaling(E4M3/E5M2)MXFP8 E4M3
梯度裁剪全域權重梯度全域權重梯度 + activation 梯度
影像純文字原生 VLM

訓練在 2026 年 8 到 9 月進行。MFU 以 GB200 的 BF16 理論峰值 2.5 PFLOPS 為分母,不計重算與 optimizer 的 FLOPs。

跟誰比、怎麼比

基礎模型沒有經過後訓練,不適合用一般的問答基準比較,所以技術文改比 held-out 文字上的 bits-per-byte(BPB,越低越好),這個指標不受各家 tokenizer 影響。四組評測資料:

所有評測集另外對訓練資料做過去污染。比較對象是 2026 年 10 月以前可取得、總參數至少一兆的模型。其中 DeepSeek-V4-Pro-Base 是唯一公開的基礎模型 checkpoint;Kimi K3、Qwen3.8-2.4T-A95B、MiMo-V2.6-Pro-RL 只有後訓練版本,技術文也註明後訓練可能影響這些分數。

圖 1 的數值(BPB):

模型程式碼問答STEM數學
MAI-Base-1(35B/1T)0.24010.44400.40260.3939
MiMo-V2.6-Pro-RL(42B/1.02T)*0.24460.45130.39750.3803
Qwen3.8-2.4T-A95B(95B/2.4T)*0.24000.44270.39780.3752
DeepSeek-V4-Pro-Base(49B/1.6T)0.23700.43510.39680.3862
Kimi K3(104B/2.8T)*0.22640.42690.37610.3589
MAI-Base-2(80B/2.8T)0.22030.41620.36530.3561

*為後訓練模型。

四組都是 MAI-Base-2 最低,但評測集全部由微軟持有或委製,外界無法重跑。差距最小的是數學,MAI-Base-2 與 Kimi K3 相差 0.0028。

算力效率的 4.2 到 12.3 倍怎麼來

直接比較模型分數,會把配方好壞跟模型大小、訓練算力混在一起。技術文的做法是:

  1. 對四組評測各自擬合一條「BPB 對訓練 FLOPs」的 scaling law,擬合點包含 MAI-Base-2 本身,讓對照點落在擬合範圍內。
  2. 推算微軟的配方要花多少算力,才能達到 DeepSeek-V4-Pro-Base 的 BPB。
  3. 拿這個數字跟 DeepSeek-V4-Pro-Base 實際的訓練算力相比。

結果依評測不同落在 4.2 到 12.3 倍。用同樣方法跟今年 6 月的 MAI-Base-1 比,是 6.3 到 12.7 倍。微軟把效率歸功於資料品質與建模上的選擇。技術文另外說明,DeepSeek-V4-Pro-Base 的訓練 FLOPs 明顯少於 MAI-Base-2,微軟把省下的算力拿去訓練更大的模型。

Scaling ladder 與 source scaling

微軟用 scaling ladder 評估每項改動:一系列只差在大小、配方固定的模型。每一階以層數 L 命名,其餘架構都由 L 決定,例如 hidden 維度是 (256/3)×L,query head 數是 L 取最接近的 16 的倍數。各階之間維持每個啟用參數分到的訓練 token 數(TPP)不變。

問題出在資料重複。完整訓練跑 50T token 時,部分資料集會重複很多遍,小規模實驗如果用全部資料,重複次數就對不上。微軟的 source scaling 在 ladder 低階刻意縮小不重複樣本的範圍,讓每個資料集在任何訓練長度下都重複相同次數。技術文的例子是:某資料集在 50T 的完整訓練中重複 8 次,5T 的實驗就只取其中 1/10,同樣繞 8 圈。

在固定 TPP 下,這等於讓不重複資料量與參數量的比例固定。技術文說,少了這一步,小規模實驗會偏好正則化不足、放大後會過擬合的架構和超參數;source scaling 是 2.8T 模型能順利放大的關鍵。用 L18 到 L42 擬合的 ladder 外推,預測值與 MAI-Base-2 最終結果吻合(圖 2)。

超參數也改用 scaling law 決定:

讓每一層都放 MoE

微軟過去的實驗顯示,以 FLOPs 計算,每層都放 MoE 一直比 dense 與 MoE 交替更有效率。MAI-Base-1 仍採交替,原因是當時的 MoE 實作慢,換算成實際訓練時間不划算。這次找出兩個拖慢的原因:

處理後,MoE 層的 MFU 從 24% 升到 35%,MAI-Base-2 才改成每層都是 MoE。

多輪 dropless MoE:直接丟掉超出容量的 token 會讓效果變差,按最壞情況預留記憶體又浪費。微軟把每張卡上的專家分成若干組,每組以固定的單一專家容量跑「派送 → 計算 → 收回」。某個專家收到的 token 超過容量時,多出來的部分在額外幾輪處理,同組其他專家沒用完的容量也可以借用。反向傳播照同樣的多輪排程,逐組重算。這樣不管路由多不平均,記憶體用量都固定,不會碎片化,也不會 OOM。容量設定成平常只需要一輪。

降低單一專家的開銷,技術文列了四項:

Muon 與 activation 梯度裁剪

微軟發現放大到 MAI-Base-2 時,Muon 比 AdamW 訓練更穩定,效率也高很多。Muon 要用 Newton–Schulz 迭代把整個動量矩陣正交化,這跟 ZeRO 把動量切散在各張卡上的做法衝突。做法是在 optimizer 步驟時重新分配:各卡更新自己那一片 FP32 動量並轉成 BF16,經一次 all-to-all 讓每張卡拿到幾個完整矩陣,在本地正交化,再經第二次 all-to-all 送回原本的分片。

最佳化之後,Muon 佔每次訓練迭代平均時間的 2.1%(165 毫秒)。微軟估計理論下限是 55 毫秒,卡在通訊上,差距來自落後的節點、kernel 派送開銷、負載不均與重新分片時的記憶體配置。

Activation 梯度裁剪針對的是一般梯度裁剪的缺點:全域裁剪看的是整個 batch 的權重梯度範數,少數異常樣本就能讓整個 batch 的梯度一起被縮小。微軟改在網路中選定的位置裁剪 activation 的梯度。配合 Muon,MAI-Base-2 的梯度範數比 MAI-Base-1 集中得多。

技術文說,多數改動在不同規模下帶來的效率提升大致固定,activation 梯度裁剪的效益卻隨規模變大:在 ladder 的 L42 階,held-out 評測上的 FLOPs 效率提升約 28%。它也被拿來檢查資料:一筆被裁剪的程式碼訓練樣本裡,授權條款在句子中間被截斷、直接接上 PR 標題,循線追查找到前處理流程的 bug。

資料

50T token 的配比:

類別佔比
程式碼45%
純影像14%
網頁12%
PDF9%
STEM9%
書籍5%
數學4%
其他2%

還不清楚的部分

參考資料

註釋

編輯本頁

Podcast

關於語音摘要

此語音摘要由 Google NotebookLM 提供,不完全代表作者本人對文章的理解。