Learned Subspace Compression for Communication-Efficient Pipeline Parallelism

Sintesi Tecnica: Compressione di Sottospazio Appresa per il Parallelismo di Pipeline Efficiente nella Comunicazione

Definizione del Problema
L'addestramento di grandi modelli linguistici (LLM) su hardware eterogenei e geograficamente distribuiti affronta significativi colli di bottiglia nella comunicazione, in particolare nel parallelismo di pipeline. Sebbene il parallelismo di pipeline partizioni i parametri del modello tra i dispositivi per superare i limiti di memoria, esso richiede lo scambio di attivazioni intermedie tra gli stadi durante sia la fase di forward che quella di backward. Su reti a bassa larghezza di banda, questa comunicazione inter-stadio diventa il collo di bottiglia dominante.

Gli approcci esistenti, come le Subspace Networks (SSN), tentano di mitigare questo problema proiettando le attivazioni su un sottospazio a basso rango fisso e condiviso definito da una matrice ortogonale. Tuttavia, questo metodo richiede modifiche architettoniche intrusive, inclusa l'ottimizzazione vincolata per mantenere i pesi all'interno di un sottospazio condiviso e l'uso di ottimizzatori modificati. Empiricamente, questi sottospazi globali fissi portano a un significativo degrado delle prestazioni, specialmente all'aumentare dei rapporti di compressione, poiché non riescono ad adattarsi alla geometria evolutiva delle attivazioni durante l'addestramento. Inoltre, l'apprendimento ingenuo delle matrici di proiezione senza vincoli causa la loro fuga dalla varietà di Stiefel (l'insieme delle matrici ortogonali), distruggendo le proprietà isometriche necessarie per un'efficace compressione e portando a gravi cali di prestazioni.

Metodologia: Manifold Aware Projection Learning (MAPL)
Gli autori propongono il Manifold Aware Projection Learning (MAPL), un metodo che tratta la compressione inter-stadio come una proiezione geometrica apprendibile piuttosto che come un vincolo architettonico fisso. I componenti principali di MAPL sono:

1. Proiettori Ortogonali Apprendibili per Stadio: A differenza delle SSN che impongono un unico sottospazio globale, MAPL equipaggia ogni confine di stadio della pipeline con il proprio proiettore ortogonale apprendibile $A_p \in \mathbb{R}^{d \times r}$. Questi proiettori sono vincolati alla varietà di Stiefel ($St(d, r)$), garantendo che $A_p^\top A_p = I_r$. Ciò consente a ogni stadio di scoprire e adattare continuamente il proprio sottospazio di compressione ottimale per il compito.
2. Ottimizzazione Vincolata alla Varietà (SPEL): Per evitare che i proiettori escano dalla varietà durante gli aggiornamenti del gradiente standard, gli autori utilizzano la Spectral Steepest Descent on the Stiefel Manifold (SPEL). Questo algoritmo aggiorna i proiettori utilizzando proiezioni tangenti e operazioni di ritrazione (tramite il metodo Polar Express) ad ogni passo di ottimizzazione, garantendo che il vincolo di ortogonalità sia rigorosamente mantenuto.
3. Embedding di Ancoraggio Fattorizzati: Per gestire gli offset ad alto rango guidati dai token nei flussi residui senza esaurire la capacità del proiettore, MAPL introduce embedding di ancoraggio fattorizzati. L'offset è decomposto in una matrice a basso rango addestrabile $E_{small}^p$ e una matrice ortogonale congelata $P_p$. Solo gli ID interi dei token vengono trasmessi per ricostruire questi offset nello stadio ricevente, aggiungendo un carico di comunicazione trascurabile.
4. Quantizzazione Vettoriale (VQ): Per un'ulteriore compressione, il metodo supporta la quantizzazione vettoriale residua (RVQ) sulle rappresentazioni proiettate a basso rango. Viene utilizzato un protocolo di sincronizzazione del codicebook in streaming per ammortizzare la comunicazione del dizionario, consentendo alti rapporti di compressione (fino a 16×) con un costo di larghezza di banda minimo.

Contributi Chiave

• Apprendimento di Sottospazio Per-Stadio: Il documento dimostra che imporre un sottospazio globale condiviso degrada l'apprendimento nei contesti di corrispondenza dei token. Al contrario, MAPL permette a ogni stadio di apprendere un proprio sottospazio di compressione distinto e vincolato alla varietà di Stiefel, che si allinea meglio con la struttura intrinseca a basso rango delle attivazioni a diverse profondità.
• Ottimizzazione Vincolata alla Varietà: Gli autori identificano la "fuga dalla varietà" come il principale modo di fallimento dell'apprendimento ingenuo dei proiettori. Introducono SPEL per mantenere l'ortogonalità, dimostrando che l'apprendimento non vincolato è strettamente peggiore rispetto all'uso di baseline fisse.
• Ricostruzione Efficiente: Combinando i proiettori apprendibili con gli embedding di ancoraggio fattorizzati, MAPL ottiene la ricostruzione delle attivazioni a pieno rango con un carico di comunicazione trascurabile, eliminando la necessità delle intrusive decomposizioni architettoniche richieste dalle SSN.
• Miglioramento della Frontiera di Pareto: Il metodo traccia la frontiera di Pareto tra la cross-entropy di validazione e il costo di comunicazione, raggiungendo un'alta compressione con un degrado delle prestazioni trascurabile rispetto ai metodi precedenti.

Risultati Sperimentali
Gli autori hanno valutato MAPL su modelli LLaMA decoder-only con parametri compresi tra 150M e 1B, addestrati sul corpus DCLM con configurazioni di pipeline di $P=4$ e $P=8$ stadi.

• Prestazioni vs. Compressione: MAPL supera costantemente le baseline SSN. Per un modello da 150M con compressione 4×, MAPL ha raggiunto una loss di validazione entro lo 0,84% della baseline non compressa, mentre le baseline SSN sono degradate dell'8–12%.
• Scalabilità: Alla scala di 1 miliardo di parametri con compressione 8×, MAPL ha mantenuto un degrado della loss di soli 1,38–2,02%, mentre le baseline SSN hanno subito degradazioni fino al 26%.
• Task Downstream: Nelle valutazioni zero-shot (HellaSwag, PIQA, ARC), MAPL ha seguito da vicino le baseline non compresse (entro circa 1 punto di accuratezza media alle scale 500M/1B), mentre le baseline SSN hanno subito cali di accuratezza fino a 9 punti.
• Compressione Estrema: Quando combinato con la quantizzazione vettoriale (MAPL+VQ), il metodo ha raggiunto una compressione di 16× con un modesto degrado aggiuntivo (ad esempio, ~2,3% di aumento della loss a 1 miliardo di parametri), sebbene l'accuratezza dei task downstream abbia visto un calo più pronunciato rispetto a MAPL non quantizzato.

Significatività e Rivendicazioni
Il documento afferma che le attivazioni dei transformer inter-stadio ammettono naturalmente strutture a bassa dimensione adattive che possono essere apprese congiuntamente con l'obiettivo del compito, piuttosto che essere imposte tramite sottospazi globali fissi. Trattando la compressione come una proiezione geometrica apprendibile sulla varietà di Stiefel, MAPL consente un'aggressiva riduzione della comunicazione (da 4× a 16×) preservando al contempo la qualità dell'addestramento e le prestazioni sui task. Il lavoro evidenzia l'apprendimento di rappresentazioni vincolate alla varietà come un meccanismo pratico per ridurre i costi di comunicazione nell'addestramento di modelli fondazionali distribuiti, offrendo un compromesso superiore tra prestazioni e compressione rispetto agli approcci basati su sottospazi esistenti.

Limitazioni
Gli autori notano che gli esperimenti sono attualmente limitati a modelli fino a 1 miliardo di parametri, e sono necessarie ulteriori valutazioni su scale più grandi e in condizioni di rete eterogenee reali. Inoltre, sebbene MAPL riduca significativamente il degrado, le prestazioni subiscono comunque un calo sotto rapporti di compressione estremi, in particolare quando viene applicata la quantizzazione vettoriale.