Команда Olmo выпустила Olmo-core 3: открытый фреймворк для обучения больших MoE

Команда Olmo выпустила Olmo-core 3, серьёзное обновление своего фреймворка для разработки больших языковых моделей, в котором переработана открытая система обучения моделей со смесью экспертов (mixture-of-experts, MoE). Авторы называют её одной из ключевых систем для следующего поколения Olmo. Задача, масштабировать обучение MoE до диапазона в триллионы параметров, сохраняя вычислительную эффективность.

Проблема, которую решают авторы: MoE-модели могут содержать много параметров, не задействуя их все на каждом входе, но вся модель всё равно должна храниться в памяти GPU и обновляться при обучении, а маршрутизация данных к нужным экспертам по кластеру создаёт издержки на связь и координацию. С ростом модели эти издержки способны съесть значительную часть выигрыша. В одном из тестов команда увеличила число экспертов с 8 до 128, по-прежнему выбирая четыре эксперта на токен; число активных параметров на токен осталось примерно прежним (около 3,2 млрд), общее число параметров выросло с 4,6 до 47 млрд, а пропускная способность обучения упала менее чем на 5%. Ту же инфраструктуру проверяли и на масштабе свыше триллиона параметров.

По архитектуре: прежняя реализация MoE в Olmo-core использовала полностью шардированный параллелизм данных (FSDP) со сбором и повторным разбиением весов на каждый малый пакет данных. Olmo-core 3 переходит на систему на основе распределённого параллелизма данных (DDP): эксперты постоянно лежат на GPU, а нужные данные направляются к ним, без повторного сбора весов. Предыдущая модель Olmo 3 была плотной, а OlmoE использовала 64 маршрутизируемых эксперта. Авторы отмечают, что Megatron-Core от NVIDIA, устоявшийся вариант для обучения больших MoE, а Olmo-core 3 добавляет в фреймворк Olmo встроенный стек для MoE. В предварительном тесте на восьми GPU NVIDIA B300 MoE на 47 млрд параметров обрабатывала 52 000 токенов в секунду на GPU против 19 400 у прежней реализации, примерно в 2,7 раза больше.

Используются три приёма распределения модели: параллелизм экспертов (эксперты разнесены по GPU), конвейерный параллелизм (слои модели разбиты по группам GPU) и распределённый оптимизатор (состояние оптимизатора разнесено по GPU вместо полной копии на каждом). Чтобы удешевить маршрутизацию и вычисления, добавлены rowwise-параллелизм экспертов (данные сразу попадают во входные буферы экспертов), маршрутизация на стороне GPU (метаданные остаются на GPU, и процессору не нужно ждать их копирования) и сгруппированное умножение матриц (grouped GEMM), объединяющее много мелких вычислений экспертов. Кроме того, поддержан формат пониженной точности MXFP8. В контролируемом тесте на четырёх GPU B300 при равномерной загрузке экспертов и включении MXFP8 там, где он помогал сильнее всего, пропускная способность обучения была примерно на 21% выше, чем с базовым BF16, а пиковая активная память снизилась со 103 до 95 GiB. Основной прирост дали вычисления в слоях прямого распространения (feed-forward) и обмен данными между экспертами, а не одно только внимание (attention).

На масштабе: среди протестированных на B300 конфигураций была модель на 1,2 трлн параметров с 58,36 млрд активных на токен на 512 GPU; наибольшая наблюдавшаяся производительность составила 858 TFLOP/s на GPU. Эти тесты шли со случайной маршрутизацией и измеряли производительность системы, а не качество обученной модели. С альтернативным способом обмена между экспертами DeepEP v2 удалось достичь конфигурации на 2,38 трлн параметров, но это короткий тест ёмкости, а не полноценное обучение: он показывает достижимый масштаб, а не устойчивую скорость обучения.

В техническом отчёте описаны и эксперименты. Метрика, призванная поощрять сбалансированную маршрутизацию, могла улучшаться, пока реальная нагрузка становилась менее сбалансированной, авторы называют этот сбой «token gerrymandering». Снижение скорости обучения (learning rate) для экспертов из-за того, что они обрабатывают меньше токенов, не улучшило результат на проверенном семействе моделей. Время GPU-вычислений зависело от самих значений, а не только от размерностей матриц, поэтому сравнивать производительность нужно на одинаковых входных значениях. Перекрытие связи и вычислений на отдельных потоках GPU не всегда ускоряло обучение, а в части тестов замедляло его.

Следующее поколение Olmo будет на архитектуре MoE; авторы стремятся сделать его самым способным Olmo на самом большом датасете и с самым длинным контекстным окном. Olmo-core 3 полностью открыт: исследователи и разработчики могут обучать на нём собственные MoE, адаптировать его под другое железо и экспериментировать с маршрутизацией и параллелизмом. Технический отчёт и код опубликованы, код, на GitHub.

Ключевые факты

• Olmo-core 3, открытый фреймворк с новой системой обучения MoE; рассчитан на масштаб порядка триллиона параметров, следующий Olmo будет на MoE.
• Переход с FSDP на DDP с постоянным размещением экспертов на GPU: в предварительном тесте на восьми B300 модель на 47 млрд параметров дала 52 000 против 19 400 токенов/с на GPU, около 2,7 раза.
• При росте числа экспертов с 8 до 128 (общее число параметров с 4,6 до 47 млрд, активных около 3,2 млрд) пропускная способность упала менее чем на 5%.
• MXFP8 на четырёх B300 дал около 21% прироста пропускной способности против BF16 и снизил пиковую память со 103 до 95 GiB.
• Бенчмарк на 1,2 трлн параметров (58,36 млрд активных, 512 GPU) достиг 858 TFLOP/s на GPU при случайной маршрутизации; 2,38 трлн параметров с DeepEP v2, лишь короткий тест ёмкости.

Почему это важно

Обучение больших моделей требует огромных вычислений, что, по словам авторов, повышает затраты и расход энергии и отодвигает разработку передовых моделей от академических исследователей и небольших лабораторий. MoE-подход эффективнее, но на больших масштабах издержки связи и координации могут съедать его выигрыш. Olmo-core 3 нацелен именно на этот разрыв и публикует весь стек обучения открытым, а не только веса моделей. Авторы подчёркивают: веса моделей полезнее, когда открыты инфраструктура и решения об обучении.

Кому это важно

Исследователям и командам, которые обучают собственные MoE-модели или хотят понять, как устроено обучение таких моделей в масштабе. Авторы прямо называют целевой аудиторией исследователей и разработчиков, которые могут адаптировать стек под другое железо и экспериментировать с маршрутизацией и параллелизмом. Интересен и тем, кто выбирает между готовыми стеками: авторы сами упоминают Megatron-Core от NVIDIA как устоявшийся вариант.

Как это применить

Код доступен на GitHub, к нему приложены технический отчёт и интерактивное объяснение того, как параллелизм данных, экспертов и конвейерный параллелизм вместе масштабируют обучение MoE от одного GPU до многих. Из отчёта можно взять практические наблюдения: не снижать learning rate экспертов только из-за меньшего числа обработанных токенов (на проверенном семействе моделей это не помогло), сравнивать производительность GPU на одинаковых входных значениях и проверять, не замедляет ли перекрытие связи и вычислений обучение. Условия лицензии в тексте материала не приведены.

Можно ли доверять

Это публикация самой команды Olmo о собственном релизе, а не независимое исследование: все цифры измерены и опубликованы авторами. Плюс, они сами указывают ограничения: 2,7-кратный прирост получен в предварительном тесте на восьми GPU B300, результаты на масштабе триллиона получены со случайной маршрутизацией, а рекорд 2,38 трлн параметров, короткий тест ёмкости. Прямого сравнения по скорости с Megatron-Core в материале нет.

Риски и подводные камни

Бенчмарки масштаба триллиона параметров измеряют производительность системы, а не качество обученной модели; обученной модели такого размера в материале не заявлено. Значение 858 TFLOP/s на GPU, наибольшее наблюдавшееся, а не типичное. Выигрыш MXFP8 измерен на четырёх GPU при равномерной загрузке экспертов и только там, где формат помогал больше всего; авторы оговаривают, что он окупается лишь когда экономия превышает стоимость преобразования форматов. Ускорение одной части обучения может создавать издержки в другой, поэтому результаты в другой конфигурации и на другом железе могут отличаться. Стоимость, энергопотребление и время обучения в бенчмарках не приведены.

«Веса модели полезнее, когда открыты и инфраструктура, и решения об обучении, стоящие за ними.»

— Команда Olmo, блог-пост о релизе Olmo-core 3