TACO: Efficient Communication Compression of Intermediate Tensors for Scalable Tensor-Parallel LLM Training

1. Enunciato del Problema

L'addestramento su larga scala dei Large Language Model (LLM) dipende fortemente dal Parallelismo 3D (Parallelismo dei Dati, Parallelismo della Pipeline e Parallelismo dei Tensori). Sebbene il Parallelismo dei Dati (DP) e il Parallelismo della Pipeline (PP) abbiano visto tecniche di compressione delle comunicazioni di successo, il Parallelismo dei Tensori (TP) rimane un collo di bottiglia critico.

• Sovraccarico di Comunicazione: Il TP richiede una comunicazione frequente e strettamente sincronizzata di tensori intermedi (attivazioni e gradienti) a ogni layer e iterazione. Nell'addestramento su larga scala, la comunicazione TP rappresenta il 50–60% del tempo totale di comunicazione e si trova sul percorso di esecuzione critico, rendendo difficile sovrapporla al calcolo.
• La Sfida della Compressione:
  • Accumulo di Errori: A differenza di DP/PP, il TP implica una sincronizzazione ad alta frequenza. Minori errori di quantizzazione nei tensori intermedi si propagano attraverso le connessioni di attenzione e residuali, amplificandosi durante la retropropagazione e causando divergenza nell'addestramento.
  • Distribuzione dei Dati: I tensori intermedi TP mostrano una distribuzione densa, vicina allo zero, con una lunga coda. I metodi di quantizzazione standard (come INT8) falliscono qui perché utilizzano una spaziatura uniforme, portando a una massiccia perdita di informazioni per i valori di piccola magnitudine che dominano il tensore.
  • Inefficienza dei Metodi Esistenti: La compressione orientata all'inferenza precedente (ad es. SmoothQuant) o i metodi focalizzati su DP (ad es. SDP4bit) non possono essere applicati direttamente all'addestramento TP senza un fallimento catastrofico dovuto all'instabilità numerica.

2. Metodologia: Framework TACO

Gli autori propongono TACO (Tensor-parallel Adaptive COmmunication compression), un framework robusto basato su FP8 progettato specificamente per i tensori intermedi TP.

A. Perché FP8?

Il paper analizza le proprietà statistiche dei tensori TP e conclude che FP8 è superiore a INT8:

• Limitazioni di INT8: La quantizzazione uniforme causa "collisioni" in cui valori distinti e piccoli mappano sullo stesso intero, portando a un alto errore relativo nella regione densa centrata sullo zero.
• Vantaggio FP8: FP8 (in particolare il formato E4M3) utilizza una scala esponenziale, fornendo alta precisione per i valori piccoli (vicini allo zero) e un ampio intervallo dinamico per i valori anomali. Questo corrisponde perfettamente alla distribuzione statistica dei tensori TP.

B. Tecniche Principali

TACO introduce tre innovazioni chiave per rendere la compressione FP8 praticabile per l'addestramento TP:

1. Trasformata Adattiva Scala–Hadamard (ASH):

   • Problema: Le trasformate Hadamard standard sono isometriche (preservano l'energia) e non possono disperdere efficacemente i densi cluster di zero dei tensori TP; rimangono intrappolati vicino allo zero, sottoutilizzando la precisione di FP8.
   • Soluzione: TACO impiega un ridimensionamento adattivo a blocchi prima della trasformata Hadamard. Calcola il momento grezzo del secondo ordine (energia) di ogni blocco e applica un fattore di scala per normalizzare l'energia a un livello target.
   • Effetto: Questo "srotola" i densi cluster di zero, distribuendoli sull'intervallo rappresentabile di FP8, seguito da una rotazione Hadamard ortogonale per creare una distribuzione simile a una Gaussiana, adatta alla quantizzazione.

2. Quantizzazione a Doppia Scala (DS):

   • Problema: Anche con ASH, i valori estremi possono causare overflow, o i valori piccoli underflow, portando a instabilità.
   • Soluzione: TACO utilizza due scalari distinti per blocco:
     • Scala pre-rotazione ($\alpha_k$): Controlla l'energia del blocco prima della trasformazione.
     • Scala post-rotazione ($s_k$): Assicura che i valori trasformati rientrino strettamente nell'intervallo dinamico FP8 ($Q_{max}$).
   • Ricostruzione: Durante la decompressione, queste due scale vengono utilizzate per invertire il processo, garantendo una ricostruzione ad alta fedeltà e prevenendo l'amplificazione degli errori.

3. Fusione di Kernel a Livello di Sistema e Sovrapposizione:

   • Operatore Fuso: Per minimizzare il sovraccarico computazionale della compressione (che è critico poiché il TP è sul percorso critico), TACO fonde il calcolo della varianza, la trasformata ASH e la quantizzazione DS in un singolo kernel GPU. Questo elimina accessi ridondanti alla memoria globale e il sovraccarico di avvio del kernel.
   • Sovrapposizione Comunicazione-Calcolo: TACO si integra con la COCCL (Compressible Collective Communication Library) e utilizza una strategia AllReduce "a due scatti" (ReduceScatter + AllGather). Sovrappone i kernel di compressione/decompressione alla comunicazione, nascondendo efficacemente la latenza.

3. Contributi Chiave

1. Analisi Sistematica: Ha identificato il fenomeno del "Collasso dello Zero" nei tensori TP, dimostrando che INT8 standard e trasformate fisse sono inadeguati, e stabilendo FP8 come formato ottimale.
2. Framework TACO: Il primo sistema di compressione delle comunicazioni 3D-parallele a raggiungere una convergenza stabile nell'addestramento TP su larga scala. Introduce la trasformata ASH e la quantizzazione a Doppia Scala per gestire la distribuzione unica dei tensori TP.
3. Implementazione ad Alte Prestazioni: Ha sviluppato un operatore di compressione altamente fuso che riduce il traffico di memoria e abilita una sovrapposizione efficiente con la comunicazione, mitigando il sovraccarico della compressione frequente.
4. Integrazione 3D Completa: Ha integrato con successo TACO con compressori DP (SDP4bit) e PP (TahQuant) all'avanguardia per creare un sistema di addestramento parallelo 3D completamente abilitato alla compressione.

4. Risultati Sperimentali

Gli autori hanno valutato TACO su modelli GPT (2.7B, 6.7B, 13B) e Qwen2.5-7B utilizzando GPU NVIDIA H100.

• Accuratezza (Convergenza):

  • TACO raggiunge un'accuratezza quasi senza perdita. Su GPT-350M, il degrado della perdita di validazione è solo +0.25% rispetto alla baseline BF16, mentre i metodi esistenti (TahQuant) degradano di circa il 2.9%.
  • Su modelli più grandi (GPT-6.7B, Qwen-2.5-7B), il degrado rimane trascurabile (<0.7%), dimostrando stabilità su tutte le scale.
  • Studi di ablazione confermano che la rimozione di ASH o DS porta a una divergenza catastrofica o a una significativa perdita di accuratezza.

• Prestazioni (Throughput):

  • Compressione TP: TACO raggiunge fino a un miglioramento del throughput end-to-end di 1.87× per la comunicazione TP rispetto alla baseline non compressa.
  • Parallelismo 3D Completo: Quando applicato a un setup parallelo 3D completo (TP=4, PP=2, DP=2), TACO offre fino a un speedup di 1.53× rispetto alla baseline.
  • Scalabilità: I guadagni di prestazioni aumentano con il grado di Parallelismo dei Tensori (TP), poiché la comunicazione TP diventa il collo di bottiglia dominante.

5. Significato

• Risoluzione del Collo di Bottiglia TP: TACO affronta l'aspetto più difficile dell'addestramento distribuito di LLM: comprimere i tensori intermedi ad alta frequenza e numericamente sensibili nel Parallelismo dei Tensori senza sacrificare la qualità del modello.
• Abilitazione di Modelli Più Grandi: Riducendo il sovraccarico di comunicazione di quasi il 50% (tramite FP8) mantenendo la stabilità, TACO permette di addestrare modelli più grandi su cluster hardware esistenti o di addestrare modelli esistenti significativamente più velocemente.
• Indipendenza dall'Hardware: Sebbene ottimizzato per FP8 (supportato da GPU moderne come H100), il framework è progettato per degradare elegantemente a INT8 su hardware privo di supporto nativo FP8, garantendo ampia applicabilità.
• Impatto Futuro: Questo lavoro apre la strada all'addestramento exascale dimostrando che una compressione aggressiva delle comunicazioni è fattibile anche sul percorso critico dell'addestramento, purché la compressione sia matematicamente adattata alla distribuzione dei dati.