📄 Технический отчет | 💻 Code | 🧩 Интерактивная демонстрация
Сегодня мы выпускаем — значительное обновление нашего фреймворка для разработки больших языковых моделей, включающее переработанную открытую систему обучения с использованием смеси экспертов (MoE).
Olmo-core 3 разработан для масштабирования обучения MoE-моделей до уровня триллиона параметров при сохранении вычислительной эффективности. Это одна из ключевых систем, лежащих в основе следующего поколения Olmo, и часть нашего постоянного стремления открывать инструменты и инфраструктуру обучения для каждой новой модели.
Обучение больших моделей ИИ требует значительных вычислительных мощностей, что увеличивает затраты и энергопотребление, делая разработку передовых моделей недоступной для многих академических исследователей и небольших лабораторий. MoE-модели предлагают более эффективный подход: они могут содержать гораздо больше обученных компонентов (параметров), не требуя использования всех их для каждого входного сигнала. Однако полная модель всё равно должна храниться в памяти GPU и обновляться во время обучения, а направление входных данных к нужным экспертам — специализированным компонентам внутри MoE — по всему кластеру создает собственные затраты на коммуникацию и координацию. По мере роста MoE эти затраты могут нивелировать большую часть вычислительного преимущества использования только части модели для каждого входного сигнала.
Olmo-core 3 создан, чтобы устранить этот разрыв. В одном из тестов мы увеличили пул экспертов с 8 до 128, при этом выбирая только четырех экспертов на токен (небольшие единицы текста, обрабатываемые языковой моделью), что позволило сохранить количество активных параметров на токен примерно на уровне 3,2 млрд. Общая емкость параметров выросла с 4,6 млрд до 47 млрд, в то время как пропускная способность обучения снизилась менее чем на 5%.
Та же инфраструктура была протестирована на моделях с общим количеством параметров более одного триллиона.
Создание стека обучения с учетом принципов работы MoE
Olmo-core развивался вместе с каждым поколением Olmo.
Наша работа над разреженными моделями восходит к OlmoE, где использовалась архитектура MoE с 64 маршрутизируемыми экспертами. Olmo 3, напротив, использовала плотную архитектуру, что означает, что почти вся модель была активна для каждого токена, и стек обучения был построен вокруг этого дизайна. Olmo-core 3 расширяет фреймворк системой обучения, разработанной для гораздо более крупных MoE-моделей.
Наша ранняя реализация MoE в Olmo-core использовала полностью шардированный параллелизм данных (FSDP), настроенный на сборку и повторное шардирование весов модели для каждого небольшого пакета обучающих данных. Olmo-core 3 переходит на систему, основанную на распределенном параллелизме данных (DDP). Она сохраняет экспертов в памяти GPU и направляет к ним соответствующие данные, избегая повторной сборки весов.
NVIDIA Megatron-Core является признанным решением для обучения больших MoE. Olmo-core 3 привносит интегрированный стек обучения MoE во фреймворк, лежащий в основе Olmo, с переработкой, которая повышает пропускную способность по сравнению с нашей предыдущей реализацией на базе FSDP. В предварительном тесте на восьми GPU NVIDIA B300 MoE с 47 миллиардами параметров обрабатывала 52 000 токенов в секунду на GPU с новым стеком, по сравнению с 19 400 при использовании нашей предыдущей реализации — это примерно в 2,7 раза выше пропускная способность.
Масштабирование и оптимизация обучения MoE
Olmo-core 3 сочетает в себе несколько методов распределения больших MoE по кластерам GPU с оптимизациями, которые делают маршрутизацию и вычисления более эффективными.
Три метода определяют, как модель и состояние её обучения распределяются между аппаратным обеспечением:
- Параллелизм экспертов распределяет экспертов по GPU, поэтому каждый GPU хранит только часть полного пула экспертов.
- Конвейерный параллелизм разделяет слои модели — последовательные этапы, преобразующие входные данные — между группами GPU, уменьшая объем модели, который каждый GPU должен хранить в памяти.
- Распределенный оптимизатор распределяет состояние оптимизатора — дополнительные данные, используемые для вычисления и применения обновлений во время обучения — между GPU вместо хранения полной копии на каждом из них.
Вместе эти методы позволяют масштабировать MoE, не требуя от каждого GPU хранить всю модель и состояние её обучения в памяти.
Olmo-core 3 также снижает затраты на маршрутизацию данных к нужным экспертам и выполнение их вычислений. Построчный параллелизм экспертов размещает маршрутизируемые данные непосредственно в буферах ввода экспертов, минимизируя дополнительную работу, необходимую для их переупорядочивания. Маршрутизация, выполняемая на GPU, сохраняет метаданные маршрутизации на GPU, поэтому CPU может ставить задачи в очередь, не дожидаясь копирования этой информации обратно. А групповой GEMM объединяет множество небольших вычислений экспертов, чтобы GPU могли выполнять их более эффективно.
Наконец, Olmo-core 3 поддерживает MXFP8, формат чисел с пониженной точностью, который представляет некоторые значения меньшим количеством бит. Это может сократить вычисления и объем данных, передаваемых между GPU, при условии, что эта экономия перевешивает затраты на преобразование между форматами чисел.
Мы измерили влияние MXFP8 на сквозную пропускную способность обучения в контролируемом тесте на четырех GPU NVIDIA B300, где работа была распределена равномерно между экспертами. При включении MXFP8 в тех частях системы, где это принесло наибольшую пользу, пропускная способность обучения была примерно на 21% выше, чем при использовании BF16 — формата с более высокой точностью, который мы использовали в качестве базового, в то время как пиковая активная память снизилась со 103 ГБ до 95 ГБ. Большая часть прироста пришлась на прямое распространение вычислений и перемещение данных между экспертами, а не только на механизм внимания.
Эти методы и оптимизации должны работать вместе. Ускорение одной части обучения может привести к затратам в других; более быстрые вычисления могут потребовать большего перемещения данных, а передача меньшего количества бит может не помочь, если преобразование данных занимает слишком много времени. Olmo-core 3 построен с учетом этих компромиссов на протяжении всего процесса обучения, предоставляя нам — и исследователям, использующим открытый стек — контроль над тем, как эти части сочетаются друг с другом.
Изучите наш интерактивный гид, чтобы увидеть, как параллелизм данных, экспертов и конвейерный параллелизм работают вместе для масштабирования обучения MoE — от одного GPU до множества.
Масштабирование до уровня триллиона параметров
Мы протестировали 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, а не устойчивую производительность обучения.
При таких масштабах производительность системы — это лишь часть картины. Наш технический отчет также документирует эксперименты, которые легли в основу того, как мы обучаем MoE и измеряем их производительность. Например:
- Показатель, призванный стимулировать сбалансированную маршрутизацию, может улучшаться даже тогда, когда фактическая рабочая нагрузка становится менее сбалансированной. Мы называем этот сбой «токеновым джерримендерингом».
- Снижение скорости обучения экспертов — размера их обновлений при обучении — из-за того, что они обрабатывают меньше токенов, не улучшило результаты в протестированном нами семействе моделей.
- Вычисления на GPU занимали разное время при изменении обрабатываемых значений, даже при одинаковых размерностях матриц. Поэтому для сравнения производительности необходимо использовать как одинаковые входные значения, так и одинаковые формы.
- Совмещение передачи данных и вычислений в отдельных потоках GPU не всегда ускоряло обучение. В некоторых тестах это замедляло сквозное выполнение — напоминание о том, что большее совмещение не обязательно означает более высокую пропускную способность.
В отчете объясняются эти выводы, а также подходы, которые мы протестировали и решили не использовать.
Создано для следующего поколения Olmo, открыто для всех
Olmo-core 3 — это фундамент для того, что мы создаем дальше. Наш Olmo следующего поколения будет использовать архитектуру MoE, и мы стремимся к тому, чтобы он стал нашим самым мощным Olmo на сегодняшний день, обученным на нашем крупнейшем наборе данных и с самым длинным контекстным окном.
Новый стек позволяет нам масштабироваться за пределы наших предыдущих работ с MoE, предоставляя при этом больше гибкости для адаптации обучения по мере развития моделей и оборудования. И он полностью открыт — исследователи и разработчики могут использовать Olmo-core 3 для обучения собственных MoE, адаптации к различному оборудованию и экспериментов с маршрутизацией, параллелизмом и другими частями системы.
Это часть нашего подхода к разработке открытых моделей — веса моделей полезнее, когда инфраструктура и решения по обучению, стоящие за ними, также открыты.
Для более глубокого ознакомления с проектированием систем, экспериментами, абляциями и подходами, которые мы протестировали в процессе работы, прочитайте наш технический отчет и .
