Суть обновления: от инференса к обучению
Компания Prime Intellect выпустила Prime Mega MoE — расширение открытого ядра DeepGEMM от DeepSeek. Оригинальное ядро DeepSeek поддерживало только прямой проход (forward pass), что ограничивало его применение инференсом. Prime Intellect реализовала быстрый, перекрывающийся по времени обратный проход (backward pass), что позволило использовать это ядро для полноценного обучения моделей типа Mixture-of-Experts (MoE).
Ключевая проблема, которую решает решение: в стандартных настройках MoE GPU простаивает, ожидая передачи данных между экспертами, расположенными на разных картах. Новое ядро объединяет диспетчеризацию, линейные слои, SwiGLU и комбинирование в один вызов, вычисляя данные на уже прибывших токенах, пока остальные находятся в пути.
Производительность: цифры и бенчмарки
Тестирование проводилось на кластерах NVIDIA H200 (8 узлов) с использованием FP8 и последовательностей длиной 16 384 токенов. Результаты показывают значительное сокращение времени одного шага обучения по сравнению с базовой конфигурацией prime-rl и TorchTitan.
| Модель (Конфигурация) | prime-rl baseline (сек) | TorchTitan (сек) | Prime Mega MoE (сек) | Ускорение vs baseline | Ускорение vs 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% |
На уровне отдельных слоев MoE (1 узел, 8x H200, 128 экспертов) прирост скорости прямого прохода для Qwen3-30B-A3B составил 3.11x (снижение с 12.0 мс до 3.9 мс). Однако обратный проход демонстрирует меньший прирост: для Qwen3-235B-A22B ускорение составило 1.24x (30.5 мс против 24.6 мс), что объясняется сложностью вычислений градиентов в больших моделях.
Технические детали и поддержка
Prime Mega MoE поддерживает следующие типы данных и архитектуры:
- NVIDIA Hopper (SM90): BF16 и blockwise FP8.
- NVIDIA Blackwell (SM100): BF16 и MXFP8.
- Режим fp8_bwd: Прямой проход в BF16, обратный в FP8 (доступно в prime-rl).
Это решение напрямую конкурирует с другими оптимизациями, такими как AMD Primus MegaMoE и стандартный TorchTitan + DeepEP. В отличие от TorchTitan, где диспетчеризация и вычисления часто разделены, Prime Mega MoE объединяет вычисления и коммуникации, что критически важно для снижения задержек в распределенных системах.
Почему это важно
Для команд, обучающих большие MoE-модели, каждое сокращение времени шага обучения переводится в прямую экономию GPU-часов. Ускорение полного шага на 32-34% означает, что проекты, занимавшие месяцы, могут быть завершены значительно быстрее. Особенно это актуально для моделей с большим количеством экспертов (как Qwen3-235B), где накладные расходы на коммуникацию традиционно являются «узким горлышком».
Источник: MarkTechPost ↗