Son 9 aydır koks yaptığımız şey şu: MoEs eğitimini ~2 kat daha hızlı ve ~2 kat daha az hafıza yapmak! Öne Çıkanlar: - MoE genellikle modern modellerde en fazla zaman ve bellek alan kişidir. Meğerse MoE'yi matematiksel olarak geri geçiş olarak yeniden yazarak ön çekimde depolama gereken aktivasyon mem'ini ~2x azaltabiliyor, böylece ekstra matmul yeniden hesaplaması olmadan aynı gradyanlar elde ediliyor. Bu sonucu gerçekten çok beğendim, çünkü hem algoritmik hem de sistem içgörülerini birleştiriyor. - MoE katmanındaki darboğazları analiz etmek, doğal bir optimizasyon stratejisine yol açar: mem okuma/yazma sayısını mümkün olduğunca azaltın! Ön çekimli ve bwd için çıktı gradusunu toplamak bazen gruplanmış GEMM'ler kadar zaman alabilir. Gruplanmış GEMM + örtüşme mem erişimi ve hesaplama ile birleştiriyoruz, böylece tüm katman ~2 kat daha hızlı gidiyor. - Uzman yönlendirme için top-k hesaplaması şaşırtıcı derecede uzun sürebilir, ~MoE katmanının %15-20'si! Standart top-k impl, büyük k için harika ama küçük k için optimal olmayan radix top-k algoritmini kullanır. Top-k'yı bitonik top-k algoritmiyle yeniden yazdık ve bazen pytorch'un top-k'sından 20-30 kat daha hızlı! Tüm ana çekirdekler Cute-DSL ile yazılmış, bu yüzden genişletilmesi (ve :D kurması kolay olmalı. Hopper çekirdekleri yok, Blackwell çekirdekleri neredeyse hazır. MoE modelleri eskiden donanım açısından 2 kat daha az verimliydi, umarım Sonic-MOE bunu değiştirir.