Inilah yang telah kami kokas selama 9 bulan terakhir: membuat pelatihan MoE berjalan ~2x lebih cepat dan ~2x lebih sedikit memori! Sorotan: - MoE biasanya membutuhkan waktu dan memori paling banyak dalam model modern. Ternyata seseorang dapat secara matematis menulis ulang MoE backward pass untuk mengurangi mem aktivasi yang perlu Anda simpan di fwd sebesar ~2x, menghasilkan gradien yang sama tanpa komputasi ulang matmul tambahan. Saya sangat menyukai hasil ini, karena menggabungkan wawasan algoritmik dan sistem. - Menganalisis kemacetan di lapisan MoE mengarah pada strategi pengoptimalan alami: kurangi pembacaan/penulisan mem sebanyak mungkin! Mengumpulkan input untuk fwd dan output grad untuk bwd terkadang dapat memakan waktu sebanyak GEMM yang dikelompokkan. Kami menggabungkan kumpulkan dengan GEMM yang dikelompokkan + tumpang tindih akses mem dan komputasi untuk membuat seluruh lapisan berjalan ~2x lebih cepat. - Komputasi top-k untuk perutean ahli dapat memakan waktu yang sangat lama, ~15-20% dari seluruh lapisan MoE! Impl top-k standar menggunakan raga top-k algo, bagus untuk k besar tetapi suboptimal untuk k kecil. Kami menulis ulang top-k menggunakan top-k bg bitonik, dan terkadang 20-30x lebih cepat daripada top-k pytorch! Semua kernel utama ditulis dalam Cute-DSL sehingga harus mudah diperluas (dan menginstal :D). Kernel hopper sudah keluar, kernel Blackwell hampir siap. Model MoE dulunya 2x kurang efisien perangkat keras untuk dilatih, semoga Sonic-MOE akan mengubahnya.