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

Эффективность моделей MoE

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

Преимущества Olmo-core 3

В одном из тестов мы увеличили пул экспертов с 8 до 128, при этом выбирая только четырех экспертов на токен — малую единицу текста, обрабатываемую языковой моделью. Общее количество параметров увеличилось с 4.6 миллиарда до 47 миллиардов, в то время как производительность обучения снизилась менее чем на 5%. Эта же инфраструктура была протестирована на более чем одном триллионе параметров.

Эволюция Olmo-core

Olmo-core развивался с каждым новым поколением Olmo. Наша работа над разреженными моделями началась с OlmoE, использовавшей архитектуру MoE с 64 маршрутизируемыми экспертами. В отличие от этого, Olmo 3 использовал плотную архитектуру, где почти вся модель была активна для каждого токена. Olmo-core 3 расширяет фреймворк с системой обучения, предназначенной для гораздо больших моделей MoE. Ранее реализация MoE в Olmo-core использовала полностью шардированную параллельность данных (FSDP), которая собирала и перераспределяла веса модели для каждой небольшой партии обучающих данных. Olmo-core 3 переключается на систему, основанную на распределенной параллельности данных (DDP), которая сохраняет экспертов на GPU и направляет соответствующие данные к ним, избегая повторного сбора весов.

Оптимизация обучения MoE

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

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

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

Распределенный оптимизатор: состояние оптимизатора распределяется по GPU, а не хранится в полном объеме на каждом GPU.

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

Поддержка MXFP8

Olmo-core 3 поддерживает MXFP8, формат чисел с пониженной точностью, который представляет некоторые значения с меньшим количеством бит. Это может снизить вычислительные затраты и объем данных, перемещаемых между GPU, если эти экономии превышают затраты на преобразование между форматами чисел. Мы измерили влияние MXFP8 на общую производительность обучения в контролируемом тесте на четырех GPU NVIDIA B300, где работа распределялась равномерно между экспертами. При включенном MXFP8 производительность обучения была примерно на 21% выше, чем с использованием BF16, более точного формата, который мы использовали в качестве базового.

Масштабирование до триллиона параметров

Мы протестировали Olmo-core 3 в различных конфигурациях на GPU NVIDIA B300, включая модель с 1.2 триллиона параметров, где 58.36 миллиарда параметров были активны на токен, распределенных по 512 GPU. Наивысшая зафиксированная производительность составила 858 TFLOP/s/GPU — мера полезных вычислений модели в секунду на каждом GPU. Эти тесты использовали случайную маршрутизацию для измерения производительности системы, а не качества обученной модели. Мы также экспериментировали с DeepEP v2, альтернативным способом обработки коммуникации между экспертами на GPU, достигнув конфигурации с 2.38 триллиона параметров.

Открытость и доступность Olmo-core 3

Olmo-core 3 является основой для того, что мы строим дальше. Наша следующая версия Olmo будет использовать архитектуру MoE, и мы стремимся сделать ее нашей самой мощной моделью, обученной на нашем самом большом наборе данных и с самым длинным контекстом. Новый стек позволяет нам масштабироваться за пределы нашей предыдущей работы с MoE, предоставляя больше гибкости для адаптации обучения по мере эволюции моделей и аппаратного обеспечения. И он полностью открыт — исследователи и разработчики могут использовать Olmo-core 3 для обучения своих собственных MoE, адаптировать его к различному оборудованию и экспериментировать с маршрутизацией, параллелизмом и другими частями системы. Это часть нашего подхода к открытому развитию моделей — веса моделей становятся более полезными, когда инфраструктура и решения по обучению за ними также открыты.