Dit is waar we de afgelopen 9 maanden aan hebben gewerkt: maak de training van MoEs ~2x sneller en ~2x minder geheugen! Hoogtepunten: - MoE neemt doorgaans de meeste tijd en geheugen in moderne modellen. Het blijkt dat je de MoE achterwaartse doorgang wiskundig kunt herschrijven om het activatiemegeheugen dat je moet opslaan in de voorwaartse doorgang met ~2x te verminderen, wat resulteert in dezelfde gradiënten zonder extra matmul-herberekening. Ik vind dit resultaat echt leuk, omdat het zowel algoritmische als systeeminzichten combineert. - Het analyseren van knelpunten in de MoE-laag leidt tot een natuurlijke optimalisatiestrategie: verminder geheugenlezen/schrijven zoveel mogelijk! Het verzamelen van de invoer voor de voorwaartse doorgang en de uitvoergradiënt voor de achterwaartse doorgang kan soms evenveel tijd kosten als de gegroepeerde GEMMs. We fuseren verzamelen met gegroepeerde GEMM + overlappen geheugen toegang en berekening om de hele laag ~2x sneller te maken. - Het berekenen van top-k voor expert routing kan verrassend lang duren, ~15-20% van de hele MoE-laag! Standaard top-k implementatie gebruikt radix top-k algoritme, geweldig voor grote k maar suboptimaal voor kleine k. We hebben top-k herschreven met behulp van bitonisch top-k algoritme, en het is soms 20-30x sneller dan pytorch's top-k! Alle belangrijkste kernels zijn geschreven in Cute-DSL, dus ze zouden gemakkelijk uit te breiden (en te installeren :D) moeten zijn. Hopper-kernels zijn uit, Blackwell-kernels zijn bijna klaar. MoE-modellen waren vroeger 2x minder hardware-efficiënt om te trainen, hopelijk zal Sonic-MOE dat veranderen.