Questo è ciò per cui abbiamo lavorato negli ultimi 9 mesi: rendere l'addestramento di MoE ~2 volte più veloce e ~2 volte meno pesante in termini di memoria! Punti salienti: - MoE richiede tipicamente il maggior tempo e memoria nei modelli moderni. Si scopre che si può riscrivere matematicamente il passaggio all'indietro di MoE per ridurre la memoria di attivazione che è necessario memorizzare nel passaggio in avanti di ~2 volte, risultando negli stessi gradienti senza ulteriore ricalcolo di matmul. Mi piace molto questo risultato, poiché combina sia intuizioni algoritmiche che di sistema. - Analizzare i colli di bottiglia nel layer MoE porta a una strategia di ottimizzazione naturale: ridurre le letture/scritture di memoria il più possibile! Raccogliere l'input per il passaggio in avanti e il gradiente di output per il passaggio all'indietro può a volte richiedere tanto tempo quanto le GEMM raggruppate. Fonde il gathering con GEMM raggruppato + sovrappone l'accesso alla memoria e il calcolo per rendere l'intero layer ~2 volte più veloce. - Calcolare il top-k per il routing degli esperti può richiedere sorprendentemente molto tempo, ~15-20% dell'intero layer MoE! L'implementazione standard del top-k utilizza l'algoritmo radix top-k, ottimo per grandi k ma subottimale per piccoli k. Abbiamo riscritto il top-k utilizzando l'algoritmo bitonico top-k, e a volte è 20-30 volte più veloce del top-k di pytorch! Tutti i principali kernel sono scritti in Cute-DSL, quindi dovrebbero essere facili da estendere (e installare :D). I kernel Hopper sono disponibili, i kernel Blackwell sono quasi pronti. I modelli MoE erano 2 volte meno efficienti in termini di hardware per l'addestramento, speriamo che Sonic-MOE cambierà questa situazione.