Prime Intellect a publié Prime Mega MoE, une extension d’entraînement du noyau GPU fusionné Mega MoE de DeepSeek. Là où le noyau de DeepSeek ne couvrait que la passe avant, Prime Intellect y ajoute une passe arrière rapide, qui se chevauche avec les autres calculs, si bien que le noyau devient utilisable pour l’entraînement et plus seulement pour l’inférence. Selon l’entreprise, dans son entraîneur prime-rl, les étapes d’entraînement complètes sur des modèles Qwen3 MoE tournent jusqu’à 34 % plus vite, et jusqu’à 55 % plus vite face à TorchTitan.
- Prime Intellect ajoute une passe arrière au noyau Mega MoE de DeepSeek, qui ne couvrait que la passe avant.
- Selon l’entreprise, les étapes d’entraînement de prime-rl sont jusqu’à 34 % plus rapides, et 55 % face à TorchTitan.
- La passe avant gagne un facteur 3,11, mais la passe arrière ne progresse que de 1,24.
- Le noyau est sous licence Apache 2.0 et s’active via le paramètre type = «mega_moe».
Dans un modèle MoE (mixture of experts, ou « mélange d’experts »), chaque token n’est envoyé qu’à un petit sous-ensemble d’experts, répartis sur différents GPU. Chaque couche doit donc envoyer les tokens, faire calculer les experts, puis renvoyer les résultats, lors de la passe avant comme de la passe arrière, la phase de rétropropagation utilisée pour l’apprentissage. Dans une configuration standard, ces opérations sont séparées : le GPU trie les tokens, attend le transport, calcule, puis attend à nouveau. Ces temps d’attente sont autant de temps de calcul inactif.
Un noyau fusionné, complété pour l’apprentissage
Le Mega MoE de DeepSeek, ajouté à la bibliothèque DeepGEMM en avril 2026, remplace cette chaîne par un seul appel de noyau. Un noyau (kernel) est un programme optimisé qui s’exécute directement sur le GPU. Celui-ci fusionne le dispatch expert-parallèle, les deux couches linéaires des experts, SwiGLU et l’étape de combinaison. Pendant qu’un groupe de tokens est en transit, le noyau calcule sur un groupe déjà arrivé, et il évite les tenseurs intermédiaires grâce à des tampons de travail situés uniquement sur le GPU. Ce noyau n’incluait toutefois qu’une implémentation de la passe avant.
Prime Intellect a écrit la passe arrière manquante, elle aussi en chevauchement. Les types de données pris en charge sont BF16, MXFP8, réservé aux GPU Blackwell, et FP8 par blocs, réservé aux GPU Hopper. La configuration de prime-rl propose également un mode fp8_bwd : passe avant en BF16, passe arrière en FP8.
Des chiffres auto-déclarés
Tous les résultats sont publiés par Prime Intellect elle-même. Les étapes d’entraînement utilisent des nœuds de 8 H200, en FP8, avec des séquences de 16 384 tokens et un checkpointing d’activation complet. Sur Qwen3-30B-A3B, réparti sur 32 H200, une étape complète dure 2,78 secondes avec la base prime-rl, 3,07 secondes avec TorchTitan (FP8) et 2,08 secondes avec Prime Mega MoE, soit des gains de 34 % et de 48 %. Sur Qwen3-235B-A22B, réparti sur 128 H200, les durées sont de 12,72 secondes, 14,98 secondes et 9,65 secondes, pour des gains de 32 % et de 55 %. Le débit passe, sur le modèle 30B, de 5 337 à 7 886 tokens par seconde et par GPU face à TorchTitan, et sur le modèle 235B, de 1 094 à 1 697.
Mesuré couche par couche sur un seul nœud, le gain est plus marqué à l’avant qu’à l’arrière. La passe avant du modèle Qwen3-30B-A3B passe de 12,0 ms à 3,9 ms, soit un facteur de 3,11. La passe arrière du modèle Qwen3-235B-A22B passe de 30,5 ms à 24,6 ms, soit 1,24. Le chiffre « jusqu’à 3,11x » ne concerne donc que la passe avant. TorchTitan revendique de son côté un gain allant jusqu’à 41 %, mais dans une configuration différente (B200, modèle de 671 milliards de paramètres, base BF16), que la source juge non directement comparable.
Prime Mega MoE est diffusé sous licence Apache 2.0. Dans prime-rl, il s’active comme backend de dispatch : il suffit de définir type = «mega_moe» sous [trainer.model.moe.dispatch], puis de choisir la précision parmi bf16, mxfp8, fp8 ou fp8_bwd. Les noyaux d’entraînement sont fournis dans une roue DeepGEMM précompilée, disponible sur le dépôt prime-kernels, car la version standard de DeepGEMM ne les contient pas.
