Sintesi Tecnica: Fusione di Modelli per l'Aggregazione in DiLoCo

Definizione del Problema

L'addestramento Distributed Low-Communication (DiLoCo) è emerso come un metodo per addestrare grandi modelli linguistici attraverso cluster di calcolo scarsamente connessi ("isole"), minimizzando la sincronizzazione globale. In DiLoCo, i worker eseguono molteplici passi di ottimizzazione locale (passi interni) prima di aggregare i loro aggiornamenti di parametri (pseudo-gradienti) tramite un ottimizzatore globale. Sebbene efficace nel ridurre i costi di comunicazione, DiLoCo soffre di un degrado delle prestazioni all'aumentare del numero di worker ($R$) o del numero di passi locali ($H$). Questo degrado avviene perché le repliche locali divergono significativamente, causando un disallineamento tra gli pseudo-gradienti mediati e gli aggiornamenti prodotti dall'addestramento data-parallel strettamente sincronizzato.

Contemporaneamente, le tecniche di Model Merging (fusione di modelli) sono state sviluppate per combinare modelli addestrati indipendentemente in un unico modello. Una sfida centrale nella fusione è l'interferenza tra task ("task interference"), dove gli aggiornamenti dei parametri provenienti da diversi task entrano in conflitto, portando a cali di prestazioni, specialmente all'aumentare del numero di modelli da fondere.

Gli autori identificano un'analogia strutturale tra questi due campi: il passo di aggregazione dello pseudo-gradiente in DiLoCo (con un ottimizzatore esterno SGD) è matematicamente equivalente alla Task Arithmetic del model merging. In questa analogia:

• Il modello globale condiviso all'inizio di un passo esterno corrisponde all'inizializzazione pre-addestrata.
• Le repliche locali dei worker corrispondono ai modelli "esperti" sottoposti a fine-tuning.
• Gli aggiornamenti locali (pseudo-gradienti) corrispondono ai "task vectors".
• Il numero di worker corrisponde al numero di task.

Il problema centrale affrontato è se tecniche avanzate di model merging, progettate specificamente per mitigare l'interferenza tra task, possano essere riutilizzate per migliorare il passo di aggregazione in DiLoCo, ripristinando così le prestazioni all'aumentare del numero di worker o dei passi locali.

Metodologia

Ponte Teorico

Il documento stabilisce che l'aggiornamento esterno di DiLoCo con SGD è un'applicazione diretta della task arithmetic. Visualizzando l'aggregazione degli pseudo-gradienti come un problema di model merging, gli autori propongono di sostituire l'operazione standard di media con algoritmi di fusione sofisticati progettati per ridurre l'interferenza negativa.

Valutazione Sperimentale dei Metodi di Fusione

Gli autori hanno testato sistematicamente diversi metodi di merging allo stato dell'arte come sostituti diretti per il passo di aggregazione di DiLoCo. I metodi valutati includevano:

• Basati sulla sparsità: TIES (TrIm, Elect Sign) e DARE (Drop And REscale).
• Basati sull'ortogonalizzazione: TSV (Task Singular Vector), Iso-C (Isotropic merging in Common subspace) e Iso-CTS.

Risultati sui Metodi di Fusione:

• Metodi di sparsità (TIES, DARE): Hanno performato in modo simile al DiLoCo SGD standard solo a livelli di sparsità molto bassi. Le prestazioni sono degradate all'aumentare della sparsità, suggerendo che non offrissero un vantaggio distinto rispetto alla semplice media in questo contesto.
• TSV: Ha performato scarsamente (loss di validazione più alta) ed era computazionalmente proibitivo, richiedendo molteplici decomposizioni in valori singolari (SVD) per ogni passo.
• Iso-CTS: Ha superato il baseline ma era computazionalmente costoso e non offriva un vantaggio significativo rispetto a Iso-C nel contesto di DiLoCo (dove i task locali sono brevi e simili).
• Iso-C: È emerso come il candidato più promettente. Migliora le prestazioni sostituendo lo spettro dei valori singolari dello pseudo-gradiente aggregato con la sua media, appiattendo efficacemente le direzioni ad alta energia per ottenere un sottospazio comune bilanciato. Richiede solo una SVD per matrice di parametri, offrendo un rapporto costo-prestazioni favorevole.

Algoritmo Proposto: IsoLoCo

Basandosi sul successo di Iso-C, gli autori propongono IsoLoCo (Isotropic Distributed Low-Communication). IsoLoCo adatta Iso-C per l'addestramento distribuito tramite:

1. Correzione Isotropa: Per le matrici di parametri 2D (layer di peso), calcola la SVD dello pseudo-gradiente mediato, sostituisce i valori singolari con la loro media e ricostruisce l'aggiornamento. Questo regolarizza lo spettro preservando la struttura del sottospazio.
2. Momentum di Nesterov: A differenza del normale Iso-C, IsoLoCo incorpora il momentum di Nesterov nell'ottimizzatore esterno. Ciò è motivato da studi precedenti che mostrano il beneficio del momentum nei contesti FedOpt/DiLoCo. Lo pseudo-gradiente corretto viene inserito in un buffer di momentum per generare l'aggiornamento finale.

Ottimizzazione dell'Implementazione

Per affrontare il costo computazionale della SVD, gli autori hanno sviluppato un'implementazione veloce di IsoLoCo utilizzando le iterazioni di Newton–Schulz. Questa sostituisce la SVD standard con l'ortogonalizzazione iterativa e utilizza un proxy RMS per la media dei valori singolari. Questa ottimizzazione produce un'ortogonalizzazione fino a due ordini di grandezza più veloce con una degradazione trascurabile delle prestazioni.

Risultati Chiave

Prestazioni in base al Numero di Worker

Gli esperimenti sono stati condotti su modelli di tipo Llama (178M, 512M e 1B di parametri) utilizzando il dataset DCLM.

• Scalabilità con i Worker: IsoLoCo supera significativamente il DiLoCo standard all'aumentare del numero di worker.
  • Alla scala di 178M con $R=128$, IsoLoCo ha ridotto il gap della loss di validazione rispetto al baseline data-parallel del 5,8% rispetto a DiLoCo.
  • A $R=128$, IsoLoCo ha mostrato un degrado del 9,7% rispetto all'addestramento data-parallel, mentre DiLoCo ha mostrato un degrado del 16,5%.
  • IsoLoCo ha persino superato il baseline data-parallel con un basso numero di worker ($R \le 4$).
• Dimensione del Modello: Il vantaggio di IsoLoCo rispetto a DiLoCo persiste attraverso le dimensioni dei modelli (da 178M a 1B), sebbene il miglioramento relativo sia leggermente minore per i modelli più grandi (1,2% a 1B rispetto a ~2,4% a 178M), probabilmente perché i modelli più grandi scalano naturalmente meglio con il numero di worker.

Impatto dei Passi Interni

IsoLoCo è stato valutato con intervalli di sincronizzazione crescenti ($H = 30$ a $240$). Ha superato costantemente DiLoCo in tutti gli intervalli, dimostrando robustezza a periodi più lunghi di divergenza locale.

Confronto con DiLoCo basato su Momentum

In particolare, la combinazione di DiLoCo SGD con l'aggregazione Iso-C (senza momentum) ha già superato il DiLoCo standard (che usa il momentum di Nesterov). L'aggiunta del momentum di Nesterov a Iso-C (creando IsoLoCo) ha ulteriormente ampliato il divario di prestazioni, indicando che l'aggregazione basata sul merging e i meccanismi di momentum sono complementari.

Compatibilità con Muon

Il documento nota che IsoLoCo può essere composto con l'ottimizzatore interno Muon (MuLoCo). In questa configurazione, IsoLoCo ha ottenuto la loss di validazione più bassa tra tutti i metodi per $R \ge 32$, mentre IsoLoCo con un ottimizzatore interno AdamW ha eguagliato o superato MuLoCo con un numero elevato di worker, con costi di ortogonalizzazione significativamente inferiori.

Significato e Rivendicazioni

Il documento sostiene che vedere l'addestramento a bassa comunicazione attraverso la lente del model merging fornisce un potente framework per migliorare la scalabilità. I principali contributi sono:

1. Unificazione Teorica: Stabilire un'analogia diretta tra l'aggregazione dello pseudo-gradiente di DiLoCo e la task-arithmetic del model merging.
2. Validazione Empirica: Dimostrare che le tecniche di model merging, specificamente Iso-C, possono servire come regole di aggregazione superiori per l'addestramento distribuito, superando la semplice media e persino il DiLoCo basato su momentum.
3. IsoLoCo: Proporre un algoritmo pratico, potenziato dal momentum, che riduce significativamente il gap di prestazioni tra l'addestramento a bassa comunicazione e gli standard gold del data-parallel, specialmente in regimi con molti worker o lunghi intervalli di sincronizzazione.
4. Efficienza: Fornire un'implementazione basata su Newton-Schulz che rende il processo di merging, pesante in termini di SVD, computazionalmente fattibile per l'addestramento distribuito iterativo.

Gli autori concludono che l'aggregazione ispirata al merging è una strategia efficace per l'addestramento a bassa comunicazione, offrendo una via per scalare l'addestramento su cluster geograficamente distribuiti o scarsamente connessi senza sacrificare la qualità del modello.