L’altra metà del muro della memoria: servire MoE da 35B da SSD con una predizione dell’instradamento addestrata

Abstract degli autori

L’inferenza con modelli mixture-of-experts (MoE) su hardware di consumo è limitata dalla memoria necessaria per i pesi: un modello della classe 35B occupa 19,5 GB a 4 bit, e la sparsità riduce il calcolo per token, non i byte da mantenere in memoria. Il semplice trasferimento dei pesi su SSD, da solo, non basta: gli esperti del livello N+1 devono essere scelti prima che sia disponibile l’output del livello N, quindi le letture non possono iniziare abbastanza presto da essere mascherate dal calcolo. Presentiamo Edge0, un motore di inferenza MoE in streaming che colma questa lacuna con un prerouter: un modulo per ciascun livello predice l’instradamento del livello successivo con un token di anticipo, e la predizione viene usata come instradamento stesso. Così, l’insieme degli esperti precaricati coincide con quello degli esperti selezionati e nulla viene scartato. Una LoRA di recupero non fusa, addestrata sul percorso del modello studente, recupera la qualità persa a causa della quantizzazione int4 e della sostituzione dell’instradamento. Su una sola macchina da 24 GB, Edge0 serve un MoE da 35B a 20tok/s con un picco di memoria attiva entro 3 GiB, con uno scarto medio di pochi punti rispetto al suo modello insegnante fp16 su cinque benchmark pubblici. Una configurazione 8B funziona sullo stesso framework; il framework, i checkpoint e gli adattatori sono open source.