Ai2 представила Olmo-core 3 — переработанную систему обучения больших языковых моделей на архитектуре mixture-of-experts (MoE). Релиз рассчитан на масштабирование MoE до триллиона параметров с сохранением вычислительной эффективности. Это один из ключевых компонентов следующего поколения Olmo, и он полностью открыт: код, технический отчёт и интерактивная демонстрация доступны публично.
Зачем понадобилась переработка
Обучение крупных моделей требует огромных вычислений — отсюда растущие расходы и энергопотребление, из-за которых разработка передовых моделей недоступна многим академическим группам и небольшим лабораториям. MoE-модели предлагают более экономный путь: они содержат гораздо больше обучаемых компонентов (параметров), но не задействуют все сразу для каждого входа. Однако полная модель всё равно должна храниться в памяти GPU и обновляться в процессе обучения, а маршрутизация входов к нужным экспертам по кластеру создаёт собственные расходы на коммуникацию и координацию. По мере роста MoE эти издержки могут съедать значительную часть выигрыша от того, что на каждый вход работает лишь часть модели.
Olmo-core 3 создана, чтобы закрыть этот разрыв. В одном из бенчмарков пул экспертов увеличили с 8 до 128, при этом на каждый токен по-прежнему выбирались только четыре эксперта, а число активных параметров на токен оставалось примерно на уровне 3,2 млрд. Общая ёмкость параметров выросла с 4,6 млрд до 47 млрд, а пропускная способность обучения упала менее чем на 5%. Та же инфраструктура прошла тесты на более чем одном триллионе суммарных параметров.
От FSDP к DDP: эксперты остаются на GPU
История стека развивалась вместе с поколениями Olmo. Работа над разреженными моделями началась с OlmoE, где использовалась MoE-архитектура с 64 маршрутизируемыми экспертами. В Olmo 3, напротив, применили плотную архитектуру — почти вся модель активна для каждого токена, и стек обучения строился под этот дизайн. Olmo-core 3 расширяет фреймворк системой, рассчитанной на куда более крупные MoE.
Прежняя реализация MoE в Olmo-core опиралась на fully sharded data parallelism (FSDP): веса модели собирались и заново шардировались для каждого небольшого батча обучающих данных. Olmo-core 3 переходит на систему на базе distributed data parallelism (DDP): эксперты постоянно живут на GPU, а к ним направляются нужные данные — без повторной сборки весов.
NVIDIA Megatron-Core — устоявшийся вариант для обучения больших MoE. Olmo-core 3 приносит интегрированный MoE-стек во фреймворк, стоящий за Olmo, и за счёт переработки повышает пропускную способность по сравнению с прежней реализацией на FSDP. В предварительном тесте на восьми GPU NVIDIA B300 MoE на 47 млрд параметров обрабатывала 52 000 токенов в секунду на GPU с новым стеком против 19 400 со старой реализацией — примерно в 2,7 раза выше.
Параллелизм и оптимизации маршрутизации
Olmo-core 3 сочетает несколько приёмов распределения крупных MoE по кластерам GPU с оптимизациями, ускоряющими маршрутизацию и вычисления.
Три техники определяют, как модель и её состояние обучения разбиваются по железу:
- Expert parallelism распределяет экспертов по GPU, так что каждая карта хранит лишь часть полного пула.
- Pipeline parallelism делит слои модели — последовательные стадии преобразования входа — между группами GPU, снижая объём модели, который каждой карте нужно держать в памяти.
- Распределённый оптимизатор разносит состояние оптимизатора — дополнительные данные для расчёта и применения обновлений при обучении — по GPU вместо хранения полной копии на каждой карте.
Вместе эти приёмы позволяют масштабировать MoE без требования держать всю модель и её состояние обучения в памяти каждого GPU.
Отдельно снижается стоимость маршрутизации данных к нужным экспертам и выполнения их вычислений. Rowwise expert parallelism размещает маршрутизированные данные прямо во входных буферах экспертов, минимизируя лишнюю работу по их перестановке. GPU-resident routing держит метаданные маршрутизации на GPU, чтобы CPU мог ставить задачи в очередь, не дожидаясь копирования этой информации обратно. А grouped GEMM объединяет множество мелких вычислений экспертов, чтобы GPU выполняли их эффективнее.
Наконец, Olmo-core 3 поддерживает MXFP8 — формат чисел пониженной точности, представляющий часть значений меньшим числом бит. Он может сократить вычисления и объём данных, перемещаемых между GPU, — при условии, что экономия перевешивает затраты на конвертацию между форматами.
Эффект MXFP8 на сквозную пропускную способность измерили в контролируемом бенчмарке на четырёх GPU NVIDIA B300 с равномерным распределением работы по экспертам. При включении MXFP8 в тех частях системы, где он помогал больше всего, пропускная способность обучения оказалась примерно на 21% выше, чем с BF16 — более точным форматом, взятым за базу, — а пиковая активная память снизилась со 103 ГиБ до 95 ГиБ. Основной выигрыш дала feed-forward-часть вычислений и перемещение данных между экспертами, а не только внимание.
Эти техники и оптимизации обязаны работать сообща. Ускорение одной части обучения создаёт издержки в другой: более быстрые вычисления могут потребовать большего перемещения данных, а передача меньшего числа бит не поможет, если конвертация занимает слишком много времени. Olmo-core 3 выстроена вокруг этих компромиссов на всём протяжении обучения.
Триллион параметров и что показали тесты
Olmo-core 3 протестировали в разных конфигурациях на GPU NVIDIA B300, включая модель на 1,2 трлн параметров с 58,36 млрд активных параметров на токен на 512 GPU. Максимальная наблюдённая пропускная способность составила 858 TFLOP/s на GPU — мера полезных вычислений модели в секунду на каждой карте. В этих тестах применялась случайная маршрутизация для измерения производительности системы, а не качества обученной модели.
Также экспериментировали с DeepEP v2 — альтернативным способом организации коммуникации между экспертами на разных GPU — и достигли конфигурации с 2,38 трлн суммарных параметров. Это был краткий тест ёмкости, а не полный прогон обучения: он демонстрирует масштаб, до которого способна дотянуться Olmo-core 3, а не устойчивую производительность обучения.
На таких масштабах производительность системы — лишь часть картины. В техническом отчёте описаны эксперименты, повлиявшие на то, как Ai2 обучает MoE и измеряет их работу. Например:
- Оценка, призванная поощрять сбалансированную маршрутизацию, могла улучшаться, тогда как фактическая нагрузка становилась менее равномерной. Этот сбой называют token gerrymandering.
- Снижение скорости обучения экспертов — размера их обучающих обновлений — из-за обработки меньшего числа токенов не улучшило результаты в протестированном семействе моделей.
- Вычисления на GPU занимали разное время при изменении обрабатываемых значений, даже при одинаковых размерах матриц. Поэтому для сравнения производительности нужно совпадение не только форм, но и входных значений.
- Перекрытие коммуникации и вычислений на отдельных потоках GPU не всегда ускоряло обучение. В некоторых тестах оно замедляло сквозное выполнение — напоминание, что больший параллелизм не обязательно означает более высокую пропускную способность.
Что дальше и кому это доступно
Olmo-core 3 — фундамент для следующего поколения Olmo. Будущая модель будет использовать MoE-архитектуру, и Ai2 рассчитывает сделать её самой способной версией Olmo: с самым большим датасетом и самым длинным контекстным окном. Новый стек позволяет выйти за пределы прежних MoE-наработок и гибче адаптировать обучение по мере эволюции моделей и железа.
Стек полностью открыт: исследователи и разработчики могут обучать на Olmo-core 3 собственные MoE, адаптировать его под другое железо и экспериментировать с маршрутизацией, параллелизмом и другими частями системы. Логика Ai2 проста: веса моделей полезнее, когда инфраструктура и решения об обучении, стоящие за ними, тоже открыты.
Для команд, которые уже используют открытые модели в продакшене, это означает возможность воспроизвести путь обучения, а не только скачать готовые веса. Если вы разворачиваете агентные системы на открытых моделях, полезно посмотреть, как устроен учёт ИИ-агентов и контроль вызовов инструментов и какие есть открытые computer-use модели.
Частые вопросы
Что такое Olmo-core 3 и кто её выпустил?
Это открытый фреймворк для обучения больших языковых моделей на архитектуре mixture-of-experts, выпущенный Ai2. Это существенное обновление предыдущей версии Olmo-core с переработанной системой обучения MoE.
Что такое MoE и почему это экономичнее обычной модели?
Mixture-of-experts содержит много обучаемых компонентов (экспертов), но на каждый вход задействует лишь часть из них. Это снижает вычисления на токен, хотя полная модель всё равно должна храниться в памяти GPU и обновляться при обучении.
Насколько выросла эффективность по сравнению с прежней реализацией?
В предварительном тесте на восьми GPU NVIDIA B300 MoE на 47 млрд параметров обрабатывала 52 000 токенов в секунду на GPU с новым стеком против 19 400 с прежней реализацией на FSDP — примерно в 2,7 раза выше.
Что даёт переход с FSDP на DDP?
При FSDP веса модели собирались и заново шардировались для каждого небольшого батча. DDP-подход держит экспертов постоянно на GPU и направляет к ним нужные данные, избавляя от повторной сборки весов.
Что такое MXFP8 и сколько он экономит?
Это формат чисел пониженной точности. В контролируемом бенчмарке на четырёх GPU NVIDIA B300 его включение дало примерно на 21% более высокую пропускную способность, чем BF16, а пиковая активная память снизилась со 103 ГиБ до 95 ГиБ.
До какого размера модели протестирована Olmo-core 3?
Тесты включали модель на 1,2 трлн параметров с 58,36 млрд активных параметров на токен на 512 GPU с пиком 858 TFLOP/s на GPU. В экспериментах с DeepEP v2 дошли до конфигурации с 2,38 трлн суммарных параметров, но это был краткий тест ёмкости, а не полный прогон обучения.
Что такое token gerrymandering?
Так в Ai2 назвали сбой, при котором оценка, поощряющая сбалансированную маршрутизацию, улучшалась, тогда как фактическая нагрузка на экспертов становилась менее равномерной.
Можно ли использовать Olmo-core 3 для собственных моделей?
Да. Стек полностью открыт: код и технический отчёт опубликованы, исследователи и разработчики могут обучать на нём собственные MoE, адаптировать под другое железо и экспериментировать с маршрутизацией и параллелизмом.
Источник: huggingface.co