ParaBlock: Communication-Computation Parallel Block Coordinate Federated Learning for Large Language Models

Sintesi Tecnica: ParaBlock: Comunicazione-Calcolo Parallela per il Block Coordinate Federated Learning di Grandi Modelli Linguistici

Problematica
L'Apprendimento Federato (Federated Learning - FL) è emerso come un paradigma di preservazione della privacy per l'addestramento di modelli su dati decentralizzati. Recentemente, gli schemi di discesa del coordinamento a blocchi federati (Federated Block Coordinate Descent - BCD) hanno guadagnato terreno per l'addestramento di modelli su larga scala, consentendo ai client di aggiornare solo un sottoinsieme (blocco) dei parametri del modello anziché l'intero modello. Tuttavia, nell'era dei Grandi Modelli Linguistici (LLM), anche un singolo blocco può contenere decine o centinaia di milioni di parametri. Questa scala crea un collo di bottiglia critico: la latenza di comunicazione necessaria per trasmettere questi grandi blocchi tra i client e il server diventa significativa, particolarmente per i client edge con risorse limitate.

I metodi tradizionali di federated BCD operano con un design a thread singolo (comunicazione-calcolo-comunicazione), dove i client devono attendere che la trasmissione del modello si completi prima di avviare il passaggio di calcolo successivo. Sebbene questo tempo di comunicazione fosse trascurabile per modelli più piccoli, ora ostacola gravemente la scalabilità e l'efficienza del fine-tuning federato per gli LLM, causando tempi di inattività sostanziali per i client.

Metodologia: ParaBlock
Per affrontare questa latenza, gli autori propongono ParaBlock, un nuovo metodo di federated BCD che introduce un'architettura di comunicazione-calcolo parallela. A differenza degli approoli standard a thread singolo, ParaBlock stabilisce due thread paralleli sui client locali:

1. Thread di Calcolo: Esegue gli aggiornamenti del blocco di coordinamento locale (fine-tuning del blocco assegnato) utilizzando i dati locali.
2. Thread di Comunicazione: Si sincronizza con il server per inviare gli aggiornamenti locali del round precedente e ricevere gli aggiornamenti globali aggregati.

Meccanismi Chiave:

• Sovrapposizione (Overlapping): I client eseguono il calcolo locale mentre comunicano simultaneamente con il server. Ciò elimina il tempo di inattività tipicamente causato dall'attesa della trasmissione di rete.
• Staleness (Obsolescenza) e Correzione: Poiché il thread di comunicazione è in ritardo di un round (i client ricevono gli aggiornamenti globali dal round $t-1$ mentre calcolano per il round $t$), i parametri del modello locale possono diventare incoerenti con il modello globale. ParaBlock affronta questo problema tramite un meccanismo di correzione. Quando un client aggiorna il proprio modello locale per il round successivo, non applica solo il nuovo aggiornamento locale al blocco corrente, ma corregge anche i parametri del blocco precedente utilizzando l'aggiornamento globale ricevuto ($\Delta_{t-1}$) per sostituire l'aggiornamento locale obsoleto ($\Delta^i_{t-1}$).
• Consistenza Ricorsiva: L'analisi teorica dimostra che questa correzione assicura che, dopo un round globale, i parametri del modello locale per il blocco precedentemente aggiornato siano esattamente allineati con i parametri del modello globale, mitigando efficacemente l'impatto del ritardo di un round.

Contributi Chiave

1. Architettura Innovativa: La proposta di ParaBlock, che sostituisce il tradizionale design FL a thread singolo con uno schema parallelo a due thread specificamente progettato per ridurre la latenza di comunicazione nel fine-tuning federato di LLM.
2. Garanzie Teoriche: Gli autori forniscono un'analisi rigorosa della convergenza per obiettivi non convessi. Dimostrano che ParaBlock raggiunge un tasso di convergenza di $O(1/\sqrt{T})$, che è coerente con i metodi standard di federated block coordinate descent. Ciò conferma che i guadagni di efficienza derivanti dal parallelismo non compromettono le prestazioni di convergenza teorica.
3. Validazione Empirica: Sono stati condotti estesi esperimenti per il fine-tuning di modelli Llama 3-8B e Llama 3.2-3B su compiti di istruzione generale (Alpaca-GPT4) e ragionamento matematico (MathInstruct/GSM8K).

Risultati

• Performance: ParaBlock mantiene prestazioni elevate nei task downstream, ottenendo punteggi comparabili o migliori rispetto ai baseline standard di federated BCD (ad es., FedBCD) e superando significativamente diversi metodi basati su Parameter-Efficient Fine-Tuning (PEFT) (ad es., FedIT, FLoRA) nei compiti di istruzione e ragionamento matematico.
• Efficienza: Il metodo riduce significativamente il tempo di esecuzione (wall-clock runtime). In scenari con bassa larghezza di banda di rete (50M/s) e piccoli batch, ParaBlock ha ridotto il tempo di esecuzione di oltre il 30% rispetto ai baseline di block-coordinate come FedCyBGD. Rispetto ai metodi basati su LoRA come FedIT, ParaBlock ha ottenuto una riduzione del runtime di circa il 40% nella maggior parte delle configurazioni.
• Robustezza: Gli studi di ablazione hanno mostrato che il metodo rimane efficace sotto distribuzioni di dati eterogenee (non-IID) e attraverso diverse architetture di modelli (inclusi Qwen-2.5). Il metodo ha inoltre dimostrato compatibilità con scenari di partecipazione parziale in contesti cross-silo.

Significato
Il documento sostiene che ParaBlock offra una soluzione semplice ed efficace al collo di bottiglia della comunicazione inerente al federated BCD per gli LLM. Sovrapponendo comunicazione e calcolo, esso consente il deployment scalabile del fine-tuning federato per grandi modelli su client con risorse limitate senza sacrificare le garanzie di convergenza o l'utilità del modello. Il lavoro evidenzia come la riduzione della latenza di comunicazione attraverso il parallelismo architettonico sia una strategia vitale e necessaria per la prossima generazione di addestramento federato di LLM, distinta da e ortogonale alle tecniche esistenti come la quantizzazione o il pruning.