Quad Length Codes for Lossless Compression of e4m3

Titolo: Quad Length Codes per la Compressione Lossless del formato e4m3

1. Il Problema

L'addestramento e l'inferenza dei Large Language Models (LLM) dipendono fortemente dalla parallelizzazione e dalle operazioni collettive (come AllReduce, AllGather), che sono spesso limitate dalla larghezza di banda della rete. Per mitigare questo collo di bottiglia, si utilizza la compressione lossless.
Tuttavia, le soluzioni esistenti presentano compromessi significativi:

• Codici di Huffman: Sfruttano la distribuzione delle frequenze dei simboli per ottenere una compressione ottimale (entropia), ma richiedono una traversata di un albero binario per la decodifica. Questo processo è sequenziale a livello di bit, lento e comporta un'alta complessità hardware a causa della profondità dell'albero.
• Codici Universali (es. Exponential-Golomb): Sono più veloci da decodificare perché non sono completamente sequenziali, ma non sfruttano la distribuzione specifica delle frequenze dei simboli, risultando meno efficienti in termini di compressione.

L'obiettivo è sviluppare un codice che combini l'efficienza della compressione (basata sulle frequenze) con la velocità di decodifica e la semplicità hardware dei codici universali.

2. Metodologia: Quad Length Codes

Gli autori propongono i Quad Length Codes, un approccio ibrido progettato per bilanciare efficienza e velocità. La metodologia si basa sui seguenti principi:

• Partizionamento in Aree: Invece di assegnare un codice unico a ciascuno dei 256 simboli possibili (come fa Huffman), i simboli vengono divisi in 8 aree distinte utilizzando 3 bit di prefisso.
• Lunghezze di Codice Variabili ma Limitate: Ogni area contiene un numero diverso di simboli e ha una lunghezza di codice fissa specifica per quell'area.
  • Le prime 5 aree contengono 8 simboli ciascuna (richiedono 3 bit di prefisso + 3 bit di simbolo = 6 bit totali).
  • La 6ª area contiene 16 simboli (3 bit di prefisso + 4 bit = 7 bit totali).
  • La 7ª area contiene 32 simboli (3 bit di prefisso + 5 bit = 8 bit totali).
  • L'8ª area contiene i restanti 168 simboli (3 bit di prefisso + 8 bit = 11 bit totali).
• Implementazione tramite Lookup Table (LUT):
  • Codifica: Gli input vengono mappati su una LUT di 256 voci. I simboli vengono ordinati in base alla probabilità decrescente e assegnati a un indice da 0 a 255. L'encoder cerca semplicemente il codice nella tabella.
  • Decodifica: Il decoder legge i primi 3 bit (codice dell'area) per determinare istantaneamente la lunghezza totale del codice (6, 7, 8 o 11 bit). Successivamente, legge i bit rimanenti, applica un offset basato sull'area e utilizza una seconda LUT per recuperare il simbolo originale.
• Adattabilità: Lo schema può essere adattato a diverse distribuzioni di probabilità modificando il numero di simboli per area e le lunghezze dei codici, mantenendo comunque solo 4 lunghezze di codice uniche.

3. Risultati Sperimentali

Lo studio è stato condotto analizzando i tensori di attivazione e gradiente del modello Gemma 2B (formato dati e4m3) durante il Supervised Fine Tuning (SFT).

• Caso 1: Distribuzione FFN1 Activation

  • Huffman: Compressione del 15,9%. Lunghezze dei codici da 6 a 18 bit (albero molto profondo).
  • Quad Length Codes (Schema Base): Compressione del 13,9%.
  • Trade-off: Una perdita di compressione di circa il 2%, ma con una drastica semplificazione hardware (solo 4 lunghezze di codice invece di 13) e decodifica non sequenziale.

• Caso 2: Distribuzione FFN2 Activation (con picco di frequenza per lo zero)

  • Huffman: Compressione del 23,2%. Lunghezze dei codici da 3 a 39 bit.
  • Quad Length Codes (Adattato): Modificando lo schema per adattarsi alla distribuzione (es. area 1 con 2 simboli), si ottiene una compressione del 19,0%.
  • Risultato: Miglioramento del 2,3% rispetto allo schema base, con un gap di circa il 4% rispetto a Huffman, ma con una complessità hardware gestibile.

4. Contributi Chiave

1. Ibridazione Ottimizzata: Introduzione di uno schema che cattura l'essenza della compressione basata sulle frequenze (come Huffman) ma elimina la necessità di una traversata ad albero complessa.
2. Semplificazione Hardware: Riduzione della complessità del decoder da una struttura ad albero profondo a una semplice logica di lettura di prefissi e accesso a tabelle di ricerca (LUT) di dimensioni fisse (256 voci).
3. Velocità di Decodifica: Eliminazione della natura "bit-sequenziale" della decodifica, permettendo di determinare la lunghezza del codice immediatamente dai primi 3 bit.
4. Flessibilità: Dimostrazione che lo schema può essere ri-addestrato o adattato empiricamente per diverse distribuzioni di dati (attivazioni vs gradienti) senza cambiare l'architettura di base del decoder.

5. Significato e Impatto

Questo lavoro è significativo per l'infrastruttura di calcolo ML (Machine Learning) su larga scala.

• Riduzione del Traffico di Rete: Anche se la compressione è leggermente inferiore a quella di Huffman, il guadagno in termini di latenza di decodifica e semplificazione hardware è cruciale per le operazioni collettive su acceleratori (TPU/GPU), dove il tempo di decodifica può diventare un collo di bottiglia quanto la larghezza di banda.
• Scalabilità: L'uso di LUT fisse e la logica di decodifica prevedibile rendono i Quad Length Codes ideali per l'implementazione su hardware dedicato (ASIC/FPGA) o in kernel ad alte prestazioni, superando i limiti di velocità dei decoder Huffman tradizionali.
• Future Directions: Gli autori suggeriscono che un'ulteriore ottimizzazione matematica del numero di aree e delle lunghezze dei codici potrebbe migliorare ulteriormente il rapporto di compressione mantenendo i vantaggi di velocità.

In sintesi, i Quad Length Codes offrono un compromesso pragmatico e altamente efficace per l'era dei LLM, privilegiando la velocità di elaborazione e la fattibilità hardware senza sacrificare eccessivamente l'efficienza della compressione.