Esto es lo que hemos estado cocinando durante los últimos 9 meses: ¡hacer que el entrenamiento de MoEs sea ~2x más rápido y ~2x menos memoria! Destacados: - MoE suele llevar más tiempo y memoria en los modelos modernos. Resulta que se puede reescribir matemáticamente el paso hacia atrás de MoE para reducir la memoria de activación que necesitas almacenar en el paso hacia adelante en ~2x, resultando en los mismos gradientes sin recomputación adicional de matmul. Me gusta mucho este resultado, ya que combina tanto conocimientos algorítmicos como de sistemas. - Analizar los cuellos de botella en la capa de MoE conduce a una estrategia de optimización natural: ¡reducir las lecturas/escrituras de memoria tanto como sea posible! Reunir la entrada para el paso hacia adelante y el gradiente de salida para el paso hacia atrás a veces puede llevar tanto tiempo como los GEMMs agrupados. Fusionamos la recolección con GEMM agrupado + superponemos el acceso a la memoria y el cálculo para hacer que toda la capa sea ~2x más rápida. - Calcular el top-k para el enrutamiento de expertos puede llevar sorprendentemente mucho tiempo, ¡~15-20% de toda la capa de MoE! La implementación estándar de top-k utiliza el algoritmo radix top-k, excelente para k grandes pero subóptimo para k pequeños. Reescribimos el top-k utilizando el algoritmo bitónico top-k, ¡y a veces es 20-30x más rápido que el top-k de pytorch! Todos los núcleos principales están escritos en Cute-DSL, por lo que deberían ser fáciles de extender (y de instalar :D). Los núcleos de Hopper están listos, los núcleos de Blackwell están casi listos. Los modelos de MoE solían ser 2x menos eficientes en hardware para entrenar, ¡esperemos que Sonic-MOE cambie eso.