MoEBlaze: Efficienza nel Training di Modelli MoE

Il training di modelli Mixture-of-Experts (MoE) di grandi dimensioni è spesso limitato dal collo di bottiglia della memoria. L'architettura sparsa dei MoE introduce overhead significativi, legati alla gestione dei buffer di routing dei token e alla necessità di materializzare tensori intermedi. Questo limita la dimensione massima dei batch e la lunghezza delle sequenze gestibili dalle GPU, impattando negativamente sulle prestazioni e sulla scalabilità del modello.

MoEBlaze è un framework progettato per affrontare queste sfide. Ottimizza l'utilizzo della memoria durante il training dei modelli MoE attraverso un approccio di co-design che include:

  • Un metodo end-to-end per il dispatch dei token e il training MoE, con strutture dati ottimizzate per eliminare buffer intermedi e la materializzazione degli activation.
  • Kernel co-progettati con tecniche di activation checkpointing per ridurre l'impronta di memoria, migliorando al contempo le prestazioni.

I risultati preliminari indicano che MoEBlaze può raggiungere un incremento di velocità di oltre 4 volte e un risparmio di memoria superiore al 50% rispetto ai framework MoE esistenti. Questo rappresenta un passo avanti significativo verso un training più efficiente e scalabile di modelli MoE su hardware moderno.