Ai2 випускає Olmo-Core 3, відкритий стек навчання для трильйонних параметрів MoE

Allen Institute for AI випустив Olmo-core 3 1 жовтня 2026 року, оновлення до своєї платформи розробки великих мовних моделей, побудованої навколо перепроектованої відкритої системи навчання mixture-of-experts, яку, за словами Ai2, було протестовано з більш ніж одним трильйоном загальних параметрів.

Ai2 заявив, що Olmo-core 3 розроблений для масштабування навчання MoE до діапазону трильйонних параметрів при збереженні обчислювальної ефективності, і описав його як одну з ключових систем, що стоять за наступним поколінням Olmo. Публікація супроводжується технічний звіт, інтерактивною демонстрацією та відкритим кодом.

Моделі MoE можуть містити набагато більше параметрів, не вимагаючи, щоб кожен вхід використовував їх усі, проте повна модель все одно повинна зберігатися у GPU‑пам’яті і оновлюватися під час навчання, а маршрутизація входів до потрібних експертів у кластері створює власні витрати на комунікацію та координацію. Ai2 зазначив, що ці витрати можуть знизити значну частину обчислювальної переваги зі зростанням MoE, і що Olmo-core 3 створений, щоб усунути цей розрив. У одному бенчмарку Ai2 пул експертів збільшився з 8 до 128, при цьому все ще вибираючи чотирьох експертів на токен, зберігаючи активні параметри приблизно на рівні 3,2 мільярда, тоді як загальна місткість параметрів зросла з 4,6 мільярда до 47 мільярдів, а пропускна здатність навчання впала менше ніж на 5%.

Від FSDP до стеку навчання на базі DDP

Попередня реалізація MoE в Olmo-core від Ai2 використовувала повністю розподілену паралельність даних, налаштовану на збір та повторне розподілення ваг моделі для кожної невеликої партії навчальних даних. Olmo-core 3 переходить до системи, заснованої на розподіленій паралельності даних, яка тримає експертів на GPU та направляє відповідні дані до них, уникаючи повторного збору ваг. У попередньому тесті на восьми GPU NVIDIA B300 Ai2 повідомляє, що MoE з 47‑мільярдними параметрами обробляв 52 000 токенів за секунду на GPU за новим стеком, у порівнянні з 19 400 при використанні попередньої реалізації, що, за словами Ai2, становить приблизно 2,7‑разовий приріст пропускної здатності.

Робота Ai2 над розрідженими моделями простежується до OlmoE, який використовував архітектуру MoE з 64 маршрутизованими експертами, тоді як Olmo 3 застосовував щільну архітектуру, у якій майже вся модель була активною для кожного токену. Olmo-core 3 розширює платформу системою навчання, розрахованою на значно більші моделі MoE. Ai2 зазначив, що Megatron-Core від NVIDIA є усталеним варіантом для навчання великих MoE, і описав Olmo-core 3 як інтегрований стек навчання MoE у платформі, що стоїть за Olmo.

Паралелізм, точність та методи управління пам’яттю

Три техніки визначають, як модель і її стан навчання розподіляються по апаратурі: експертна паралельність розподіляє експертів між GPU, конвеєрна (pipeline) паралельність розбиває шари моделі на групи GPU, а розподілений оптимізатор розподіляє стан оптимізатора між GPU замість зберігання повної копії на кожному GPU. Olmo-core 3 також зменшує витрати на маршрутизацію даних до потрібних експертів: рядкова експертна паралельність розміщує маршрутизовані дані безпосередньо у буфери входу експертів, маршрутизація, що перебуває на GPU, тримає метадані маршрутизації на GPU, дозволяючи процесору чергувати роботу без очікування копіювання назад, а згруповані GEMM об’єднують багато малих обчислень експертів, щоб GPU могли виконувати їх ефективніше. Технічний звіт описує дизайн рядкової експертної паралельності на базі NVSHMEM, який записує кожен токен безпосередньо у буфер відповідного експерта, з пристроєм‑запланованими згрупованими множеннями матриць, що робить експертну паралельність повністю без синхронізації.

Olmo-core 3 підтримує MXFP8, формат чисел нижчої точності. У контрольованому бенчмарку на чотирьох GPU NVIDIA B300 з рівномірним розподілом роботи між експертами Ai2 повідомляє, що пропускна здатність навчання була приблизно на 21 % вища, ніж при базовому BF16, при цьому пікова активна пам’ять знизилася з 103 GiB до 95 GiB, причому більша частина приросту походить від обчислень feed‑forward та переміщення даних між експертами. У звіті зазначається, що топологічно‑агностичні контрольні точки записують глобальні тензори FP32 незалежно від паралельної розкладки, що дозволяє продовжити запуск на іншій топології, а у контрольованому порівнянні повторний обчислення активацій усунуло майже дві третини пам’яті активацій і зменшило пікову пам’ять приблизно на чверть за рахунок близько однієї п’ятої пропускної здатності.

Трильйонні параметричні бенчмарки та заявлені обмеження

Ai2 зазначив, що провів бенчмарки Olmo-core 3 у різних конфігураціях на GPU NVIDIA B300, включаючи модель з 1,2 трильйона параметрів, у якій 58,36 мільярда параметрів активні на токен на 512 GPU, де найвища зафіксована пропускна здатність склала 858 TFLOP/s на GPU. Ai2 заявив, що ці тести використовували випадкову маршрутизацію для вимірювання продуктивності системи, а не якість навченого моделі. Технічний звіт перераховує виміряні робочі точки від моделі з 12,9 мільярда параметрів на 16 GPU до моделі з 1,2 трильйона параметрів на 512 GPU, і зазначає, що наведені показники є посиланнями на системи, виміряними при випадковій маршрутизації, а не контрольованою кривою масштабування або заявою про час до досягнення якості.

Ai2 також експериментував з DeepEP v2, альтернативним бекендом для комунікації між експертами через GPU, досягнувши конфігурації з 2,38 трильйонам загальної кількості параметрів. Ai2 описує цей результат як тест короткої місткості, який демонструє масштаб, до якого може дістатися Olmo-core 3, а не як стабільну продуктивність навчання. Технічний звіт під назвою “Supercharging Olmo-core for Efficient and Scalable MoE Training” датований жовтнем 2026 року; його автори — представники Allen Institute for AI та University of Washington, серед яких Tianhua Tao зазначений як головний автор і головний розробник.

Документовані результати та плани щодо наступного покоління Olmo

У звіті документовано патерн збою, який Ai2 називає “token gerrymandering”, коли оцінка, призначена для заохочення збалансованого маршрутування, могла покращуватись, хоча реальне навантаження ставало менш збалансованим. Інші задокументовані висновки включають: зниження швидкості навчання експертів, оскільки вони обробляють менше токенів, не покращувало результати у тестованій сімейсті моделей; обчислення на GPU займали різний час, коли змінювалися оброблювані значення, навіть при однакових розмірах матриць; та накладання комунікації та обчислень на окремих потоках GPU не завжди прискорювало навчання і в деяких тестах уповільнювало загальне виконання. У звіті також описані підходи, які були випробувані, але не прийняті, включаючи вивантаження активацій на CPU, яке не могло бути передане через хостове з’єднання, та два схеми накладання, що конкурували з обчисленнями експертів за ресурсами GPU.

Ai2 заявила, що її наступне покоління Olmo використовуватиме архітектуру MoE і що вона прагне зробити його найпотужнішим Olmo на сьогодні, навченим на найбільшому наборі даних та з найдовшим контекстним вікном. Організація повідомила, що Olmo-core 3 є повністю відкритим, тому дослідники та розробники можуть використовувати його для навчання власних MoEs, адаптувати його до різного обладнання та експериментувати з маршрутизацією, паралелізмом та іншими частинами системи. Olmo-core GitHub репозиторій описує проєкт як набори будівельних блоків PyTorch для екосистеми OLMo, має ліцензію Apache-2.0 і може бути встановлений з вихідного коду або з PyPI як ai2-olmo-core.