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 використовує all-to-all у torch і експерти 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. Ядра для навчання постачаються у попередньо зібраному колесі 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, АБО релізу продукту, АБО вебінару тощо? Зв’яжіться з нами
Перекладено автоматично з англійської. Оригінал статті — за посиланням нижче.