IT
FlashSO(2) accelera gli strati che preservano la simmetria con l’emulazione intera
Accelerating SO(2)-Equivariant Materials Models through Low-Precision Emulation · arXiv:2610.35164
In parole semplici
Alcuni modelli per i materiali sono progettati in modo che, ruotando una struttura, ruotino coerentemente anche le loro previsioni. Ma lo strato che gestisce questa simmetria rotazionale richiede una parte consistente del calcolo, e i normali calcoli a precisione ridotta possono compromettere tale coerenza.
L’articolo presenta FlashSO(2), che suddivide i calcoli delle matrici in piccoli calcoli interi e combina operazioni vicine per risparmiare tempo e memoria. Nei test sugli output di MALOQ e UMA, la configurazione Emu 2 ha rispettato il limite dichiarato per l’errore rotazionale; il risultato riguarda soltanto gli output verificati.
La ricerca
I modelli per i materiali che mantengono coerenti le previsioni al variare della rotazione svolgono gran parte dei calcoli in uno strato costoso, mentre i normali calcoli a precisione ridotta possono introdurre errori dipendenti dall’orientamento.
Questo strato può usare hardware più veloce per i calcoli interi e preservare comunque la simmetria su cui si basano questi modelli?
FlashSO(2) emula prodotti di matrici precisi su unità intere veloci e fonde le operazioni SO(2) circostanti per accelerare modelli per i materiali che preservano la simmetria.
Gli autori hanno verificato se lo strato SO(2) potesse usare hardware per il calcolo di matrici intere senza perdere l’equivarianza rotazionale, la proprietà per cui gli output si trasformano coerentemente quando ruotano gli input.
Hanno misurato il massimo errore degli output al variare degli angoli di rotazione per gli hamiltoniani di MALOQ e le forze di UMA; per entrambi i risultati considerati, Emu 2 è rimasta entro la tolleranza di 10⁻⁴.
Hanno poi confrontato lo strato con il riferimento FP32 e misurato carichi di lavoro a livello di modello. Emu 2 fusa ha aumentato il throughput dello strato in entrambe le configurazioni testate; anche l’addestramento di MALOQ e l’inferenza di UMA-M sono stati più veloci dei rispettivi confronti FP32, nelle condizioni specificate.
Nel complesso, questi risultati mostrano che l’approccio può preservare la simmetria verificata e migliorare le prestazioni di questi modelli.
I risultati provengono da modelli rappresentativi MALOQ e UMA; la maggior parte delle valutazioni è stata eseguita su un unico superchip NVIDIA GH200.
Gli aumenti di velocità riportati per i modelli valgono per specifiche dimensioni di input e configurazioni, quindi non dimostrano che gli stessi miglioramenti si ottengano con altri modelli, hardware o carichi di lavoro.
Cosa mostrano le prove
- Velocità e simmetria
FlashSO(2) è un’implementazione progettata congiuntamente per software e hardware dello strato SO(2) nei modelli di materiali basati su reti neurali a grafo. Questi modelli preservano l’equivarianza rotazionale: gli output cambiano in modo coerente quando ruota l’input.
In molti modelli di questo tipo, lo strato rappresenta gran parte del tempo di inferenza, ma la precisione ridotta dei formati nativi può introdurre errori dipendenti dalla rotazione e limitare l’uso di hardware più veloce per il calcolo delle matrici.
- Come funziona FlashSO(2)
Per affrontare questo problema, FlashSO(2) emula i prodotti di matrici usando sezioni INT8 riscalate, accumulandone esattamente i risultati in INT32 e ricostruendo poi il risultato. La configurazione Emu 2 usa due sezioni.
Il ridimensionamento a blocchi assegna una scala a ciascun gruppo di valori, aiutando a mantenere la precisione delle attivazioni piccole e a evitare overflow nelle somme lunghe. Inoltre, l’implementazione fonde le operazioni attorno ai prodotti di matrici, mantenendo i risultati intermedi nei registri invece di spostarli tra kernel separati.
- Verifica della simmetria
Gli autori hanno verificato se l’implementazione preservava l’equivarianza rotazionale misurando il massimo scarto tra le singole componenti al variare degli angoli di rotazione. Hanno confrontato gli output hamiltoniani di MALOQ e gli output di forza di UMA con una tolleranza di 10⁻⁴.
Tra i formati nativi, solo FP32 rispettava quel limite; TF32 restava appena fuori, mentre Emu 2 lo rispettava per entrambi i risultati considerati. Emu 3 era indistinguibile da FP32.
- Da dove arrivano i miglioramenti di velocità
Un confronto in due fasi sullo strato ha separato l’effetto della fusione dei kernel da quello dell’emulazione intera: gli autori hanno confrontato lo strato di riferimento FP32 con la versione FP32 fusa, poi FP32 fusa con Emu 2 fusa.
Rispetto al riferimento, Emu 2 fusa ha raggiunto un throughput pari a 40x in MALOQ e 7x in UMA-L su GH200. Sono risultati del singolo strato per le configurazioni testate, non miglioramenti dell’intero modello.
- Verifica dei miglioramenti nei modelli
Gli autori hanno misurato anche le prestazioni a livello di modello. Nell’addestramento di MALOQ con torch.compile, Emu 2 ha reso i passi fino a 3.93x più veloci alla dimensione di input massima che rientrava nella memoria FP32 e ha ridotto la memoria di picco del 58% alla dimensione per cui FP32 esauriva la memoria.
Nell’inferenza di UMA-M, ha eseguito un input specificato di rame fcc 3.33x più velocemente di FP32. Questi risultati sostengono la presenza di miglioramenti oltre il singolo strato nelle condizioni testate.
Dettagli
- Autori del paper originale
- Denghui Lu, Alexander Maeder, Alexandros Nikolaos Ziogas, Mathieu Luisier, Manasa Kaniselvan
- Inviato ad arXiv
- Licenza
- arXiv non-exclusive