Вышло Prime Mega MoE — fused-ядро для обучения MoE с ускорением forward pass до 3.11x

Prime Intellect открыла исходный код Prime Mega MoE — обучения расширения fused-ядра Mega MoE из DeepSeek. DeepSeek реализовал только прямой проход; Prime Intellect добавил обратный проход с перекрытием (overlapping backward pass), что делает fused-ядро пригодным для полноценного обучения, а не только инференса. В своём фреймворке prime-rl полные шаги обучения на моделях Qwen3 MoE выполняются на 32–34% быстрее baseline, а в сравнении с TorchTitan прирост достигает 48–55%.

Ядро тестировали на двух архитектурах Qwen3: 30B-A3B и 235B-A22B (128 экспертов, top-8 routing). Лучший результат — forward pass на слое Qwen3-30B-A3B ускорился с 12,0 до 3,9 мс (в 3,11 раза). Даже самый скромный прирост — backward pass на Qwen3-235B-A22B: с 30,5 до 24,6 мс (в 1,24 раза). Ядро работает на NVIDIA Hopper (блочный FP8) и Blackwell (MXFP8), поддерживает BF16 на обеих архитектурах и экспериментальный режим fp8_bwd (BF16 forward + FP8 backward). Измерения проводили на 8 узлах H200 (от 32 до 128 GPU).

Проблема MoE-обучения — узкое место по коммуникации: каждый токен отправляется к малой группе экспертов, разбросанных по разным GPU. DeepSeek в ядре Mega MoE решил это слиянием операций диспатча, вычислений и комбайна в один вызов, но только для forward. Prime Intellect закрыл пробел, написав перекрывающий backward. Для команд, тренирующих большие MoE-модели, это прямой выигрыш в GPU-часах — никакой магии, просто чистая инженерная оптимизация.