Das ist es, wonach wir in den letzten 9 Monaten gesucht haben: MoEs-Training soll ~2x schneller und ~2x weniger Speicher benötigen! Höhepunkte: - MoE benötigt typischerweise die meiste Zeit und den meisten Speicher in modernen Modellen. Es stellt sich heraus, dass man den MoE-Rückwärtsdurchlauf mathematisch umschreiben kann, um den Aktivierungsspeicher, den man im Vorwärtsdurchlauf speichern muss, um ~2x zu reduzieren, was zu denselben Gradienten ohne zusätzliche Matmul-Neuberechnung führt. Ich mag dieses Ergebnis wirklich, da es sowohl algorithmische als auch systemtechnische Erkenntnisse kombiniert. - Die Analyse von Engpässen in der MoE-Schicht führt zu einer natürlichen Optimierungsstrategie: Reduziere Speicherlese-/schreibvorgänge so weit wie möglich! Das Sammeln der Eingaben für den Vorwärtsdurchlauf und der Ausgabegradienten für den Rückwärtsdurchlauf kann manchmal genauso viel Zeit in Anspruch nehmen wie die gruppierten GEMMs. Wir fusionieren das Sammeln mit gruppiertem GEMM + überlappen den Speicherzugriff und die Berechnung, um die gesamte Schicht um ~2x schneller zu machen. - Die Berechnung von top-k für die Expertenweiterleitung kann überraschend lange dauern, ~15-20% der gesamten MoE-Schicht! Die Standard-Top-k-Implementierung verwendet den Radix-Top-k-Algorithmus, der großartig für großes k, aber suboptimal für kleines k ist. Wir haben top-k mit dem bitonischen Top-k-Algorithmus umgeschrieben, und es ist manchmal 20-30x schneller als Pytorchs top-k! Alle Hauptkerne sind in Cute-DSL geschrieben, sodass sie leicht erweiterbar (und installierbar :D) sein sollten. Hopper-Kerne sind draußen, Blackwell-Kerne sind fast bereit. MoE-Modelle waren früher 2x weniger hardwareeffizient im Training, hoffentlich wird Sonic-MOE das ändern.