Folding Tensor and Sequence Parallelism for Memory-Efficient Transformer Training & Inference

1. Enunciato del Problema

L'addestramento e l'inferenza di modelli linguistici su larga scala (LLM) affrontano due principali colli di bottiglia nella memoria:

1. Memoria dei Parametri: Lo spazio di archiviazione richiesto per pesi del modello, gradienti e stati dell'ottimizzatore.
2. Memoria delle Attivazioni: Lo spazio di archiviazione richiesto per le attivazioni intermedie, che scala quadraticamente con la lunghezza della sequenza ($S^2$) nei meccanismi di attenzione.

Le strategie di parallelismo esistenti affrontano questi problemi in isolamento o con compromessi:

• Tensor Parallelism (TP): Frammenta i pesi del modello tra i dispositivi, riducendo la memoria dei parametri ma lasciando la memoria delle attivazioni completamente replicata. Diventa inefficiente a lunghe lunghezze di sequenza a causa dell'elevata memoria delle attivazioni e del volume di comunicazione.
• Sequence Parallelism (SP): Frammenta la sequenza di input tra i dispositivi, riducendo la memoria delle attivazioni ma lasciando i pesi del modello completamente replicati. È inefficace quando la memoria dei parametri è il collo di bottiglia.
• TP+SP Convenzionale: Combina entrambi assegnandoli ad assi ortogonali in una griglia di dispositivi (ad esempio, una griglia $T \times \Sigma$). Sebbene efficace, questo consuma un gran numero di dispositivi per il parallelismo del modello ($T \times \Sigma$), lasciando meno dispositivi per il parallelismo dei dati. Inoltre, spesso costringe i gruppi paralleli a estendersi su collegamenti inter-nodo più lenti, sottoutilizzando le interconnessioni intra-nodo ad alta larghezza di banda (come NVLink o AMD Infinity Fabric).

La Sfida Principale: Come ridurre simultaneamente sia la memoria dei parametri che quella delle attivazioni su un singolo asse di dispositivo per massimizzare l'efficienza della memoria e liberare dispositivi per il parallelismo dei dati, senza incorrere in costi di comunicazione proibitivi.

2. Metodologia: Tensor and Sequence Parallelism (TSP)

Gli autori propongono TSP, una strategia che "piega" il parallelismo tensoriale e sequenziale su un singolo asse di dispositivo. Invece di assegnare TP e SP a dimensioni ortogonali, TSP assegna a ogni rank (GPU) sia un frammento dei pesi del modello che un frammento della sequenza.

Componenti Chiave del Design:

• Frammentazione Unificata: Ogni dispositivo detiene $1/D$ dei parametri del modello e $1/D$ dei token della sequenza. Questo riduce sia la memoria proporzionale ai parametri che la memoria delle attivazioni di un fattore $D$ su un singolo asse.
• Pianificazioni di Runtime: Per gestire il fatto che pesi e attivazioni sono frammentati diversamente, TSP impiega schemi di comunicazione specifici per diversi livelli:
  • Livelli di Attenzione: Utilizza uno schema di broadcast.
    • I rank iterano sui frammenti di peso (pacchettizzati $W_Q, W_K, W_V, W_O$).
    • In ogni iterazione, un rank trasmette in broadcast il proprio frammento di peso a tutti gli altri.
    • Tutti i rank calcolano proiezioni locali sui propri frammenti di sequenza.
    • Un'operazione All-Gather ricostruisce il contesto completo Chiave/Valore lungo la dimensione della sequenza per eseguire FlashAttention.
    • Viene utilizzato uno schema di partizionamento a zigzag per bilanciare il carico di lavoro dell'attenzione causale (assicurando che i token successivi non comportino un carico di lavoro sproporzionatamente maggiore).
  • Livelli MLP con Gate: Utilizza uno schema di comunicazione ad anello.
    • I frammenti di peso ($W_1, W_2, W_3$) circolano punto-a-punto attorno all'anello.
    • Ogni rank accumula parzialmente gli output localmente.
    • Questo elimina la necessità di un'operazione All-Reduce (standard in TP), poiché la sequenza rimane locale e si muovono solo i pesi.
• Consapevolezza dell'Hardware: Il design è ottimizzato per topologie come AMD MI300X (Infinity Fabric), dove la larghezza di banda collettiva scala con il numero di GPU partecipanti. TSP garantisce che i primitivi di comunicazione (broadcast e anello) coinvolgano l'intero gruppo o siano sovrapposti al calcolo per massimizzare l'efficienza.

3. Contributi Chiave

1. Strategia di Parallelismo Innovativa: Introduzione di TSP, che collassa due assi di parallelismo ortogonali in uno, permettendo la riduzione simultanea della memoria dei parametri e delle attivazioni.
2. Analisi Teorica: Un rigoroso modello matematico che confronta TSP con DP, TP, SP e TP+SP riguardo a:
   • Memoria: TSP è l'unico schema che riduce sia la memoria proporzionale ai pesi che quella delle attivazioni di $1/D$ senza richiedere $D^2$ dispositivi.
   • Comunicazione: Sebbene TSP introduca comunicazione di spostamento dei pesi (che TP+SP evita), gli autori dimostrano che per sequenze lunghe o dimensioni di batch moderate, la comunicazione proporzionale alle attivazioni domina, rendendo TSP competitivo.
   • FLOPs: TSP mantiene la stessa efficienza computazionale per dispositivo di TP e SP.
3. Implementazione e Sovrapposizione: Implementazione dettagliata dei blocchi di attenzione e MLP di TSP con strategie esplicite di sovrapposizione comunicazione/calcolo (pipelining dei trasferimenti di pesi dietro i calcoli GEMM).
4. Ottimizzazione Topologica: Dimostrazione che TSP permette all'intera replica del modello di rientrare in un singolo nodo ad alta larghezza di banda (intra-nodo), evitando le penalità di prestazioni derivanti dall'estensione su collegamenti inter-nodo più lenti spesso richiesti dalle mesh 2D TP+SP.

4. Risultati

Gli autori hanno valutato TSP su un cluster di fino a 1.024 GPU AMD MI300X (128 nodi) con un modello transformer denso da 7 miliardi di parametri.

• Efficienza della Memoria:
  • TSP ha costantemente raggiunto la memoria di picco più bassa su tutte le lunghezze di sequenza testate (da 16K a 128K token).
  • A contesti brevi, TSP ha eguagliato TP (dominato dalla memoria dei parametri).
  • A contesti lunghi, TSP ha superato significativamente TP e TP+SP (dominati dalla memoria delle attivazioni), prevenendo errori Out-Of-Memory (OOM) dove altre strategie fallivano.
• Throughput:
  • Nonostante un volume di comunicazione teorico più elevato (dovuto allo spostamento dei pesi), TSP ha mantenuto un throughput competitivo con TP+SP.
  • La comunicazione aggiuntiva è stata nascosta con successo dietro il calcolo tramite pipelining.
  • I vantaggi del throughput sono cresciuti con l'aumento dei gradi di parallelismo ($D$) e delle lunghezze di sequenza.
• Scalabilità:
  • TSP ha permesso lotti micro più grandi grazie alla ridotta impronta di memoria, migliorando ulteriormente il throughput.
  • Ha permesso di adattare intere repliche del modello all'interno di un singolo nodo (8 GPU), liberando rank per il Parallelismo dei Dati (DP) in cluster più grandi.

5. Significato

• Scalabilità Consapevole dell'Hardware: TSP fornisce una soluzione pratica per i moderni cluster HPC in cui la larghezza di banda intra-nodo è elevata ma quella inter-nodo è un collo di bottiglia. Massimizza l'utilità delle costose interconnessioni intra-nodo.
• Fattibilità per Contesti Lunghi: Man mano che i modelli si spostano verso finestre di contesto più lunghe (100k+ token), la memoria delle attivazioni diventa il collo di bottiglia principale. TSP è posizionato in modo unico per gestire questo senza l'enorme sovraccarico di dispositivi del parallelismo 2D.
• Flessibilità: TSP può essere composto ortogonalmente con altri schemi di parallelismo (Pipeline Parallelism, Expert Parallelism, Data Parallelism), offrendo un "asse" flessibile per i progettisti di sistemi da ottimizzare per topologie hardware specifiche.
• Efficienza dei Costi: Riducendo il numero di dispositivi necessari per il parallelismo del modello, TSP permette di utilizzare più dispositivi per il parallelismo dei dati, potenzialmente riducendo la dimensione totale del cluster richiesta per un dato lavoro di addestramento o aumentando la dimensione del batch effettiva.

In conclusione, TSP rappresenta un'evoluzione significativa nelle strategie di addestramento distribuito, scambiando un aumento gestibile del volume di comunicazione con guadagni sostanziali nell'efficienza della memoria e nell'utilizzo dell'hardware, rendendolo uno strumento critico per l'addestramento di LLM di prossima generazione con contesti lunghi.