Prime Intellect выпустила Prime Mega MoE: обучаемое объединённое MoE-ядро, прямой проход в 3,11 раза быстрее
Prime Intellect выпустила Prime Mega MoE — расширение для обучения на базе объединённого ядра Mega MoE от DeepSeek. Ядро DeepSeek выполняло только прямой проход. Prime Intellect добавила быстрый обратный проход с перекрытием операций. Благодаря этому объединённое ядро можно использовать для обучения, а не только для инференса. В своём обучающем модуле prime-rl полные шаги обучения на моделях Qwen3 MoE выполняются до 34% быстрее. По сравнению с TorchTitan прирост достигает 55%. Для команд, обучающих модели со смесью экспертов (MoE), это означает прямое сокращение времени работы на GPU.
Кратко
- Размер: Ядро GPU, а не модель. Протестировано на Qwen3-30B-A3B и Qwen3-235B-A22B (128 экспертов, маршрутизация top-8).
- Работает на: GPU NVIDIA Hopper (блочный FP8) и Blackwell (MXFP8), BF16 поддерживается на обеих платформах. Опубликованные результаты получены на узлах с 8×H200, от 32 до 128 GPU.
- Производительность: полные шаги обучения на 32–34% быстрее базовой версии prime-rl. На 48–55% быстрее TorchTitan.
- Лучший результат: прямой проход на слое MoE Qwen3-30B-A3B выполняется в 3,11 раза быстрее (12,0 мс против 3,9 мс).
- Худший результат: обратный проход на слое MoE Qwen3-235B-A22B выполняется в 1,24 раза быстрее (30,5 мс против 24,6 мс).
Что такое Prime Mega MoE?
Prime Mega MoE — это объединённое ядро MoE с перекрытием операций для обучения от Prime Intellect. Оно основано на ядре Mega MoE из библиотеки DeepSeek с открытым исходным кодом DeepGEMM.
Почему обучение MoE ограничено пропускной способностью коммуникаций?
В модели MoE каждый токен направляется к небольшому подмножеству экспертов. Эти эксперты работают на разных GPU. Поэтому каждый слой отправляет токены, выполняет вычисления экспертов, а затем отправляет результаты обратно. Это происходит как в прямом, так и в обратном проходе.
В стандартной конфигурации это отдельные операции. GPU сортирует токены и ждёт транспортный уровень. Затем он выполняет вычисления, возвращает результаты и снова ждёт. Эти ожидания означают простой GPU.
Как работает объединённое ядро?
Mega MoE от DeepSeek, добавленное в DeepGEMM в апреле 2026 года, заменяет эту цепочку одним вызовом ядра. Оно объединяет диспетчеризацию параллельных экспертов, 2 линейных слоя экспертов, SwiGLU и объединение результатов. Пока одна группа токенов находится в пути, ядро выполняет вычисления над уже прибывшей группой. Кроме того, оно пропускает промежуточные тензоры, используя буферы временного хранения только на GPU. Однако изначально была реализована только версия для прямого прохода.
Prime Intellect написала недостающее ядро обратного прохода, также с перекрытием операций. Поддерживаются типы BF16, MXFP8 (только Blackwell) и блочный FP8 (только Hopper). Конфигурация prime-rl также предоставляет режим fp8_bwd: прямой проход в BF16, обратный проход в FP8.
Насколько быстро работает Prime Mega MoE?
Все приведённые ниже показатели предоставлены самой Prime Intellect. Шаги обучения выполнялись на узлах с 8×H200, в FP8, с последовательностями длиной 16 384 токена и полным контрольным сохранением активаций. Каждый результат представляет собой среднее значение для шагов с 6-го по 20-й.
Полный шаг обучения, секунды (меньше — лучше):
| Модель (кластер) | Базовая версия prime-rl | TorchTitan (FP8) | Prime Mega MoE | Прирост относительно базовой версии | Прирост относительно TorchTitan |
|---|---|---|---|---|---|
| Qwen3-30B-A3B (32 H200) | 2.78 | 3.07 | 2.08 | 34% | 48% |
| Qwen3-235B-A22B (128 H200) | 12.72 | 14.98 | 9.65 | 32% | 55% |
По сравнению с TorchTitan пропускная способность выросла с 5 337 до 7 886 токенов/с/GPU на модели 30B. На модели 235B она выросла с 1 094 до 1 697. Базовая версия prime-rl использует torch all-to-all и экспертов DeepGEMM FP8.
Один слой MoE, миллисекунды (1 узел, 8×H200, EP=8, 128 экспертов, top-8, случайная маршрутизация):
| Слой | Прямой проход | Обратный проход | Полный слой (2x прямой проход + обратный проход) |
|---|---|---|---|
| Qwen3-30B-A3B | 12.0–3.9 (3.11x) | 15.4–8.2 (1.87x) | 39.3–15.9 (2.47x) |
| Qwen3-235B-A22B | 20.7–9.5 (2.17x) | 30.5–24.6 (1.24x) | 71.9–43.7 (1.64x) |
Картина ясна. Прирост в прямом проходе велик, а в обратном — меньше. Заголовок «до 3,11 раза» относится только к прямому проходу.
Как это решение сравнивается с другими объединёнными вариантами MoE?
| Характеристика | Prime Mega MoE | DeepGEMM Mega MoE | TorchTitan + DeepEP | AMD Primus MegaMoE |
|---|---|---|---|---|
| Разработчик | Prime Intellect | DeepSeek | PyTorch, Nebius | AMD |
| Обучение (обратный проход) | Да, объединённый и с перекрытием операций | Нет, только прямой проход | Да (отдельные ядра диспетчеризации) | Да, согласно проектной документации |
| Объединение вычислений и коммуникаций | Да | Да | Нет | Да (диспетчеризация+GEMM1, GEMM2+объединение) |
| GPU | Hopper, Blackwell | SM90, SM100 (библиотека) | B200 в опубликованных запусках | MI355X (gfx950) |
| Типы данных | BF16, MXFP8, блочный FP8 | FP8xFP8, FP8xFP4 | BF16, сгруппированные GEMM в MXFP8 | BF16, MXFP8 |
| Опубликованный прирост скорости | До 55% по времени шага по сравнению с TorchTitan | Не раскрыт | До 41% на DeepSeek-V3 671B | Не раскрыт |
| Лицензия | Apache 2.0 (prime-rl) | MIT | BSD-3-Clause | MIT |
Прирост TorchTitan на 41% получен в другой конфигурации (B200, модель 671B, базовая версия BF16). Эти показатели нельзя напрямую сравнивать.
Как разработчики используют это решение?
Prime Mega MoE — это бэкенд диспетчеризации в prime-rl. Установите type = "mega_moe" в разделе [trainer.model.moe.dispatch]. Выберите precision из вариантов bf16, mxfp8, fp8 или fp8_bwd. Ядра для обучения поставляются в предварительно собранном wheel-пакете DeepGEMM из prime-kernels. В стандартной сборке DeepGEMM их нет.
Основные выводы
- Prime Mega MoE добавляет к Mega MoE от DeepSeek объединённый обратный проход с перекрытием операций.
- Полные шаги обучения выполняются на 32–34% быстрее базовой версии prime-rl.
- По сравнению с TorchTitan шаги на моделях Qwen3 MoE выполняются на 48–55% быстрее.
- Прирост в прямом проходе достигает 3,11 раза, но в обратном проходе снижается до 1,24 раза на больших слоях.
- Решение поставляется в prime-rl как бэкенд диспетчеризации
mega_moe.
Ознакомьтесь с публикацией в блоге о Prime Mega MoE, объявлением в X и репозиторием prime-rl. Также подписывайтесь на нас в Twitter и не забудьте присоединиться к нашему сабреддиту по машинному обучению с более чем 150 тыс. участников и подписаться на нашу рассылку. Стоп! Вы пользуетесь Telegram? теперь к нам можно присоединиться и в Telegram.
Хотите сотрудничать с нами для продвижения своего репозитория на GitHub, страницы на Hugging Face, выпуска продукта, вебинара и т. д.? Свяжитесь с нами
Переведено автоматически с английского. Оригинал статьи — по ссылке ниже.