ZorBA: Zeroth-order Federated Fine-tuning of LLMs with Heterogeneous Block Activation

Il paper propone ZorBA, un framework di fine-tuning federato per grandi modelli linguistici che utilizza ottimizzazione di ordine zero e attivazione eterogenea dei blocchi per ridurre drasticamente l'uso di VRAM e l'overhead di comunicazione mantenendo un'alta efficienza di convergenza.

L'essenziale

Immagina di voler insegnare a un gruppo di amici molto diversi tra loro (alcuni con computer potenti, altri con laptop vecchi) come parlare una nuova lingua usando un "super-linguista" digitale (un modello di intelligenza artificiale gigante). Questo è il cuore del Federated Learning: tutti imparano insieme senza condividere i loro appunti privati, ma solo le loro lezioni apprese.

Il problema? Il "super-linguista" è enorme. Per aggiornarlo, normalmente serve un computer potentissimo e tantissima memoria video (VRAM), cosa che molti non hanno. Inoltre, inviare tutti i dati aggiornati via internet è come cercare di spedire un camion intero per posta: costa troppo e ci vuole troppo tempo.

Ecco che entra in scena ZorBA, la soluzione proposta in questo articolo.

1. Il Problema: Il Gigante che non entra nella stanza

Immagina che il modello di intelligenza artificiale sia un elefante (il modello LLM).

• Il problema della memoria (VRAM): Per allenare l'elefante, di solito devi tenerlo in una stanza piena di specchi (i gradienti) per vedere come si muove. Ma se la stanza è piccola (il computer del cliente), l'elefante non ci sta e gli specchi si rompono.
• Il problema della comunicazione: Se ogni elefante deve inviare una foto di sé stesso al capo ogni volta che si muove, la posta si intasa completamente.

2. La Soluzione ZorBA: Tre Trucchi Geniali

ZorBA risolve questi problemi con tre idee semplici ma brillanti:

A. Il "Metodo del Tocco" (Ottimizzazione di Ordine Zero)

Invece di usare gli specchi (che richiedono molta memoria), ZorBA usa un metodo più semplice: il tocco.
Immagina di dover trovare la cima di una montagna al buio. Il metodo normale ti dà una mappa dettagliata (gradiente) che richiede di memorizzare tutto il terreno. ZorBA invece dice: "Fai un piccolo passo a caso, vedi se sali o scendi, e poi torna indietro".

• Vantaggio: Non devi memorizzare la mappa intera (nessuna memoria extra!). Ti basta solo vedere cosa succede dopo un piccolo passo. Questo riduce drasticamente lo spazio necessario sul computer.

B. Il "Dizionario Condiviso" (Attivazione Eterogenea dei Blocchi)

Il modello è come un libro con 24 capitoli (blocchi). Non tutti hanno bisogno di leggere tutti i capitoli per imparare.

• L'idea: Il capo (il server centrale) assegna a ogni amico solo alcuni capitoli da leggere e aggiornare, in base a quanto è grande il loro computer.
  • Chi ha un computer potente legge 10 capitoli.
  • Chi ha un computer debole ne legge solo 3.
• Il trucco: Anche se ognuno legge capitoli diversi, alla fine si mettono d'accordo per creare un unico libro completo. Questo permette a tutti di partecipare senza schiacciare i computer più piccoli.

C. Il "Seme Magico" (Condivisione dei Semi Casuali)

Per fare i "piccoli passi a caso" (il metodo del tocco), serve un generatore di numeri casuali. Di solito, inviare questi numeri casuali via internet consuma molta banda.

• L'idea: Invece di inviare i numeri, il capo invia un seme (una chiave segreta). Tutti usano lo stesso seme per generare esattamente la stessa sequenza di numeri casuali.
• Risultato: Si risparmia un'enorme quantità di dati da inviare. È come dire a tutti: "Usate il seme numero 42", invece di inviare 1 milione di numeri.

3. L'Algoritmo: Il Cuoco che bilancia gli ingredienti

Il team ha creato un algoritmo intelligente (chiamato "algoritmo lessicografico") che decide quali capitoli dare a chi.
Immagina di dover distribuire ingredienti per una torta tra 50 cuochi:

• Alcuni cuochi hanno forni piccoli (poca memoria).
• Altri hanno forni grandi.
• L'algoritmo deve assicurarsi che:
  1. La torta venga fatta bene (convergenza veloce).
  2. Nessuno bruci il forno (memoria non superata).
  3. Nessuno si lamenti di non aver ricevuto abbastanza ingredienti (equità).

L'algoritmo prova diverse combinazioni finché non trova il perfetto equilibrio: dare più capitoli a chi può permetterselo e meno a chi è limitato, garantendo che l'elefante (il modello) impari comunque velocemente.

I Risultati: Cosa hanno scoperto?

Gli esperimenti hanno mostrato che ZorBA è un vincitore:

1. Risparmio di memoria: Rispetto ai metodi tradizionali, ZorBA riduce l'uso della memoria video (VRAM) fino al 62%. È come se un computer da 100 euro potesse fare il lavoro di uno da 250 euro.
2. Velocità: Nonostante usi meno risorse, impara quasi alla stessa velocità (o talvolta più velocemente) dei metodi che usano tutti i capitoli.
3. Comunicazione: Riduce il traffico internet necessario per scambiare i dati, rendendo il processo molto più leggero.

In sintesi

ZorBA è come un'orchestra dove ogni musicista suona solo lo strumento che può permettersi, ma tutti seguono lo stesso spartito segreto (i semi casuali) e il direttore d'orchestra (il server) assegna le parti in modo che la musica (l'intelligenza artificiale) suoni perfetta senza che nessuno debba comprare uno strumento costoso o inviare file pesanti. È un modo intelligente per rendere l'intelligenza artificiale accessibile a tutti, anche con computer modesti.

Sommario tecnico

1. Il Problema

L'addestramento federato (Federated Learning - FL) di grandi modelli linguistici (LLM) presenta sfide critiche legate alle risorse e alla comunicazione:

• Consumo di VRAM: I metodi convenzionali di fine-tuning utilizzano la retropropagazione (backpropagation), che richiede lo stoccaggio dei gradienti e delle attivazioni forward-pass. Per modelli di grandi dimensioni (miliardi di parametri), questo supera spesso la capacità di memoria video (VRAM) dei client con risorse limitate.
• Overhead di Comunicazione: Lo scambio frequente di aggiornamenti del modello o dei gradienti tra client e server centrale genera un carico di comunicazione significativo, specialmente in scenari con modelli ad alta dimensionalità.
• Limitazioni dell'Ottimizzazione del Terzo Ordine (Zeroth-Order): Sebbene l'ottimizzazione del terzo ordine (che stima i gradienti tramite differenze finite senza calcolare le derivate esatte) elimini la necessità di memorizzare i gradienti, le implementazioni esistenti attivano tutti i blocchi del modello. Questo porta a:
  1. Un consumo di VRAM elevato a causa delle attivazioni forward-pass (che crescono linearmente con il numero di blocchi).
  2. Tassi di convergenza più lenti rispetto ai metodi del primo ordine a causa dell'alta varianza introdotta dai vettori di perturbazione in spazi ad alta dimensionalità.
  3. Un overhead di comunicazione ancora elevato se si trasmettono gradienti stimati per l'intero modello.

2. Metodologia: Il Framework ZorBA

Gli autori propongono ZorBA, un framework di fine-tuning federato basato sull'ottimizzazione del terzo ordine con attivazione eterogenea dei blocchi.

Componenti Chiave:

• Ottimizzazione del Terzo Ordine (Zeroth-Order): Sostituisce la retropropagazione con un approccio basato solo sul forward-pass. I gradienti vengono stimati utilizzando differenze finite di valori di perdita generati da vettori di perturbazione casuali. Questo elimina lo stoccaggio dei gradienti, riducendo drasticamente l'uso di VRAM.
• Attivazione Eterogenea dei Blocchi (Heterogeneous Block Activation):
  • Invece di aggiornare tutti i blocchi del Transformer su tutti i client, il server centrale assegna a ciascun client un sottoinsieme diverso di blocchi da attivare.
  • Questa decisione è guidata dai vincoli di VRAM di ciascun client e dall'obiettivo di massimizzare la velocità di convergenza.
  • I blocchi non attivati vengono "congelati" (frozen), riducendo il numero di attivazioni forward-pass da memorizzare.
• Semi Casuali Condivisi (Shared Random Seeds):
  • Per evitare lo scambio di vettori di perturbazione ad alta dimensionalità, il server e i client condividono un insieme di semi casuali.
  • Questo permette a tutte le parti di generare indipendentemente gli stessi vettori di perturbazione, riducendo l'overhead di comunicazione a zero per quanto riguarda i vettori stessi.
• Comunicazione Efficiente: Invece di inviare gradienti stimati completi, i client inviano solo le differenze finite delle stime dei gradienti. Il server, avendo gli stessi semi e lo stesso modello, può ricostruire i gradienti aggregati.

Formulazione del Problema e Algoritmo:

Gli autori formulano un problema di ottimizzazione multi-obiettivo per:

1. Minimizzare il termine di bias nella bound di convergenza (rappresentato da $\Lambda(A)$, che dipende dalla distribuzione dei blocchi attivati).
2. Minimizzare l'uso totale di VRAM.

Per risolvere questo problema NP-hard, propongono un algoritmo lessicografico basato sul vincolo $\epsilon$:

1. Massimizzazione della Popolarità Minima: Si determina il massimo numero minimo di client che possono attivare un blocco (popolarità minima) rispettando i vincoli di VRAM.
2. Regolazione Greedy: Si attivano blocchi aggiuntivi in modo greedy per minimizzare il numero di client che rimangono al limite minimo di popolarità, bilanciando convergenza e risorse.
3. Si costruisce un fronte di Pareto per selezionare la matrice di attivazione ottimale che bilancia velocità di convergenza e consumo di memoria.

3. Contributi Principali

1. Framework ZorBA: Introduzione di un nuovo paradigma che combina ottimizzazione del terzo ordine e attivazione selettiva dei blocchi per LLM federati.
2. Analisi Teorica:
   • Derivazione di un limite di convergenza per ZorBA in contesti non convessi.
   • Dimostrazione che la velocità di convergenza è governata dalla distribuzione della "popolarità minima" dei blocchi tra i client, non solo dal numero totale di blocchi attivati.
   • Dimostrazione che è possibile ottenere una bound di convergenza indipendente dalla dimensionalità del modello ($d$) sotto certe condizioni di rango effettivo dell'Hessiano.
3. Algoritmo di Ottimizzazione: Sviluppo di un algoritmo efficiente ($\epsilon$-constraint lexicographic) per determinare le decisioni di attivazione dei blocchi, riducendo la complessità computazionale rispetto alla soluzione diretta del problema.
4. Validazione Sperimentale: Test su dataset reali (AG-News, SST-2, SNLI) con modelli OPT-125M e OPT-1.3B.

4. Risultati Sperimentali

I risultati mostrano che ZorBA supera significativamente le baseline (FedIT, FedZO, DeComFL):

• Riduzione della VRAM: ZorBA riduce l'uso totale di VRAM fino al 62.41% rispetto alle altre metodologie, permettendo il fine-tuning su dispositivi con risorse limitate.
• Velocità di Convergenza: ZorBA converge più velocemente rispetto a FedZO e DeComFL (fino al 23.76% in meno di round di addestramento per raggiungere la stessa accuratezza), dimostrando che l'attivazione eterogenea ottimizzata è superiore all'attivazione completa casuale o uniforme.
• Overhead di Comunicazione: ZorBA riduce drasticamente l'overhead di comunicazione rispetto ai metodi che scambiano gradienti completi, mantenendo un carico comparabile o inferiore a DeComFL, ma con una convergenza superiore.

5. Significato e Impatto

Questo lavoro è significativo perché risolve il collo di bottiglia principale nell'addestramento federato di LLM: la scarsità di memoria GPU.

• Accessibilità: Abilita il fine-tuning collaborativo di modelli di grandi dimensioni su dispositivi edge con risorse limitate, senza richiedere l'invio di dati grezzi.
• Efficienza: Dimostra che non è necessario aggiornare l'intero modello per ottenere buone prestazioni; una strategia intelligente di attivazione parziale può accelerare la convergenza e risparmiare risorse.
• Privacy e Sicurezza: L'uso di semi condivisi e la trasmissione solo di differenze finite riducono il rischio di leakage di gradienti e proteggono meglio la privacy rispetto allo scambio di parametri completi.

In sintesi, ZorBA offre una soluzione pratica e teoricamente fondata per scalare il fine-tuning federato di LLM in ambienti reali con vincoli di risorse stringenti.