FlashFFTConv: Convoluzioni Efficienti per Sequenze Lunghe con Tensor Core

FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores

November 10, 2023

Autori: Daniel Y. Fu, Hermann Kumbong, Eric Nguyen, Christopher Ré

cs.AI

Abstract

I modelli convoluzionali con filtri lunghi hanno dimostrato capacità di ragionamento all'avanguardia in molti compiti che coinvolgono sequenze lunghe, ma rimangono indietro rispetto ai Transformer più ottimizzati in termini di tempo di esecuzione. Un collo di bottiglia principale è la Trasformata Rapida di Fourier (FFT), che consente alle convoluzioni lunghe di essere eseguite in tempo O(N logN) rispetto alla lunghezza della sequenza N, ma ha una scarsa utilizzazione hardware. In questo articolo, studiamo come ottimizzare la convoluzione FFT. Identifichiamo due colli di bottiglia chiave: la FFT non utilizza efficacemente le unità specializzate per la moltiplicazione di matrici e comporta costose operazioni di I/O tra i livelli della gerarchia di memoria. In risposta, proponiamo FlashFFTConv. FlashFFTConv utilizza una decomposizione matriciale che calcola la FFT utilizzando le unità di moltiplicazione di matrici e abilita la fusione dei kernel per sequenze lunghe, riducendo l'I/O. Presentiamo inoltre due algoritmi di convoluzione sparsa: 1) convoluzioni parziali e 2) convoluzioni sparse in frequenza, che possono essere implementati semplicemente saltando blocchi nella decomposizione matriciale, offrendo ulteriori opportunità di risparmio di memoria e calcolo. FlashFFTConv accelera le convoluzioni FFT esatte fino a 7,93 volte rispetto a PyTorch e raggiunge un miglioramento end-to-end fino a 4,4 volte. Con lo stesso budget di calcolo, FlashFFTConv consente a Hyena-GPT-s di ottenere 2,3 punti in meno di perplessità su PILE e a M2-BERT-base di ottenere 3,3 punti in più nel punteggio GLUE, eguagliando modelli con il doppio del numero di parametri. FlashFFTConv raggiunge anche il 96,1% di accuratezza su Path-512, un compito di visione ad alta risoluzione in cui nessun modello aveva precedentemente superato il 50%. Inoltre, le convoluzioni parziali abilitano modelli per sequenze più lunghe, producendo il primo modello di DNA in grado di elaborare i geni umani più lunghi (2,3 milioni di coppie di basi), e le convoluzioni sparse in frequenza accelerano i modelli pre-addestrati mantenendo o migliorando la qualità del modello.

English

Convolution models with long filters have demonstrated state-of-the-art reasoning abilities in many long-sequence tasks but lag behind the most optimized Transformers in wall-clock time. A major bottleneck is the Fast Fourier Transform (FFT)--which allows long convolutions to run in O(N logN) time in sequence length N but has poor hardware utilization. In this paper, we study how to optimize the FFT convolution. We find two key bottlenecks: the FFT does not effectively use specialized matrix multiply units, and it incurs expensive I/O between layers of the memory hierarchy. In response, we propose FlashFFTConv. FlashFFTConv uses a matrix decomposition that computes the FFT using matrix multiply units and enables kernel fusion for long sequences, reducing I/O. We also present two sparse convolution algorithms--1) partial convolutions and 2) frequency-sparse convolutions--which can be implemented simply by skipping blocks in the matrix decomposition, enabling further opportunities for memory and compute savings. FlashFFTConv speeds up exact FFT convolutions by up to 7.93times over PyTorch and achieves up to 4.4times speedup end-to-end. Given the same compute budget, FlashFFTConv allows Hyena-GPT-s to achieve 2.3 points better perplexity on the PILE and M2-BERT-base to achieve 3.3 points higher GLUE score--matching models with twice the parameter count. FlashFFTConv also achieves 96.1% accuracy on Path-512, a high-resolution vision task where no model had previously achieved better than 50%. Furthermore, partial convolutions enable longer-sequence models--yielding the first DNA model that can process the longest human genes (2.3M base pairs)--and frequency-sparse convolutions speed up pretrained models while maintaining or improving model quality.