Olmo-core 3 開源混合專家訓練基礎設施
TL;DR
Olmo-core 3 是一款新發布的開源訓練堆疊，可高效訓練達到千億參數規模的混合專家（MoE）模型，相比先前實作，每秒處理的 token 數最多可提升 2.7 倍，並支援 MXFP8 精度等先進優化技術。
Olmo‑core 3 是什麼以及它的重要性
Olmo‑core 3 是 Allen Institute 對其 MoE 訓練框架的重新設計，旨在隨著專家池從數十個擴展至數百個時，仍能保持計算效率，同時每 token 的活躍參數數量大致維持不變。透過消除傳統上會侵蝕 MoE 優勢的通訊與路由開銷，此堆疊讓學術實驗室與小型研究團隊也能負擔得起千億參數的稀疏模型訓練。
核心架構變更
從 FSDP 改為 DDP
- 早期的 Olmo‑core 使用完全分片資料平行（FSDP），需針對每個 mini‑batch 反覆收集並重新共享模型權重。
- Olmo‑core 3 採用分散式資料平行（DDP），讓專家權重保留在 GPU 上，僅傳輸必要的 token 資料。這消除了耗時的權重收集步驟。
整合式 MoE 堆疊 vs. Megatron‑Core
- NVIDIA 的 Megatron‑Core 仍是大型 MoE 的參考實作。
- Olmo‑core 3 提供端到端、緊密整合的 MoE 堆疊，性能超越先前基於 FSDP 的版本。在八張 NVIDIA B300 GPU 的基準測試中，一個 47 B 參數的 MoE 每 GPU 可達 52,000 tokens / s，相比先前的 19,400 tokens / s，提升了 2.7× 的吞吐量。
採用的擴展技術
平行策略
| 技術 | 目的 |
|---|---|
| 專家平行 | 將專家池分散至各 GPU，使每張 GPU 僅儲存部分專家。 |
| 流水線平行 | 將模型層切分至 GPU 群組，降低每 GPU 的記憶體負載。 |
| 分散式優化器 | 將優化器狀態分散至各 GPU，避免每台設備都複製完整狀態。 |
這三種機制共同作用，使 MoE 模型得以擴展，而無需每張 GPU 都儲存完整模型或優化器狀態。
路由與計算優化
- 逐列專家平行 – 直接將路由後的 token 寫入專家輸入緩衝區，減少資料重排。
- GPU 保留路由 – 將路由元資料保留在 GPU 上，避免 CPU 成為瓶頸。
- 群組 GEMM – 將來自不同專家的許多小型矩陣乘法合併為單一大型 GEMM 呼叫，提升 GPU 使用率。
使用 MXFP8 的精度優化
- MXFP8 是一種低精度數值格式，可減少計算與 GPU 之間的通訊流量。
- 在四張 NVIDIA B300 GPU 的控制基準測試中，啟用 MXFP8 後，端到端 吞吐量相比 BF16 基準提升了 ~21 %，同時峰值活躍記憶體從 103 GiB 降低至 95 GiB。
- 大部分效能提升來自更快的前饋計算與減少專家間的資料搬移。
展示的規模
| 配置 | 總參數 | 每 token 活躍參數 | GPU | 峰值吞吐量 |
|---|---|---|---|---|
| 1.2 T 參數 MoE | 1.2 T | 58.36 B | 512 | 858 TFLOP/s / GPU |
| 2.38 T 參數（DeepEP v2） | 2.38 T | — | — | — |
1.2 T 參數的執行使用隨機路由以壓力測試系統效能；此結果不代表訓練後模型的品質。
超越原始速度的實證發現
- token 畫界 – 路由平衡分數可能提升，但實際工作負載平衡反而惡化，突顯指標設計的陷阱。
- 專家學習率縮放 – 在測試的模型家族中，降低使用頻率較低的專家學習率並未帶來更好的結果。
- 輸入值依賴的執行時間 – 當輸入值不同時，相同矩陣形狀的 GPU 核心執行時間會有差異，強調需使用值匹配的基準測試。
- 通訊與計算重疊 – 有時重疊通訊與計算流會降低整體執行速度，顯示並非所有情況下重疊越多越好。
這些觀察詳述於附帶的技術報告中，並指引未來 MoE 訓練研究的方向。
開源可用性與社群影響
- 程式碼 – 完全開源於 GitHub： https://github.com/allenai/olmo-core
- 技術報告 – 提供深入的系統設計、消融實驗與實驗方法論： https://allenai.org/papers/olmocore3
- 互動式示範 – 資料、專家與流水線平行的視覺化導覽： https://narrative.allen.ai/scaling-up-training
透過以開源授權釋出此堆疊，Allen Institute 使研究人員能訓練自己的千億參數 MoE，將系統適應至其他硬體，並實驗新型路由或精度方案。
未來展望
Olmo‑core 3 將成為下一代 Olmo 模型的骨幹，這些模型將結合新的 MoE 架構、更大的資料集與更長的上下文視窗。其模組化設計確保能隨著硬體進步持續演進，讓研究社群始終處於可擴展稀疏模型訓練的最前線。
重點摘要
- Olmo‑core 3 在保持每 token 活躍參數不變的情況下，將 MoE 訓練擴展至千億參數規模。
- 透過 DDP、專家/流水線平行與路由優化，其每秒處理的 token 數比先前基於 FSDP 的堆疊最高提升 2.7×。
- MXFP8 精度進一步帶來 ~21 % 的速度提升並降低記憶體使用。
- 此框架完全開源，為社群提供生產級工具，用於大規模稀疏模型研究。