今天我們發布 Olmo-core 3，這是我們開發大型語言模型框架的一項重大升級，其特色是重新設計的開放式專家混合模型（MoE）訓練系統。

Olmo-core 3 旨在將 MoE 訓練擴展至兆級參數規模，同時保持運算效率。它是下一代 Olmo 的核心系統之一，也是我們持續致力於開放每個新模型背後工具與訓練基礎設施的承諾。

訓練大型 AI 模型需要大量的運算資源，導致成本和能源消耗增加，使許多學術研究人員和小型實驗室難以進行先進的模型開發。MoE 模型提供了一種更有效率的方法，它們可以包含更多的學習組件或參數，而無需每個輸入都使用所有這些組件。然而，完整的模型仍需儲存在 GPU 記憶體中並在訓練期間更新，而且在叢集內將輸入導向正確的專家（MoE 中的專用組件）會產生額外的通訊和協調成本。

隨著 MoE 模型規模的擴大，這些成本可能會侵蝕掉每個輸入僅使用部分模型所帶來的運算優勢。

Olmo-core 3 的建構宗旨就是彌補這個差距。在一項基準測試中，我們將專家池從 8 個增加到 128 個，同時每個 token（語言模型處理的最小文本單位）仍僅選擇四個專家，使每個 token 的活躍參數數量大致維持在約 32 億。總參數容量從 46 億增長到 470 億，而訓練吞吐量下降不到 5%。相同的基礎設施已在超過一兆個總參數的規模下進行了基準測試。

Olmo-core 隨著每一代 Olmo 的發展而演進。我們在稀疏模型方面的工作可追溯到 OlmoE，它採用了包含 64 個路由專家的 MoE 架構。相較之下，Olmo 3 則使用了密集架構，這意味著幾乎所有模型對於每個 token 都是活躍的，其訓練堆疊也是圍繞該設計構建的。Olmo-core 3 透過專為更大規模 MoE 模型設計的訓練系統，擴展了該框架。

我們在 Olmo-core 中較早的 MoE 實作使用了完全分片資料平行處理（FSDP），其配置是為每個小批次的訓練資料收集並重新分片模型權重。Olmo-core 3 則改用基於分散式資料平行處理（DDP）的系統。它讓專家常駐於 GPU 上，並將相關資料路由給它們，從而避免了重複的權重收集。

NVIDIA 的 Megatron-Core 是訓練大型 MoE 模型的成熟選項。Olmo-core 3 為 Olmo 背後的框架帶來了整合的 MoE 訓練堆疊，其重新設計改善了相較於我們早期基於 FSDP 實作的吞吐量。在八個 NVIDIA B300 GPU 上的初步測試中，一個 470 億參數的 MoE 模型使用新堆疊每秒每個 GPU 處理 52,000 個 token，而使用我們早期實作則為 19,400 個，吞吐量約提高了 2.7 倍。

Olmo-core 3 結合了多種技術，用於將大型 MoE 模型分散到 GPU 叢集，並透過最佳化來提高路由和運算的效率。

有三種技術決定了模型及其訓練狀態如何在硬體上分佈：專家平行處理（Expert parallelism）將專家分散到多個 GPU 上，因此每個 GPU 只儲存部分完整的專家池。管線平行處理（Pipeline parallelism）將模型的層（轉換輸入的連續階段）分散到多個 GPU 群組，減少每個 GPU 需要保留在記憶體中的模型量。

分散式最佳化器（Distributed optimizer）將最佳化器狀態（用於在訓練期間計算和應用更新的額外資料）分散到多個 GPU 上，而不是在每個 GPU 上儲存完整副本。

這些技術共同使得 MoE 模型能夠擴展，而無需每個 GPU 都將整個模型及其訓練狀態保留在記憶體中。

Olmo-core 3 還降低了將資料路由到正確專家並執行其運算的成本。行式專家平行處理（Rowwise expert parallelism）將路由後的資料直接放入專家輸入緩衝區，最大限度地減少了重新排列所需的額外工作。GPU 常駐路由（GPU-resident routing）將路由中繼資料保留在 GPU 上，因此 CPU 可以排隊工作，而無需等待該資訊被複製回來。

而分組 GEMM（Grouped GEMM）則將許多小型專家運算結合起來，使 GPU 能夠更有效率地執行它們。

最後，Olmo-core 3 支援 MXFP8，這是一種低精度數字格式，用較少的位元表示某些數值。這可以減少運算和 GPU 之間移動的資料量，前提是這些節省的成本超過了數字格式轉換的成本。

我們在四個 NVIDIA B300 GPU 上進行了受控基準測試，評估 MXFP8 對端到端訓練吞吐量的影響，其中工作均勻分佈於各個專家。在系統中最有幫助的部分啟用 MXFP8 後，訓練吞吐量比我們作為基準的更高精度格式 BF16 高出約 21%，而峰值活躍記憶體從 103 GiB 下降到 95 GiB。大部分的增益來自前饋運算和專家之間的資料移動，而非僅僅注意力機制。

這些技術和最佳化必須協同運作。加速訓練的某個部分可能會在其他地方產生額外成本；更快的運算可能需要更多的資料移動，而減少位元移動如果資料轉換時間過長也可能沒有幫助。Olmo-core 3 圍繞著整個訓練過程中的這些權衡而建構，讓我們以及使用這個開放堆疊的研究人員能夠控制各個組件如何協同運作。

探索我們的互動式導覽，了解資料、專家和管線平行處理如何協同運作，以擴展 MoE 訓練，從單一 GPU 到多個 GPU。

我們已在 NVIDIA B300 GPU 上對 Olmo-core 3 進行了一系列配置的基準測試，其中包括一個 1.2 兆參數的模型，該模型在 512 個 GPU 上每個 token 活躍參數為 583.6 億。其觀察到的最高吞吐量為每 GPU 每秒 858 TFLOP/s，這是衡量每個 GPU 每秒有用模型運算量的指標。這些測試使用隨機路由來衡量系統性能，而非訓練模型的品質。

我們也實驗了 DeepEP v2，這是一種處理 GPU 之間專家通訊的替代方法，達到了一個總參數為 2.38 兆的配置。這是一個短期容量測試，而非完整的訓練運行，因此它展示了 Olmo-core 3 可以達到的規模，而非持續的訓練性能。

在這些規模下，系統性能只是全貌的一部分。我們的技術報告也記錄了指導我們如何訓練 MoE 模型並衡量其性能的實驗。例如：旨在鼓勵平衡路由的分數可能會提高，即使實際工作負載變得不那麼平衡，我們稱之為「token 不公平劃分」。降低專家的學習率（即其訓練更新的大小），因為它們處理的 token 較少，並未改善我們測試的模型系列的結果。

即使矩陣維度相同，當處理的值改變時，GPU 運算所需的時間也不同，因此性能比較需要匹配的輸入值和匹配的形狀。在單獨的 GPU 串流上重疊通訊和運算並不總是能加快訓練速度，在某些測試中，它反而減慢了端到端執行，這提醒我們更多的重疊不一定意味著更高的吞吐量。

該報告解釋了這些發現，以及我們測試過但未採用的方法。

Olmo-core 3 是我們下一步建構的基礎。我們的下一代 Olmo 將採用 MoE 架構，我們的目標是使其成為迄今為止最強大的 Olmo，將在我們最大的資料集上進行訓練，並擁有最長的上下文視窗。

這個新堆疊使我們能夠超越之前在 MoE 方面的工作，同時提供更大的靈活性，以隨著模型和硬體的演進來調整訓練。它也完全開放，研究人員和開發者可以使用 Olmo-core 3 訓練自己的 MoE 模型，將其適應不同的硬體，並實驗路由、平行處理和系統的其他部分。

這也是我們對開放模型開發的看法之一，當模型權重背後的基礎設施和訓練決策也開放時，模型權重會更有用。

如需更深入了解系統設計、實驗、消融研究以及我們在此過程中測試的方法，請閱讀我們的技術報告並在 GitHub 上探索 Olmo-core 3。