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

Институт Аллена в области ИИ выпустил Olmo-core 3 1 октября 2026 года, обновление своей платформы разработки больших языковых моделей, построенной вокруг переработанной открытой системы обучения смеси экспертов, которую, по словам 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

Ранее реализованная Ai2 система MoE в Olmo-core использовала полностью шардинг‑параллелизм данных, настроенный на сбор и повторное распределение весов модели для каждой небольшой партии обучающих данных. 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, конвейерный параллелизм делит слои модели между группами GPU, а распределённый оптимизатор распределяет состояние оптимизатора по GPU вместо хранения полной копии на каждом GPU. Olmo-core 3 также снижает затраты на маршрутизацию данных к нужным экспертам: построчный параллелизм экспертов помещает маршрутизируемые данные непосредственно в буферы ввода экспертов, маршрутизация, находящаяся на GPU, сохраняет метаданные маршрутизации на GPU, позволяя CPU ставить задачи в очередь без ожидания их копирования обратно, а групповой GEMM объединяет множество небольших вычислений экспертов, чтобы GPU могли выполнять их более эффективно. Технический доклад описывает дизайн построчного параллелизма экспертов на основе NVSHMEM, который записывает каждый токен напрямую в буфер соответствующего эксперта, с планированными на устройстве групповыми умножениями матриц, делающими параллелизм экспертов полностью свободным от синхронизации.

Olmo-core 3 поддерживает MXFP8, формат чисел с пониженной точностью. В контролируемом бенчмарке на четырёх GPU NVIDIA B300 с равномерным распределением работы между экспертами Ai2 сообщает, что пропускная способность обучения увеличилась примерно на 21 % по сравнению с базовым BF16, при этом пиковая активная память упала с 103 ГиБ до 95 ГиБ, причём большинство прироста пришлось на вычисления прямого распространения и перемещение данных между экспертами. В докладе добавлено, что топологически независимые контрольные точки фиксируют глобальные тензоры FP32 независимо от параллельного расположения, позволяя возобновлять запуск на иной топологии, а в контролируемом сравнении перевычисление активаций устранило почти две трети памяти активаций и сократило пиковую память примерно на четверть ценой потери около пятой части пропускной способности.

Бенчмарки с триллионными параметрами и заявленные ограничения

Ai2 заявила, что провела бенчмарки Olmo-core 3 в различных конфигурациях на GPU NVIDIA B300, включая модель с 1,2 триллиона параметров, у которой на каждый токен активно 58,36 миллиарда параметров на 512 GPU, где максимальная зафиксированная пропускная способность составила 858 TFLOP/с на 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, при этом Тянхуа Тао указан как ведущий автор и основной разработчик.

Документированные выводы и планы по следующему поколению Olmo

В отчёте описан паттерн отказа, который Ai2 называет «токеновым джерримендерингом», при котором оценка, предназначенная для поощрения сбалансированной маршрутизации, могла улучшаться, даже когда реальная нагрузка становилась менее сбалансированной. Другие зафиксированные выводы включают: снижение скорости обучения экспертов, поскольку они обрабатывают меньше токенов, не улучшало результаты в тестируемом семействе моделей; расчёты на GPU занимали разное время при изменении обрабатываемых значений, даже при одинаковых размерах матриц; а наложение коммуникаций и вычислений на отдельные потоки GPU не всегда ускоряло обучение и в некоторых тестах замедляло конечное выполнение. В отчёте также описаны подходы, протестированные, но не принятые, включая выгрузку активаций на CPU, которую не могла поддержать связь хоста, и две схемы наложения, конкурирующие с вычислениями экспертов за ресурсы GPU.

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