Prime Intellect пусна Prime Mega MoE: обучаемо слято MoE ядро, 3,11 пъти по-бърз преден проход

Prime Intellect пусна Prime Mega MoE — разширение за обучение на слетия Mega MoE kernel на DeepSeek. Ядрото на DeepSeek изпълняваше само предния проход. Prime Intellect добави бърз, припокриващ се обратен проход. Това прави слетото ядро приложимо за обучение, а не само за инференс. В своя prime-rl обучител пълните стъпки на обучение при моделите Qwen3 MoE се изпълняват с до 34% по-висока скорост. Спрямо TorchTitan подобрението достига 55%. За екипите, които обучават модели от типа mixture-of-experts (MoE), това означава пряко намаляване на GPU часовете.

Накратко

• Размер: GPU ядро, а не модел. Тествано с Qwen3-30B-A3B и Qwen3-235B-A22B (128 експерта, маршрутизиране top-8).
• Работи на: NVIDIA Hopper (blockwise FP8) и Blackwell (MXFP8) GPU, BF16 и на двете платформи. Публикуваните резултати са получени с 8×H200 възела и от 32 до 128 GPU.
• Производителност: От 32% до 34% по-бързи пълни стъпки на обучение спрямо базовата версия на prime-rl. От 48% до 55% по-бързи спрямо TorchTitan.
• Най-добър резултат: 3.11 пъти по-бърз преден проход на MoE слой Qwen3-30B-A3B (от 12.0 ms до 3.9 ms).
• Най-слаб резултат: 1.24 пъти по-бърз обратен проход на MoE слой Qwen3-235B-A22B (от 30.5 ms до 24.6 ms).

Какво представлява Prime Mega MoE?

Prime Mega MoE е слето, припокриващо се MoE ядро за обучение от Prime Intellect. То се основава на ядрото Mega MoE в библиотеката DeepSeek с отворен код DeepGEMM.

Защо обучението на MoE е ограничено от комуникацията?

В един MoE модел всеки токен се насочва към малък набор от експерти. Тези експерти се намират на различни GPU. Затова всеки слой изпраща токените, изпълнява изчисленията на експертите и връща резултатите. Това се случва в предния проход и отново в обратния проход.

При стандартна конфигурация това са отделни операции. GPU сортира токените и изчаква транспортния слой. След това извършва изчисленията, връща резултатите и отново изчаква. Тези изчаквания са неизползвано GPU време.

Как работи слетото ядро?

Mega MoE на DeepSeek, добавено към DeepGEMM през април 2026 г., заменя тази последователност с едно извикване на ядрото. То слива dispatch при паралелизъм между експертите, двата линейни слоя на експертите, SwiGLU и combine. Докато една група токени се предава, ядрото изчислява върху група, която вече е пристигнала. То също така пропуска междинните тензори, като използва scratch буфери, достъпни само за GPU. Но първоначално включваше реализация само на предния проход.

Prime Intellect създаде липсващото ядро за обратния проход, също с припокриване. Поддържаните типове са BF16, MXFP8 (само Blackwell) и blockwise FP8 (само Hopper). Конфигурацията на prime-rl също предоставя режим fp8_bwd: BF16 преден проход, FP8 обратен проход.

Колко бърз е Prime Mega MoE?

Всички числа по-долу са предоставени от Prime Intellect. Стъпките на обучение използват 8×H200 възела, FP8, последователности от 16 384 токена и пълно checkpointing на активациите. Всеки резултат е средната стойност от стъпки 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 токена/s/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.11x“ се отнася само до предния проход.

Как се сравнява с други слети MoE опции?

Характеристика	Prime Mega MoE	DeepGEMM Mega MoE	TorchTitan + DeepEP	AMD Primus MegaMoE
Разработчик	Prime Intellect	DeepSeek	PyTorch, Nebius	AMD
Обучение (обратен проход)	Да, слето и с припокриване	Не, само преден проход	Да (отделни dispatch ядра)	Да, според проектната документация
Слети изчисления и комуникации	Да	Да	Не	Да (dispatch+GEMM1, GEMM2+combine)
GPU	Hopper, Blackwell	SM90, SM100 (библиотека)	B200 при публикуваните тестове	MI355X (gfx950)
Типове данни	BF16, MXFP8, blockwise FP8	FP8xFP8, FP8xFP4	BF16, MXFP8 grouped GEMMs	BF16, MXFP8
Публикувано ускорение	До 55% по-кратко време за стъпка спрямо TorchTitan	Не е оповестено	До 41% при DeepSeek-V3 671B	Не е оповестено
Лиценз	Apache 2.0 (prime-rl)	MIT	BSD-3-Clause	MIT

Подобрението от 41% при TorchTitan идва от различна конфигурация (B200, модел 671B, BF16 базова версия). То не е пряко сравнимо.

Как го използват разработчиците?

Prime Mega MoE е dispatch бекенд в prime-rl. Задайте type = "mega_moe" под [trainer.model.moe.dispatch]. Изберете precision измежду bf16, mxfp8, fp8 или fp8_bwd. Ядрата за обучение се предоставят в предварително изграден DeepGEMM wheel от prime-kernels. Стандартната компилация на DeepGEMM не ги включва.

Основни изводи

• Prime Mega MoE добавя слет, припокриващ се обратен проход към Mega MoE на DeepSeek.
• Пълните стъпки на обучение се изпълняват с 32% до 34% по-висока скорост спрямо базовата версия на prime-rl.
• Спрямо TorchTitan стъпките се изпълняват с 48% до 55% по-висока скорост при моделите Qwen3 MoE.
• Подобрението при предния проход достига 3.11 пъти, но при обратния проход спада до 1.24 пъти при големите слоеве.
• Доставя се в prime-rl като dispatch бекенда mega_moe.

---

Разгледайте блог публикацията за Prime Mega MoE, съобщението в X и хранилището на prime-rl. Също така можете да ни последвате в Twitter и не забравяйте да се присъедините към нашия ML SubReddit с над 150 хил. членове и да се абонирате за нашия бюлетин. Чакайте! В Telegram ли сте? Вече можете да се присъедините към нас и в Telegram.

Имате нужда от партньор за популяризиране на вашето GitHub хранилище, страница в Hugging Face, продуктово представяне, уебинар и др.? Свържете се с нас