Prime Intellect представила Prime Mega MoE — расширение для обучения объединённого GPU-ядра Mega MoE от DeepSeek. Если ядро DeepSeek охватывало только прямой проход, то Prime Intellect добавила к нему быстрый обратный проход, который перекрывается с другими вычислениями, благодаря чему ядро стало пригодным для обучения, а не только для вывода. По заявлению компании, в её фреймворке prime-rl полные шаги обучения на моделях Qwen3 MoE выполняются до 34% быстрее, а по сравнению с TorchTitan — до 55% быстрее.
В модели MoE (mixture of experts, или «смесь экспертов») каждый токен отправляется лишь в небольшое подмножество экспертов, распределённых по разным GPU. Каждый слой должен отправлять токены, выполнять вычисления экспертов, а затем возвращать результаты как при прямом, так и при обратном проходе — фазе обратного распространения ошибки, используемой для обучения. В стандартной конфигурации эти операции разделены: GPU сортирует токены, ожидает передачи, вычисляет, затем снова ожидает. Эти периоды ожидания представляют собой простои вычислительных мощностей.
Объединённое ядро, дополненное для обучения
Mega MoE от DeepSeek, добавленный в библиотеку DeepGEMM в апреле 2026 года, заменяет эту цепочку единым вызовом ядра. Ядро (kernel) — это оптимизированная программа, выполняемая непосредственно на GPU. Оно объединяет параллельную диспетчеризацию экспертов, два линейных слоя экспертов, SwiGLU и этап комбинирования. Пока одна группа токенов передаётся, ядро вычисляет для уже прибывшей группы и избегает промежуточных тензоров благодаря рабочим буферам, расположенным исключительно на GPU. Однако это ядро изначально включало только реализацию прямого прохода.
Prime Intellect написала недостающий обратный проход, также с перекрытием вычислений. Поддерживаемые типы данных: BF16, MXFP8 (только для GPU Blackwell) и блочный FP8 (только для GPU Hopper). Конфигурация prime-rl также предлагает режим fp8_bwd: прямой проход в BF16, обратный — в FP8.
Самостоятельно заявленные показатели
Все результаты опубликованы самой Prime Intellect. Шаги обучения используют узлы из 8 H200, в FP8, с последовательностями в 16 384 токена и полным чекпоинтингом активаций. На Qwen3-30B-A3B, распределённом на 32 H200, полный шаг занимает 2,78 секунды в базовом prime-rl, 3,07 секунды в TorchTitan (FP8) и 2,08 секунды в Prime Mega MoE, что даёт прирост в 34% и 48% соответственно. На Qwen3-235B-A22B, распределённом на 128 H200, длительности составляют 12,72 секунды, 14,98 секунды и 9,65 секунды, что даёт прирост в 32% и 55%. Пропускная способность на модели 30B возрастает с 5 337 до 7 886 токенов в секунду на GPU по сравнению с TorchTitan, а на модели 235B — с 1 094 до 1 697.
При послойном измерении на одном узле прирост более заметен на прямом проходе, чем на обратном. Прямой проход модели Qwen3-30B-A3B сокращается с 12,0 мс до 3,9 мс, то есть в 3,11 раза. Обратный проход модели Qwen3-235B-A22B сокращается с 30,5 мс до 24,6 мс, то есть в 1,24 раза. Таким образом, показатель «до 3,11x» относится только к прямому проходу. TorchTitan, со своей стороны, заявляет о приросте до 41%, но в другой конфигурации (B200, модель на 671 миллиард параметров, база BF16), которую источник считает не напрямую сопоставимой.
Prime Mega MoE распространяется под лицензией Apache 2.0. В prime-rl он активируется как бэкенд диспетчеризации: достаточно задать type = «mega_moe» в разделе [trainer.model.moe.dispatch], а затем выбрать точность среди bf16, mxfp8, fp8 или fp8_bwd. Ядра обучения поставляются в предварительно скомпилированном пакете DeepGEMM (wheel), доступном в репозитории prime-kernels, поскольку стандартная версия DeepGEMM их не содержит.
