Ai2 выпустила Olmo-core 3 для обучения MoE-моделей до триллиона параметров
Ai2 открыла Olmo-core 3, фреймворк для обучения больших MoE-моделей. В тесте на восьми GPU он дал в 2,7 раза большую пропускную способность, чем прежняя реализация.
Ai2 выпустила Olmo-core 3, открытый фреймворк для обучения больших MoE-моделей. MoE (mixture of experts, смесь экспертов) — архитектура, в которой для каждого токена (небольшого фрагмента текста) задействуется только часть параметров. В одном тесте число экспертов выросло с 8 до 128, общее число параметров модели увеличилось с 4,6 до 47 млрд, а пропускная способность снизилась менее чем на 5%.
Новый стек рассчитан на обучение MoE с триллионом параметров и больше. Прежняя реализация в Olmo-core собирала и пересобирала веса при обработке каждого небольшого батча через FSDP. В Olmo-core 3 применена распределённая параллелизация данных (DDP): эксперты постоянно хранятся на GPU, а данные направляются к нужным из них. На восьми NVIDIA B300 модель на 47 млрд параметров обрабатывала 52 тыс. токенов в секунду на GPU против 19,4 тыс. у прежней реализации, примерно в 2,7 раза больше.
Чтобы распределять модель по кластеру, Olmo-core 3 использует параллелизм экспертов, конвейерный параллелизм и распределённый оптимизатор. Первый разносит экспертов по GPU, второй делит слои модели между группами GPU, третий хранит состояние оптимизатора распределённо. Дополнительно rowwise expert parallelism размещает данные сразу во входных буферах экспертов, GPU-resident routing держит метаданные маршрутизации на видеокартах, а grouped GEMM объединяет мелкие вычисления экспертов.
Фреймворк поддерживает формат MXFP8 с меньшей точностью. На четырёх B300 при равномерном распределении работы он дал примерно на 21% большую пропускную способность, чем BF16, а пиковая активная память снизилась со 103 до 95 GiB. В тесте на 512 GPU модель на 1,2 трлн параметров с 58,36 млрд активных на токен показала до 858 TFLOP/s на GPU. Тесты использовали случайную маршрутизацию: проверялась работа системы, качество обученной модели не измерялось.
Есть и эксперимент с DeepEP v2 на 2,38 трлн параметров. Его провели как короткий тест ёмкости, без полного обучения. Olmo-core 3 полностью открыт: код доступен на GitHub, исследователи и разработчики могут обучать на нём свои MoE и адаптировать его под другое железо. Следующую Olmo компания планирует строить на MoE-архитектуре с самым большим набором данных и самым длинным контекстом.