Voici ce que nous avons préparé au cours des 9 derniers mois : faire en sorte que l'entraînement des MoE soit ~2x plus rapide et nécessite ~2x moins de mémoire ! Points forts : - Les MoE prennent généralement le plus de temps et de mémoire dans les modèles modernes. Il s'avère qu'on peut réécrire mathématiquement le passage arrière des MoE pour réduire la mémoire d'activation que vous devez stocker dans le passage avant d'environ 2x, ce qui permet d'obtenir les mêmes gradients sans recomputation supplémentaire de matmul. J'aime vraiment ce résultat, car il combine à la fois des idées algorithmiques et des perspectives systèmes. - L'analyse des goulets d'étranglement dans la couche MoE conduit à une stratégie d'optimisation naturelle : réduire les lectures/écritures de mémoire autant que possible ! Rassembler l'entrée pour le passage avant et le gradient de sortie pour le passage arrière peut parfois prendre autant de temps que les GEMM groupés. Nous fusionnons le rassemblement avec le GEMM groupé + chevauchons l'accès mémoire et le calcul pour rendre l'ensemble de la couche ~2x plus rapide. - Le calcul du top-k pour le routage des experts peut prendre étonnamment longtemps, ~15-20 % de l'ensemble de la couche MoE ! L'implémentation standard du top-k utilise l'algorithme top-k radix, excellent pour de grands k mais sous-optimal pour de petits k. Nous avons réécrit le top-k en utilisant l'algorithme top-k bitonique, et il est parfois 20-30x plus rapide que le top-k de pytorch ! Tous les principaux noyaux sont écrits en Cute-DSL, donc ils devraient être faciles à étendre (et à installer :D). Les noyaux Hopper sont sortis, les noyaux Blackwell sont presque prêts. Les modèles MoE étaient auparavant 2x moins efficaces en termes de matériel pour l'entraînement, espérons que Sonic-MOE changera cela.

Nvidia continue de proposer certains des modèles ouverts les plus puissants et les plus rapides. Les données de pré-entraînement et de post-entraînement sont également publiées, ce que très peu d'organisations ont fait.