Efficient, VRAM-Constrained xLM Inference on Clients

1. Enunciato del Problema

Il documento affronta la sfida critica di eseguire modelli AI su larga scala ad alta accuratezza (collettivamente denominati xLM, inclusi LLM Densi, LLM Mixture-of-Experts (MoE) e Modelli Vision-Language (VLM)) su sistemi client (desktop e laptop) con Video RAM (VRAM) limitata.

• Il Vincolo: Le GPU client spesso dispongono di una VRAM significativamente inferiore rispetto alla memoria totale richiesta dai moderni modelli ad alta accuratezza (ad esempio, un modello da 77 GB su una GPU da 16 GB).
• I Limiti delle Soluzioni Esistenti:
  • Quantizzazione/Potatura: Queste tecniche riducono le dimensioni del modello ma sacrificano l'accuratezza, cosa inaccettabile per applicazioni ad alta fedeltà.
  • Offloading Statico: La pianificazione ibrida esistente (ad esempio in llama.cpp) spesso si basa su configurazioni manuali (ad esempio, specificare il numero di livelli GPU) e non riesce ad adattarsi dinamicamente alle condizioni variabili del sistema (thread CPU, larghezza di banda PCIe) o alle fasi di inferenza (prefill vs decode).
  • Mancanza di Universalità: Nessun singolo prodotto esistente gestisce in modo efficiente le modalità interattive, l'elaborazione batch, i VLM ad alta risoluzione e le architetture sia dense che MoE simultaneamente sotto vincoli rigorosi di VRAM.

2. Metodologia

Gli autori propongono una nuova tecnica di pianificazione ibrida guidata da profili di benchmark chiamata Pipelined Sharding, combinata con un insieme di ottimizzazioni specifiche per la visione chiamate VLMOpt.

A. Pipelined Sharding (Innovazione Principale)

Questa è una tecnica di pianificazione completamente automatica che interrompe il grafo di inferenza a livello di sotto-strato (ad esempio, dividendo i blocchi Attention e FFN) piuttosto che solo a livello di strato. Opera in tre fasi:

1. Fase di Installazione (Profiling):

   • Una suite di benchmark eseguita sull'hardware client specifico analizza le prestazioni del kernel (FLOPS) per varie dimensioni dei tensori, livelli di quantizzazione e conteggi di thread CPU.
   • Crucialmente, analizza sotto carico concorrente (simulando il traffico PCIe e la contesa del controller di memoria) per generare stime temporali accurate per l'esecuzione ibrida CPU-GPU.
   • Questo crea un piccolo database di profili (~170 KB) utilizzato per la stima dei costi.

2. Fase di Pianificazione (Generazione della Programmazione):

   • Il sistema definisce un insieme di "Livelli di Token" (ad esempio, 1, 4, 16, 64 token per iterazione) che rappresentano diversi carichi di inferenza.
   • Per ogni livello, genera tre piani di programmazione distinti basati sul budget VRAM corrente e sulle condizioni del sistema:
     • Solo GPU: Tutti i sotto-strati non bloccati vengono eseguiti sulla GPU; i pesi vengono trasmessi via PCIe. Ottimale per carichi di calcolo elevati (batch grandi/contesti ampi) dove la latenza PCIe è nascosta dal calcolo.
     • Statico: I sotto-strati ad alta priorità (Attention, KV Cache) sono bloccati nella VRAM; quelli a bassa priorità (FFN) vengono eseguiti sulla CPU. Minimizza il traffico PCIe ma si affida al calcolo della CPU.
     • Dinamico: Un approccio ibrido in cui i pesi per i sotto-strati GPU vengono trasmessi dalla RAM di sistema mentre la CPU esegue i sotto-strati precedenti. Questo sovrappone i trasferimenti PCIe al calcolo.
   • Un stimatore temporale basato su Roofline utilizza il database di profili per prevedere la latenza di ogni piano e seleziona quello ottimale per il livello di token specifico.

3. Fase di Inferenza (Esecuzione):

   • Durante l'inferenza, il sistema determina il conteggio dei "nuovi token" per il batch corrente.
   • Cerca la programmazione ottimale precalcolata per quel livello di token e la esegue.
   • Ciò permette al sistema di adattarsi senza soluzione di continuità tra la modalità interattiva (Dimensione Batch = 1, pochi token) e la modalità batch (Dimensione Batch > 1, molti token).

B. VLMOpt (Ottimizzazione Vision-Language)

Per supportare VLM ad alta risoluzione (in particolare il modello Cosmos-Reason1), gli autori hanno implementato tre ottimizzazioni all'interno di llama.cpp:

1. Offloading CPU dei Tensori Vision: I pesi del modello vision sono bloccati nella memoria CPU e trasmessi alla GPU solo quando necessario, riducendo l'impronta VRAM.
2. Flash Attention per Vision: Implementa Flash Attention nel codificatore vision per ridurre la complessità di memoria $O(N^2)$ dei punteggi di attenzione per immagini ad alta risoluzione.
3. Evitamento della Sovrapposizione VRAM: Ristruttura l'esecuzione per garantire che la codifica vision completi e liberi la VRAM prima che inizi la decodifica linguistica, assicurando che l'uso massimo della VRAM sia il massimo delle due parti piuttosto che la loro somma.

3. Contributi Chiave

1. Algoritmo Pipelined Sharding: Un pianificatore ibrido automatico e innovativo che si adatta alle condizioni del sistema (thread CPU, larghezza di banda PCIe) e agli stati di inferenza (contesto vs decode) senza sintonizzazione manuale.
2. Implementazione VLMOpt: Una suite di ottimizzazioni che abilita l'inferenza VLM ad alta risoluzione su hardware client con requisiti VRAM drasticamente ridotti.
3. Implementazione Open-Source in C++: La soluzione è implementata in llama.cpp (un framework C++-first), rendendola accessibile agli sviluppatori di giochi e agli ingegneri di AI fisica.
4. Valutazione Completa: Test rigorosi su LLM Densi, LLM MoE e VLM su hardware client diversificato (dai laptop ai desktop di fascia alta).

4. Risultati

La valutazione dimostra significativi miglioramenti delle prestazioni rispetto a baseline aggressive (llama.cpp e vLLM sintonizzati manualmente):

• Modalità Interattiva (Dimensione Batch = 1):
  • TTFT (Time-to-First-Token): Miglioramento fino a 6,7×.
  • TPS (Tokens Per Second): Miglioramento fino a 30× (ad esempio, per il modello Qwen235B da 77 GB).
  • Fattibilità: Abilita l'esecuzione interattiva di modelli massicci (ad esempio, 77 GB) su GPU con VRAM minima di 2 GB (raggiungendo >5 TPS, la soglia per la velocità di lettura umana).
• Modalità Batch (Dimensione Batch > 1):
  • Miglioramenti del throughput fino a 8,2× rispetto alle baseline.
  • Il pianificatore sposta efficacemente il lavoro sulla GPU all'aumentare della dimensione del batch, evitando colli di bottiglia della CPU.
• Prestazioni VLM (Cosmos-Reason1):
  • Riduzione VRAM: Riduzione della domanda di VRAM di 10× (da ~20 GB a ~2 GB), abilitando inferenza ad alta risoluzione (1440p) su sistemi client dove era precedentemente impossibile.
  • Latenza: Raggiunto un aumento di velocità fino a 1,78× nella Latenza End-to-End (E2EL) rispetto alle baseline con budget VRAM equivalenti.
• Co-Esecuzione: In scenari che eseguono un gioco (Cyberpunk 2077) e un LLM contemporaneamente, la soluzione permette al gioco di trattenere più asset nella VRAM, migliorando sia gli FPS del gioco che i TPS dell'LLM evitando la contesa delle risorse.

5. Significato

Questo lavoro colma il divario tra i requisiti di alta accuratezza dell'AI moderna e i vincoli hardware dei dispositivi client consumer.

• Inferenza Senza Perdite: Dimostra che modelli ad alta accuratezza possono essere eseguiti su hardware client senza sacrificare la fedeltà del modello attraverso quantizzazione o potatura.
• Democratizzazione: Abilitando VLM ad alta risoluzione e grandi modelli MoE su laptop e desktop standard, permette agli sviluppatori di costruire avanzate "AI Fisiche" e assistenti AI in-game senza dipendere da infrastrutture cloud costose.
• Adattabilità: L'approccio guidato dai livelli di token e dai profili fornisce un framework robusto per l'evoluzione futura dell'hardware, garantendo che le prestazioni di inferenza scalino con le risorse disponibili (thread CPU, larghezza di banda PCIe) invece di essere staticamente vincolate a una configurazione fissa.

In sintesi, il documento presenta una soluzione pronta per la produzione, automatizzata, che rende la prossima generazione di AI lato client fattibile, efficiente e senza perdita di qualità.