Vai al contenuto
AI.info

The Pulse

TurboPairFormer accelera fino a 1,73× l’addestramento di OpenFold3 nei test

I ricercatori riferiscono che TurboPairFormer ha ridotto il tempo di calcolo su GPU per l’addestramento di OpenFold3 rispetto al backend Triton originale. Nelle condizioni testate, i kernel open source hanno inoltre prodotto gradienti più accurati e riproducibili.

TurboPairFormer accelera fino a 1,73× l’addestramento di OpenFold3 nei test

Secondo i ricercatori, in un preprint TurboPairFormer ha accelerato fino a 1,73× l’addestramento di OpenFold3, misurato in tempo di calcolo su GPU per passaggio dell’ottimizzatore, rispetto al backend Triton originale del modello. I kernel open source sono stati sviluppati da ricercatori di Lambda, dello Stevens Institute of Technology e del team OpenFold per accelerare le operazioni usate nella previsione della struttura biomolecolare.

Lambda ha descritto i risultati in un post del 13 novembre, successivo al preprint datato 5 ottobre. I ricercatori riferiscono anche che nei test i gradienti erano più accurati e i risultati riproducibili, sebbene queste conclusioni si applichino alle configurazioni specifiche che hanno valutato.

Dove si impiega il tempo

OpenFold3 supporta l’addestramento e il fine-tuning di modelli che prevedono strutture composte da proteine, acidi nucleici e piccole molecole. TurboPairFormer implementa l’attenzione triangolare e la moltiplicazione triangolare, operazioni usate in questi modelli.

L’attenzione triangolare permette a un modello di usare informazioni su una terza componente molecolare quando ragiona su una coppia. Secondo i ricercatori, il suo costo computazionale cresce in modo cubico con il numero di token: raddoppiare la dimensione del crop di addestramento rende l’operazione circa otto volte più costosa.

Durante il passaggio all’indietro, l’addestramento calcola i gradienti per query, chiavi, valori e bias di coppia. TurboPairFormer ricalcola una volta ogni tile delle probabilità di attenzione e lo riutilizza per tutti e quattro i gradienti, invece di ricalcolarlo più volte per i diversi gradienti.

Benchmark di addestramento

I ricercatori hanno integrato TurboPairFormer in OpenFold3 insieme ai propri kernel per la moltiplicazione triangolare e lo hanno confrontato con altre sette configurazioni di backend. I test sono stati eseguiti su 16 GPU NVIDIA H100 distribuite su due nodi, con dimensioni del crop di 384, 640 e 768 token.

Con una dimensione del crop di 768, TurboPairFormer ha impiegato 12,58 secondi di calcolo su GPU per passaggio dell’ottimizzatore, contro i 21,72 secondi di OpenFold3 Triton e i 14,25 secondi di NVIDIA cuEquivariance, secondo i risultati riportati. I ricercatori hanno rilevato che, tra le configurazioni testate, aveva il tempo di calcolo per passaggio più basso a tutte e tre le dimensioni del crop.

Le misurazioni hanno riguardato i passaggi dal 21 al 200 ed escluso il caricamento dei dati, i trasferimenti dall’host al dispositivo e gli aggiornamenti EMA. Su GPU H200 con una dimensione del crop di 768, Lambda ha riportato un’accelerazione di 1,54× rispetto a OpenFold3 Triton, con il tempo per passaggio sceso da 22,53 a 14,67 secondi.

Accuratezza e riproducibilità

I kernel affrontano anche l’errore di arrotondamento che si verifica quando i valori vengono convertiti dalla precisione FP32 a quella BF16. TurboPairFormer salva un’approssimazione del residuo di arrotondamento e la usa nel calcolo all’indietro senza modificare l’output BF16.

Secondo il preprint, in esperimenti controllati questa compensazione ha ridotto del 28–47% l’errore medio dei gradienti di query, chiave e bias di coppia. Rispetto a un riferimento FP64, tra le otto implementazioni dell’attenzione valutate l’implementazione ha registrato gli errori medi più bassi per quei gradienti; secondo Lambda, l’accuratezza del gradiente dei valori era sostanzialmente alla pari.

Il team ha anche verificato se chiamate ripetute producessero gradienti identici. Su una singola H100, nelle condizioni di test, tutti e quattro i gradienti sono risultati identici bit per bit in cinque esecuzioni per ciascuno dei 600 casi con input fisso.

Secondo i ricercatori, una prosecuzione dell’addestramento con una versione precedente dei kernel ha mantenuto finite le loss, mentre l’esecuzione di OpenFold3 Triton ha prodotto loss NaN entro 10 aggiornamenti. TurboPairFormer è disponibile su PyPI e attualmente supporta le GPU H100 e H200.

Fonti

Notizie correlate

Esplora

Altri articoli