Sintesi Tecnica: DASH – Shampoo più veloce tramite il precondizionamento a blocchi batch e solver di radice inversa efficienti

Problema

Shampoo è un noto ottimizzatore di secondo ordine approssimato, noto per produrre modelli con minori outlier di attivazione e una migliore comprimibilità. Tuttavia, la sua adozione diffusa è ostacolata da un overhead computazionale significativo, dovuto principalmente all'operazione costosa del calcolo delle radici di matrici inverse ($A^{-1/4}$) per le matrici di precondizionamento. Il Distributed Shampoo tenta di mitigare questo problema dividendo i precondizionatori in blocchi $B \times B$ per ridurre la complessità da $O(n^3)$ a $O(Bn^2)$, ma le implementazioni esistenti soffrono di critiche lacune di efficienza:

1. Elaborazione Sequenziale: Le implementazioni standard elaborano i blocchi in modo sequenziale, fallendo nel sfruttare appieno il parallelismo delle moderne GPU.
2. Solver Inefficienti: La dipendenza predefinita dalla Decomposizione in Autovettori (EVD) è difficile da parallelizzare sulle GPU. Metodi iterativi alternativi come il Coupled-Newton (CN) sono spesso disabilitati di default a causa di preoccupazioni sulla stabilità numerica e sulla velocità di convergenza.
3. Scalabilità Subottimale: I metodi attuali spesso scalano le matrici utilizzando la norma di Frobenius, che è computazionalmente economica ma può spingere gli autovalori verso lo zero, richiedendo significativamente più iterazioni per convergere a una precisione target.

Metodologia

Il documento propone DASH (Distributed Accelerated SHampoo), un'implementazione ad alte prestazioni progettata per colmare il divario tra l'efficienza teorica e le capacità dell'hardware. La metodologia si basa su due pilastri principali:

1. Precondizionamento a Blocchi Batch (Ottimizzazione a livello di Sistema)

Invece di elaborare i blocchi di precondizionamento sequenzialmente, DASH impila i blocchi indipendenti in tensori 3D.

• Impilamento di Tensori 3D: I blocchi di gradiente $G_{ij}$ vengono impilati in una matrice di forma $(N, B, B)$, dove $N$ è il numero totale di blocchi.
• Esecuzione Parallela: Ciò consente all'ottimizzatore di applicare le procedure di radice inversa (EVD, CN o nuovi metodi iterativi) all'intero batch di blocchi simultaneamente, utilizzando primitive di moltiplicazione di matrici batch (bmm).
• Utilizzo della GPU: Questo cambiamento architettonico massimizza il throughput dei Tensor Core della GPU, riducendo la frammentazione della memoria e consentendo l'uso efficiente dell'aritmetica a mezza precisione (FP16) per determinati solver.

2. Solver di Radice Inversa Efficienti (Ottimizzazione Algoritmica)

Il documento introduce e analizza approcci di algebra lineare avanzata per sostituire o integrare l'EVD e il CN standard:

• Iterazione Newton-Denman-Beavers (NDB): Un approccio innovativo per Shampoo che calcola sia le radici quadrate che le radici quadrate inverse. Gli autori ottimizzano la prima iterazione per eliminare le moltiplicazioni di matrici ridondanti, riducendo l'overhead a 3 matmul per step (effettivamente meno in pratica).
• Approssimazioni con Polinomi di Chebyshev: Utilizzo dell'algoritmo di Clenshaw per approssimare le potenze delle matrici, offrendo un metodo iterativo alternativo.
• Multi-Power-Iteration per lo Scaling: Per affrontare i problemi di convergenza, il documento sostituisce lo scaling standard della norma di Frobenius con un metodo Multi-Power-Iteration.
  • Meccanismo: Stima il raggio spettrale ($\lambda_{max}$) utilizzando un pool di vettori di partenza (fino a 32) elaborati in parallelo tramite Power-Iteration.
  • Vantaggio: Fornisce un fattore di scala più accurato rispetto alla norma di Frobenius, garantendo che gli autovalori rimangano in un regime in cui i solver iterativi convergono rapidamente (vicini a 1), riducendo così il numero di iterazioni richieste e migliorando la stabilità numerica.

Contributi Chiave

1. Implementazione DASH: Un framework efficiente per GPU che sostituisce i loop sequenziali con operazioni tensoriali 3D parallele, ottenendo un incremento di velocità fino a 5,6× rispetto al Distributed Shampoo ben ottimizzato.
2. Nuovi Solver: L'introduzione dell'iterazione Newton-DB (NDB) e delle approssimazioni con polinomi di Chebyshev come alternative più veloci a EVD e CN. Il documento dimostra che NDB produce una perplexity di validazione inferiore rispetto a Coupled-Newton e spesso supera l'EVD.
3. Analisi della Convergenza: Un'analisi comportamentale dettagliata che mostra come lo scaling standard della norma di Frobenius sia subottimale, portando spesso a una convergenza più lenta. Il documento fornisce spiegazioni intuitive per i distinti comportamenti di convergenza di CN rispetto a NDB.
4. Multi-Power-Iteration: Un'implementazione efficiente in mezza precisione della Power-Iteration che stima robustamente i raggi spettrali, consentendo uno scaling ottimale delle matrici che soddisfa i criteri di convergenza per NDB e CN con meno step.

Risultati Sperimentali

Gli autori hanno valutato DASH su un modello Llama con 953M di parametri (dimensione embedding 2048, lunghezza sequenza 1024) addestrato sul dataset C4.

• Velocità:
  • DASH riduce il tempo di esecuzione dello step dell'ottimizzatore di 4× - 5,6× rispetto al Distributed Shampoo, a seconda del metodo di radice inversa e della precisione utilizzata.
  • Il maggior incremento di velocità (5,6×) è stato osservato confrontando l'EVD (FP32) di Distributed Shampoo contro il CN (FP16) di DASH con stacking a livello di GPU.
  • Anche con aggiornamenti frequenti del precondizionatore ($f=1$), DASH mantiene sostanziali vantaggi di velocità.
• Accuratezza (Perplexity di Validazione):
  • DASH eguaglia il Distributed Shampoo in quasi tutte le configurazioni stabili.
  • NDB produce costantemente una perplexity di validazione inferiore rispetto a CN e spesso supera l'EVD.
  • L'uso dello scaling con Power-Iteration con NDB ha prodotto la perplexity di validazione più bassa (11,68) rispetto allo scaling Frobenius (11,76), validando l'analisi teorica.
• Precisione:
  • L'implementazione supporta FP16 per CN ed EVD senza degradare la perplexity, riducendo ulteriormente i tempi di esecuzione.
  • È stato riscontrato che NDB in FP16/BF16 è numericamente instabile nei loro esperimenti, sebbene gli autori lo indichino come un'area per lavori futuri.

Significato e Rivendicazioni

Il documento sostiene che DASH riduce significativamente la barriera computazionale all'uso di ottimizzatori di secondo ordine come Shampoo, rendendo pratico l'uso di precondizionamenti di alta qualità su larga scala.

• Colmare il Divario tra Teoria e Hardware: Allineando il design algoritmico (solver iterativi, scaling delle matrici) con le capacità dell'hardware (tensor core, operazioni batch), DASH consente ai ricercatori di investigare le proprietà degli ottimizzatori senza tempi di esecuzione proibitivi.
• Stabilità Numerica vs Velocità: Il lavoro sfida l'assunto che i metodi iterativi siano intrinsecamente meno stabili o più lenti dell'EVD. Introducendo un migliore scaling (Multi-Power-Iteration) e iterazioni ottimizzate (NDB), gli autori dimostrano che i metodi iterativi possono raggiungere una convergenza superiore e una perplexity inferiore pur essendo significativamente più veloci.
• Direzioni Future: Gli autori suggeriscono che le prestazioni superiori di Newton-DB aprono la strada alla selezione dinamica del solver basata sui numeri di condizionamento dei blocchi e al potenziale di stabilizzare NDB in formati a precisione inferiore utilizzando metodi stocastici o di correzione dell'errore.

Il documento conclude che il collo di bottiglia primario di Shampoo non era la complessità matematica, ma l'esecuzione frammentata e l'instabilità numerica, entrambi aspetti che DASH affronta efficacemente.