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

Что выпустил Ai2

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

Заявленная цель — масштабировать обучение MoE до триллиона параметров, не теряя вычислительной эффективности; это одна из базовых систем следующего поколения Olmo.

Обучение больших моделей требует много вычислений, что повышает стоимость и энергопотребление и делает передовые разработки недоступными для академических исследователей и небольших лабораторий.

MoE экономнее устроены: в модели может быть гораздо больше параметров, но каждый вход использует не все из них. При этом полную модель всё равно приходится хранить в памяти GPU и обновлять при обучении, а маршрутизация входов к нужным экспертам по кластеру создаёт собственные расходы на коммуникацию и координацию.

По мере роста MoE эти издержки способны съесть значительную часть выигрыша — Olmo-core 3 и создавался, чтобы закрыть этот разрыв.

Замеры: больше экспертов без потери скорости

В одном из бенчмарков Ai2 увеличил пул экспертов с 8 до 128, по-прежнему выбирая для каждого токена только четыре эксперта, и держал число активных параметров на токен примерно постоянным — около 3,2 млрд.

При этом общая ёмкость выросла с 4,6 до 47 млрд параметров, а пропускная способность обучения упала меньше чем на 5 %.

Ту же инфраструктуру тестировали и на конфигурации более чем с триллионом суммарных параметров.

От FSDP к DDP

Olmo-core развивается вместе с поколениями Olmo: работа над разреженными моделями началась с OlmoE с 64 routed experts, а Olmo 3 был плотной архитектурой, где почти вся модель работала на каждом токене.

Прежняя реализация 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 и по предварительному тесту на восьми GPU NVIDIA B300 модель с 47 млрд параметров обрабатывала 52 000 токенов в секунду на GPU против 19 400 у прежней реализации — примерно в 2,7 раза больше.

Как модель распределяется по кластеру

Три техники определяют, как модель и её состояние обучения раскладываются по железу: expert parallelism распределяет экспертов по GPU, так что каждая хранит лишь часть полного пула.

Pipeline parallelism делит слои модели между группами GPU и снижает объём модели, который каждой из них нужно держать в памяти.

Distributed optimizer распределяет состояние оптимизатора по GPU вместо полной копии на каждой.

Вместе эти приёмы позволяют масштабировать MoE, не требуя, чтобы каждая GPU держала всю модель и её состояние обучения в памяти.

Маршрутизация и MXFP8

Отдельно снижена цена маршрутизации данных к экспертам: rowwise expert parallelism кладёт маршрутизированные данные прямо в буферы входов экспертов, GPU-resident routing держит метаданные маршрутизации на GPU, чтобы CPU не ждал их копирования назад, а grouped GEMM объединяет множество мелких вычислений экспертов.

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

В контролируемом замере на четырёх GPU NVIDIA B300 с равномерным распределением работы по экспертам включённый там, где он помогал, MXFP8 дал примерно на 21 % более высокую пропускную способность обучения, чем BF16, взятый за базу, а пиковая активная память упала с 103 до 95 GiB.

Основную часть выигрыша дали вычисления feed-forward и пересылка данных между экспертами, а не только внимание.

Триллион параметров и оговорки

Ai2 протестировал Olmo-core 3 на ряде конфигураций с GPU NVIDIA B300, включая модель на 1,2 трлн параметров, из которых 58,36 млрд активны на токен, на 512 GPU: максимальная наблюдённая пропускная способность составила 858 TFLOP/s/GPU.

Эти тесты использовали случайную маршрутизацию, чтобы измерить производительность системы, а не качество обученной модели.

Отдельно экспериментировали с DeepEP v2 — альтернативным способом организации коммуникации между экспертами на разных GPU — и дошли до конфигурации с 2,38 трлн суммарных параметров; это был короткий тест ёмкости, а не полный прогон обучения, поэтому он показывает достижимый масштаб, а не устойчивую производительность.

Что нашлось в техническом отчёте

В техническом отчёте описаны эксперименты, которые повлияли на то, как Ai2 обучает MoE и измеряет их качество.

Оценка, призванная поощрять сбалансированную маршрутизацию, могла улучшаться, пока реальная нагрузка становилась менее сбалансированной — этот провал в Ai2 называют token gerrymandering.

Понижение learning rate экспертов из-за того, что они обрабатывают меньше токенов, не улучшило результаты в протестированном семействе моделей.

Вычисления на GPU занимали разное время при разных обрабатываемых значениях даже при одинаковых размерах матриц, поэтому сравнивать производительность нужно при совпадающих входных значениях, а не только формах.

Совмещение коммуникации и вычислений на разных GPU-стримах не всегда ускоряло обучение, а в некоторых тестах замедляло end-to-end исполнение.

Что дальше и что это за цифры

Olmo-core 3 — фундамент следующего шага Ai2: следующее поколение Olmo будет использовать архитектуру MoE и должно стать самым способным Olmo, обученным на самом большом датасете и с самым длинным контекстным окном.

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

Технический отчёт, код на GitHub и интерактивный разбор доступны по ссылкам в самой публикации.

Все приведённые числа — данные Ai2 из её собственной публикации, без независимых замеров: токены в секунду на GPU, TFLOP/s/GPU и объём памяти получены в тестах на стендах института.