Riepilogo Tecnico: PagedWeight

Problema

I modelli Large Language Models (LLM) Mixture-of-Experts (MoE) offrono un'elevata efficienza e precisione attivando solo un sottoinsieme di esperti per ogni token. Tuttavia, negli scenari di serving, i modelli MoE affrontano una tensione critica nella memoria: la GPU deve memorizzare sia i massicci pesi del modello che la crescente cache Key-Value (KV). Sebbene tecniche come PagedAttention gestiscano la frammentazione della cache KV, esse non affrontano l'imponente impronta di memoria dei pesi MoE, che possono occupare oltre il 60% della memoria GPU.

I metodi di quantizzazione esistenti sono ampiamente statici, fissando la precisione prima dell'inizio dell'inferenza. Sebbene esistano alcuni metodi di adattamento della precisione a runtime, essi non utilizzano i cambiamenti di precisione per riallocare dinamicamente la memoria GPU tra i pesi e la cache KV. Di conseguenza, manca un sistema in grado di liberare dinamicamente memoria dai pesi degli esperti meno critici man mano che la cache KV si espande, senza sacrificare l'accuratezza del compito o interrompere l'inferenza.

Metodologia: PagedWeight

Gli autori propongono PagedWeight, un nuovo sistema di gestione per il serving di MoE LLM che quantizza dinamicamente i pesi del modello a runtime per bilanciare la precisione dei pesi degli esperti con le dimensioni della cache KV. Il sistema tratta i pesi degli esperti in modo simile ai blocchi di cache KV paginati, consentendo l'offloading e il ricaricamento dinamico.

Componenti Core

1. Rappresentazione Paged Weight:

   • PagedWeight opera alla granularità di blocchi lineari per esperto (identificati da layer, esperto e tipo di modulo: gate_up o down).
   • Utilizza la tecnologia Any-Precision LLM (APL), rappresentando i pesi in un formato a piani di bit sovrapposti con relative tabelle di lookup (LUT).
   • Una Weight Page Table traccia la larghezza di bit desiderata ($d_i$) e la larghezza di bit impegnata ($q_i$) per ogni blocco lineare. Le pagine possono esistere in stati residenti in GPU, residenti in CPU o in stato di trasferimento.

2. Planner di Runtime Consapevole della Qualità:
   Il planner determina quali pesi quantizzare (ridurre la larghezza di bit) basandosi su tre fattori per minimizzare la perdita di accuratezza:

   • Sensibilità Globale Offline: Utilizza punteggi pesati tramite l'Hessiana per stimare la sensibilità intrinseca di ogni blocco lineare alla quantizzazione.
   • Statistiche di Routing Online: Monitora la "massa di routing" (frequenza e peso della selezione) degli esperti. Gli esperti frequentemente instradati ("hot") sono protetti con soglie di larghezza di bit più elevate e moltiplicatori di danno.
   • Residui del Prompt: Un termine di correzione specifico per il prompt predetto tramite teste di regressione lineare. Questo regola la stima della sensibilità globale in base alla sequenza di input corrente, riconoscendo che l'impatto della quantizzazione varia a seconda del prompt.

   Il planner calcola un punteggio di "danno" per le potenziali riduzioni della larghezza di bit. Quando la pressione della cache KV attiva un obiettivo di memoria, il planner seleziona greedymente le riduzioni con il minor danno (massima qualità per byte risparmiato) per raggiungere l'obiettivo.

3. Esecuzione Asincrona e Ottimizzazione dei Kernel:

   • Movimento di Pagine Asincrono: Per nascondere la latenza, il sistema scarica le pagine dei pesi nella RAM della CPU e le ricarica in modo asincrono. Gli impegni sulla larghezza di bit vengono aggiornati solo in confini sicuri (ad esempio, dopo che un piano è stato completamente eseguito), garantendo che l'inferenza non venga interrotta.
   • Kernel a Precisione Mista Fuso: Un kernel CUDA personalizzato legge direttamente i piani di bit Any-Precision e le LUT, supportando diverse larghezze di bit per ogni blocco lineare all'interno di una singola operazione fusa per ridurre l'overhead.

Contributi Chiave

1. Design del Sistema: La proposta di PagedWeight, un sistema di pesi paginati per l'online Any-Precision MoE che gestisce larghezze di bit impegnate/desiderate, stati delle pagine e consumo di memoria dinamicamente.
2. Algoritmo di Planning: Un planner di runtime consapevole della qualità che sintetizza la sensibilità offline, le statistiche di routing online e i residui del prompt per selezionare strategie di riduzione dei pesi a basso danno sotto pressione di memoria.
3. Strategia di Esecuzione: Implementazione del movimento di pagine asincrono per nascondere la latenza di offload/reload con una perdita minima di throughput, insieme a un kernel MoE a precisione mista fuso.
4. Validazione Empirica: Valutazione completa su tre modelli MoE (Qwen1.5-MoE-A2.7B, Mixtral-8×7B, Gemma-4-26B-A4B) che dimostra prestazioni superiori rispetto ai baseline esistenti.

Risultati

Gli autori hanno valutato PagedWeight rispetto a FP16, Any-Precision LLM (APL), DP-LLM e MxMoE in compiti di modellazione del linguaggio, ragionamento e contesto lungo.

• Tradeoff Qualità-Memoria: PagedWeight raggiunge un'accuratezza equivalente a FP16 con un risparmio di memoria GPU fino al 72,0% e un miglioramento del throughput di 1,94× rispetto ai baseline.
• Qualità a Budget Fisso: A parità di budget di memoria, PagedWeight migliora la qualità del compito rispetto ai metodi di quantizzazione fino al 39,3%, con una perdita di throughput di al massimo il 4,1%.
• Prestazioni a Lungo Contesto: Sotto budget di memoria ristretti (es. 10 GB), PagedWeight mantiene una qualità a lungo contesto (Passage Retrieval, NarrativeQA) paragonabile a FP16, mentre i baseline di quantizzazione statica (come APL) subiscono un degrado significativo.
• Throughput: PagedWeight eguaglia il throughput dei baseline di quantizzazione uniforme, con la più grande riduzione osservata del 4,1% con batch size 4.
• Ablazione: Rimuovere componenti come le statistiche di routing, i residui del prompt o il movimento delle pagine degrada costantemente le prestazioni, confermando la necessità del design completo del sistema.

Significato

Il paper afferma che PagedWeight naviga efficacementamente il complesso tradeoff tra l'accuratezza del compito del modello, il consumo di memoria e il throughput/latenza. Trattando i pesi degli esperti come unità gestibili e paginate anziché come asset statici, PagedWeight sblocca uno spazio di tradeoff inesplorato tra l'accuratezza del compito e l'allocazione della memoria GPU tra i pesi e la cache KV. Il sistema dimostra che la quantizzazione dinamica e consapevole della qualità è una strategia viabile e superiore per il serving di MoE rispetto alla quantizzazione statica o ai metodi che non si adattano alla pressione di memoria a runtime.