To jest to, na co czekaliśmy przez ostatnie 9 miesięcy: sprawić, aby trening MoE był ~2x szybszy i ~2x mniej pamięciożerny! Najważniejsze punkty: - MoE zazwyczaj zajmuje najwięcej czasu i pamięci w nowoczesnych modelach. Okazuje się, że można matematycznie przekształcić wsteczny przebieg MoE, aby zmniejszyć pamięć aktywacji, którą trzeba przechowywać w przebiegu do ~2x, co skutkuje tymi samymi gradientami bez dodatkowej rekalkulacji matmul. Naprawdę podoba mi się ten wynik, ponieważ łączy zarówno wnioski algorytmiczne, jak i systemowe. - Analiza wąskich gardeł w warstwie MoE prowadzi do naturalnej strategii optymalizacji: zredukować odczyty/zapisy pamięci tak bardzo, jak to możliwe! Zbieranie danych wejściowych do przebiegu i gradientu wyjściowego do wstecznego przebiegu może czasami zajmować tyle samo czasu, co grupowe GEMM. Łączymy zbieranie z grupowym GEMM + nakładamy dostęp do pamięci i obliczenia, aby cała warstwa działała ~2x szybciej. - Obliczanie top-k dla routingu ekspertów może zająć zaskakująco dużo czasu, ~15-20% całej warstwy MoE! Standardowa implementacja top-k używa algorytmu radix top-k, świetnego dla dużego k, ale suboptymalnego dla małego k. Przepisaliśmy top-k używając algorytmu bitonic top-k, i czasami jest on 20-30x szybszy niż top-k w pytorch! Wszystkie główne jądra są napisane w Cute-DSL, więc powinny być łatwe do rozszerzenia (i zainstalowania :D). Jądra Hopper są już dostępne, jądra Blackwell są prawie gotowe. Modele MoE były 2x mniej wydajne sprzętowo w treningu, miejmy nadzieję, że Sonic-MOE to zmieni.

Nvidia nadal wydaje jedne z najsilniejszych i najszybszych otwartych modeli. Udostępniane są również dane do wstępnego i po treningu, co zrobiło bardzo niewiele organizacji.