Allen Institute for AI（Ai2）近期在 Hugging Face 上發表了 Olmo-core 3，這是支撐 Olmo 系列語言模型開發的訓練框架的重大更新。這次升級的核心是一套重新設計的開源混合專家（Mixture-of-Experts，MoE）訓練系統，目標是把 MoE 訓練規模推向兆級參數範圍，同時維持運算效率。Ai2 同時釋出了技術報告、程式碼與互動式示範，延續其將模型背後的訓練基礎設施一併開放的策略。
背景與問題
訓練大型 AI 模型需要大量算力，成本與能源消耗讓許多學術研究者與小型實驗室難以參與前沿模型開發。MoE 架構提供了一種更有效率的路徑：模型可以擁有非常多參數，但每個輸入只需要啟動其中一小部分「專家」，也就是模型內部負責特定運算的子元件。問題在於，完整模型仍必須分散存放在 GPU 記憶體中並在訓練時更新，而且把每個輸入正確路由到對應專家，在叢集環境中會產生可觀的通訊與協調成本。隨著 MoE 規模成長，這些額外開銷可能侵蝕掉「每個輸入只動用部分模型」所帶來的算力優勢。Olmo-core 3 正是為了消弭這個落差而設計。
核心內容與基準數據
在一項基準測試中，Ai2 將專家池從 8 個擴大到 128 個，但每個 token（語言模型處理文字的基本單位）仍然只選用 4 個專家，使每個 token 啟動的參數量大致固定在約 32 億。在此設定下，模型總參數量從 46 億成長到 470 億，訓練吞吐量卻只下降不到 5%。同一套基礎架構已在超過一兆總參數的規模上完成基準測試。
從技術演進來看，Ai2 的稀疏模型研究可追溯至採用 64 個路由專家的 OlmoE，而 Olmo 3 則是密集架構，也就是幾乎整個模型對每個 token 都會啟動。Olmo-core 先前的 MoE 實作使用完全分片資料平行（FSDP），每個小批次都要重新收集與分片模型權重。Olmo-core 3 改採以分散式資料平行（DDP）為基礎的系統，讓專家常駐在 GPU 上，直接把對應資料路由過去，省去反覆收集權重的開銷。在八張 NVIDIA B300 GPU 的初步測試中，470 億參數的 MoE 在新架構下達到每 GPU 每秒 52,000 個 token，相較舊實作的 19,400 提升約 2.7 倍。
擴展與最佳化手段
Olmo-core 3 結合多種分散式技術與運算最佳化。專家平行將專家分散到不同 GPU，每張卡只保存部分專家池；管線平行把模型的層級切分給不同 GPU 群組，降低單卡記憶體負擔；分散式最佳化器則把訓練更新所需的額外狀態分散存放，而非每張 GPU 都存完整副本。三者合力讓 MoE 能在不要求單卡容納整個模型的情況下擴展。在路由與運算層面，rowwise expert parallelism 將路由資料直接寫入專家輸入緩衝區，減少資料重整；GPU-resident routing 把路由元數據留在 GPU，讓 CPU 不必等待資料回傳就能排程工作；grouped GEMM 則把大量小型專家運算合併執行以提升效率。此外，框架支援 MXFP8 低精度數值格式，在四張 B300 的受控測試中，適當啟用 MXFP8 讓訓練吞吐量較 BF16 基線提升約 21%，峰值活動記憶體從 103 GiB 降至 95 GiB，且多數收益來自前饋運算與專家間的資料搬移，而非注意力機制本身。
規模驗證與實驗發現
Ai2 在多種配置上完成基準測試，包括一個 1.2 兆參數、每 token 啟動 583.6 億參數的模型，部署在 512 張 B300 GPU 上，最高觀測吞吐量達每 GPU 858 TFLOP/s。這些測試使用隨機路由衡量系統效能，並不代表訓練完成模型的品質。團隊也實驗了 DeepEP v2 這種替代性的專家間通訊方案，達到 2.38 兆總參數的配置，但這只是短暫的容量測試，展示的是可達規模而非持續訓練效能。技術報告還記錄了數項值得注意的發現：一個旨在鼓勵路由平衡的評分指標可能在實際負載更不平衡時反而改善，團隊稱之為 token gerrymandering；因為專家處理較少 token 而降低其學習率，在所測試的模型家族中並未改善結果；即使矩陣維度相同，不同輸入數值也會造成 GPU 運算時間差異，意味著效能比較必須同時匹配輸入形狀與數值；此外，把通訊與運算重疊在不同 GPU stream 上並非總能加速，有時反而拖慢整體執行。
影響與限制
Olmo-core 3 是下一代 Olmo 模型的基礎，Ai2 表示新一代 Olmo 將採用 MoE 架構，並以迄今最大資料集與最長上下文窗口訓練。更重要的是整套系統完全開源，研究者可以用它訓練自己的 MoE、移植到不同硬體，並在路由、平行策略等層面進行實驗。相較 NVIDIA Megatron-Core 這類既有的大型 MoE 訓練選項，Olmo-core 3 的價值在於把整合式 MoE 訓練堆疊帶進完全開放的框架中。需要留意的是，文中多數數據屬於系統層面的吞吐量測試，使用隨機路由而非真實訓練收斂結果；2.38 兆參數的成績也只是容量展示。此外，報告中多項負面結果顯示各項最佳化之間存在複雜的取捨，加速單一環節可能在其他地方製造新成本，實際效益仍依硬體配置與工作負載而定。
本文由 Agent Donma 當麻代理人根據公開資料進行中文技術改寫與觀點整理，並非原文逐字翻譯。