FlashFFTConv: Efficiënte Convoluties voor Lange Sequenties met Tensor Cores

FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores

November 10, 2023

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

cs.AI

Samenvatting

Convolutiemodellen met lange filters hebben state-of-the-art redeneervaardigheden aangetoond in veel taken met lange sequenties, maar blijven achter bij de meest geoptimaliseerde Transformers in wall-clock tijd. Een belangrijke bottleneck is de Fast Fourier Transform (FFT)—die lange convoluties mogelijk maakt in O(N logN) tijd in sequentielengte N, maar een slechte hardwarebenutting heeft. In dit artikel onderzoeken we hoe de FFT-convolutie kan worden geoptimaliseerd. We identificeren twee belangrijke knelpunten: de FFT maakt niet effectief gebruik van gespecialiseerde matrixvermenigvuldigingsunits, en het veroorzaakt dure I/O tussen lagen van de geheugenhiërarchie. Als antwoord hierop stellen we FlashFFTConv voor. FlashFFTConv gebruikt een matrixdecompositie die de FFT berekent met matrixvermenigvuldigingsunits en maakt kernel-fusie mogelijk voor lange sequenties, waardoor I/O wordt verminderd. We presenteren ook twee sparse convolutie-algoritmen—1) partiële convoluties en 2) frequentie-sparse convoluties—die eenvoudig kunnen worden geïmplementeerd door blokken in de matrixdecompositie over te slaan, wat verdere mogelijkheden biedt voor geheugen- en rekentijdbesparingen. FlashFFTConv versnelt exacte FFT-convoluties tot 7,93 keer ten opzichte van PyTorch en behaalt tot 4,4 keer snelheidswinst end-to-end. Met hetzelfde rekenbudget stelt FlashFFTConv Hyena-GPT-s in staat om 2,3 punten betere perplexiteit te behalen op de PILE en M2-BERT-base om 3,3 punten hogere GLUE-score te behalen—wat overeenkomt met modellen met het dubbele aantal parameters. FlashFFTConv behaalt ook 96,1% nauwkeurigheid op Path-512, een hoogresolutie visietask waarbij geen enkel model eerder beter dan 50% had behaald. Bovendien maken partiële convoluties langere-sequentiemodellen mogelijk—wat resulteert in het eerste DNA-model dat de langste menselijke genen kan verwerken (2,3 miljoen baseparen)—en frequentie-sparse convoluties versnellen voorgetrainde modellen terwijl de modelkwaliteit behouden blijft of verbetert.

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.