Riepilogo Tecnico: Riduzione del Collo di Bottiglia della Memoria GPU con la Compressione Lossless per il ML

Definizione del Problema

I moderni carichi di lavoro di machine learning (ML), inclusi i Graph Neural Networks (GNN), i Deep Learning Recommendation Models (DLRM) e i Large Language Models (LLM), elaborano frequentemente dataset che vanno da centinaia di gigabyte a terabyte. Questi volumi di dati superano la capacità della memoria GPU disponibile, rendendo necessari trasferimenti di dati on-demand dalla memoria CPU, dal disco o dall'archiviazione di rete tramite l'interconnessione PCIe.

Il principale vincolo di prestazioni in questi scenari è la larghezza di banda PCIe, che è di oltre un ordine di grandezza inferiore alla larghezza di banda della memoria GPU. Di conseguenza, l'esecuzione è spesso limitata dalla latenza di trasferimento dei dati piuttosto che dalla capacità di calcolo della GPU. Sebbene tecniche come il caching e la sovrapposizione tra calcolo e trasferimenti offrano un parziale sollievo, non possono mitigare completamente il collo di bottiglia, specialmente quando si verificano mancanze di cache (cache misses).

Le soluzioni esistenti si affidano spesso alla compressione lossy (ad esempio, la quantizzazione) per ridurre il volume dei dati. Tuttavia, i metodi lossy introducono una degradazione dell'accuratezza imprevedibile che varia a seconda del modello e del carico di lavoro. Nelle implementazioni commerciali, anche perdite di accuratezza minime (ad esempio, lo 0,1%) sono spesso inaccettabili a causa dei rischi per i ricavi, rendendo difficile l'adozione della compressione lossy come soluzione general-purpose.

La compressione lossless offre un'alternativa adatta al deployment che preserva l'accuratezza del modello, ma affronta sfide implementative significative nelle pipeline di ML:

1. Competizione sulla Latenza: Le latenze di trasferimento CPU-GPU sono competitive con le latenze di decompressione della GPU; gli algoritmi di compressione tradizionali spesso aggiungono un overhead che annulla i risparmi nel trasferimento.
2. Vincoli di Memoria: La decompressione richiede spesso grandi buffer intermedi o strutture di metadati complesse che competono con gli input del modello per la scarsa memoria GPU.
3. Discrepanza di Parallelismo: Molti algoritmi lossless standard (ad esempio, Huffman coding, LZ4) si affidano a lookup di dizionari o pattern di accesso alla memoria irregolari che ostacolano il parallelismo dei warp della GPU e gli accessi alla memoria coalescenti.

Metodologia: Invariant Bit Packing (IBP)

Gli autori propongono Invariant Bit Packing (IBP), un nuovo algoritmo di compressione lossless progettato specificamente per i trasferimenti di tensori ML. L'intuizione fondamentale è che i tensori ML, anche quando numericamente diversificati, presentano spesso bit invarianti a bassa entropia attraverso gruppi di tensori. Ad esempio, nelle rappresentazioni in virgola mobile, i bit del segno e dell'esponente rimangono spesso costanti tra molti tensori a causa della regolarizzazione e della dinamica di addestramento.

Algoritmo Core

IBP opera in tre fasi principali:

1. Preprocessing (Generazione dei Metadati):

   • Il sistema analizza un batch di tensori di input per identificare le posizioni dei bit che sono invarianti (ovvero hanno lo stesso valore) attraverso una percentuale elevata ($T$) dei tensori.
   • Genera una Mask a dimensione fissa (che indica quali posizioni di bit sono invarianti) e un Bitval (che memorizza i valori invarianti).
   • Questi metadati sono memorizzati nella memoria GPU (tipicamente di dimensioni nell'ordine dei kilobyte) e vengono riutilizzati per l'intero dataset, evitando l'overhead dei metadati per ogni singolo tensore.

2. Compressione:

   • I tensori di input vengono elaborati in chunk. Per ogni chunk, l'algoritmo controlla se i bit nelle posizioni invarianti corrispondono al Bitval memorizzato.
   • Se corrispondono, i bit invarianti vengono rimossi (eliminati) e i bit rimanenti vengono traslati per colmare il vuoto.
   • Un bit di partecipazione viene preposto a ogni chunk per indicare se è stato compresso o lasciato non compresso (per gestire i casi in cui la compressione non produce risparmi di spazio).
   • La compressione viene eseguita sulla GPU utilizzando la memoria virtuale unificata zero-copy, consentendo alla GPU di accedere direttamente alla memoria della CPU senza la necessità di una configurazione DMA esplicita avviata dalla CPU per piccoli trasferimenti.

3. Decompressione:

   • La decompressione è ottimizzata per il parallelismo dei warp della GPU. Un warp di 32 thread legge cooperativamente un segmento contiguo di dati compressi dalla memoria della CPU in uno spazio di lavoro di memoria condivisa (shared memory).
   • I thread utilizzano primitive di warp (comunicazione a bassa latenza all'interno di un warp) per eseguire uno scan prefisso, calcolando l'offset iniziale dei bit per il chunk assegnato a ciascun thread.
   • I thread ricostruiscono quindi indipendentemente i propri dati reinserendo i bit invarianti dal Bitval globale nelle corrette posizioni.
   • Il processo utilizza accessi alla memoria asincroni per sovrapporre il recupero dei dati alla logica di decompressione, nascondendo la latenza PCIe.

Integrazione del Sistema

Gli autori integrano IBP nelle pipeline di ML in due modi primari:

• Cache Compression: Comprimere i tensori memorizzati nelle cache software della GPU (ad esempio, nei framework GNN) per aumentare la capacità effettiva della cache.
• PCIe Transfer Compression: Comprimere i tensori nella memoria della CPU prima del trasferimento. Per gestire i tensori compressi di dimensioni variabili, il sistema utilizza la compressione in-place (allocando buffer di dimensione fissa ma leggendo solo la porzione compressa) e mantiene una bitmask nella memoria GPU per tracciare quali tensori sono compressi.

Contributi Chiave

1. Analisi dell'Invarianza dei Dati ML: Il documento identifica e quantifica la presenza di bit invarianti attraverso diversi carichi di lavoro ML (GNN, DLRM, LLM) e tipi di dati (FP32, FP16, BF16), dimostrando che le distribuzioni strutturate permettono una significativa compressione lossless senza perdita di accuratezza.
2. Algoritmo Invariant Bit Packing (IBP): Una tecnica di compressione lossless ottimizzata per GPU che rimuove i bit invarianti, memorizza essi in metadati globali minimi e utilizza la decompressione in parallelo ai warp per ottenere un alto throughput.
3. Integrazione del Sistema: Implementazione di IBP in popolari framework di ML (DGL, Legion, Colossal-AI, FlexGen, InfiniGen) tramite API facili da usare (estensioni PyTorch e librerie CUDA).
4. Valutazione Esaustiva: Benchmarking esteso rispetto alle librerie allo stato dell'arte accelerate da GPU (nvCOMP, ndzip-gpu) e sistemi nativi su dataset reali.

Risultati

La valutazione è stata condotta su una GPU NVIDIA A100 con PCIe Gen 4.

• Prestazioni di Compressione: IBP ottiene risparmi di spazio del 10–12% per dataset GNN densi, dell'8% per i pesi DLRM e fino al 27% per i pesi e le KV-cache degli LLM (particolarmente con tipi di dati BF16).
• Throughput: IBP supera costantemente le librerie di compressione accelerate da GPU esistenti. Per i dataset GNN sparsi, raggiunge un'accelerazione del 9,7x nel throughput di trasferimento rispetto ai trasferimenti non compressi.
• Accelerazioni dei Carichi di Lavoro:
  • Addestramento GNN: Accelerazione media del 74% (fino a 2,7x in casi specifici) riducendo il tempo di trasferimento PCIe e aumentando i tassi di hit della cache.
  • Inferenza DLRM: Accelerazione media del 180% per il lookup degli embedding, guidata principalmente da trasferimenti efficienti avviati dalla GPU e dalla riduzione dell'overhead della CPU.
  • Inferenza LLM: Accelerazione del 24% per l'inferenza offloaded. Per l'offloading dei pesi (BF16), la latenza è stata ridotta del 22%; per l'offloading della KV-cache, la latenza è stata ridotta fino al 27%.
• Overhead: L'overhead di decompressione è minimo; per risparmi di spazio inferiori al 50% (tipici per i dati ML densi), il throughput rimane superiore al 95% del tasso di trasferimento ideale non compresso. Il tempo di preprocessing è moderato e avviene fuori dal percorso critico.

Significato e Rivendicazioni

Il documento sostiene che IBP fornisca una soluzione pratica e lossless al collo di bottiglia della memoria GPU senza i rischi di accuratezza associati alla quantizzazione o alla compressione lossy. Sfruttando la struttura intrinseca dei dati ML (bit invarianti) e ottimizzando per l'architettura GPU (parallelismo dei warp, accesso coalescente, zero-copy), IBP riesce efficacementamente a nascondere l'overhead di decompressione dietro la latenza di trasferimento PCIe.

Gli autori sottolineano che IBP è un algoritmo general-purpose applicabile a vari tipi di dati (float32, float16, bfloat16, vettori sparsi) e carichi di lavoro (addestramento e inferenza). La capacità di integrare IBP con modifiche minime al codice (spesso meno di 35 righe di Python) nei framework esistenti dimostra il suo potenziale per essere adottato in implementazioni ML commerciali per scalare i modelli oltre i limiti della memoria GPU, mantenendo al contempo prestazioni e accuratezza.