Per anni la spiegazione ufficiale della lentezza dei Transformer sui testi lunghi è stata una sola: l’attention costa il quadrato della lunghezza della frase. Raddoppi le parole, quadruplichi il lavoro. Vero, ma incompleto — e l’incompletezza nascondeva un’occasione. Nel 2022 Tri Dao, insieme a colleghi di Stanford e della State University di New York a Buffalo, pubblicò un paper che spostava lo sguardo dal quanto si calcola al quanto si sposta. Il titolo mette subito la parola chiave in chiaro: FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. La tesi è che l’attention, sulle GPU vere, non è limitata dai calcoli. È limitata dal traffico di dati dentro la scheda. E una volta capito questo, si può renderla molto piu veloce senza toccare di una virgola la matematica.
Perché il problema dell’attenzione è la memoria e non i conti, e come il calcolo a tessere lo aggira: nel libro.
Leggi «FlashAttention: l’attenzione che non spreca memoria» nel libro →
Il collo di bottiglia è il traffico, non il conto
Dentro una GPU convivono due tipi di memoria molto diversi. C’è la HBM, ampia — decine di gigabyte — dove sta tutto: i dati, i pesi, i risultati intermedi. E c’è la SRAM, la memoria on-chip, vicinissima alle unità di calcolo: velocissima ma minuscola, dell’ordine dei kilobyte per blocco. Il rapporto è quello tra un magazzino sterminato in periferia e un piano di lavoro accanto alla scrivania. Puoi tenere tutto in magazzino, ma ogni volta che ti serve qualcosa devi mandarci un furgone, e il furgone è lento.
Qui sta il problema dell’attention classica. Per calcolare a quali parole guardare, costruisce una matrice enorme di punteggi: una casella per ogni coppia di parole. Con mille parole sono un milione di caselle; con ottomila, sessantaquattro milioni. Questa matrice è troppo grande per la SRAM, quindi viene scritta per intero nella HBM, riletta per applicare la softmax, riscritta, e riletta di nuovo per combinare i valori. Il calcolo vero e proprio — le moltiplicazioni — le GPU lo sbrigano in un lampo. È il continuo andirivieni di quella matrice tra magazzino e scrivania a mangiarsi il tempo. In gergo si dice che l’operazione è memory-bound: a frenarla non sono le operazioni aritmetiche, ma le letture e scritture in memoria.
Il calcolo era gia veloce. Era il pendolarismo dei dati tra le due memorie a costare caro.
L’idea: non scrivere mai la matrice intera
La proposta di FlashAttention è tanto semplice da enunciare quanto delicata da realizzare: non materializzare mai la matrice completa dei punteggi nella HBM. Invece di calcolare tutti i punteggi, poi tutta la softmax, poi tutta la combinazione, si lavora a blocchi (in inglese tiling). Si carica in SRAM un pezzetto delle query, un pezzetto delle key e dei value, si fa tutto il lavoro su quel riquadro — punteggi, pesatura, accumulo nel risultato — e si scrive nella HBM soltanto il pezzo di output finale. La matrice gigante non viene mai scritta per intero da nessuna parte: esiste solo un riquadro alla volta, e vive e muore dentro la SRAM.
L’immagine è quella di uno chef con un piano di lavoro piccolo. Non può stendere tutti gli ingredienti insieme: prepara una porzione, la manda in forno, pulisce il piano, passa alla successiva. Il piatto finale è identico a quello che uscirebbe da una cucina enorme — stessi ingredienti, stesse dosi — ma non ha mai avuto bisogno di un piano grande quanto la sala. Per questo il paper insiste su una parola: exact. FlashAttention non approssima, non taglia, non salta coppie di parole. Produce esattamente lo stesso risultato dell’attention classica, bit per bit nella logica del calcolo. Cambia solo il percorso dei dati.
Le matrici Q, K e V restano nella HBM ampia ma lenta. Un blocco alla volta viene portato nella SRAM, dove si calcolano punteggi, softmax e accumulo dell’output; solo l’output torna in HBM. La matrice completa dei punteggi non viene mai scritta per intero, ed è da qui che nasce il risparmio.
Il trucco della softmax che non ha bisogno di vedere tutto
C’è un ostacolo che rende l’idea meno banale di quanto sembri. La softmax — la funzione che trasforma i punteggi in pesi che sommano a uno — per definizione ha bisogno dell’intera riga: per normalizzare deve conoscere il totale, e per essere numericamente stabile deve conoscere il valore massimo. Ma se lavoriamo un blocco alla volta, un blocco non vede l’intera riga. Come si fa una somma di cose che non abbiamo ancora tutte davanti?
La risposta è la softmax online: si tiene un totale corrente e un massimo corrente, e a ogni nuovo blocco li si aggiorna, correggendo all’indietro quanto accumulato finora. È lo stesso principio con cui si calcola una media camminando: non serve aspettare l’ultimo dato per avere una stima, basta aggiornarla a ogni passo. A fine giro, il risultato è identico a quello che si otterrebbe avendo visto tutta la riga in una volta. FlashAttention aggiunge un secondo accorgimento nella fase di apprendimento all’indietro: invece di conservare in memoria la matrice dei punteggi per riusarla, la ricalcola al volo. Sembra uno spreco — rifare un conto già fatto — ma sulle GPU rifare una moltiplicazione veloce costa meno che andare a ripescare un dato lento dalla HBM. È il capovolgimento che regge tutto il paper.
Cosa mostravano i numeri
I risultati sono interessanti proprio perché non riguardano la qualità del modello — che resta identica, trattandosi di attention esatta — ma la sua economia. Nell’addestramento, FlashAttention rendeva BERT-large piu veloce del 15% rispetto al record MLPerf 1.1 dell’epoca, accelerava GPT-2 di 3 volte e i benchmark Long-Range Arena di 2,4 volte. Il consumo di memoria, che nell’attention classica cresce col quadrato della lunghezza, diventava lineare: un cambiamento di regime, non una limatura. Ed è questo a sbloccare i risultati piu vistosi. Su Path-X, un compito con sequenze da 16.000 elementi, i Transformer erano fino ad allora fermi al caso; con FlashAttention ottenevano 61,4% di accuratezza. La variante block-sparse spingeva fino a Path-256, sequenze da 64.000 elementi, con il 63,1%. Sequenze che prima erano semplicemente fuori portata diventavano trattabili.
Perché rileggerlo oggi
Se oggi diamo per scontato di poter incollare un intero PDF, un contratto lungo o un’ora di trascrizione dentro un assistente, lo dobbiamo in buona parte a questo filone. Le finestre di contesto da centomila token e oltre, diventate normali nei modelli piu recenti, poggiano su varianti e discendenti di FlashAttention: senza rendere l’attention lineare in memoria, quelle lunghezze resterebbero economicamente proibitive. Non è un caso che la tecnica sia entrata in fretta nelle librerie standard su cui si addestra e si serve praticamente ogni modello.
Ma la lezione piu duratura è metodologica, e va oltre l’attention. Per anni l’ottimizzazione del deep learning ha guardato quasi solo al numero di operazioni. Dao e colleghi hanno ricordato che un algoritmo non gira nel vuoto: gira su un chip con una gerarchia di memoria precisa, e ignorarla lascia performance sul tavolo. Progettare tenendo conto di dove vivono i dati — non solo di quanti conti si fanno — è un principio che vale ben oltre questo paper. A volte la strada per andare piu veloci non è calcolare di meno. È smettere di spostare i dati avanti e indietro.
Il paper
Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022. Disponibile su arXiv: arxiv.org/abs/2205.14135.
I commenti sono riservati agli iscritti.
Accedi per commentare