Riepilogo Tecnico: Parallelismo Eterogeneo per l'Addestramento di Modelli Linguistici Multimodali di Grande Dimensione

1. Enunciato del Problema

Mentre l'addestramento dei modelli fondazionali evolve dai transformer basati solo su testo verso sistemi multimodali (che coprono immagini, audio, video e documenti), il grafico di addestramento diventa sempre più eterogeneo in termini di profondità, larghezza e dimensione del contesto dei moduli. Tuttavia, gli stack di parallelismo attuali (Megatron-LM, DeepSpeed, ecc.) impongono tipicamente un layout di parallelismo unico e condiviso (Parallelismo di Tensori, Contesto, Pipeline, Dati ed Esperti) su tutto il grafico end-to-end.

Questo approccio "taglia unica" crea inefficienze significative quando i requisiti degli encoder e dei Modelli Linguistici di Grande Dimensione (LLM) divergono:

• Comunicazione Non Necessaria: I grandi LLM richiedono spesso gradi elevati di Parallelismo di Tensori (TP). Costringere piccoli encoder a ereditare questi gradi di TP introduce collettivi all-reduce non necessari attorno ai calcoli locali dell'encoder, aggiungendo un sovraccarico che supera i risparmi di memoria.
• Disallineamento del Parallelismo di Contesto: Il Parallelismo di Contesto (CP) è essenziale per i LLM per gestire la memoria dell'attenzione su sequenze lunghe. Tuttavia, le lunghezze delle sequenze degli encoder sono limitate dalla risoluzione della modalità (ad esempio, patch di immagini), non dalla finestra di contesto del LLM. Ereditare shard di CP del LLM per sequenze di encoder a dimensione fissa non offre sollievo alla memoria, mentre quei rank potrebbero funzionare in modo più efficiente come repliche indipendenti di Parallelismo di Dati (DP) per elaborare più campioni.
• Vincoli della Pipeline: Il Parallelismo di Pipeline (PP) confina spesso gli encoder alla prima fase della pipeline del LLM. Questo costringe l'encoder a ereditare il grado di TP del LLM e può sovraccaricare rank specifici, impedendo al LLM di utilizzare un layout di pipeline più efficiente.

La sfida principale è che il punto operativo ottimale per calcolo, memoria e comunicazione si sposta in base ai parametri del carico di lavoro (scala dell'encoder, lunghezza del contesto, densità dei token, dimensione del mondo). Un layout statico e condiviso non può adattarsi a questi spostamenti.

2. Metodologia

Gli autori propongono il Parallelismo Eterogeneo, un'astrazione runtime che permette ai moduli all'interno di un singolo grafico multimodale end-to-end di scegliere layout logici indipendenti e posizionamenti fisici dei rank. Il sistema supporta due modalità di esecuzione distinte:

A. Astrazione Principale: Comunicatori di Confine

Per disaccoppiare i layout dei moduli preservando al contempo la semantica dell'addestramento, il sistema introduce comunicatori di confine. Questi sono operatori consapevoli dell'autograd attaccati agli archi inter-modulo che gestiscono la trasformazione dei tensori tra diverse rappresentazioni distribuite.

• Passo in Avanti: Materializza l'attivazione sorgente nel layout del modulo di destinazione.
• Passo all'Indietro: Applica la trasformazione inversa per instradare i gradienti al layout del modulo sorgente.
• Correttezza: Il sistema garantisce che la semantica dei tensori sia preservata attraverso griglie indipendenti. Mappature errate sono prevenute definendo rigorosamente trasformazioni logiche di shard di batch (Equal-DP, Fan-in e Fan-out) per ogni arco.

B. Modalità di Esecuzione

1. Esecuzione Collocata:

   • Meccanismo: I moduli condividono le stesse GPU fisiche ma le interpretano secondo griglie logiche diverse (ad esempio, un encoder con alto DP e basso TP che condivide i rank con un LLM con basso DP e alto TP).
   • Implementazione: Il comunicatore agisce come operatore di layout a rank condiviso. Esegue operazioni collettive locali (ad esempio, all-gather o selezione di intervalli) per reinterpretare il layout dei dati senza spostare tensori tra insiemi di rank.
   • Pianificazione: Utilizza una pianificazione a tre fasi per evitare conflitti tra collettivi dell'encoder e la pipeline scalinata del LLM (1F1B).
     • Fase 1: Forward dell'encoder + trasformazione forward collocata.
     • Fase 2: Esecuzione della pipeline LLM staccata (1F1B) sul tensore trasformato.
     • Fase 3: Trasformazione backward collocata + backward dell'encoder.
   • Vantaggio: Permette di scaricare i parametri dell'encoder durante la fase della pipeline LLM, migliorando l'efficienza della memoria.

2. Esecuzione Non Collocata:

   • Meccanismo: I moduli occupano insiemi disgiunti di rank (ad esempio, un "isola encoder" e un "isola LLM").
   • Implementazione: Utilizza comunicazione esplicita Punto-a-Punto (P2P) tra leader di confine deterministici (rank con coordinate canoniche).
     • Forward: I leader sorgente inviano ai leader di destinazione; broadcast locali materializzano i dati.
     • Backward: I gradienti sono instradati indietro attraverso i percorsi inversi dei leader.
   • Pianificazione: Utilizza un dispatcher 1F1B consapevole del grafico che risolve gli archi in un grafico multistadio, gestendo giunzioni e rami dove più flussi di encoder alimentano il LLM.
   • Vantaggio: Previene che lo stato dell'encoder (parametri, attivazioni, stati dell'ottimizzatore) competano con il LLM per la memoria sugli stessi rank, permettendo al LLM di utilizzare un layout di pipeline più efficiente.

3. Contributi Chiave

1. Astrazione Runtime di Parallelismo Eterogeneo: Un sistema che permette layout TP/CP/PP/DP/EP indipendenti e posizionamenti fisici dei rank per moduli in un singolo grafico MLLM, supportando sia l'esecuzione collocata che non collocata.
2. Comunicatori di Confine e Pianificazione: L'introduzione di comunicatori di confine consapevoli dell'autograd che gestiscono le trasformazioni di layout, accoppiati a scheduler specializzati (a tre fasi per l'esecuzione collocata, 1F1B consapevole del grafico per quella non collocata) per mantenere l'efficienza della pipeline.
3. Valutazione del Regime Operativo: Una valutazione completa che caratterizza quando l'eterogeneità è benefica. Lo studio identifica che l'eterogeneità collocata è ottimale quando c'è spazio sufficiente nella memoria ma i layout logici differiscono, mentre l'eterogeneità non collocata diventa superiore quando lo stato dell'encoder vincola la memoria o il layout della pipeline del LLM.
4. Correttezza e Open Source: Validazione della parità di convergenza della perdita rispetto a baseline omogenee tramite controlli FP32 a livello di passo e scansioni di addestramento in stile LLaVA. Il sistema è rilasciato come estensione open source a Megatron-LM.

4. Risultati Sperimentali

Il sistema è stato valutato su GPU NVIDIA H100 da 80GB attraverso vari carichi di lavoro multimodali (encoder da 1B a 12B, LLM da 14B a 120B) e lunghezze di contesto (8K–128K).

• Eterogeneità Collocata:
  • Migliora i TFLOPS/GPU fino al 49,3% (osservato in scenari a contesto lungo, ad esempio 64K–128K, dove i LLM richiedono CP ma gli encoder no).
  • Migliora il throughput del 12,1%–21,6% in varie configurazioni di encoder 1B/3B + LLM 14B/70B/120B permettendo all'encoder di utilizzare un DP più alto e un TP più basso.
• Eterogeneità Non Collocata:
  • Migliora il throughput aggregato di token fino al 13,0% e i TFLOPS/GPU fino al 9,6% per carichi di lavoro con encoder grandi (6B–12B) + LLM da 120B.
  • Diventa il punto operativo preferito quando la scala dell'encoder (ad esempio, 6B+) crea pressione sulla memoria che limita il layout della pipeline del LLM nelle impostazioni collocate.
• Convergenza: Validato che i layout eterogenei producono curve di perdita identiche alle baseline omogenee, confermando che i comunicatori di confine preservano correttamente la semantica dei gradienti.

5. Significato e Affermazioni

Il documento afferma che il parallelismo eterogeneo è un'evoluzione necessaria per l'addestramento multimodale, allontanandosi da vincoli rigidi di layout a livello di modello. Espandendo gradi di libertà aggiuntivi, il sistema permette all'infrastruttura di addestramento di sintonizzare dinamicamente il punto operativo in base alle caratteristiche del carico di lavoro (dimensione dell'encoder, lunghezza del contesto, densità dei token).

Gli autori sottolineano che questo non è meramente un cambiamento di configurazione, ma un substrato runtime fondamentale che preserva la semantica dell'addestramento end-to-end consentendo:

• Esecuzione collocata per estrarre il massimo throughput da un budget GPU fisso quando i moduli possono condividere l'hardware ma necessitano di griglie logiche diverse.
• Esecuzione non collocata per disaccoppiare le pressioni di memoria e pianificazione quando lo stato dell'encoder altrimenti costituirebbe un collo di bottiglia per il LLM.

Il lavoro è presentato come un'estensione open source a Megatron-LM, fornendo una base riutilizzabile per futuri framework MLLM modulari e sistemi di addestramento disaggregati. Il documento non afferma di risolvere tutte le sfide multimodali, ma affronta specificamente l'inefficienza causata dal forzare moduli eterogenei in un singolo layout di parallelismo.