BitMTP: quando la predizione multi-token incontra i grandi modelli linguistici a pochi bit

Abstract degli autori

Per rendere efficiente l’impiego dei grandi modelli linguistici (LLM), i metodi basati su backbone a bassissimo numero di bit riducono lo spazio necessario per memorizzare i pesi, il traffico di memoria e il costo di esecuzione dei livelli lineari. I metodi di predizione multi-token (MTP), invece, riducono i passaggi di decodifica sequenziale prevedendo in parallelo più token futuri. Pur essendo complementari, la loro combinazione diretta non riesce ad accelerare la decodifica a parità di qualità: nella generazione preliminare con mascheramento manca l’input dei token intermedi disponibile nell’addestramento alla predizione del token successivo, mentre un calcolo più ampio e gli aggiornamenti causali key-value (KV) possono annullare i risparmi ottenuti riducendo il numero di passaggi. Per affrontare queste difficoltà proponiamo BitMTP, un framework di progettazione congiunta di algoritmo e sistema per una MTP efficiente su backbone LLM a bassissimo numero di bit. Un obiettivo congiunto di predizione del token successivo e MTP permette a un unico modello sia di generare in modo autoregressivo (AR) sia di produrre sequenze preliminari con causalità a blocchi. La decodifica basata soltanto sulla confidenza, quella esclusivamente speculativa e quella ibrida offrono diversi compromessi tra qualità e throughput, senza ricorrere a modelli separati per la generazione preliminare e la verifica. L’allineamento del prefisso accettato mantiene coerenti il prefisso confermato, il confine causale KV materializzato e lo stato di riesecuzione dei CUDA Graph per proseguire la generazione. BitMTP sfrutta questa garanzia per selezionare modalità di esecuzione a pochi bit specializzate in base alla larghezza e riutilizzare i grafi acquisiti su uno spazio di archiviazione KV condiviso. Su NVIDIA H20, BitMTP con decodifica ibrida raggiunge accelerazioni medie su cinque attività di 2,2 volte con BitNet b1.58-2B e di 2,1-2,4 volte con i modelli Bonsai da 1,7B-4B, a parità di qualità nelle attività. Su RTX 3090 e Jetson Orin raggiunge inoltre, rispettivamente, un throughput medio su cinque attività fino a 2,0 volte superiore e un’accelerazione di 3,3 volte su MBPP+. Questi risultati dimostrano che BitMTP accelera la generazione degli LLM a bassissimo numero di bit su diversi backbone, dimensioni dei modelli e tipi di GPU. Il codice sorgente è disponibile all’indirizzo https://github.com/PolyU-VCLab/BitMTP