FlashFFTConv: Convoluciones Eficientes para Secuencias Largas con Tensor Cores

FlashFFTConv: Efficient Convolutions for Long Sequences with Tensor Cores

November 10, 2023

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

cs.AI

Resumen

Los modelos de convolución con filtros largos han demostrado capacidades de razonamiento de vanguardia en muchas tareas de secuencias largas, pero se quedan atrás de los Transformers más optimizados en términos de tiempo de ejecución. Un cuello de botella importante es la Transformada Rápida de Fourier (FFT), que permite que las convoluciones largas se ejecuten en tiempo O(N logN) para una secuencia de longitud N, pero tiene una mala utilización del hardware. En este artículo, estudiamos cómo optimizar la convolución FFT. Identificamos dos cuellos de botella clave: la FFT no utiliza eficientemente las unidades especializadas de multiplicación de matrices, y genera costosas operaciones de E/S entre las capas de la jerarquía de memoria. En respuesta, proponemos FlashFFTConv. FlashFFTConv utiliza una descomposición matricial que calcula la FFT empleando unidades de multiplicación de matrices y permite la fusión de kernels para secuencias largas, reduciendo las operaciones de E/S. También presentamos dos algoritmos de convolución dispersa: 1) convoluciones parciales y 2) convoluciones esparcidas en frecuencia, que pueden implementarse simplemente omitiendo bloques en la descomposición matricial, lo que ofrece oportunidades adicionales para ahorrar memoria y cómputo. FlashFFTConv acelera las convoluciones FFT exactas hasta 7.93 veces en comparación con PyTorch y logra una aceleración de hasta 4.4 veces en el proceso completo. Con el mismo presupuesto de cómputo, FlashFFTConv permite que Hyena-GPT-s alcance 2.3 puntos mejor en perplejidad en PILE y que M2-BERT-base obtenga 3.3 puntos más en la puntuación GLUE, igualando modelos con el doble de parámetros. FlashFFTConv también logra un 96.1% de precisión en Path-512, una tarea de visión de alta resolución donde ningún modelo había superado previamente el 50%. Además, las convoluciones parciales permiten modelos de secuencias más largas, dando lugar al primer modelo de ADN que puede procesar los genes humanos más largos (2.3 millones de pares de bases), y las convoluciones esparcidas en frecuencia aceleran modelos preentrenados manteniendo o mejorando la calidad del modelo.

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.